Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add a size limit to outputs from mustache #114002

Merged
merged 11 commits into from
Oct 8, 2024
5 changes: 5 additions & 0 deletions docs/changelog/114002.yaml
Original file line number Diff line number Diff line change
@@ -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: []
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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();
}

Expand Down Expand Up @@ -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++) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 = """
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
/*
* 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 {
thecoop marked this conversation as resolved.
Show resolved Hide resolved

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) {
if (getBuffer().length() + additionalChars > sizeLimit) {
throw new SizeLimitExceededException(
Strings.format("String [%s...] has exceeded the size limit [%s]", getBuffer().substring(0, 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
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
);
Expand Down Expand Up @@ -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
);
Expand Down Expand Up @@ -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
);
Expand Down Expand Up @@ -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
);
Expand Down Expand Up @@ -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
);
Expand All @@ -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
);
Expand Down Expand Up @@ -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
);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -434,7 +434,7 @@ public void testRealmWithTemplatedRoleMapping() throws Exception {

final ScriptService scriptService = new ScriptService(
settings,
Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
ScriptModule.CORE_CONTEXTS,
() -> 1L
);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -533,7 +533,7 @@ public void testLdapRealmWithTemplatedRoleMapping() throws Exception {

final ScriptService scriptService = new ScriptService(
defaultGlobalSettings,
Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
ScriptModule.CORE_CONTEXTS,
() -> 1L
);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ public class ClusterStateRoleMapperTests extends ESTestCase {
public void setup() {
scriptService = new ScriptService(
Settings.EMPTY,
Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
ScriptModule.CORE_CONTEXTS,
() -> 1L
);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ public class NativeRoleMappingStoreTests extends ESTestCase {
public void setup() {
scriptService = new ScriptService(
Settings.EMPTY,
Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine(Settings.EMPTY)),
ScriptModule.CORE_CONTEXTS,
() -> 1L
);
Expand Down
Loading