Merge pull request #20 from peternied/pie-petern

Fix tests and fully implement named writeable for PrincipalIdentifierToken

server/src/main/java/org/opensearch/extensions/rest/RestExecuteOnExtensionRequest.java

@@ -30,10 +30,10 @@ public class RestExecuteOnExtensionRequest extends TransportRequest {
 
     private PrincipalIdentifierToken requestIssuerIdentity;
 
-    public RestExecuteOnExtensionRequest(Method method, String uri, PrincipalIdentifierToken token) {
+    public RestExecuteOnExtensionRequest(Method method, String uri, PrincipalIdentifierToken requestorIdentifier) {
         this.method = method;
         this.uri = uri;
-        this.requestIssuerIdentity = token;
+        this.requestIssuerIdentity = requestorIdentifier;
     }
 
     public RestExecuteOnExtensionRequest(StreamInput in) throws IOException {
@@ -44,7 +44,7 @@ public RestExecuteOnExtensionRequest(StreamInput in) throws IOException {
             throw new IOException(e);
         }
         uri = in.readString();
-        requestIssuerIdentity = in.readNamedWriteable(PrincipalIdentifierToken.class, PrincipalIdentifierToken.NAME);
+        requestIssuerIdentity = in.readNamedWriteable(PrincipalIdentifierToken.class);
     }
 
     @Override

server/src/main/java/org/opensearch/extensions/rest/RestSendToExtensionAction.java

@@ -15,7 +15,6 @@
 import org.opensearch.extensions.DiscoveryExtension;
 import org.opensearch.extensions.ExtensionsOrchestrator;
 import org.opensearch.identity.ExtensionTokenProcessor;
-import org.opensearch.identity.Principal;
 import org.opensearch.identity.PrincipalIdentifierToken;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.BytesRestResponse;
@@ -29,6 +28,7 @@
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.security.Principal;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -48,6 +48,13 @@ 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";
+    // To replace with user identity see https://github.com/opensearch-project/OpenSearch/pull/4247
+    private static final Principal DEFAULT_PRINCIPAL = new Principal() {
+        @Override
+        public String getName() {
+            return "OpenSearchUser";
+        }
+    };
 
     private final List<Route> routes;
     private final String uriPrefix;
@@ -100,21 +107,10 @@ public List<Route> routes() {
     public RestChannelConsumer prepareRequest(final RestRequest request, final NodeClient client) throws IOException {
         Method method = request.getHttpRequest().method();
         String uri = request.getHttpRequest().uri();
-
-        // TODO: should be replaced with MyShiro calls (fetch user/identity from shiro)
-        // Principal principal = getCurrentSubject().getPrincipal();
-        /* assuming admin principal for now until shiro realm is implemented */
-        Principal principal = new Principal("admin");
-        /* getting extension id for this request which is unique to every extension */
-        String discoveryExtensionId = this.discoveryExtension.getId();
-        ExtensionTokenProcessor extensionTokenProcessor = new ExtensionTokenProcessor(discoveryExtensionId);
-
-        PrincipalIdentifierToken token = extensionTokenProcessor.generateRequestIssuerIdentity(principal);
-
         if (uri.startsWith(uriPrefix)) {
             uri = uri.substring(uriPrefix.length());
         }
-        String message = "Forwarding the request " + method + " " + uri + " with owner " + token.getToken() + " to " + discoveryExtension;
+        String message = "Forwarding the request " + method + " " + uri + " to " + discoveryExtension;
         logger.info(message);
         // Initialize response. Values will be changed in the handler.
         final RestExecuteOnExtensionResponse restExecuteOnExtensionResponse = new RestExecuteOnExtensionResponse(
@@ -167,12 +163,15 @@ public String executor() {
             }
         };
         try {
+            final ExtensionTokenProcessor extensionTokenProcessor = new ExtensionTokenProcessor(discoveryExtension.getId());
+            final PrincipalIdentifierToken requestIssuerIdentity = extensionTokenProcessor.generateToken(DEFAULT_PRINCIPAL);
+
             transportService.sendRequest(
                 discoveryExtension,
                 ExtensionsOrchestrator.REQUEST_REST_EXECUTE_ON_EXTENSION_ACTION,
                 // HERE BE DRAGONS - DO NOT INCLUDE HEADERS
                 // SEE https://github.com/opensearch-project/OpenSearch/issues/4429
-                new RestExecuteOnExtensionRequest(method, uri, token),
+                new RestExecuteOnExtensionRequest(method, uri, requestIssuerIdentity),
                 restExecuteOnExtensionResponseHandler
             );
             try {

server/src/main/java/org/opensearch/identity/ExtensionTokenProcessor.java

@@ -11,6 +11,8 @@
 
 package org.opensearch.identity;
 
+import java.security.Principal;
+
 /**
  * Token processor class to handle token encryption/decryption
  */
@@ -31,7 +33,7 @@ public String getExtensionUniqueId() {
      * Create a two-way encrypted access token for given principal for this extension
      * @return token generated from principal
      */
-    public PrincipalIdentifierToken generateRequestIssuerIdentity(Principal principal) {
+    public PrincipalIdentifierToken generateToken(Principal principal) {
         String token = principal.getName() + ":" + extensionUniqueId;
         return new PrincipalIdentifierToken(token);
     }
@@ -49,7 +51,10 @@ public Principal extractPrincipal(PrincipalIdentifierToken token) {
             return null;
         }
 
-        return new Principal(parts[0]);
+        return new Principal() {
+            public String getName() {
+                return parts[0];
+            };
+        };
     }
-
 }

server/src/main/java/org/opensearch/identity/PrincipalIdentifierToken.java

@@ -14,6 +14,7 @@
 
 import java.io.IOException;
 import java.io.Serializable;
+import java.util.Objects;
 
 /**
  * Requester Token for requests to/from an extension
@@ -44,4 +45,17 @@ public String getWriteableName() {
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(token);
     }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj) return true;
+        if (!(obj instanceof PrincipalIdentifierToken)) return false;
+        PrincipalIdentifierToken other = (PrincipalIdentifierToken) obj;
+        return Objects.equals(token, other.token);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(token);
+    }
 }

server/src/test/java/org/opensearch/extensions/rest/RestExecuteOnExtensionTests.java

@@ -8,13 +8,13 @@
 
 package org.opensearch.extensions.rest;
 
-import org.opensearch.identity.ExtensionTokenProcessor;
-import org.opensearch.identity.Principal;
 import org.opensearch.identity.PrincipalIdentifierToken;
 import org.opensearch.rest.RestStatus;
 import org.opensearch.common.bytes.BytesReference;
 import org.opensearch.common.io.stream.BytesStreamInput;
 import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.common.io.stream.NamedWriteableRegistry;
+import org.opensearch.common.io.stream.NamedWriteableAwareStreamInput;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestRequest.Method;
 import org.opensearch.test.OpenSearchTestCase;
@@ -28,27 +28,30 @@ public class RestExecuteOnExtensionTests extends OpenSearchTestCase {
     public void testRestExecuteOnExtensionRequest() throws Exception {
         Method expectedMethod = Method.GET;
         String expectedUri = "/test/uri";
+        PrincipalIdentifierToken expectedRequestIssuerIdentity = new PrincipalIdentifierToken("requestor-identifier");
+        NamedWriteableRegistry registry = new NamedWriteableRegistry(
+            org.opensearch.common.collect.List.of(
+                new NamedWriteableRegistry.Entry(PrincipalIdentifierToken.class, PrincipalIdentifierToken.NAME, PrincipalIdentifierToken::new)
+            )
+        );
 
-        Principal p1 = new Principal("admin");
-        String extensionId = "ext_1";
-        ExtensionTokenProcessor extensionTokenProcessor = new ExtensionTokenProcessor(extensionId);
-        PrincipalIdentifierToken expectedRequesterToken = extensionTokenProcessor.generateRequestIssuerIdentity(p1);
-
-        RestExecuteOnExtensionRequest request = new RestExecuteOnExtensionRequest(expectedMethod, expectedUri, expectedRequesterToken);
+        RestExecuteOnExtensionRequest request = new RestExecuteOnExtensionRequest(expectedMethod, expectedUri, expectedRequestIssuerIdentity);
 
         assertEquals(expectedMethod, request.getMethod());
         assertEquals(expectedUri, request.getUri());
-        assertEquals(expectedRequesterToken, request.getRequestIssuerIdentity());
+        assertEquals(expectedRequestIssuerIdentity, request.getRequestIssuerIdentity());
 
         try (BytesStreamOutput out = new BytesStreamOutput()) {
             request.writeTo(out);
             out.flush();
             try (BytesStreamInput in = new BytesStreamInput(BytesReference.toBytes(out.bytes()))) {
-                request = new RestExecuteOnExtensionRequest(in);
+                try (NamedWriteableAwareStreamInput nameWritableAwareIn = new NamedWriteableAwareStreamInput(in, registry)) {
+                    request = new RestExecuteOnExtensionRequest(nameWritableAwareIn);
+                }
 
                 assertEquals(expectedMethod, request.getMethod());
                 assertEquals(expectedUri, request.getUri());
-                assertEquals(expectedRequesterToken, request.getRequestIssuerIdentity());
+                assertEquals(expectedRequestIssuerIdentity, request.getRequestIssuerIdentity());
             }
         }
     }