From a5a6a847b498b4b38b575a679d72ada4278b7dae Mon Sep 17 00:00:00 2001
From: Simon Cooper <simon.cooper@elastic.co>
Date: Mon, 9 Sep 2024 11:33:54 +0100
Subject: [PATCH] Add out-of-order object/array checks

---
 .../xcontent/ChunkedToXContentBuilder.java    | 44 +++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/server/src/main/java/org/elasticsearch/common/xcontent/ChunkedToXContentBuilder.java b/server/src/main/java/org/elasticsearch/common/xcontent/ChunkedToXContentBuilder.java
index 8fccfa2b26288..64bbb0740aa7e 100644
--- a/server/src/main/java/org/elasticsearch/common/xcontent/ChunkedToXContentBuilder.java
+++ b/server/src/main/java/org/elasticsearch/common/xcontent/ChunkedToXContentBuilder.java
@@ -9,11 +9,14 @@
 package org.elasticsearch.common.xcontent;
 
 import org.elasticsearch.common.CheckedBiConsumer;
+import org.elasticsearch.core.Assertions;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.Deque;
 import java.util.Iterator;
 import java.util.Map;
 import java.util.function.BiConsumer;
@@ -26,10 +29,14 @@
  */
 public class ChunkedToXContentBuilder implements Iterator<ToXContent> {
 
+    private enum ElementType { OBJECT, ARRAY };
+
     private final ToXContent.Params params;
     private final Stream.Builder<ToXContent> builder = Stream.builder();
     private Iterator<ToXContent> iterator;
 
+    private final Deque<ElementType> elementTracker = Assertions.ENABLED ? new ArrayDeque<>() : null;
+
     public ChunkedToXContentBuilder(ToXContent.Params params) {
         this.params = params;
     }
@@ -40,16 +47,25 @@ private void addChunk(ToXContent content) {
     }
 
     public ChunkedToXContentBuilder startObject() {
+        if (elementTracker != null) {
+            elementTracker.push(ElementType.OBJECT);
+        }
         addChunk((b, p) -> b.startObject());
         return this;
     }
 
     public ChunkedToXContentBuilder startObject(String name) {
+        if (elementTracker != null) {
+            elementTracker.push(ElementType.OBJECT);
+        }
         addChunk((b, p) -> b.startObject(name));
         return this;
     }
 
     public ChunkedToXContentBuilder endObject() {
+        if (elementTracker != null && elementTracker.poll() != ElementType.OBJECT) {
+            throw new IllegalStateException("Out-of-order object close");
+        }
         addChunk((b, p) -> b.endObject());
         return this;
     }
@@ -59,16 +75,25 @@ public ChunkedToXContentBuilder object(String name, Consumer<ChunkedToXContentBu
     }
 
     public ChunkedToXContentBuilder startArray() {
+        if (elementTracker != null) {
+            elementTracker.push(ElementType.ARRAY);
+        }
         addChunk((b, p) -> b.startArray());
         return this;
     }
 
     public ChunkedToXContentBuilder startArray(String name) {
+        if (elementTracker != null) {
+            elementTracker.push(ElementType.ARRAY);
+        }
         addChunk((b, p) -> b.startArray(name));
         return this;
     }
 
     public ChunkedToXContentBuilder endArray() {
+        if (elementTracker != null && elementTracker.poll() != ElementType.ARRAY) {
+            throw new IllegalStateException("Out-of-order array close");
+        }
         addChunk((b, p) -> b.endArray());
         return this;
     }
@@ -109,11 +134,18 @@ public ChunkedToXContentBuilder execute(Consumer<ChunkedToXContentBuilder> consu
         return this;
     }
 
+    /**
+     * Adds chunks from an iterator. Each item is passed to {@code create} to add chunks to this builder.
+     */
     public <T> ChunkedToXContentBuilder forEach(Iterator<T> items, BiConsumer<? super T, ChunkedToXContentBuilder> create) {
         items.forEachRemaining(t -> create.accept(t, this));
         return this;
     }
 
+    /**
+     * Adds chunks from an iterator. Each item is passed to {@code create}, and the resulting {@code ToXContent}-like objects
+     * are added to this builder in order.
+     */
     public <T> ChunkedToXContentBuilder forEach(
         Iterator<T> items,
         Function<? super T, CheckedBiConsumer<XContentBuilder, ToXContent.Params, IOException>> create
@@ -122,14 +154,23 @@ public <T> ChunkedToXContentBuilder forEach(
         return this;
     }
 
+    /**
+     * Each entry in {@code map} is added to the builder as a separate field, named with the entry key.
+     */
     public <T> ChunkedToXContentBuilder appendEntries(Map<String, ?> map) {
         return forEach(map.entrySet().iterator(), (e, b) -> b.field(e.getKey(), e.getValue()));
     }
 
+    /**
+     * Each entry in {@code map} is added to the builder as a separate object, named with the entry key.
+     */
     public <T> ChunkedToXContentBuilder appendXContentObjects(Map<String, ? extends ToXContent> map) {
         return forEach(map.entrySet().iterator(), (e, b) -> b.startObject(e.getKey()).appendXContent(e.getValue()).endObject());
     }
 
+    /**
+     * Each entry in {@code map} is added to the builder as a separate XContent field, named with the entry key.
+     */
     public <T> ChunkedToXContentBuilder appendXContentFields(Map<String, ? extends ToXContent> map) {
         return forEach(map.entrySet().iterator(), (e, b) -> b.field(e.getKey()).appendXContent(e.getValue()));
     }
@@ -183,6 +224,9 @@ public ChunkedToXContentBuilder appendIfPresent(@Nullable ChunkedToXContent chun
 
     private Iterator<ToXContent> checkCreateIterator() {
         if (iterator == null) {
+            if (elementTracker != null && elementTracker.isEmpty() == false) {
+                throw new IllegalStateException("Unclosed XContent elements present: " + elementTracker);
+            }
             iterator = builder.build().iterator();
         }
         return iterator;