Call ensureNoSelfReferences() on _agg state variable after scripted m…

…etric agg script executions (#31044)

Previously this was called for the combine script only. This change checks for self references for
init, map, and reduce scripts as well, and adds unit test coverage for the init, map, and combine cases.

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

@@ -157,7 +157,7 @@ public Iterator<Object> iterator() {
 
     @Override
     public String stringify(Object object) {
-        CollectionUtils.ensureNoSelfReferences(object);
+        CollectionUtils.ensureNoSelfReferences(object, "CustomReflectionObjectHandler stringify");
         return super.stringify(object);
     }
 }

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

@@ -29,6 +29,7 @@
 import java.util.IdentityHashMap;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.RandomAccess;
@@ -40,6 +41,7 @@
 import org.apache.lucene.util.BytesRefBuilder;
 import org.apache.lucene.util.InPlaceMergeSorter;
 import org.apache.lucene.util.IntroSorter;
+import org.elasticsearch.common.Strings;
 
 /** Collections-related utility methods. */
 public class CollectionUtils {
@@ -225,10 +227,17 @@ public static int[] toArray(Collection<Integer> ints) {
         return ints.stream().mapToInt(s -> s).toArray();
     }
 
-    public static void ensureNoSelfReferences(Object value) {
+    /**
+     * Deeply inspects a Map, Iterable, or Object array looking for references back to itself.
+     * @throws IllegalArgumentException if a self-reference is found
+     * @param value The object to evaluate looking for self references
+     * @param messageHint A string to be included in the exception message if the call fails, to provide
+     *                    more context to the handler of the exception
+     */
+    public static void ensureNoSelfReferences(Object value, String messageHint) {
         Iterable<?> it = convert(value);
         if (it != null) {
-            ensureNoSelfReferences(it, value, Collections.newSetFromMap(new IdentityHashMap<>()));
+            ensureNoSelfReferences(it, value, Collections.newSetFromMap(new IdentityHashMap<>()), messageHint);
         }
     }
 
@@ -247,13 +256,15 @@ private static Iterable<?> convert(Object value) {
         }
     }
 
-    private static void ensureNoSelfReferences(final Iterable<?> value, Object originalReference, final Set<Object> ancestors) {
+    private static void ensureNoSelfReferences(final Iterable<?> value, Object originalReference, final Set<Object> ancestors,
+                                               String messageHint) {
         if (value != null) {
             if (ancestors.add(originalReference) == false) {
-                throw new IllegalArgumentException("Iterable object is self-referencing itself");
+                String suffix = Strings.isNullOrEmpty(messageHint) ? "" : String.format(Locale.ROOT, " (%s)", messageHint);
+                throw new IllegalArgumentException("Iterable object is self-referencing itself" + suffix);
             }
             for (Object o : value) {
-                ensureNoSelfReferences(convert(o), o, ancestors);
+                ensureNoSelfReferences(convert(o), o, ancestors, messageHint);
             }
             ancestors.remove(originalReference);
         }

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

@@ -21,6 +21,7 @@
 
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.Script;
@@ -97,7 +98,11 @@ public InternalAggregation doReduce(List<InternalAggregation> aggregations, Redu
             ExecutableScript.Factory factory = reduceContext.scriptService().compile(
                 firstAggregation.reduceScript, ExecutableScript.AGGS_CONTEXT);
             ExecutableScript script = factory.newInstance(vars);
-            aggregation = Collections.singletonList(script.run());
+
+            Object scriptResult = script.run();
+            CollectionUtils.ensureNoSelfReferences(scriptResult, "reduce script");
+
+            aggregation = Collections.singletonList(scriptResult);
         } else if (reduceContext.isFinalReduce())  {
             aggregation = Collections.singletonList(aggregationObjects);
         } else {

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

@@ -69,6 +69,7 @@ public void collect(int doc, long bucket) throws IOException {
                 assert bucket == 0 : bucket;
                 leafMapScript.setDocument(doc);
                 leafMapScript.run();
+                CollectionUtils.ensureNoSelfReferences(params, "Scripted metric aggs map script");
             }
         };
     }
@@ -78,7 +79,7 @@ public InternalAggregation buildAggregation(long owningBucketOrdinal) {
         Object aggregation;
         if (combineScript != null) {
             aggregation = combineScript.run();
-            CollectionUtils.ensureNoSelfReferences(aggregation);
+            CollectionUtils.ensureNoSelfReferences(aggregation, "Scripted metric aggs combine script");
         } else {
             aggregation = params.get("_agg");
         }

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

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.search.aggregations.metrics.scripted;
 
+import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.SearchScript;
@@ -89,6 +90,7 @@ public Aggregator createInternal(Aggregator parent, boolean collectsFromSingleBu
         final Script reduceScript = deepCopyScript(this.reduceScript, context);
         if (initScript != null) {
             initScript.run();
+            CollectionUtils.ensureNoSelfReferences(aggParams.get("_agg"), "Scripted metric aggs init script");
         }
         return new ScriptedMetricAggregator(name, mapScript,
                 combineScript, reduceScript, aggParams, context, parent,

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

@@ -458,7 +458,7 @@ public boolean advanceExact(int doc) throws IOException {
                         final BytesRef value = bytesValues.nextValue();
                         script.setNextAggregationValue(value.utf8ToString());
                         Object run = script.run();
-                        CollectionUtils.ensureNoSelfReferences(run);
+                        CollectionUtils.ensureNoSelfReferences(run, "ValuesSource.BytesValues script");
                         values[i].copyChars(run.toString());
                     }
                     sort();

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

@@ -45,7 +45,7 @@ private void set(int i, Object o) {
         if (o == null) {
             values[i].clear();
         } else {
-            CollectionUtils.ensureNoSelfReferences(o);
+            CollectionUtils.ensureNoSelfReferences(o, "ScriptBytesValues value");
             values[i].copyChars(o.toString());
         }
     }

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

@@ -65,7 +65,7 @@ public void hitsExecute(SearchContext context, SearchHit[] hits) throws IOExcept
                 final Object value;
                 try {
                     value = leafScripts[i].run();
-                    CollectionUtils.ensureNoSelfReferences(value);
+                    CollectionUtils.ensureNoSelfReferences(value, "ScriptFieldsFetchSubPhase leaf script " + i);
                 } catch (RuntimeException e) {
                     if (scriptFields.get(i).ignoreException()) {
                         continue;

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

@@ -343,7 +343,7 @@ public boolean advanceExact(int doc) throws IOException {
                             @Override
                             public BytesRef binaryValue() {
                                 final Object run = leafScript.run();
-                                CollectionUtils.ensureNoSelfReferences(run);
+                                CollectionUtils.ensureNoSelfReferences(run, "ScriptSortBuilder leaf script");
                                 spare.copyChars(run.toString());
                                 return spare.get();
                             }

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

@@ -183,13 +183,14 @@ public void testPerfectPartition() {
     }
 
     public void testEnsureNoSelfReferences() {
-        CollectionUtils.ensureNoSelfReferences(emptyMap());
-        CollectionUtils.ensureNoSelfReferences(null);
+        CollectionUtils.ensureNoSelfReferences(emptyMap(), "test with empty map");
+        CollectionUtils.ensureNoSelfReferences(null, "test with 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"));
+        IllegalArgumentException e = expectThrows(IllegalArgumentException.class,
+            () ->  CollectionUtils.ensureNoSelfReferences(map, "test with self ref"));
+        assertThat(e.getMessage(), containsString("Iterable object is self-referencing itself (test with self ref)"));
     }
 }

server/src/test/java/org/elasticsearch/common/xcontent/BaseXContentTestCase.java

@@ -843,8 +843,8 @@ public void testEnsureNotNull() {
     }
 
     public void testEnsureNoSelfReferences() throws IOException {
-        CollectionUtils.ensureNoSelfReferences(emptyMap());
-        CollectionUtils.ensureNoSelfReferences(null);
+        builder().map(emptyMap());
+        builder().map(null);
 
         Map<String, Object> map = new HashMap<>();
         map.put("field", map);

server/src/test/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorTests.java

@@ -73,6 +73,13 @@ public class ScriptedMetricAggregatorTests extends AggregatorTestCase {
             Collections.singletonMap("divisor", 4));
     private static final String CONFLICTING_PARAM_NAME = "initialValue";
 
+    private static final Script INIT_SCRIPT_SELF_REF = new Script(ScriptType.INLINE, MockScriptEngine.NAME, "initScriptSelfRef",
+            Collections.emptyMap());
+    private static final Script MAP_SCRIPT_SELF_REF = new Script(ScriptType.INLINE, MockScriptEngine.NAME, "mapScriptSelfRef",
+            Collections.emptyMap());
+    private static final Script COMBINE_SCRIPT_SELF_REF = new Script(ScriptType.INLINE, MockScriptEngine.NAME, "combineScriptSelfRef",
+            Collections.emptyMap());
+
     private static final Map<String, Function<Map<String, Object>, Object>> SCRIPTS = new HashMap<>();
 
     @BeforeClass
@@ -127,6 +134,25 @@ public static void initMockScripts() {
             int divisor = ((Integer) params.get("divisor"));
             return ((List<Integer>) agg.get("collector")).stream().mapToInt(Integer::intValue).map(i -> i / divisor).sum();
         });
+
+        SCRIPTS.put("initScriptSelfRef", params -> {
+            Map<String, Object> agg = (Map<String, Object>) params.get("_agg");
+            agg.put("collector", new ArrayList<Integer>());
+            agg.put("selfRef", agg);
+            return agg;
+        });
+
+        SCRIPTS.put("mapScriptSelfRef", params -> {
+            Map<String, Object> agg = (Map<String, Object>) params.get("_agg");
+            agg.put("selfRef", agg);
+            return agg;
+        });
+
+        SCRIPTS.put("combineScriptSelfRef", params -> {
+           Map<String, Object> agg = (Map<String, Object>) params.get("_agg");
+           agg.put("selfRef", agg);
+           return agg;
+        });
     }
 
     @SuppressWarnings("unchecked")
@@ -257,6 +283,60 @@ public void testConflictingAggAndScriptParams() throws IOException {
         }
     }
 
+    public void testSelfReferencingAggStateAfterInit() throws IOException {
+        try (Directory directory = newDirectory()) {
+            try (RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) {
+                // No need to add docs for this test
+            }
+            try (IndexReader indexReader = DirectoryReader.open(directory)) {
+                ScriptedMetricAggregationBuilder aggregationBuilder = new ScriptedMetricAggregationBuilder(AGG_NAME);
+                aggregationBuilder.initScript(INIT_SCRIPT_SELF_REF).mapScript(MAP_SCRIPT);
+
+                IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () ->
+                    search(newSearcher(indexReader, true, true), new MatchAllDocsQuery(), aggregationBuilder)
+                );
+                assertEquals("Iterable object is self-referencing itself (Scripted metric aggs init script)", ex.getMessage());
+            }
+        }
+    }
+
+    public void testSelfReferencingAggStateAfterMap() throws IOException {
+        try (Directory directory = newDirectory()) {
+            Integer numDocs = randomInt(100);
+            try (RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) {
+                for (int i = 0; i < numDocs; i++) {
+                    indexWriter.addDocument(singleton(new SortedNumericDocValuesField("number", i)));
+                }
+            }
+            try (IndexReader indexReader = DirectoryReader.open(directory)) {
+                ScriptedMetricAggregationBuilder aggregationBuilder = new ScriptedMetricAggregationBuilder(AGG_NAME);
+                aggregationBuilder.initScript(INIT_SCRIPT).mapScript(MAP_SCRIPT_SELF_REF);
+
+                IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () ->
+                    search(newSearcher(indexReader, true, true), new MatchAllDocsQuery(), aggregationBuilder)
+                );
+                assertEquals("Iterable object is self-referencing itself (Scripted metric aggs map script)", ex.getMessage());
+            }
+        }
+    }
+
+    public void testSelfReferencingAggStateAfterCombine() throws IOException {
+        try (Directory directory = newDirectory()) {
+            try (RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) {
+                // No need to add docs for this test
+            }
+            try (IndexReader indexReader = DirectoryReader.open(directory)) {
+                ScriptedMetricAggregationBuilder aggregationBuilder = new ScriptedMetricAggregationBuilder(AGG_NAME);
+                aggregationBuilder.initScript(INIT_SCRIPT).mapScript(MAP_SCRIPT).combineScript(COMBINE_SCRIPT_SELF_REF);
+
+                IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () ->
+                    search(newSearcher(indexReader, true, true), new MatchAllDocsQuery(), aggregationBuilder)
+                );
+                assertEquals("Iterable object is self-referencing itself (Scripted metric aggs combine script)", ex.getMessage());
+            }
+        }
+    }
+
     /**
      * We cannot use Mockito for mocking QueryShardContext in this case because
      * script-related methods (e.g. QueryShardContext#getLazyExecutableScript)