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

Backport #114002 to 8.16

Author: thecoop

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: []

libs/logstash-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);
  }

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

modules/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);
  }

modules/lang-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);

modules/lang-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++) {

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 = """

qa/smoke-test-ingest-with-all-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);
  }

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
+}

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));
+ }
+}