mustache: Improve the mustache script engine

* Added a `content_type` option at compile time to decide how variable values are encoded. Possible values are `application/json` and `plain/text`. Defaults to `application/json`.
* Added support for variable placeholders to lookup values from specific slots in arrays/lists.

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

@@ -0,0 +1,153 @@
+/*
+ * 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.reflect.ReflectionObjectHandler;
+import org.elasticsearch.common.util.iterable.Iterables;
+
+import java.lang.reflect.Array;
+import java.util.AbstractMap;
+import java.util.Collection;
+import java.util.Set;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.HashMap;
+
+final class CustomReflectionObjectHandler extends ReflectionObjectHandler {
+
+    @Override
+    public Object coerce(Object object) {
+        if (object == null) {
+            return null;
+        }
+
+        if (object.getClass().isArray()) {
+            return new ArrayMap(object);
+        } else if (object instanceof Collection) {
+            @SuppressWarnings("unchecked")
+            Collection<Object> collection = (Collection<Object>) object;
+            return new CollectionMap(collection);
+        } else {
+            return super.coerce(object);
+        }
+    }
+
+    final static class ArrayMap extends AbstractMap<Object, Object> implements Iterable<Object> {
+
+        private final Object array;
+        private final int length;
+
+        public ArrayMap(Object array) {
+            this.array = array;
+            this.length = Array.getLength(array);
+        }
+
+        @Override
+        public Object get(Object key) {
+            if (key instanceof Number) {
+                return Array.get(array, ((Number) key).intValue());
+            }
+            try {
+                int index = Integer.parseInt(key.toString());
+                return Array.get(array, index);
+            } catch (NumberFormatException nfe) {
+                // if it's not a number it is as if the key doesn't exist
+                return null;
+            }
+        }
+
+        @Override
+        public boolean containsKey(Object key) {
+            return get(key) != null;
+        }
+
+        @Override
+        public Set<Entry<Object, Object>> entrySet() {
+            Map<Object, Object> map = new HashMap<>(length);
+            for (int i = 0; i < length; i++) {
+                map.put(i, Array.get(array, i));
+            }
+            return map.entrySet();
+        }
+
+        @Override
+        public Iterator<Object> iterator() {
+            return new Iterator<Object>() {
+
+                int index = 0;
+
+                @Override
+                public boolean hasNext() {
+                    return index < length;
+                }
+
+                @Override
+                public Object next() {
+                    return Array.get(array, index++);
+                }
+            };
+        }
+
+    }
+
+    final static class CollectionMap extends AbstractMap<Object, Object> implements Iterable<Object> {
+
+        private final Collection<Object> col;
+
+        public CollectionMap(Collection<Object> col) {
+            this.col = col;
+        }
+
+        @Override
+        public Object get(Object key) {
+            if (key instanceof Number) {
+                return Iterables.get(col, ((Number) key).intValue());
+            }
+            try {
+                int index = Integer.parseInt(key.toString());
+                return Iterables.get(col, index);
+            } catch (NumberFormatException nfe) {
+                // if it's not a number it is as if the key doesn't exist
+                return null;
+            }
+        }
+
+        @Override
+        public boolean containsKey(Object key) {
+            return get(key) != null;
+        }
+
+        @Override
+        public Set<Entry<Object, Object>> entrySet() {
+            Map<Object, Object> map = new HashMap<>(col.size());
+            int i = 0;
+            for (Object item : col) {
+                map.put(i++, item);
+            }
+            return map.entrySet();
+        }
+
+        @Override
+        public Iterator<Object> iterator() {
+            return col.iterator();
+        }
+    }
+
+}

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

@@ -28,13 +28,12 @@ import java.io.Writer;
 /**
  * A MustacheFactory that does simple JSON escaping.
  */
-public final class JsonEscapingMustacheFactory extends DefaultMustacheFactory {
-    
+final class JsonEscapingMustacheFactory extends DefaultMustacheFactory {
+
     @Override
     public void encode(String value, Writer writer) {
         try {
-            JsonStringEncoder utils = new JsonStringEncoder();
-            writer.write(utils.quoteAsString(value));;
+            writer.write(JsonStringEncoder.getInstance().quoteAsString(value));
         } catch (IOException e) {
             throw new MustacheException("Failed to encode value: " + value);
         }

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

@@ -22,6 +22,7 @@ import java.lang.ref.SoftReference;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 import java.util.Collections;
+import java.io.Reader;
 import java.util.Map;
 
 import org.elasticsearch.SpecialPermission;
@@ -40,6 +41,7 @@ import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import com.github.mustachejava.Mustache;
+import com.github.mustachejava.DefaultMustacheFactory;
 
 /**
  * Main entry point handling template registration, compilation and
@@ -49,9 +51,12 @@ import com.github.mustachejava.Mustache;
  * process: First compile the string representing the template, the resulting
  * {@link Mustache} object can then be re-used for subsequent executions.
  */
-public class MustacheScriptEngineService extends AbstractComponent implements ScriptEngineService {
+public final class MustacheScriptEngineService extends AbstractComponent implements ScriptEngineService {
 
     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<>();
@@ -86,8 +91,21 @@ public class MustacheScriptEngineService extends AbstractComponent implements Sc
      * */
     @Override
     public Object compile(String template, Map<String, String> params) {
-        /** Factory to generate Mustache objects from. */
-        return (new JsonEscapingMustacheFactory()).compile(new FastStringReader(template), "query-template");
+        String contentType = params.getOrDefault(CONTENT_TYPE_PARAM, JSON_CONTENT_TYPE);
+        final DefaultMustacheFactory mustacheFactory;
+        switch (contentType){
+            case PLAIN_TEXT_CONTENT_TYPE:
+                mustacheFactory = new NoneEscapingMustacheFactory();
+                break;
+            case JSON_CONTENT_TYPE:
+            default:
+                // assume that the default is json encoding:
+                mustacheFactory = new JsonEscapingMustacheFactory();
+                break;
+        }
+        mustacheFactory.setObjectHandler(new CustomReflectionObjectHandler());
+        Reader reader = new FastStringReader(template);
+        return mustacheFactory.compile(reader, "query-template");
     }
 
     @Override

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

@@ -0,0 +1,40 @@
+/*
+ * 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.DefaultMustacheFactory;
+import com.github.mustachejava.MustacheException;
+
+import java.io.IOException;
+import java.io.Writer;
+
+/**
+ * A MustacheFactory that does no string escaping.
+ */
+final class NoneEscapingMustacheFactory extends DefaultMustacheFactory {
+
+    @Override
+    public void encode(String value, Writer writer) {
+        try {
+            writer.write(value);
+        } catch (IOException e) {
+            throw new MustacheException("Failed to encode value: " + value);
+        }
+    }
+}

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

@@ -48,12 +48,13 @@ public class MustacheScriptEngineTests extends ESTestCase {
     }
 
     public void testSimpleParameterReplace() {
+        Map<String, String> compileParams = Collections.singletonMap("content_type", "application/json");
         {
             String template = "GET _search {\"query\": " + "{\"boosting\": {" + "\"positive\": {\"match\": {\"body\": \"gift\"}},"
                     + "\"negative\": {\"term\": {\"body\": {\"value\": \"solr\"}" + "}}, \"negative_boost\": {{boost_val}} } }}";
             Map<String, Object> vars = new HashMap<>();
             vars.put("boost_val", "0.3");
-            BytesReference o = (BytesReference) qe.executable(new CompiledScript(ScriptService.ScriptType.INLINE, "", "mustache", qe.compile(template, Collections.emptyMap())), vars).run();
+            BytesReference o = (BytesReference) qe.executable(new CompiledScript(ScriptService.ScriptType.INLINE, "", "mustache", qe.compile(template, compileParams)), vars).run();
             assertEquals("GET _search {\"query\": {\"boosting\": {\"positive\": {\"match\": {\"body\": \"gift\"}},"
                     + "\"negative\": {\"term\": {\"body\": {\"value\": \"solr\"}}}, \"negative_boost\": 0.3 } }}",
                     new String(o.toBytes(), Charset.forName("UTF-8")));
@@ -64,7 +65,7 @@ public class MustacheScriptEngineTests extends ESTestCase {
             Map<String, Object> vars = new HashMap<>();
             vars.put("boost_val", "0.3");
             vars.put("body_val", "\"quick brown\"");
-            BytesReference o = (BytesReference) qe.executable(new CompiledScript(ScriptService.ScriptType.INLINE, "", "mustache", qe.compile(template, Collections.emptyMap())), vars).run();
+            BytesReference o = (BytesReference) qe.executable(new CompiledScript(ScriptService.ScriptType.INLINE, "", "mustache", qe.compile(template, compileParams)), vars).run();
             assertEquals("GET _search {\"query\": {\"boosting\": {\"positive\": {\"match\": {\"body\": \"gift\"}},"
                     + "\"negative\": {\"term\": {\"body\": {\"value\": \"\\\"quick brown\\\"\"}}}, \"negative_boost\": 0.3 } }}",
                     new String(o.toBytes(), Charset.forName("UTF-8")));

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

@@ -18,36 +18,142 @@
  */
 package org.elasticsearch.script.mustache;
 
-import com.github.mustachejava.DefaultMustacheFactory;
 import com.github.mustachejava.Mustache;
-import com.github.mustachejava.MustacheFactory;
+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.script.ScriptService;
 import org.elasticsearch.test.ESTestCase;
 
-import java.io.StringReader;
-import java.io.StringWriter;
+
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.HashSet;
+
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonMap;
+import static org.elasticsearch.script.mustache.MustacheScriptEngineService.CONTENT_TYPE_PARAM;
+import static org.elasticsearch.script.mustache.MustacheScriptEngineService.JSON_CONTENT_TYPE;
+import static org.elasticsearch.script.mustache.MustacheScriptEngineService.PLAIN_TEXT_CONTENT_TYPE;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.both;
+import static org.hamcrest.Matchers.containsString;
 
-/**
- * Figure out how Mustache works for the simplest use case. Leaving in here for now for reference.
- * */
 public class MustacheTests extends ESTestCase {
-    public void test() {
-        HashMap<String, Object> scopes = new HashMap<>();
-        scopes.put("boost_val", "0.2");
 
+    private ScriptEngineService engine = new MustacheScriptEngineService(Settings.EMPTY);
+
+    public void testBasics() {
         String template = "GET _search {\"query\": " + "{\"boosting\": {"
-                + "\"positive\": {\"match\": {\"body\": \"gift\"}},"
-                + "\"negative\": {\"term\": {\"body\": {\"value\": \"solr\"}"
-                + "}}, \"negative_boost\": {{boost_val}} } }}";
-        MustacheFactory f = new DefaultMustacheFactory();
-        Mustache mustache = f.compile(new StringReader(template), "example");
-        StringWriter writer = new StringWriter();
-        mustache.execute(writer, scopes);
-        writer.flush();
+            + "\"positive\": {\"match\": {\"body\": \"gift\"}},"
+            + "\"negative\": {\"term\": {\"body\": {\"value\": \"solr\"}"
+            + "}}, \"negative_boost\": {{boost_val}} } }}";
+        Map<String, Object> params = Collections.singletonMap("boost_val", "0.2");
+
+        Mustache mustache = (Mustache) engine.compile(template, Collections.emptyMap());
+        CompiledScript compiledScript = new CompiledScript(ScriptService.ScriptType.INLINE, "my-name", "mustache", mustache);
+        ExecutableScript result = engine.executable(compiledScript, params);
         assertEquals(
                 "Mustache templating broken",
                 "GET _search {\"query\": {\"boosting\": {\"positive\": {\"match\": {\"body\": \"gift\"}},"
                         + "\"negative\": {\"term\": {\"body\": {\"value\": \"solr\"}}}, \"negative_boost\": 0.2 } }}",
-                writer.toString());
+                ((BytesReference) result.run()).toUtf8()
+        );
+    }
+
+    public void testArrayAccess() throws Exception {
+        String template = "{{data.0}} {{data.1}}";
+        CompiledScript mustache = new CompiledScript(ScriptService.ScriptType.INLINE, "inline", "mustache", engine.compile(template, Collections.emptyMap()));
+        Map<String, Object> vars = new HashMap<>();
+        Object data = randomFrom(
+            new String[] { "foo", "bar" },
+            Arrays.asList("foo", "bar"));
+        vars.put("data", data);
+        Object output = engine.executable(mustache, vars).run();
+        assertThat(output, notNullValue());
+        assertThat(output, instanceOf(BytesReference.class));
+        BytesReference bytes = (BytesReference) output;
+        assertThat(bytes.toUtf8(), equalTo("foo bar"));
+
+        // Sets can come out in any order
+        Set<String> setData = new HashSet<>();
+        setData.add("foo");
+        setData.add("bar");
+        vars.put("data", setData);
+        output = engine.executable(mustache, vars).run();
+        assertThat(output, notNullValue());
+        assertThat(output, instanceOf(BytesReference.class));
+        bytes = (BytesReference) output;
+        assertThat(bytes.toUtf8(), both(containsString("foo")).and(containsString("bar")));
+    }
+
+    public void testArrayInArrayAccess() throws Exception {
+        String template = "{{data.0.0}} {{data.0.1}}";
+        CompiledScript mustache = new CompiledScript(ScriptService.ScriptType.INLINE, "inline", "mustache", engine.compile(template, Collections.emptyMap()));
+        Map<String, Object> vars = new HashMap<>();
+        Object data = randomFrom(
+            new String[][] { new String[] { "foo", "bar" }},
+            Collections.singletonList(new String[] { "foo", "bar" }),
+            singleton(new String[] { "foo", "bar" })
+        );
+        vars.put("data", data);
+        Object output = engine.executable(mustache, vars).run();
+        assertThat(output, notNullValue());
+        assertThat(output, instanceOf(BytesReference.class));
+        BytesReference bytes = (BytesReference) output;
+        assertThat(bytes.toUtf8(), equalTo("foo bar"));
+    }
+
+    public void testMapInArrayAccess() throws Exception {
+        String template = "{{data.0.key}} {{data.1.key}}";
+        CompiledScript mustache = new CompiledScript(ScriptService.ScriptType.INLINE, "inline", "mustache", engine.compile(template, Collections.emptyMap()));
+        Map<String, Object> vars = new HashMap<>();
+        Object data = randomFrom(
+            new Object[] { singletonMap("key", "foo"), singletonMap("key", "bar") },
+            Arrays.asList(singletonMap("key", "foo"), singletonMap("key", "bar")));
+        vars.put("data", data);
+        Object output = engine.executable(mustache, vars).run();
+        assertThat(output, notNullValue());
+        assertThat(output, instanceOf(BytesReference.class));
+        BytesReference bytes = (BytesReference) output;
+        assertThat(bytes.toUtf8(), equalTo("foo bar"));
+
+        // HashSet iteration order isn't fixed
+        Set<Object> setData = new HashSet<>();
+        setData.add(singletonMap("key", "foo"));
+        setData.add(singletonMap("key", "bar"));
+        vars.put("data", setData);
+        output = engine.executable(mustache, vars).run();
+        assertThat(output, notNullValue());
+        assertThat(output, instanceOf(BytesReference.class));
+        bytes = (BytesReference) output;
+        assertThat(bytes.toUtf8(), both(containsString("foo")).and(containsString("bar")));
     }
+
+    public void testEscaping() {
+        // json string escaping enabled:
+        Map<String, String> params = randomBoolean() ? Collections.emptyMap() : Collections.singletonMap(CONTENT_TYPE_PARAM, JSON_CONTENT_TYPE);
+        Mustache mustache = (Mustache) engine.compile("{ \"field1\": \"{{value}}\"}", Collections.emptyMap());
+        CompiledScript compiledScript = new CompiledScript(ScriptService.ScriptType.INLINE, "name", "mustache", mustache);
+        ExecutableScript executableScript = engine.executable(compiledScript, Collections.singletonMap("value", "a \"value\""));
+        BytesReference rawResult = (BytesReference) executableScript.run();
+        String result = rawResult.toUtf8();
+        assertThat(result, equalTo("{ \"field1\": \"a \\\"value\\\"\"}"));
+
+        // json string escaping disabled:
+        mustache = (Mustache) engine.compile("{ \"field1\": \"{{value}}\"}", Collections.singletonMap(CONTENT_TYPE_PARAM, PLAIN_TEXT_CONTENT_TYPE));
+        compiledScript = new CompiledScript(ScriptService.ScriptType.INLINE, "name", "mustache", mustache);
+        executableScript = engine.executable(compiledScript, Collections.singletonMap("value", "a \"value\""));
+        rawResult = (BytesReference) executableScript.run();
+        result = rawResult.toUtf8();
+        assertThat(result, equalTo("{ \"field1\": \"a \"value\"\"}"));
+    }
+
 }