Mustache: Add {{#url}}{{/url}} function to URL encode strings (#20838)

This commit adds a new Mustache function (codename: url) and a new URLEncoder that can be used to URL encode strings in mustache templates.

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

@@ -27,41 +27,75 @@ import com.github.mustachejava.Mustache;
 import com.github.mustachejava.MustacheException;
 import com.github.mustachejava.MustacheVisitor;
 import com.github.mustachejava.TemplateContext;
+import com.github.mustachejava.codes.DefaultMustache;
 import com.github.mustachejava.codes.IterableCode;
 import com.github.mustachejava.codes.WriteCode;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentType;
 
 import java.io.IOException;
 import java.io.StringWriter;
 import java.io.Writer;
+import java.net.URLEncoder;
+import java.nio.charset.StandardCharsets;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.StringJoiner;
-import java.util.function.BiConsumer;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 public class CustomMustacheFactory extends DefaultMustacheFactory {
 
-    private final BiConsumer<String, Writer> encoder;
+    static final String CONTENT_TYPE_PARAM = "content_type";
 
-    public CustomMustacheFactory(boolean escaping) {
+    static final String JSON_MIME_TYPE = "application/json";
+    static final String PLAIN_TEXT_MIME_TYPE = "text/plain";
+    static final String X_WWW_FORM_URLENCODED_MIME_TYPE = "application/x-www-form-urlencoded";
+
+    private static final String DEFAULT_MIME_TYPE = JSON_MIME_TYPE;
+
+    private static final Map<String, Supplier<Encoder>> ENCODERS;
+    static {
+        Map<String, Supplier<Encoder>> encoders = new HashMap<>();
+        encoders.put(JSON_MIME_TYPE, JsonEscapeEncoder::new);
+        encoders.put(PLAIN_TEXT_MIME_TYPE, DefaultEncoder::new);
+        encoders.put(X_WWW_FORM_URLENCODED_MIME_TYPE, UrlEncoder::new);
+        ENCODERS = Collections.unmodifiableMap(encoders);
+    }
+
+    private final Encoder encoder;
+
+    public CustomMustacheFactory(String mimeType) {
         super();
         setObjectHandler(new CustomReflectionObjectHandler());
-        if (escaping) {
+        this.encoder = createEncoder(mimeType);
-            this.encoder = new JsonEscapeEncoder();
-        } else {
-            this.encoder = new NoEscapeEncoder();
     }
+
+    public CustomMustacheFactory() {
+        this(DEFAULT_MIME_TYPE);
     }
 
     @Override
     public void encode(String value, Writer writer) {
-        encoder.accept(value, writer);
+        try {
+            encoder.encode(value, writer);
+        } catch (IOException e) {
+            throw new MustacheException("Unable to encode value", e);
+        }
+    }
+
+    static Encoder createEncoder(String mimeType) {
+        Supplier<Encoder> supplier = ENCODERS.get(mimeType);
+        if (supplier == null) {
+            throw new IllegalArgumentException("No encoder found for MIME type [" + mimeType + "]");
+        }
+        return supplier.get();
     }
 
     @Override
@@ -83,6 +117,8 @@ public class CustomMustacheFactory extends DefaultMustacheFactory {
                 list.add(new JoinerCode(templateContext, df, mustache));
             } else if (CustomJoinerCode.match(variable)) {
                 list.add(new CustomJoinerCode(templateContext, df, mustache, variable));
+            } else if (UrlEncoderCode.match(variable)) {
+                list.add(new UrlEncoderCode(templateContext, df, mustache, variable));
             } else {
                 list.add(new IterableCode(templateContext, df, mustache, variable));
             }
@@ -253,27 +289,85 @@ public class CustomMustacheFactory extends DefaultMustacheFactory {
         }
     }
 
-    class NoEscapeEncoder implements BiConsumer<String, Writer> {
+    /**
+     * This function encodes a string using the {@link URLEncoder#encode(String, String)} method
+     * with the UTF-8 charset.
+     */
+    static class UrlEncoderCode extends DefaultMustache {
+
+        private static final String CODE = "url";
+        private final Encoder encoder;
+
+        public UrlEncoderCode(TemplateContext tc, DefaultMustacheFactory df, Mustache mustache, String variable) {
+            super(tc, df, mustache.getCodes(), variable);
+            this.encoder = new UrlEncoder();
+        }
 
         @Override
-        public void accept(String s, Writer writer) {
+        public Writer run(Writer writer, List<Object> scopes) {
-            try {
+            if (getCodes() != null) {
+                for (Code code : getCodes()) {
+                    try (StringWriter capture = new StringWriter()) {
+                        code.execute(capture, scopes);
+
+                        String s = capture.toString();
+                        if (s != null) {
+                            encoder.encode(s, writer);
+                        }
+                    } catch (IOException e) {
+                        throw new MustacheException("Exception while parsing mustache function at line " + tc.line(), e);
+                    }
+                }
+            }
+            return writer;
+        }
+
+        static boolean match(String variable) {
+            return CODE.equalsIgnoreCase(variable);
+        }
+    }
+
+    @FunctionalInterface
+    interface Encoder {
+        /**
+         * Encodes the {@code s} string and writes it to the {@code writer} {@link Writer}.
+         *
+         * @param s      The string to encode
+         * @param writer The {@link Writer} to which the encoded string will be written to
+         */
+        void encode(final String s, final Writer writer) throws IOException;
+    }
+
+    /**
+     * Encoder that simply writes the string to the writer without encoding.
+     */
+    static class DefaultEncoder implements Encoder {
+
+        @Override
+        public void encode(String s, Writer writer) throws IOException {
             writer.write(s);
-            } catch (IOException e) {
-                throw new MustacheException("Failed to encode value: " + s);
-            }
         }
     }
 
-    class JsonEscapeEncoder implements BiConsumer<String, Writer> {
+    /**
+     * Encoder that escapes JSON string values/fields.
+     */
+    static class JsonEscapeEncoder implements Encoder {
 
         @Override
-        public void accept(String s, Writer writer) {
+        public void encode(String s, Writer writer) throws IOException {
-            try {
             writer.write(JsonStringEncoder.getInstance().quoteAsString(s));
-            } catch (IOException e) {
-                throw new MustacheException("Failed to escape and encode value: " + s);
         }
     }
+
+    /**
+     * Encoder that escapes strings using HTML form encoding
+     */
+    static class UrlEncoder implements Encoder {
+
+        @Override
+        public void encode(String s, Writer writer) throws IOException {
+            writer.write(URLEncoder.encode(s, StandardCharsets.UTF_8.name()));
+        }
     }
 }

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

@@ -43,6 +43,8 @@ import java.security.PrivilegedAction;
 import java.util.Collections;
 import java.util.Map;
 
+import static org.elasticsearch.script.mustache.CustomMustacheFactory.CONTENT_TYPE_PARAM;
+
 /**
  * Main entry point handling template registration, compilation and
  * execution.
@@ -55,10 +57,6 @@ public final class MustacheScriptEngineService extends AbstractComponent impleme
 
     public static final String NAME = "mustache";
 
-    static final String CONTENT_TYPE_PARAM = "content_type";
-    static final String JSON_CONTENT_TYPE = "application/json";
-    static final String PLAIN_TEXT_CONTENT_TYPE = "text/plain";
-
     /** Thread local UTF8StreamWriter to store template execution results in, thread local to save object creation.*/
     private static ThreadLocal<SoftReference<UTF8StreamWriter>> utf8StreamWriter = new ThreadLocal<>();
 
@@ -91,13 +89,16 @@ public final class MustacheScriptEngineService extends AbstractComponent impleme
      * */
     @Override
     public Object compile(String templateName, String templateSource, Map<String, String> params) {
-        final MustacheFactory factory = new CustomMustacheFactory(isJsonEscapingEnabled(params));
+        final MustacheFactory factory = createMustacheFactory(params);
         Reader reader = new FastStringReader(templateSource);
         return factory.compile(reader, "query-template");
     }
 
-    private boolean isJsonEscapingEnabled(Map<String, String> params) {
+    private CustomMustacheFactory createMustacheFactory(Map<String, String> params) {
-        return JSON_CONTENT_TYPE.equals(params.getOrDefault(CONTENT_TYPE_PARAM, JSON_CONTENT_TYPE));
+        if (params == null || params.isEmpty() || params.containsKey(CONTENT_TYPE_PARAM) == false) {
+            return new CustomMustacheFactory();
+        }
+        return new CustomMustacheFactory(params.get(CONTENT_TYPE_PARAM));
     }
 
     @Override

modules/lang-mustache/src/test/java/org/elasticsearch/script/mustache/CustomMustacheFactoryTests.java

@@ -0,0 +1,97 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.script.mustache;
+
+import com.github.mustachejava.Mustache;
+import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.script.CompiledScript;
+import org.elasticsearch.script.ExecutableScript;
+import org.elasticsearch.script.ScriptEngineService;
+import org.elasticsearch.test.ESTestCase;
+
+import java.util.Map;
+
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.singletonMap;
+import static org.elasticsearch.script.ScriptService.ScriptType.INLINE;
+import static org.elasticsearch.script.mustache.CustomMustacheFactory.CONTENT_TYPE_PARAM;
+import static org.elasticsearch.script.mustache.CustomMustacheFactory.JSON_MIME_TYPE;
+import static org.elasticsearch.script.mustache.CustomMustacheFactory.PLAIN_TEXT_MIME_TYPE;
+import static org.elasticsearch.script.mustache.CustomMustacheFactory.X_WWW_FORM_URLENCODED_MIME_TYPE;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
+
+public class CustomMustacheFactoryTests extends ESTestCase {
+
+    public void testCreateEncoder() {
+        IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> CustomMustacheFactory.createEncoder(null));
+        assertThat(e.getMessage(), equalTo("No encoder found for MIME type [null]"));
+
+        e = expectThrows(IllegalArgumentException.class, () -> CustomMustacheFactory.createEncoder(""));
+        assertThat(e.getMessage(), equalTo("No encoder found for MIME type []"));
+
+        e = expectThrows(IllegalArgumentException.class, () -> CustomMustacheFactory.createEncoder("test"));
+        assertThat(e.getMessage(), equalTo("No encoder found for MIME type [test]"));
+
+        assertThat(CustomMustacheFactory.createEncoder(CustomMustacheFactory.JSON_MIME_TYPE),
+                instanceOf(CustomMustacheFactory.JsonEscapeEncoder.class));
+        assertThat(CustomMustacheFactory.createEncoder(CustomMustacheFactory.PLAIN_TEXT_MIME_TYPE),
+                instanceOf(CustomMustacheFactory.DefaultEncoder.class));
+        assertThat(CustomMustacheFactory.createEncoder(CustomMustacheFactory.X_WWW_FORM_URLENCODED_MIME_TYPE),
+                instanceOf(CustomMustacheFactory.UrlEncoder.class));
+    }
+
+    public void testJsonEscapeEncoder() {
+        final ScriptEngineService engine = new MustacheScriptEngineService(Settings.EMPTY);
+        final Map<String, String> params = randomBoolean() ? singletonMap(CONTENT_TYPE_PARAM, JSON_MIME_TYPE) : emptyMap();
+
+        Mustache script = (Mustache) engine.compile(null, "{\"field\": \"{{value}}\"}", params);
+        CompiledScript compiled = new CompiledScript(INLINE, null, MustacheScriptEngineService.NAME, script);
+
+        ExecutableScript executable = engine.executable(compiled, singletonMap("value", "a \"value\""));
+        BytesReference result = (BytesReference) executable.run();
+        assertThat(result.utf8ToString(), equalTo("{\"field\": \"a \\\"value\\\"\"}"));
+    }
+
+    public void testDefaultEncoder() {
+        final ScriptEngineService engine = new MustacheScriptEngineService(Settings.EMPTY);
+        final Map<String, String> params = singletonMap(CONTENT_TYPE_PARAM, PLAIN_TEXT_MIME_TYPE);
+
+        Mustache script = (Mustache) engine.compile(null, "{\"field\": \"{{value}}\"}", params);
+        CompiledScript compiled = new CompiledScript(INLINE, null, MustacheScriptEngineService.NAME, script);
+
+        ExecutableScript executable = engine.executable(compiled, singletonMap("value", "a \"value\""));
+        BytesReference result = (BytesReference) executable.run();
+        assertThat(result.utf8ToString(), equalTo("{\"field\": \"a \"value\"\"}"));
+    }
+
+    public void testUrlEncoder() {
+        final ScriptEngineService engine = new MustacheScriptEngineService(Settings.EMPTY);
+        final Map<String, String> params = singletonMap(CONTENT_TYPE_PARAM, X_WWW_FORM_URLENCODED_MIME_TYPE);
+
+        Mustache script = (Mustache) engine.compile(null, "{\"field\": \"{{value}}\"}", params);
+        CompiledScript compiled = new CompiledScript(INLINE, null, MustacheScriptEngineService.NAME, script);
+
+        ExecutableScript executable = engine.executable(compiled, singletonMap("value", "tilde~ AND date:[2016 FROM*]"));
+        BytesReference result = (BytesReference) executable.run();
+        assertThat(result.utf8ToString(), equalTo("{\"field\": \"tilde%7E+AND+date%3A%5B2016+FROM*%5D\"}"));
+    }
+}

modules/lang-mustache/src/test/java/org/elasticsearch/script/mustache/MustacheScriptEngineTests.java

@@ -49,7 +49,7 @@ public class MustacheScriptEngineTests extends ESTestCase {
     @Before
     public void setup() {
         qe = new MustacheScriptEngineService(Settings.Builder.EMPTY_SETTINGS);
-        factory = new CustomMustacheFactory(true);
+        factory = new CustomMustacheFactory();
     }
 
     public void testSimpleParameterReplace() {

modules/lang-mustache/src/test/java/org/elasticsearch/script/mustache/MustacheTests.java

@@ -30,6 +30,8 @@ import org.elasticsearch.script.ScriptEngineService;
 import org.elasticsearch.test.ESTestCase;
 import org.hamcrest.Matcher;
 
+import java.net.URLEncoder;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
@@ -43,8 +45,6 @@ import static java.util.Collections.singleton;
 import static java.util.Collections.singletonMap;
 import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.elasticsearch.script.ScriptService.ScriptType.INLINE;
-import static org.elasticsearch.script.mustache.MustacheScriptEngineService.CONTENT_TYPE_PARAM;
-import static org.elasticsearch.script.mustache.MustacheScriptEngineService.PLAIN_TEXT_CONTENT_TYPE;
 import static org.hamcrest.Matchers.both;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
@@ -144,24 +144,6 @@ public class MustacheTests extends ESTestCase {
         assertThat(bytes.utf8ToString(), both(containsString("foo")).and(containsString("bar")));
     }
 
-    public void testEscaping() {
-        // json string escaping enabled:
-        Mustache mustache = (Mustache) engine.compile(null, "{ \"field1\": \"{{value}}\"}", Collections.emptyMap());
-        CompiledScript compiledScript = new CompiledScript(INLINE, "name", "mustache", mustache);
-        ExecutableScript executableScript = engine.executable(compiledScript, Collections.singletonMap("value", "a \"value\""));
-        BytesReference rawResult = (BytesReference) executableScript.run();
-        String result = rawResult.utf8ToString();
-        assertThat(result, equalTo("{ \"field1\": \"a \\\"value\\\"\"}"));
-
-        // json string escaping disabled:
-        mustache = (Mustache) engine.compile(null, "{ \"field1\": \"{{value}}\"}",
-                Collections.singletonMap(CONTENT_TYPE_PARAM, PLAIN_TEXT_CONTENT_TYPE));
-        compiledScript = new CompiledScript(INLINE, "name", "mustache", mustache);
-        executableScript = engine.executable(compiledScript, Collections.singletonMap("value", "a \"value\""));
-        rawResult = (BytesReference) executableScript.run();
-        result = rawResult.utf8ToString();
-        assertThat(result, equalTo("{ \"field1\": \"a \"value\"\"}"));
-    }
 
     public void testSizeAccessForCollectionsAndArrays() throws Exception {
         String[] randomArrayValues = generateRandomStringArray(10, 20, false);
@@ -375,6 +357,44 @@ public class MustacheTests extends ESTestCase {
         assertScript("{{#join delimiter=' and '}}params{{/join delimiter=' and '}}", params, equalTo("1 and 2 and 3 and 4"));
     }
 
+    public void testUrlEncoder() {
+        Map<String, String> urls = new HashMap<>();
+        urls.put("https://www.elastic.co",
+                "https%3A%2F%2Fwww.elastic.co");
+        urls.put("<logstash-{now/d}>",
+                "%3Clogstash-%7Bnow%2Fd%7D%3E");
+        urls.put("?query=(foo:A OR baz:B) AND title:/joh?n(ath[oa]n)/ AND date:{* TO 2012-01}",
+                "%3Fquery%3D%28foo%3AA+OR+baz%3AB%29+AND+title%3A%2Fjoh%3Fn%28ath%5Boa%5Dn%29%2F+AND+date%3A%7B*+TO+2012-01%7D");
+
+        for (Map.Entry<String, String> url : urls.entrySet()) {
+            assertScript("{{#url}}{{params}}{{/url}}", singletonMap("params", url.getKey()), equalTo(url.getValue()));
+        }
+    }
+
+    public void testUrlEncoderWithParam() throws Exception {
+        assertScript("{{#url}}{{index}}{{/url}}", singletonMap("index", "<logstash-{now/d{YYYY.MM.dd|+12:00}}>"),
+                equalTo("%3Clogstash-%7Bnow%2Fd%7BYYYY.MM.dd%7C%2B12%3A00%7D%7D%3E"));
+
+        final String random = randomAsciiOfLength(10);
+        assertScript("{{#url}}prefix_{{s}}{{/url}}", singletonMap("s", random),
+                equalTo("prefix_" + URLEncoder.encode(random, StandardCharsets.UTF_8.name())));
+    }
+
+    public void testUrlEncoderWithJoin() {
+        Map<String, Object> params = singletonMap("emails", Arrays.asList("john@smith.com", "john.smith@email.com", "jsmith@email.com"));
+        assertScript("?query={{#url}}{{#join}}emails{{/join}}{{/url}}", params,
+                equalTo("?query=john%40smith.com%2Cjohn.smith%40email.com%2Cjsmith%40email.com"));
+
+        params = singletonMap("indices", new String[]{"<logstash-{now/d-2d}>", "<logstash-{now/d-1d}>", "<logstash-{now/d}>"});
+        assertScript("{{#url}}https://localhost:9200/{{#join}}indices{{/join}}/_stats{{/url}}", params,
+                equalTo("https%3A%2F%2Flocalhost%3A9200%2F%3Clogstash-%7Bnow%2Fd-2d%7D" +
+                        "%3E%2C%3Clogstash-%7Bnow%2Fd-1d%7D%3E%2C%3Clogstash-%7Bnow%2Fd%7D%3E%2F_stats"));
+
+        params = singletonMap("fibonacci", new int[]{1, 1, 2, 3, 5, 8, 13, 21, 34, 55});
+        assertScript("{{#url}}{{#join delimiter='+'}}fibonacci{{/join delimiter='+'}}{{/url}}", params,
+                equalTo("1%2B1%2B2%2B3%2B5%2B8%2B13%2B21%2B34%2B55"));
+    }
+
     private void assertScript(String script, Map<String, Object> vars, Matcher<Object> matcher) {
         Object result = engine.executable(new CompiledScript(INLINE, "inline", "mustache", compile(script)), vars).run();
         assertThat(result, notNullValue());

modules/lang-mustache/src/test/resources/rest-api-spec/test/lang_mustache/25_custom_functions.yaml

@@ -0,0 +1,43 @@
+---
+"Rendering using {{url}} function":
+
+  - do:
+      render_search_template:
+        body:  >
+          {
+            "inline": {
+              "query": {
+                "match": {
+                  "url": "https://localhost:9200/{{#url}}{{index}}{{/url}}/{{#url}}{{type}}{{/url}}/_search"
+                }
+              }
+            },
+            "params": {
+              "index": "<logstash-{now/d-2d}>",
+              "type" : "métriques"
+            }
+          }
+
+  - match: { template_output.query.match.url: "https://localhost:9200/%3Clogstash-%7Bnow%2Fd-2d%7D%3E/m%C3%A9triques/_search" }
+
+---
+"Rendering using {{url}} and {{join}} functions":
+
+  - do:
+      render_search_template:
+        body:  >
+          {
+            "inline": {
+              "query": {
+                "match": {
+                  "url": "{{#url}}https://localhost:9200/{{#join}}indices{{/join}}/_stats{{/url}}"
+                }
+              }
+            },
+            "params": {
+              "indices": ["<logstash-{now/d-2d}>", "<logstash-{now/d-1d}>", "<logstash-{now/d}>"]
+            }
+          }
+
+    # Decoded URL is https://localhost:9200/<logstash-{now/d-2d}>,<logstash-{now/d-1d}>,<logstash-{now/d}>/_stats
+  - match: { template_output.query.match.url: "https%3A%2F%2Flocalhost%3A9200%2F%3Clogstash-%7Bnow%2Fd-2d%7D%3E%2C%3Clogstash-%7Bnow%2Fd-1d%7D%3E%2C%3Clogstash-%7Bnow%2Fd%7D%3E%2F_stats" }