Ensure we protect Collections obtained from scripts from self-referencing

modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/CustomReflectionObjectHandler.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.script.mustache;
 
 import com.github.mustachejava.reflect.ReflectionObjectHandler;
+import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.common.util.iterable.Iterables;
 
 import java.lang.reflect.Array;
@@ -154,4 +155,9 @@ public Iterator<Object> iterator() {
         }
     }
 
+    @Override
+    public String stringify(Object object) {
+        CollectionUtils.ensureNoSelfReferences(object);
+        return super.stringify(object);
+    }
 }

modules/lang-painless/src/test/resources/rest-api-spec/test/painless/15_update.yml

@@ -137,4 +137,4 @@
 
   - match: { error.root_cause.0.type: "remote_transport_exception" }
   - match: { error.type: "illegal_argument_exception" }
-  - match: { error.reason: "Object has already been built and is self-referencing itself" }
+  - match: { error.reason: "Iterable object is self-referencing itself" }

modules/lang-painless/src/test/resources/rest-api-spec/test/painless/30_search.yml

@@ -406,3 +406,39 @@
     - match: { hits.hits.0._score: 1.0 }
     - match: { aggregations.value_agg.buckets.0.key: 2 }
     - match: { aggregations.value_agg.buckets.0.doc_count: 1 }
+
+---
+"Return self-referencing map":
+    - do:
+        indices.create:
+            index: test
+            body:
+              settings:
+                number_of_shards: "1"
+
+    - do:
+        index:
+            index: test
+            type: test
+            id: 1
+            body: { "genre": 1 }
+
+    - do:
+        indices.refresh: {}
+
+    - do:
+        catch: bad_request
+        index: test
+        search:
+            body:
+                aggs:
+                    genre:
+                        terms:
+                            script:
+                                lang: painless
+                                source: "def x = [:] ; def y = [:] ; x.a = y ; y.a = x ; return x"
+
+    - match: { error.root_cause.0.type: "illegal_argument_exception" }
+    - match: { error.root_cause.0.reason: "Iterable object is self-referencing itself" }
+    - match: { error.type: "search_phase_execution_exception" }
+    - match: { error.reason: "all shards failed" }

server/src/main/java/org/elasticsearch/common/util/CollectionUtils.java

@@ -19,16 +19,20 @@
 
 package org.elasticsearch.common.util;
 
+import java.nio.file.Path;
 import java.util.AbstractList;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.IdentityHashMap;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.RandomAccess;
+import java.util.Set;
 
 import com.carrotsearch.hppc.ObjectArrayList;
 import org.apache.lucene.util.BytesRef;
@@ -221,6 +225,40 @@ public static int[] toArray(Collection<Integer> ints) {
         return ints.stream().mapToInt(s -> s).toArray();
     }
 
+    public static void ensureNoSelfReferences(Object value) {
+        Iterable<?> it = convert(value);
+        if (it != null) {
+            ensureNoSelfReferences(it, value, Collections.newSetFromMap(new IdentityHashMap<>()));
+        }
+    }
+
+    private static Iterable<?> convert(Object value) {
+        if (value == null) {
+            return null;
+        }
+        if (value instanceof Map) {
+            return ((Map<?,?>) value).values();
+        } else if ((value instanceof Iterable) && (value instanceof Path == false)) {
+            return (Iterable<?>) value;
+        } else if (value instanceof Object[]) {
+            return Arrays.asList((Object[]) value);
+        } else {
+            return null;
+        }
+    }
+
+    private static void ensureNoSelfReferences(final Iterable<?> value, Object originalReference, final Set<Object> ancestors) {
+        if (value != null) {
+            if (ancestors.add(originalReference) == false) {
+                throw new IllegalArgumentException("Iterable object is self-referencing itself");
+            }
+            for (Object o : value) {
+                ensureNoSelfReferences(convert(o), o, ancestors);
+            }
+            ancestors.remove(originalReference);
+        }
+    }
+
     private static class RotatedList<T> extends AbstractList<T> implements RandomAccess {
 
         private final List<T> in;

server/src/main/java/org/elasticsearch/common/xcontent/XContentBuilder.java

@@ -28,6 +28,7 @@
 import org.elasticsearch.common.text.Text;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.unit.TimeValue;
+import org.elasticsearch.common.util.CollectionUtils;
 import org.joda.time.DateTimeZone;
 import org.joda.time.ReadableInstant;
 import org.joda.time.format.DateTimeFormatter;
@@ -43,7 +44,6 @@
 import java.util.Collections;
 import java.util.Date;
 import java.util.HashMap;
-import java.util.IdentityHashMap;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
@@ -780,7 +780,6 @@ private XContentBuilder values(Object[] values, boolean ensureNoSelfReferences)
         if (values == null) {
             return nullValue();
         }
-
         return value(Arrays.asList(values), ensureNoSelfReferences);
     }
 
