Make scripting media type parsing more flexible

libs/x-content/src/main/java/org/elasticsearch/common/xcontent/MediaType.java

@@ -22,8 +22,8 @@
 import org.elasticsearch.common.collect.Tuple;
 
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 /**
  * Abstracts a <a href="http://en.wikipedia.org/wiki/Internet_media_type">Media Type</a> and a query parameter <code>format</code>.
@@ -44,7 +44,7 @@ public interface MediaType {
      * Returns a set of HeaderValues - allowed media type values on Accept or Content-Type headers
      * Also defines media type parameters for validation.
      */
-    Set<HeaderValue> headerValues();
+    List<HeaderValue> headerValues();
 
     /**
      * A class to represent supported mediaType values i.e. application/json and parameters to be validated.

libs/x-content/src/main/java/org/elasticsearch/common/xcontent/MediaTypeRegistry.java

@@ -22,7 +22,6 @@
 import java.util.HashMap;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Set;
 import java.util.regex.Pattern;
 
 /**
@@ -63,12 +62,16 @@ public Map<String, Pattern> parametersFor(String typeWithSubtype) {
 
     public MediaTypeRegistry<T> register(T[] mediaTypes ) {
         for (T mediaType : mediaTypes) {
-            Set<MediaType.HeaderValue> tuples = mediaType.headerValues();
-            for (MediaType.HeaderValue headerValue : tuples) {
-                queryParamToMediaType.put(mediaType.queryParameter(), mediaType);
-                typeWithSubtypeToMediaType.put(headerValue.v1(), mediaType);
-                parametersMap.put(headerValue.v1(), convertPatterns(headerValue.v2()));
-            }
+            register(mediaType);
+        }
+        return this;
+    }
+
+    public MediaTypeRegistry<T> register(T mediaType) {
+        for (MediaType.HeaderValue headerValue : mediaType.headerValues()) {
+            queryParamToMediaType.put(mediaType.queryParameter(), mediaType);
+            typeWithSubtypeToMediaType.put(headerValue.v1(), mediaType);
+            parametersMap.put(headerValue.v1(), convertPatterns(headerValue.v2()));
         }
         return this;
     }

libs/x-content/src/main/java/org/elasticsearch/common/xcontent/ParsedMediaType.java

@@ -106,8 +106,8 @@ public static ParsedMediaType parseMediaType(String headerValue) {
         return null;
     }
 
-    public static ParsedMediaType parseMediaType(XContentType requestContentType, Map<String, String> parameters) {
-        ParsedMediaType parsedMediaType = parseMediaType(requestContentType.mediaTypeWithoutParameters());
+    public static ParsedMediaType parseMediaType(MediaType requestContentType, Map<String, String> parameters) {
+        ParsedMediaType parsedMediaType = parseMediaType(requestContentType.headerValues().get(0).v1());
         parsedMediaType.parameters.putAll(parameters);
         return parsedMediaType;
     }
@@ -159,6 +159,10 @@ public String toString() {
         return originalHeaderValue;
     }
 
+    public String responseContentTypeHeader() {
+        return responseContentTypeHeader(this.parameters);
+    }
+
     public String responseContentTypeHeader(Map<String,String> parameters) {
         return this.mediaTypeWithoutParameters() + formatParameters(parameters);
     }

libs/x-content/src/main/java/org/elasticsearch/common/xcontent/XContentBuilder.java

@@ -208,9 +208,7 @@ public XContentBuilder(XContent xContent, OutputStream os, Set<String> includes,
     }
 
     public String getResponseContentTypeString() {
-        Map<String, String> parameters = responseContentType != null ?
-            responseContentType.getParameters() : Collections.emptyMap();
-        return responseContentType.responseContentTypeHeader(parameters);
+        return responseContentType.responseContentTypeHeader();
     }
 
     public XContentType contentType() {

libs/x-content/src/main/java/org/elasticsearch/common/xcontent/XContentType.java

@@ -24,8 +24,8 @@
 import org.elasticsearch.common.xcontent.smile.SmileXContent;
 import org.elasticsearch.common.xcontent.yaml.YamlXContent;
 
+import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 /**
  * The content type of {@link org.elasticsearch.common.xcontent.XContent}.
@@ -57,8 +57,8 @@ public XContent xContent() {
         }
 
         @Override
-        public Set<HeaderValue> headerValues() {
-            return Set.of(
+        public List<HeaderValue> headerValues() {
+            return List.of(
                 new HeaderValue("application/json"),
                 new HeaderValue("application/x-ndjson"),
                 new HeaderValue("application/*"));
@@ -84,8 +84,8 @@ public XContent xContent() {
         }
 
         @Override
-        public Set<HeaderValue> headerValues() {
-            return Set.of(
+        public List<HeaderValue> headerValues() {
+            return List.of(
                 new HeaderValue("application/smile"));
         }
     },
@@ -109,8 +109,8 @@ public XContent xContent() {
         }
 
         @Override
-        public Set<HeaderValue> headerValues() {
-            return Set.of(
+        public List<HeaderValue> headerValues() {
+            return List.of(
                 new HeaderValue("application/yaml"));
         }
     },
@@ -134,8 +134,8 @@ public XContent xContent() {
         }
 
         @Override
-        public Set<HeaderValue> headerValues() {
-            return Set.of(
+        public List<HeaderValue> headerValues() {
+            return List.of(
                 new HeaderValue("application/cbor"));
         }
     },
@@ -159,8 +159,8 @@ public XContent xContent() {
         }
 
         @Override
-        public Set<HeaderValue> headerValues() {
-            return Set.of(
+        public List<HeaderValue> headerValues() {
+            return List.of(
                 new HeaderValue(VENDOR_APPLICATION_PREFIX + "json",
                     Map.of(COMPATIBLE_WITH_PARAMETER_NAME, VERSION_PATTERN)),
                 new HeaderValue(VENDOR_APPLICATION_PREFIX + "x-ndjson",
@@ -192,8 +192,8 @@ public XContent xContent() {
         }
 
         @Override
-        public Set<HeaderValue> headerValues() {
-            return Set.of(
+        public List<HeaderValue> headerValues() {
+            return List.of(
                 new HeaderValue(VENDOR_APPLICATION_PREFIX + "smile",
                     Map.of(COMPATIBLE_WITH_PARAMETER_NAME, VERSION_PATTERN)));
         }
@@ -223,8 +223,8 @@ public XContent xContent() {
         }
 
         @Override
-        public Set<HeaderValue> headerValues() {
-            return Set.of(
+        public List<HeaderValue> headerValues() {
+            return List.of(
                 new HeaderValue(VENDOR_APPLICATION_PREFIX + "yaml",
                     Map.of(COMPATIBLE_WITH_PARAMETER_NAME, VERSION_PATTERN)));
         }
@@ -254,8 +254,8 @@ public XContent xContent() {
         }
 
         @Override
-        public Set<HeaderValue> headerValues() {
-            return Set.of(
+        public List<HeaderValue> headerValues() {
+            return List.of(
                 new HeaderValue(VENDOR_APPLICATION_PREFIX + "cbor",
                     Map.of(COMPATIBLE_WITH_PARAMETER_NAME, VERSION_PATTERN)));
         }

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

@@ -30,10 +30,12 @@
 import com.github.mustachejava.codes.DefaultMustache;
 import com.github.mustachejava.codes.IterableCode;
 import com.github.mustachejava.codes.WriteCode;
-import org.apache.lucene.search.highlight.DefaultEncoder;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.xcontent.MediaType;
+import org.elasticsearch.common.xcontent.ParsedMediaType;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.script.MustacheMediaType;
 
 import java.io.IOException;
 import java.io.StringWriter;
@@ -51,28 +53,32 @@
 import java.util.regex.Pattern;
 
 public class CustomMustacheFactory extends DefaultMustacheFactory {
+    private static final MediaType DEFAULT_MIME_TYPE = XContentType.JSON;
 
-    static final String JSON_MIME_TYPE_WITH_CHARSET = "application/json;charset=utf-8";
-    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 = Map.of(
-            JSON_MIME_TYPE_WITH_CHARSET, JsonEscapeEncoder::new,
-            JSON_MIME_TYPE, JsonEscapeEncoder::new,
-            PLAIN_TEXT_MIME_TYPE, DefaultEncoder::new,
-            X_WWW_FORM_URLENCODED_MIME_TYPE, UrlEncoder::new);
+    private static final Map<MediaType, Supplier<Encoder>> ENCODERS = Map.of(
+        XContentType.JSON, JsonEscapeEncoder::new,
+        MustacheMediaType.PLAIN_TEXT, DefaultEncoder::new,
+        MustacheMediaType.X_WWW_FORM_URLENCODED, UrlEncoder::new);
 
     private final Encoder encoder;
 
-    public CustomMustacheFactory(String mimeType) {
-        super();
+    public CustomMustacheFactory(MediaType mimeType) {
         setObjectHandler(new CustomReflectionObjectHandler());
         this.encoder = createEncoder(mimeType);
     }
 
+    public CustomMustacheFactory(String mimeType) {
+        this(parseMediaType(mimeType));
+    }
+
+    private static MediaType parseMediaType(String mimeType) {
+        MediaType mustacheMediaType = ParsedMediaType.parseMediaType(mimeType).toMediaType(MustacheMediaType.REGISTRY);
+        if (mustacheMediaType == null) {
+            throw new IllegalArgumentException("No encoder found for MIME type [" + mimeType + "]");
+        }
+        return mustacheMediaType;
+    }
+
     public CustomMustacheFactory() {
         this(DEFAULT_MIME_TYPE);
     }
@@ -86,14 +92,19 @@ public void encode(String value, Writer writer) {
         }
     }
 
-    static Encoder createEncoder(String mimeType) {
-        final Supplier<Encoder> supplier = ENCODERS.get(mimeType);
+    private static Encoder createEncoder(MediaType mustacheMediaType) {
+        final Supplier<Encoder> supplier = ENCODERS.get(mustacheMediaType);
         if (supplier == null) {
-            throw new IllegalArgumentException("No encoder found for MIME type [" + mimeType + "]");
+            throw new IllegalArgumentException("No encoder found for MIME type [" + mustacheMediaType + "]");
         }
         return supplier.get();
     }
 
+    // package scope for testing
+    Encoder getEncoder() {
+        return encoder;
+    }
+
     @Override
     public MustacheVisitor createMustacheVisitor() {
         return new CustomMustacheVisitor(this);

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

@@ -19,54 +19,74 @@
 
 package org.elasticsearch.script.mustache;
 
+import org.elasticsearch.common.xcontent.ParsedMediaType;
+import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.script.MustacheMediaType;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptEngine;
 import org.elasticsearch.script.TemplateScript;
 import org.elasticsearch.test.ESTestCase;
 
+import java.util.Collections;
 import java.util.Map;
 
-import static java.util.Collections.emptyMap;
 import static java.util.Collections.singletonMap;
-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 {
+    String JSON_MIME_TYPE_WITH_CHARSET = ParsedMediaType.parseMediaType(XContentType.JSON, Map.of("charset", "utf-8"))
+        .responseContentTypeHeader();
+    String JSON_MIME_TYPE = ParsedMediaType.parseMediaType(XContentType.JSON, Collections.emptyMap())
+        .responseContentTypeHeader();
+    String PLAIN_TEXT_MIME_TYPE = ParsedMediaType.parseMediaType(MustacheMediaType.PLAIN_TEXT, Collections.emptyMap())
+        .responseContentTypeHeader();
+    String X_WWW_FORM_URLENCODED_MIME_TYPE = ParsedMediaType.parseMediaType(MustacheMediaType.X_WWW_FORM_URLENCODED, Collections.emptyMap())
+        .responseContentTypeHeader();
 
     public void testCreateEncoder() {
         {
             final IllegalArgumentException e =
-                    expectThrows(IllegalArgumentException.class, () -> CustomMustacheFactory.createEncoder("non-existent"));
-            assertThat(e.getMessage(), equalTo("No encoder found for MIME type [non-existent]"));
+                expectThrows(IllegalArgumentException.class, () -> new CustomMustacheFactory("non-existent"));
+            assertThat(e.getMessage(), equalTo("invalid media-type [non-existent]"));
         }
 
         {
-            final IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> CustomMustacheFactory.createEncoder(""));
-            assertThat(e.getMessage(), equalTo("No encoder found for MIME type []"));
+            final IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> new CustomMustacheFactory(""));
+            assertThat(e.getMessage(), equalTo("invalid media-type []"));
         }
 
         {
             final IllegalArgumentException e =
-                    expectThrows(IllegalArgumentException.class, () -> CustomMustacheFactory.createEncoder("test"));
-            assertThat(e.getMessage(), equalTo("No encoder found for MIME type [test]"));
+                expectThrows(IllegalArgumentException.class, () -> new CustomMustacheFactory("test"));
+            assertThat(e.getMessage(), equalTo("invalid media-type [test]"));
         }
 
-        assertThat(CustomMustacheFactory.createEncoder(CustomMustacheFactory.JSON_MIME_TYPE_WITH_CHARSET),
+        String contentType = ParsedMediaType.parseMediaType("application/json ; charset=UTF-8")
+            .responseContentTypeHeader();
+        assertThat(new CustomMustacheFactory(contentType).getEncoder(),
+            instanceOf(CustomMustacheFactory.JsonEscapeEncoder.class));
+        assertThat(new CustomMustacheFactory(JSON_MIME_TYPE_WITH_CHARSET).getEncoder(),
+            instanceOf(CustomMustacheFactory.JsonEscapeEncoder.class));
+        assertThat(new CustomMustacheFactory(JSON_MIME_TYPE).getEncoder(),
+            instanceOf(CustomMustacheFactory.JsonEscapeEncoder.class));
+
+        assertThat(new CustomMustacheFactory(PLAIN_TEXT_MIME_TYPE).getEncoder(),
+            instanceOf(CustomMustacheFactory.DefaultEncoder.class));
+        assertThat(new CustomMustacheFactory(X_WWW_FORM_URLENCODED_MIME_TYPE).getEncoder(),
+            instanceOf(CustomMustacheFactory.UrlEncoder.class));
+
+        //TODO PG should this be rejected?
+        contentType = ParsedMediaType.parseMediaType("application/json ; unknown=UTF-8")
+            .responseContentTypeHeader();
+        assertThat(new CustomMustacheFactory(contentType).getEncoder(),
             instanceOf(CustomMustacheFactory.JsonEscapeEncoder.class));
-        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 ScriptEngine engine = new MustacheScriptEngine();
-        final Map<String, String> params = randomBoolean() ? singletonMap(Script.CONTENT_TYPE_OPTION, JSON_MIME_TYPE) : emptyMap();
+        final Map<String, String> params = randomBoolean() ? singletonMap(Script.CONTENT_TYPE_OPTION, JSON_MIME_TYPE) :
+            Collections.emptyMap();
 
         TemplateScript.Factory compiled = engine.compile(null, "{\"field\": \"{{value}}\"}", TemplateScript.CONTEXT, params);
 

qa/rolling-upgrade/src/test/resources/rest-api-spec/test/mixed_cluster/10_basic.yml

@@ -1,9 +1,5 @@
 ---
 "Verify that we can still find things with the template":
-  - skip:
-      version: "all"
-      reason: "Awaits fix: https://github.com/elastic/elasticsearch/issues/66986"
-
   - do:
       search_template:
         rest_total_hits_as_int: true

qa/rolling-upgrade/src/test/resources/rest-api-spec/test/upgraded_cluster/10_basic.yml

@@ -1,8 +1,5 @@
 ---
 "Verify that we can still find things with the template":
-  - skip:
-      version: "all"
-      reason: "Awaits fix: https://github.com/elastic/elasticsearch/issues/66986"
 
   - do:
       search_template: