Add a size limit to outputs from mustache (elastic#114002)

thecoop · Oct 8, 2024 · c645ab1 · c645ab1
1 parent 029c683
commit c645ab1
Show file tree

Hide file tree

Showing 19 changed files with 181 additions and 24 deletions.
diff --git a/docs/changelog/114002.yaml b/docs/changelog/114002.yaml
@@ -0,0 +1,5 @@
+pr: 114002
+summary: Add a `mustache.max_output_size_bytes` setting to limit the length of results from mustache scripts
+area: Infra/Scripting
+type: enhancement
+issues: []
diff --git a/...ash-bridge/src/main/java/org/elasticsearch/logstashbridge/script/ScriptServiceBridge.java b/...ash-bridge/src/main/java/org/elasticsearch/logstashbridge/script/ScriptServiceBridge.java
@@ -53,7 +53,7 @@ private static ScriptService getScriptService(final Settings settings, final Lon
             PainlessScriptEngine.NAME,
             new PainlessScriptEngine(settings, scriptContexts),
             MustacheScriptEngine.NAME,
-            new MustacheScriptEngine()
+            new MustacheScriptEngine(settings)
         );
         return new ScriptService(settings, scriptEngines, ScriptModule.CORE_CONTEXTS, timeProvider);
     }

diff --git a/modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/MustachePlugin.java b/modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/MustachePlugin.java
@@ -44,7 +44,7 @@ public class MustachePlugin extends Plugin implements ScriptPlugin, ActionPlugin
 
     @Override
     public ScriptEngine getScriptEngine(Settings settings, Collection<ScriptContext<?>> contexts) {
-        return new MustacheScriptEngine();
+        return new MustacheScriptEngine(settings);
     }
 
     @Override

diff --git a/...s/lang-mustache/src/main/java/org/elasticsearch/script/mustache/MustacheScriptEngine.java b/...s/lang-mustache/src/main/java/org/elasticsearch/script/mustache/MustacheScriptEngine.java
@@ -14,6 +14,13 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.elasticsearch.ElasticsearchParseException;
+import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.common.settings.Setting;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.text.SizeLimitingStringWriter;
+import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.common.unit.MemorySizeValue;
 import org.elasticsearch.script.GeneralScriptException;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptContext;
@@ -47,6 +54,19 @@ public final class MustacheScriptEngine implements ScriptEngine {
 
     public static final String NAME = "mustache";
 
+    public static final Setting<ByteSizeValue> MUSTACHE_RESULT_SIZE_LIMIT = new Setting<>(
+        "mustache.max_output_size_bytes",
+        s -> "1mb",
+        s -> MemorySizeValue.parseBytesSizeValueOrHeapRatio(s, "mustache.max_output_size_bytes"),
+        Setting.Property.NodeScope
+    );
+
+    private final int sizeLimit;
+
+    public MustacheScriptEngine(Settings settings) {
+        sizeLimit = (int) MUSTACHE_RESULT_SIZE_LIMIT.get(settings).getBytes();
+    }
+
     /**
      * Compile a template string to (in this case) a Mustache object than can
      * later be re-used for execution to fill in missing parameter values.
@@ -118,10 +138,15 @@ private class MustacheExecutableScript extends TemplateScript {
 
         @Override
         public String execute() {
-            final StringWriter writer = new StringWriter();
+            StringWriter writer = new SizeLimitingStringWriter(sizeLimit);
             try {
                 template.execute(writer, params);
             } catch (Exception e) {
+                // size limit exception can appear at several places in the causal list depending on script & context
+                if (ExceptionsHelper.unwrap(e, SizeLimitingStringWriter.SizeLimitExceededException.class) != null) {
+                    // don't log, client problem
+                    throw new ElasticsearchParseException("Mustache script result size limit exceeded", e);
+                }
                 if (shouldLogException(e)) {
                     logger.error(() -> format("Error running %s", template), e);
                 }

diff --git a/...-mustache/src/test/java/org/elasticsearch/script/mustache/CustomMustacheFactoryTests.java b/...-mustache/src/test/java/org/elasticsearch/script/mustache/CustomMustacheFactoryTests.java
@@ -9,6 +9,7 @@
 
 package org.elasticsearch.script.mustache;
 
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptEngine;
 import org.elasticsearch.script.TemplateScript;
@@ -65,7 +66,7 @@ public void testCreateEncoder() {
     }
 
     public void testJsonEscapeEncoder() {
-        final ScriptEngine engine = new MustacheScriptEngine();
+        final ScriptEngine engine = new MustacheScriptEngine(Settings.EMPTY);
         final Map<String, String> params = randomBoolean() ? Map.of(Script.CONTENT_TYPE_OPTION, JSON_MEDIA_TYPE) : Map.of();
 
         TemplateScript.Factory compiled = engine.compile(null, "{\"field\": \"{{value}}\"}", TemplateScript.CONTEXT, params);
@@ -75,7 +76,7 @@ public void testJsonEscapeEncoder() {
     }
 
     public void testDefaultEncoder() {
-        final ScriptEngine engine = new MustacheScriptEngine();
+        final ScriptEngine engine = new MustacheScriptEngine(Settings.EMPTY);
         final Map<String, String> params = Map.of(Script.CONTENT_TYPE_OPTION, PLAIN_TEXT_MEDIA_TYPE);
 
         TemplateScript.Factory compiled = engine.compile(null, "{\"field\": \"{{value}}\"}", TemplateScript.CONTEXT, params);
@@ -85,7 +86,7 @@ public void testDefaultEncoder() {
     }
 
     public void testUrlEncoder() {
-        final ScriptEngine engine = new MustacheScriptEngine();
+        final ScriptEngine engine = new MustacheScriptEngine(Settings.EMPTY);
         final Map<String, String> params = Map.of(Script.CONTENT_TYPE_OPTION, X_WWW_FORM_URLENCODED_MEDIA_TYPE);
 
         TemplateScript.Factory compiled = engine.compile(null, "{\"field\": \"{{value}}\"}", TemplateScript.CONTEXT, params);

diff --git a/...g-mustache/src/test/java/org/elasticsearch/script/mustache/MustacheScriptEngineTests.java b/...g-mustache/src/test/java/org/elasticsearch/script/mustache/MustacheScriptEngineTests.java
@@ -8,8 +8,13 @@
  */
 package org.elasticsearch.script.mustache;
 
+import com.github.mustachejava.MustacheException;
 import com.github.mustachejava.MustacheFactory;
 
+import org.elasticsearch.ElasticsearchParseException;
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.text.SizeLimitingStringWriter;
 import org.elasticsearch.script.GeneralScriptException;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.TemplateScript;
@@ -24,6 +29,9 @@
 import java.util.List;
 import java.util.Map;
 
+import static org.elasticsearch.test.LambdaMatchers.transformedMatch;
+import static org.hamcrest.Matchers.allOf;
+import static org.hamcrest.Matchers.endsWith;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.startsWith;
@@ -37,7 +45,7 @@ public class MustacheScriptEngineTests extends ESTestCase {
 
     @Before
     public void setup() {
-        qe = new MustacheScriptEngine();
+        qe = new MustacheScriptEngine(Settings.builder().put(MustacheScriptEngine.MUSTACHE_RESULT_SIZE_LIMIT.getKey(), "1kb").build());
         factory = CustomMustacheFactory.builder().build();
     }
 
@@ -402,6 +410,24 @@ public void testEscapeJson() throws IOException {
         }
     }
 
+    public void testResultSizeLimit() throws IOException {
+        String vals = "\"" + "{{val}}".repeat(200) + "\"";
+        String params = "\"val\":\"aaaaaaaaaa\"";
+        XContentParser parser = createParser(JsonXContent.jsonXContent, Strings.format("{\"source\":%s,\"params\":{%s}}", vals, params));
+        Script script = Script.parse(parser);
+        var compiled = qe.compile(null, script.getIdOrCode(), TemplateScript.CONTEXT, Map.of());
+        TemplateScript templateScript = compiled.newInstance(script.getParams());
+        var ex = expectThrows(ElasticsearchParseException.class, templateScript::execute);
+        assertThat(ex.getCause(), instanceOf(MustacheException.class));
+        assertThat(
+            ex.getCause().getCause(),
+            allOf(
+                instanceOf(SizeLimitingStringWriter.SizeLimitExceededException.class),
+                transformedMatch(Throwable::getMessage, endsWith("has exceeded the size limit [1024]"))
+            )
+        );
+    }
+
     private String getChars() {
         String string = randomRealisticUnicodeOfCodepointLengthBetween(0, 10);
         for (int i = 0; i < string.length(); i++) {

diff --git a/modules/lang-mustache/src/test/java/org/elasticsearch/script/mustache/MustacheTests.java b/modules/lang-mustache/src/test/java/org/elasticsearch/script/mustache/MustacheTests.java
@@ -9,6 +9,7 @@
 package org.elasticsearch.script.mustache;
 
 import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.core.Strings;
 import org.elasticsearch.script.ScriptEngine;
@@ -39,7 +40,7 @@
 
 public class MustacheTests extends ESTestCase {
 
-    private ScriptEngine engine = new MustacheScriptEngine();
+    private ScriptEngine engine = new MustacheScriptEngine(Settings.EMPTY);
 
     public void testBasics() {
         String template = """

diff --git a/...l-dependencies/src/yamlRestTest/java/org/elasticsearch/ingest/AbstractScriptTestCase.java b/...l-dependencies/src/yamlRestTest/java/org/elasticsearch/ingest/AbstractScriptTestCase.java
@@ -31,7 +31,7 @@ public abstract class AbstractScriptTestCase extends ESTestCase {
 
     @Before
     public void init() throws Exception {
-        MustacheScriptEngine engine = new MustacheScriptEngine();
+        MustacheScriptEngine engine = new MustacheScriptEngine(Settings.EMPTY);
         Map<String, ScriptEngine> engines = Collections.singletonMap(engine.getType(), engine);
         scriptService = new ScriptService(Settings.EMPTY, engines, ScriptModule.CORE_CONTEXTS, () -> 1L);
     }

diff --git a/server/src/main/java/org/elasticsearch/common/text/SizeLimitingStringWriter.java b/server/src/main/java/org/elasticsearch/common/text/SizeLimitingStringWriter.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.common.text;
+
+import org.elasticsearch.common.Strings;
+
+import java.io.StringWriter;
+
+/**
+ * A {@link StringWriter} that throws an exception if the string exceeds a specified size.
+ */
+public class SizeLimitingStringWriter extends StringWriter {
+
+    public static class SizeLimitExceededException extends IllegalStateException {
+        public SizeLimitExceededException(String message) {
+            super(message);
+        }
+    }
+
+    private final int sizeLimit;
+
+    public SizeLimitingStringWriter(int sizeLimit) {
+        this.sizeLimit = sizeLimit;
+    }
+
+    private void checkSizeLimit(int additionalChars) {
+        int bufLen = getBuffer().length();
+        if (bufLen + additionalChars > sizeLimit) {
+            throw new SizeLimitExceededException(
+                Strings.format("String [%s...] has exceeded the size limit [%s]", getBuffer().substring(0, Math.min(bufLen, 20)), sizeLimit)
+            );
+        }
+    }
+
+    @Override
+    public void write(int c) {
+        checkSizeLimit(1);
+        super.write(c);
+    }
+
+    // write(char[]) delegates to write(char[], int, int)
+
+    @Override
+    public void write(char[] cbuf, int off, int len) {
+        checkSizeLimit(len);
+        super.write(cbuf, off, len);
+    }
+
+    @Override
+    public void write(String str) {
+        checkSizeLimit(str.length());
+        super.write(str);
+    }
+
+    @Override
+    public void write(String str, int off, int len) {
+        checkSizeLimit(len);
+        super.write(str, off, len);
+    }
+
+    // append(...) delegates to write(...) methods
+}
diff --git a/server/src/test/java/org/elasticsearch/common/text/SizeLimitingStringWriterTests.java b/server/src/test/java/org/elasticsearch/common/text/SizeLimitingStringWriterTests.java
@@ -0,0 +1,29 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.common.text;
+
+import org.elasticsearch.test.ESTestCase;
+
+public class SizeLimitingStringWriterTests extends ESTestCase {
+    public void testSizeIsLimited() {
+        SizeLimitingStringWriter writer = new SizeLimitingStringWriter(10);
+
+        writer.write("a".repeat(10));
+
+        // test all the methods
+        expectThrows(SizeLimitingStringWriter.SizeLimitExceededException.class, () -> writer.write('a'));
+        expectThrows(SizeLimitingStringWriter.SizeLimitExceededException.class, () -> writer.write("a"));
+        expectThrows(SizeLimitingStringWriter.SizeLimitExceededException.class, () -> writer.write(new char[1]));
+        expectThrows(SizeLimitingStringWriter.SizeLimitExceededException.class, () -> writer.write(new char[1], 0, 1));
+        expectThrows(SizeLimitingStringWriter.SizeLimitExceededException.class, () -> writer.append('a'));
+        expectThrows(SizeLimitingStringWriter.SizeLimitExceededException.class, () -> writer.append("a"));
+        expectThrows(SizeLimitingStringWriter.SizeLimitExceededException.class, () -> writer.append("a", 0, 1));
+    }
+}
diff --git a/...ava/org/elasticsearch/xpack/core/security/authc/support/mapper/TemplateRoleNameTests.java b/...ava/org/elasticsearch/xpack/core/security/authc/support/mapper/TemplateRoleNameTests.java
@@ -89,7 +89,7 @@ public void testEqualsAndHashCode() throws Exception {
     public void testEvaluateRoles() throws Exception {
         final ScriptService scriptService = new ScriptService(
             Settings.EMPTY,
-            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
+            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
             ScriptModule.CORE_CONTEXTS,
             () -> 1L
         );
@@ -145,7 +145,7 @@ public void tryEquals(TemplateRoleName original) {
     public void testValidate() {
         final ScriptService scriptService = new ScriptService(
             Settings.EMPTY,
-            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
+            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
             ScriptModule.CORE_CONTEXTS,
             () -> 1L
         );
@@ -173,7 +173,7 @@ public void testValidate() {
     public void testValidateWillPassWithEmptyContext() {
         final ScriptService scriptService = new ScriptService(
             Settings.EMPTY,
-            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
+            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
             ScriptModule.CORE_CONTEXTS,
             () -> 1L
         );
@@ -204,7 +204,7 @@ public void testValidateWillPassWithEmptyContext() {
     public void testValidateWillFailForSyntaxError() {
         final ScriptService scriptService = new ScriptService(
             Settings.EMPTY,
-            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
+            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
             ScriptModule.CORE_CONTEXTS,
             () -> 1L
         );
@@ -268,7 +268,7 @@ public void testValidationWillFailWhenInlineScriptIsNotEnabled() {
         final Settings settings = Settings.builder().put("script.allowed_types", ScriptService.ALLOW_NONE).build();
         final ScriptService scriptService = new ScriptService(
             settings,
-            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
+            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
             ScriptModule.CORE_CONTEXTS,
             () -> 1L
         );
@@ -285,7 +285,7 @@ public void testValidateWillFailWhenStoredScriptIsNotEnabled() {
         final Settings settings = Settings.builder().put("script.allowed_types", ScriptService.ALLOW_NONE).build();
         final ScriptService scriptService = new ScriptService(
             settings,
-            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
+            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
             ScriptModule.CORE_CONTEXTS,
             () -> 1L
         );
@@ -314,7 +314,7 @@ public void testValidateWillFailWhenStoredScriptIsNotEnabled() {
     public void testValidateWillFailWhenStoredScriptIsNotFound() {
         final ScriptService scriptService = new ScriptService(
             Settings.EMPTY,
-            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
+            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
             ScriptModule.CORE_CONTEXTS,
             () -> 1L
         );

diff --git a/.../elasticsearch/xpack/core/security/authz/support/SecurityQueryTemplateEvaluatorTests.java b/.../elasticsearch/xpack/core/security/authz/support/SecurityQueryTemplateEvaluatorTests.java
@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.core.security.authz.support;
 
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.query.TermQueryBuilder;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptService;
@@ -94,7 +95,7 @@ public void testDocLevelSecurityTemplateWithOpenIdConnectStyleMetadata() throws
             true
         );
 
-        final MustacheScriptEngine mustache = new MustacheScriptEngine();
+        final MustacheScriptEngine mustache = new MustacheScriptEngine(Settings.EMPTY);
 
         when(scriptService.compile(any(Script.class), eq(TemplateScript.CONTEXT))).thenAnswer(inv -> {
             assertThat(inv.getArguments(), arrayWithSize(2));

diff --git a/...c/test/java/org/elasticsearch/xpack/idp/saml/sp/WildcardServiceProviderResolverTests.java b/...c/test/java/org/elasticsearch/xpack/idp/saml/sp/WildcardServiceProviderResolverTests.java
@@ -95,7 +95,7 @@ public void setUpResolver() {
         final Settings settings = Settings.EMPTY;
         final ScriptService scriptService = new ScriptService(
             settings,
-            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
+            Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
             ScriptModule.CORE_CONTEXTS,
             () -> 1L
         );

diff --git a/...ml/src/test/java/org/elasticsearch/xpack/ml/inference/ltr/LearningToRankServiceTests.java b/...ml/src/test/java/org/elasticsearch/xpack/ml/inference/ltr/LearningToRankServiceTests.java
@@ -241,7 +241,7 @@ private LearningToRankService getTestLearningToRankService(TrainedModelProvider
     }
 
     private ScriptService getTestScriptService() {
-        ScriptEngine scriptEngine = new MustacheScriptEngine();
+        ScriptEngine scriptEngine = new MustacheScriptEngine(Settings.EMPTY);
         return new ScriptService(Settings.EMPTY, Map.of(DEFAULT_TEMPLATE_LANG, scriptEngine), ScriptModule.CORE_CONTEXTS, () -> 1L);
     }
 }