Add content and expand param getters

dbwiddis · dbwiddis · commit fe3d26c491db · 2022-09-29T19:39:48.000-07:00
Signed-off-by: Daniel Widdis &lt;widdis@gmail.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -104,7 +104,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
  - Adding create component extension point support for AD ([#4517](https://github.com/opensearch-project/OpenSearch/pull/4517))
  - Add getSettings support for AD([#4519](https://github.com/opensearch-project/OpenSearch/pull/4519))
  - Fixed javadoc warning for build failure([#4581](https://github.com/opensearch-project/OpenSearch/pull/4581))
- - Pass REST params to extensions ([#4633](https://github.com/opensearch-project/OpenSearch/pull/4633))
+ - Pass REST params and content to extensions ([#4633](https://github.com/opensearch-project/OpenSearch/pull/4633))
 
 ## [2.x]
 ### Added
diff --git a/server/src/main/java/org/opensearch/extensions/rest/ExtensionRestRequest.java b/server/src/main/java/org/opensearch/extensions/rest/ExtensionRestRequest.java
@@ -8,8 +8,10 @@
 
 package org.opensearch.extensions.rest;
 
+import org.opensearch.common.bytes.BytesReference;
 import org.opensearch.common.io.stream.StreamInput;
 import org.opensearch.common.io.stream.StreamOutput;
+import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.identity.PrincipalIdentifierToken;
 import org.opensearch.rest.RestRequest;
 import org.opensearch.rest.RestRequest.Method;
@@ -33,27 +35,44 @@ public class ExtensionRestRequest extends TransportRequest {
     private Method method;
     private String uri;
     private Map<String, String> params;
+    private XContentType xContentType = null;
+    private BytesReference content;
     // The owner of this request object
     private PrincipalIdentifierToken principalIdentifierToken;
-    // Tracks consumed parameters
+
+    // Tracks consumed parameters and content
     private final Set<String> consumedParams = new HashSet<>();
+    private boolean contentConsumed = false;
 
     /**
-     * This object can be instantiated given method, uri, params and identifier
+     * This object can be instantiated given method, uri, params, content and identifier
+     *
      * @param method of type {@link Method}
      * @param uri url string
      * @param params the REST params
+     * @param xContentType the content type, or null for plain text or no content
+     * @param content the REST request content
      * @param principalIdentifier the owner of this request
      */
-    public ExtensionRestRequest(Method method, String uri, Map<String, String> params, PrincipalIdentifierToken principalIdentifier) {
+    public ExtensionRestRequest(
+        Method method,
+        String uri,
+        Map<String, String> params,
+        XContentType xContentType,
+        BytesReference content,
+        PrincipalIdentifierToken principalIdentifier
+    ) {
         this.method = method;
         this.uri = uri;
         this.params = params;
+        this.xContentType = xContentType;
+        this.content = content;
         this.principalIdentifierToken = principalIdentifier;
     }
 
     /**
      * Instantiate this request from input stream
+     *
      * @param in Input stream
      * @throws IOException on failure to read the stream
      */
@@ -62,6 +81,10 @@ public ExtensionRestRequest(StreamInput in) throws IOException {
         method = in.readEnum(RestRequest.Method.class);
         uri = in.readString();
         params = in.readMap(StreamInput::readString, StreamInput::readString);
+        if (in.readBoolean()) {
+            xContentType = in.readEnum(XContentType.class);
+        }
+        content = in.readBytesReference();
         principalIdentifierToken = in.readNamedWriteable(PrincipalIdentifierToken.class);
     }
 
@@ -71,25 +94,55 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeEnum(method);
         out.writeString(uri);
         out.writeMap(params, StreamOutput::writeString, StreamOutput::writeString);
+        out.writeBoolean(xContentType != null);
+        if (xContentType != null) {
+            out.writeEnum(xContentType);
+        }
+        out.writeBytesReference(content);
         out.writeNamedWriteable(principalIdentifierToken);
     }
 
     /**
+     * Gets the REST method
+     *
      * @return This REST request {@link Method} type
      */
     public Method method() {
         return method;
     }
 
     /**
+     * Gets the REST uri
+     *
      * @return This REST request's uri
      */
     public String uri() {
         return uri;
     }
 
+    /**
+     * Gets the full map of params without consuming them. Rest Handlers should use {@link #param(String)} or {@link #param(String, String)}
+     * to get parameter values.
+     *
+     * @return This REST request's params
+     */
+    public Map<String, String> params() {
+        return params;
+    }
+
+    /**
+     * Tests whether a parameter named {@code key} exists.
+     *
+     * @param key The parameter to test.
+     * @return True if there is a value for this parameter.
+     */
+    public boolean hasParam(String key) {
+        return params.containsKey(key);
+    }
+
     /**
      * Gets the value of a parameter, consuming it in the process.
+     *
      * @param key The parameter key
      * @return The parameter value if it exists, or null.
      */
@@ -100,8 +153,9 @@ public String param(String key) {
 
     /**
      * Gets the value of a parameter, consuming it in the process.
+     *
      * @param key The parameter key
-     * @param defaultValue A value to return if the parameter value doesn't exist.
+     * @param defaultValue  A value to return if the parameter value doesn't exist.
      * @return The parameter value if it exists, or the default value.
      */
     public String param(String key, String defaultValue) {
@@ -110,13 +164,22 @@ public String param(String key, String defaultValue) {
     }
 
     /**
-     * Gets the full map of params without consuming them.
-     * Rest Handlers should use {@link #param(String)} or {@link #param(String, String)} to get parameter values.
+     * Gets the value of a parameter as a long, consuming it in the process.
      *
-     * @return This REST request's params
+     * @param key The parameter key
+     * @param defaultValue A value to return if the parameter value doesn't exist.
+     * @return The parameter value if it exists, or the default value.
      */
-    public Map<String, String> params() {
-        return params;
+    public long paramAsLong(String key, long defaultValue) {
+        String value = param(key);
+        if (value == null) {
+            return defaultValue;
+        }
+        try {
+            return Long.parseLong(value);
+        } catch (NumberFormatException e) {
+            throw new IllegalArgumentException("Unable to parse param '" + key + "' value '" + value + "' to a long.", e);
+        }
     }
 
     /**
@@ -128,6 +191,43 @@ public List<String> consumedParams() {
         return new ArrayList<>(consumedParams);
     }
 
+    /**
+     * Gets the content type, if any.
+     *
+     * @return the content type of the {@link #content()}, or null if the context is plain text or if there is no content.
+     */
+    public XContentType getXContentType() {
+        return xContentType;
+    }
+
+    /**
+     * Gets the content.
+     *
+     * @return This REST request's content
+     */
+    public BytesReference content() {
+        contentConsumed = true;
+        return content;
+    }
+
+    /**
+     * Tests whether content exists.
+     *
+     * @return True if there is non-empty content.
+     */
+    public boolean hasContent() {
+        return content.length() > 0;
+    }
+
+    /**
+     * Tests whether content has been consumed.
+     *
+     * @return True if the content was consumed.
+     */
+    public boolean isContentConsumed() {
+        return contentConsumed;
+    }
+
     /**
      * @return This REST request issuer's identity token
      */
@@ -143,6 +243,10 @@ public String toString() {
             + uri
             + ", params="
             + params
+            + ", xContentType="
+            + xContentType
+            + ", contentLength="
+            + content.length()
             + ", requester="
             + principalIdentifierToken.getToken()
             + "}";
@@ -156,11 +260,13 @@ public boolean equals(Object obj) {
         return Objects.equals(method, that.method)
             && Objects.equals(uri, that.uri)
             && Objects.equals(params, that.params)
+            && Objects.equals(xContentType, that.xContentType)
+            && Objects.equals(content, that.content)
             && Objects.equals(principalIdentifierToken, that.principalIdentifierToken);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(method, uri, params, principalIdentifierToken);
+        return Objects.hash(method, uri, params, xContentType, content, principalIdentifierToken);
     }
 }
diff --git a/server/src/main/java/org/opensearch/extensions/rest/RestSendToExtensionAction.java b/server/src/main/java/org/opensearch/extensions/rest/RestSendToExtensionAction.java
@@ -11,7 +11,9 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.client.node.NodeClient;
+import org.opensearch.common.bytes.BytesReference;
 import org.opensearch.common.io.stream.StreamInput;
+import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.extensions.DiscoveryExtension;
 import org.opensearch.extensions.ExtensionsOrchestrator;
 import org.opensearch.identity.ExtensionTokenProcessor;
@@ -48,6 +50,7 @@ public class RestSendToExtensionAction extends BaseRestHandler {
     private static final String SEND_TO_EXTENSION_ACTION = "send_to_extension_action";
     private static final Logger logger = LogManager.getLogger(RestSendToExtensionAction.class);
     private static final String CONSUMED_PARAMS_KEY = "extension.consumed.parameters";
+    private static final String CONSUMED_CONTENT_KEY = "extension.consumed.content";
     // To replace with user identity see https://github.com/opensearch-project/OpenSearch/pull/4247
     private static final Principal DEFAULT_PRINCIPAL = new Principal() {
         @Override
@@ -108,6 +111,8 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
         Method method = request.method();
         String uri = request.uri();
         Map<String, String> params = request.params();
+        XContentType contentType = request.getXContentType();
+        BytesReference content = request.content();
 
         if (uri.startsWith(uriPrefix)) {
             uri = uri.substring(uriPrefix.length());
@@ -136,12 +141,20 @@ public void handleResponse(RestExecuteOnExtensionResponse response) {
                 restExecuteOnExtensionResponse.setStatus(response.getStatus());
                 restExecuteOnExtensionResponse.setContentType(response.getContentType());
                 restExecuteOnExtensionResponse.setContent(response.getContent());
-                // Extract the consumed parameters from the header
+                // Extract the consumed parameters and content from the header
                 Map<String, List<String>> headers = response.getHeaders();
                 List<String> consumedParams = headers.get(CONSUMED_PARAMS_KEY);
                 if (consumedParams != null) {
+                    // consume each param
                     consumedParams.stream().forEach(p -> request.param(p));
                 }
+                List<String> consumedContent = headers.get(CONSUMED_CONTENT_KEY);
+                if (consumedContent != null) {
+                    // conditionally consume content
+                    if (consumedParams.stream().filter(c -> Boolean.parseBoolean(c)).count() > 0) {
+                        request.content();
+                    }
+                }
                 Map<String, List<String>> headersWithoutConsumedParams = headers.entrySet()
                     .stream()
                     .filter(e -> !e.getKey().equals(CONSUMED_PARAMS_KEY))
@@ -173,7 +186,7 @@ public String executor() {
                 ExtensionsOrchestrator.REQUEST_REST_EXECUTE_ON_EXTENSION_ACTION,
                 // HERE BE DRAGONS - DO NOT INCLUDE HEADERS
                 // SEE https://github.com/opensearch-project/OpenSearch/issues/4429
-                new ExtensionRestRequest(method, uri, params, requestIssuerIdentity),
+                new ExtensionRestRequest(method, uri, params, contentType, content, requestIssuerIdentity),
                 restExecuteOnExtensionResponseHandler
             );
             try {
diff --git a/server/src/test/java/org/opensearch/extensions/rest/ExtensionRestRequestTests.java b/server/src/test/java/org/opensearch/extensions/rest/ExtensionRestRequestTests.java
