Stop runtime script from emitting too many values

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/AbstractScriptFieldScript.java

@@ -27,6 +27,11 @@
  * {@link AggregationScript} but hopefully with less historical baggage.
  */
 public abstract class AbstractScriptFieldScript {
+    /**
+     * The maximum number of values a script should be allowed to emit.
+     */
+    public static final int MAX_VALUES = 1000;
+
     public static <F> ScriptContext<F> newContext(String name, Class<F> factoryClass) {
         return new ScriptContext<F>(
             name + "_script_field",

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScript.java

@@ -19,6 +19,11 @@
 import java.util.Map;
 
 public abstract class StringScriptFieldScript extends AbstractScriptFieldScript {
+    /**
+     * The maximum number of chars a script should be allowed to emit.
+     */
+    public static final long MAX_CHARS = 1024 * 1024;
+
     public static final ScriptContext<Factory> CONTEXT = newContext("string_script_field", Factory.class);
 
     static List<Whitelist> whitelist() {
@@ -36,6 +41,7 @@ public interface LeafFactory {
     }
 
     private final List<String> results = new ArrayList<>();
+    private long chars;
 
     public StringScriptFieldScript(Map<String, Object> params, SearchLookup searchLookup, LeafReaderContext ctx) {
         super(params, searchLookup, ctx);
@@ -49,12 +55,20 @@ public StringScriptFieldScript(Map<String, Object> params, SearchLookup searchLo
      */
     public final List<String> resultsForDoc(int docId) {
         results.clear();
+        chars = 0;
         setDocument(docId);
         execute();
         return results;
     }
 
     protected final void emitValue(String v) {
+        if (results.size() >= MAX_VALUES) {
+            throw new IllegalArgumentException("too many runtime values");
+        }
+        chars += v.length();
+        if (chars >= MAX_CHARS) {
+            throw new IllegalArgumentException("too many characters in runtime values [" + chars + "]");
+        }
         results.add(v);
     }
 

x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScriptTests.java

@@ -6,7 +6,22 @@
 
 package org.elasticsearch.xpack.runtimefields;
 
+import org.apache.lucene.document.StoredField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.RandomIndexWriter;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.script.ScriptContext;
+import org.elasticsearch.search.lookup.SearchLookup;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.startsWith;
+import static org.mockito.Mockito.mock;
 
 public class StringScriptFieldScriptTests extends ScriptFieldScriptTestCase<StringScriptFieldScript.Factory> {
     public static final StringScriptFieldScript.Factory DUMMY = (params, lookup) -> ctx -> new StringScriptFieldScript(
@@ -29,4 +44,53 @@ protected ScriptContext<StringScriptFieldScript.Factory> context() {
     protected StringScriptFieldScript.Factory dummyScript() {
         return DUMMY;
     }
+
+    public void testTooManyValues() throws IOException {
+        try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
+            iw.addDocument(List.of(new StoredField("_source", new BytesRef("{}"))));
+            try (DirectoryReader reader = iw.getReader()) {
+                StringScriptFieldScript script = new StringScriptFieldScript(
+                    Map.of(),
+                    new SearchLookup(mock(MapperService.class), (ft, lookup) -> null),
+                    reader.leaves().get(0)
+                ) {
+                    @Override
+                    public void execute() {
+                        for (int i = 0; i <= AbstractScriptFieldScript.MAX_VALUES; i++) {
+                            emitValue("test");
+                        }
+                    }
+                };
+                Exception e = expectThrows(IllegalArgumentException.class, script::execute);
+                assertThat(e.getMessage(), equalTo("too many runtime values"));
+            }
+        }
+    }
+
+    public void testTooManyChars() throws IOException {
+        try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
+            iw.addDocument(List.of(new StoredField("_source", new BytesRef("{}"))));
+            try (DirectoryReader reader = iw.getReader()) {
+                StringScriptFieldScript script = new StringScriptFieldScript(
+                    Map.of(),
+                    new SearchLookup(mock(MapperService.class), (ft, lookup) -> null),
+                    reader.leaves().get(0)
+                ) {
+                    @Override
+                    public void execute() {
+                        StringBuilder big = new StringBuilder();
+                        while (big.length() < StringScriptFieldScript.MAX_CHARS / 4) {
+                            big.append("test");
+                        }
+                        String bigString = big.toString();
+                        for (int i = 0; i <= 4; i++) {
+                            emitValue(bigString);
+                        }
+                    }
+                };
+                Exception e = expectThrows(IllegalArgumentException.class, script::execute);
+                assertThat(e.getMessage(), startsWith("too many characters in runtime values ["));
+            }
+        }
+    }
 }