@@ -865,7 +864,7 @@ private XContentBuilder map(Map<String, ?> values, boolean ensureNoSelfReference
         // checks that the map does not contain references to itself because
         // iterating over map entries will cause a stackoverflow error
         if (ensureNoSelfReferences) {
-            ensureNoSelfReferences(values);
+            CollectionUtils.ensureNoSelfReferences(values);
         }
 
         startObject();
@@ -894,9 +893,8 @@ private XContentBuilder value(Iterable<?> values, boolean ensureNoSelfReferences
             // checks that the iterable does not contain references to itself because
             // iterating over entries will cause a stackoverflow error
             if (ensureNoSelfReferences) {
-                ensureNoSelfReferences(values);
+                CollectionUtils.ensureNoSelfReferences(values);
             }
-
             startArray();
             for (Object value : values) {
                 // pass ensureNoSelfReferences=false as we already performed the check at a higher level
@@ -1067,32 +1065,4 @@ static void ensureNotNull(Object value, String message) {
             throw new IllegalArgumentException(message);
         }
     }
-
-    static void ensureNoSelfReferences(Object value) {
-        ensureNoSelfReferences(value, Collections.newSetFromMap(new IdentityHashMap<>()));
-    }
-
-    private static void ensureNoSelfReferences(final Object value, final Set<Object> ancestors) {
-        if (value != null) {
-
-            Iterable<?> it;
-            if (value instanceof Map) {
-                it = ((Map<?,?>) value).values();
-            } else if ((value instanceof Iterable) && (value instanceof Path == false)) {
-                it = (Iterable<?>) value;
-            } else if (value instanceof Object[]) {
-                it = Arrays.asList((Object[]) value);
-            } else {
-                return;
-            }
-
-            if (ancestors.add(value) == false) {
-                throw new IllegalArgumentException("Object has already been built and is self-referencing itself");
-            }
-            for (Object o : it) {
-                ensureNoSelfReferences(o, ancestors);
-            }
-            ancestors.remove(value);
-        }
-    }
 }

server/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.search.aggregations.metrics.scripted;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.SearchScript;
@@ -77,6 +78,7 @@ public InternalAggregation buildAggregation(long owningBucketOrdinal) {
         Object aggregation;
         if (combineScript != null) {
             aggregation = combineScript.run();
+            CollectionUtils.ensureNoSelfReferences(aggregation);
         } else {
             aggregation = params.get("_agg");
         }

server/src/main/java/org/elasticsearch/search/aggregations/pipeline/bucketscript/BucketScriptPipelineAggregator.java

@@ -112,10 +112,11 @@ public InternalAggregation reduce(InternalAggregation aggregation, ReduceContext
             } else {
                 ExecutableScript executableScript = factory.newInstance(vars);
                 Object returned = executableScript.run();
+                // no need to check for self references since only numbers are valid
                 if (returned == null) {
                     newBuckets.add(bucket);
                 } else {
-                    if (!(returned instanceof Number)) {
+                    if ((returned instanceof Number) == false) {
                         throw new AggregationExecutionException("series_arithmetic script for reducer [" + name()
                                 + "] must return a Number");
                     }

server/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSource.java

@@ -30,6 +30,7 @@
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.lucene.ScorerAware;
+import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.index.fielddata.AbstractSortingNumericDocValues;
 import org.elasticsearch.index.fielddata.AtomicOrdinalsFieldData;
 import org.elasticsearch.index.fielddata.IndexFieldData;
@@ -460,7 +461,9 @@ public boolean advanceExact(int doc) throws IOException {
                     for (int i = 0; i < count; ++i) {
                         final BytesRef value = bytesValues.nextValue();
                         script.setNextAggregationValue(value.utf8ToString());
-                        values[i].copyChars(script.run().toString());
+                        Object run = script.run();
+                        CollectionUtils.ensureNoSelfReferences(run);
+                        values[i].copyChars(run.toString());
                     }
                     sort();
                     return true;

server/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptBytesValues.java

@@ -20,6 +20,7 @@
 
 import org.apache.lucene.search.Scorer;
 import org.elasticsearch.common.lucene.ScorerAware;
+import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
 import org.elasticsearch.index.fielddata.SortingBinaryDocValues;
 import org.elasticsearch.script.SearchScript;
@@ -44,6 +45,7 @@ private void set(int i, Object o) {
         if (o == null) {
             values[i].clear();
         } else {
+            CollectionUtils.ensureNoSelfReferences(o);
             values[i].copyChars(o.toString());
         }
     }

server/src/main/java/org/elasticsearch/search/fetch/subphase/ScriptFieldsFetchSubPhase.java

@@ -22,6 +22,7 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.ReaderUtil;
 import org.elasticsearch.common.document.DocumentField;
+import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.fetch.FetchSubPhase;
@@ -64,6 +65,7 @@ public void hitsExecute(SearchContext context, SearchHit[] hits) throws IOExcept
                 final Object value;
                 try {
                     value = leafScripts[i].run();
+                    CollectionUtils.ensureNoSelfReferences(value);
                 } catch (RuntimeException e) {
                     if (scriptFields.get(i).ignoreException()) {
                         continue;

server/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java

@@ -32,6 +32,7 @@
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.logging.DeprecationLogger;
 import org.elasticsearch.common.logging.Loggers;
+import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.common.xcontent.ConstructingObjectParser;
 import org.elasticsearch.common.xcontent.ObjectParser.ValueType;
 import org.elasticsearch.common.xcontent.XContentBuilder;
@@ -341,7 +342,9 @@ public boolean advanceExact(int doc) throws IOException {
                             }
                             @Override
                             public BytesRef binaryValue() {
-                                spare.copyChars(leafScript.run().toString());
+                                final Object run = leafScript.run();
+                                CollectionUtils.ensureNoSelfReferences(run);
+                                spare.copyChars(run.toString());
                                 return spare.get();
                             }
                         };

server/src/test/java/org/elasticsearch/common/util/CollectionUtilsTests.java

@@ -25,16 +25,21 @@
 import org.apache.lucene.util.Counter;
 import org.elasticsearch.test.ESTestCase;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.SortedSet;
 import java.util.TreeSet;
 
+import static java.util.Collections.emptyMap;
 import static org.elasticsearch.common.util.CollectionUtils.eagerPartition;
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
 
@@ -176,4 +181,15 @@ public void testPerfectPartition() {
                 eagerPartition(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12), 6)
         );
     }
+
+    public void testEnsureNoSelfReferences() {
+        CollectionUtils.ensureNoSelfReferences(emptyMap());
+        CollectionUtils.ensureNoSelfReferences(null);
+
+        Map<String, Object> map = new HashMap<>();
+        map.put("field", map);
+
+        IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () ->  CollectionUtils.ensureNoSelfReferences(map));
+        assertThat(e.getMessage(), containsString("Iterable object is self-referencing itself"));
+    }
 }