Provide service accounts tokens to extensions (opensearch-project#9618)

Provide service accounts tokens to extensions This change adds a new transport action which passes the extension a string representation of its service account auth token. This token is created by the TokenManager interface implementation. The token is expected to be an encoded basic auth credential string which can be used by the extension to interact with its own system index. Signed-off-by: Stephen Crawford <steecraw@amazon.com> Signed-off-by: Stephen Crawford <65832608+scrawfor99@users.noreply.github.com> Signed-off-by: Peter Nied <petern@amazon.com> Co-authored-by: Owais Kazi <owaiskazi19@gmail.com> Co-authored-by: Peter Nied <petern@amazon.com> Signed-off-by: Shivansh Arora <hishiv@amazon.com>
Gaurav614 · Apr 25, 2024 · 8aa919b · 8aa919b
1 parent cae484f
commit 8aa919b
Show file tree

Hide file tree

Showing 15 changed files with 119 additions and 34 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Introduce new dynamic cluster setting to control slice computation for concurrent segment search ([#9107](https://github.com/opensearch-project/OpenSearch/pull/9107))
 - Implement on behalf of token passing for extensions ([#8679](https://github.com/opensearch-project/OpenSearch/pull/8679))
 - Implement Visitor Design pattern in QueryBuilder to enable the capability to traverse through the complex QueryBuilder tree. ([#10110](https://github.com/opensearch-project/OpenSearch/pull/10110))
+- Provide service accounts tokens to extensions ([#9618](https://github.com/opensearch-project/OpenSearch/pull/9618))
 
 ### Dependencies
 - Bump `log4j-core` from 2.18.0 to 2.19.0

diff --git a/plugins/identity-shiro/src/main/java/org/opensearch/identity/shiro/ShiroTokenManager.java b/plugins/identity-shiro/src/main/java/org/opensearch/identity/shiro/ShiroTokenManager.java
@@ -64,9 +64,9 @@ public Optional<AuthenticationToken> translateAuthToken(org.opensearch.identity.
     public AuthToken issueOnBehalfOfToken(Subject subject, OnBehalfOfClaims claims) {
 
         String password = generatePassword();
-        final byte[] rawEncoded = Base64.getEncoder().encode((claims.getAudience() + ":" + password).getBytes(UTF_8)); // Make a new
-                                                                                                                       // ShiroSubject w/
-                                                                                                                       // audience as name
+        // Make a new ShiroSubject audience as name
+        final byte[] rawEncoded = Base64.getUrlEncoder().encode((claims.getAudience() + ":" + password).getBytes(UTF_8));
+
         final String usernamePassword = new String(rawEncoded, UTF_8);
         final String header = "Basic " + usernamePassword;
         BasicAuthToken token = new BasicAuthToken(header);
@@ -75,6 +75,19 @@ public AuthToken issueOnBehalfOfToken(Subject subject, OnBehalfOfClaims claims)
         return token;
     }
 
+    @Override
+    public AuthToken issueServiceAccountToken(String audience) {
+
+        String password = generatePassword();
+        final byte[] rawEncoded = Base64.getUrlEncoder().withoutPadding().encode((audience + ":" + password).getBytes(UTF_8)); // Make a new
+        final String usernamePassword = new String(rawEncoded, UTF_8);
+        final String header = "Basic " + usernamePassword;
+
+        BasicAuthToken token = new BasicAuthToken(header);
+        shiroTokenPasswordMap.put(token, password);
+        return token;
+    }
+
     @Override
     public Subject authenticateToken(AuthToken authToken) {
         return new NoopSubject();

diff --git a/...ins/identity-shiro/src/test/java/org/opensearch/identity/shiro/AuthTokenHandlerTests.java b/...ins/identity-shiro/src/test/java/org/opensearch/identity/shiro/AuthTokenHandlerTests.java
@@ -167,6 +167,18 @@ public void testTokenNoopIssuance() {
         Subject subject = new NoopSubject();
         AuthToken token = tokenManager.issueOnBehalfOfToken(subject, claims);
         assertTrue(token instanceof AuthToken);
+        AuthToken serviceAccountToken = tokenManager.issueServiceAccountToken("test");
+        assertTrue(serviceAccountToken instanceof AuthToken);
+        assertEquals(serviceAccountToken.asAuthHeaderValue(), "noopToken");
     }
 
+    public void testShouldSucceedIssueServiceAccountToken() {
+        String audience = "testExtensionName";
+        BasicAuthToken authToken = (BasicAuthToken) shiroAuthTokenHandler.issueServiceAccountToken(audience);
+        assertTrue(authToken instanceof BasicAuthToken);
+        UsernamePasswordToken translatedToken = (UsernamePasswordToken) shiroAuthTokenHandler.translateAuthToken(authToken).get();
+        assertEquals(authToken.getPassword(), new String(translatedToken.getPassword()));
+        assertTrue(shiroAuthTokenHandler.getShiroTokenPasswordMap().containsKey(authToken));
+        assertEquals(shiroAuthTokenHandler.getShiroTokenPasswordMap().get(authToken), new String(translatedToken.getPassword()));
+    }
 }
diff --git a/server/src/main/java/org/opensearch/discovery/InitializeExtensionRequest.java b/server/src/main/java/org/opensearch/discovery/InitializeExtensionRequest.java
@@ -25,23 +25,27 @@
 public class InitializeExtensionRequest extends TransportRequest {
     private final DiscoveryNode sourceNode;
     private final DiscoveryExtensionNode extension;
+    private final String serviceAccountHeader;
 
-    public InitializeExtensionRequest(DiscoveryNode sourceNode, DiscoveryExtensionNode extension) {
+    public InitializeExtensionRequest(DiscoveryNode sourceNode, DiscoveryExtensionNode extension, String serviceAccountHeader) {
         this.sourceNode = sourceNode;
         this.extension = extension;
+        this.serviceAccountHeader = serviceAccountHeader;
     }
 
     public InitializeExtensionRequest(StreamInput in) throws IOException {
         super(in);
         sourceNode = new DiscoveryNode(in);
         extension = new DiscoveryExtensionNode(in);
+        serviceAccountHeader = in.readString();
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         sourceNode.writeTo(out);
         extension.writeTo(out);
+        out.writeString(serviceAccountHeader);
     }
 
     public DiscoveryNode getSourceNode() {
@@ -52,6 +56,10 @@ public DiscoveryExtensionNode getExtension() {
         return extension;
     }
 
+    public String getServiceAccountHeader() {
+        return serviceAccountHeader;
+    }
+
     @Override
     public String toString() {
         return "InitializeExtensionsRequest{" + "sourceNode=" + sourceNode + ", extension=" + extension + '}';
@@ -62,7 +70,9 @@ public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         InitializeExtensionRequest that = (InitializeExtensionRequest) o;
-        return Objects.equals(sourceNode, that.sourceNode) && Objects.equals(extension, that.extension);
+        return Objects.equals(sourceNode, that.sourceNode)
+            && Objects.equals(extension, that.extension)
+            && Objects.equals(serviceAccountHeader, that.getServiceAccountHeader());
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/extensions/ExtensionsManager.java b/server/src/main/java/org/opensearch/extensions/ExtensionsManager.java
@@ -41,6 +41,7 @@
 import org.opensearch.extensions.settings.CustomSettingsRequestHandler;
 import org.opensearch.extensions.settings.RegisterCustomSettingsRequest;
 import org.opensearch.identity.IdentityService;
+import org.opensearch.identity.tokens.AuthToken;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.ConnectTransportException;
 import org.opensearch.transport.TransportException;
@@ -101,14 +102,15 @@ public static enum OpenSearchRequestType {
     private Settings environmentSettings;
     private AddSettingsUpdateConsumerRequestHandler addSettingsUpdateConsumerRequestHandler;
     private NodeClient client;
+    private IdentityService identityService;
 
     /**
      * Instantiate a new ExtensionsManager object to handle requests and responses from extensions. This is called during Node bootstrap.
      *
      * @param additionalSettings  Additional settings to read in from extension initialization request
      * @throws IOException  If the extensions discovery file is not properly retrieved.
      */
-    public ExtensionsManager(Set<Setting<?>> additionalSettings) throws IOException {
+    public ExtensionsManager(Set<Setting<?>> additionalSettings, IdentityService identityService) throws IOException {
         logger.info("ExtensionsManager initialized");
         this.initializedExtensions = new HashMap<String, DiscoveryExtensionNode>();
         this.extensionIdMap = new HashMap<String, DiscoveryExtensionNode>();
@@ -123,6 +125,7 @@ public ExtensionsManager(Set<Setting<?>> additionalSettings) throws IOException
         }
         this.client = null;
         this.extensionTransportActionsHandler = null;
+        this.identityService = identityService;
     }
 
     /**
@@ -400,7 +403,7 @@ protected void doRun() throws Exception {
                 transportService.sendRequest(
                     extension,
                     REQUEST_EXTENSION_ACTION_NAME,
-                    new InitializeExtensionRequest(transportService.getLocalNode(), extension),
+                    new InitializeExtensionRequest(transportService.getLocalNode(), extension, issueServiceAccount(extension)),
                     initializeExtensionResponseHandler
                 );
             }
@@ -443,6 +446,15 @@ TransportResponse handleExtensionRequest(ExtensionRequest extensionRequest) thro
         }
     }
 
+    /**
+     * A helper method called during initialization that issues a service accounts to extensions
+     * @param extension The extension to be issued a service account
+     */
+    private String issueServiceAccount(DiscoveryExtensionNode extension) {
+        AuthToken serviceAccountToken = identityService.getTokenManager().issueServiceAccountToken(extension.getId());
+        return serviceAccountToken.asAuthHeaderValue();
+    }
+
     static String getRequestExtensionActionName() {
         return REQUEST_EXTENSION_ACTION_NAME;
     }

diff --git a/server/src/main/java/org/opensearch/extensions/NoopExtensionsManager.java b/server/src/main/java/org/opensearch/extensions/NoopExtensionsManager.java
@@ -20,6 +20,7 @@
 import org.opensearch.transport.TransportService;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 
@@ -31,7 +32,7 @@
 public class NoopExtensionsManager extends ExtensionsManager {
 
     public NoopExtensionsManager() throws IOException {
-        super(Set.of());
+        super(Set.of(), new IdentityService(Settings.EMPTY, List.of()));
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/identity/noop/NoopTokenManager.java b/server/src/main/java/org/opensearch/identity/noop/NoopTokenManager.java
@@ -37,6 +37,20 @@ public String asAuthHeaderValue() {
         };
     }
 
+    /**
+     * Issue a new Noop Token
+     * @return a new Noop Token
+     */
+    @Override
+    public AuthToken issueServiceAccountToken(final String audience) {
+        return new AuthToken() {
+            @Override
+            public String asAuthHeaderValue() {
+                return "noopToken";
+            }
+        };
+    }
+
     @Override
     public Subject authenticateToken(AuthToken authToken) {
         return null;

diff --git a/server/src/main/java/org/opensearch/identity/tokens/BasicAuthToken.java b/server/src/main/java/org/opensearch/identity/tokens/BasicAuthToken.java
@@ -23,7 +23,7 @@ public final class BasicAuthToken implements AuthToken {
 
     public BasicAuthToken(final String headerValue) {
         final String base64Encoded = headerValue.substring(TOKEN_IDENTIFIER.length()).trim();
-        final byte[] rawDecoded = Base64.getDecoder().decode(base64Encoded);
+        final byte[] rawDecoded = Base64.getUrlDecoder().decode(base64Encoded);
         final String usernamepassword = new String(rawDecoded, StandardCharsets.UTF_8);
 
         final String[] tokenParts = usernamepassword.split(":", 2);

diff --git a/server/src/main/java/org/opensearch/identity/tokens/TokenManager.java b/server/src/main/java/org/opensearch/identity/tokens/TokenManager.java
@@ -23,6 +23,14 @@ public interface TokenManager {
      */
     public AuthToken issueOnBehalfOfToken(final Subject subject, final OnBehalfOfClaims claims);
 
+    /**
+     * Create a new service account token
+     *
+     * @param audience: A string representing the unique id of the extension for which a service account token should be generated
+     * @return a new auth token
+     */
+    public AuthToken issueServiceAccountToken(final String audience);
+
     /**
      * Authenticates a provided authToken
      * @param authToken: The authToken to authenticate

diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java
@@ -496,7 +496,7 @@ protected Node(
                 for (ExtensionAwarePlugin extAwarePlugin : extensionAwarePlugins) {
                     additionalSettings.addAll(extAwarePlugin.getExtensionSettings());
                 }
-                this.extensionsManager = new ExtensionsManager(additionalSettings);
+                this.extensionsManager = new ExtensionsManager(additionalSettings, identityService);
             } else {
                 this.extensionsManager = new NoopExtensionsManager();
             }

diff --git a/server/src/test/java/org/opensearch/action/ActionModuleTests.java b/server/src/test/java/org/opensearch/action/ActionModuleTests.java
@@ -143,7 +143,7 @@ public void testSetupRestHandlerContainsKnownBuiltin() throws IOException {
             usageService,
             null,
             new IdentityService(Settings.EMPTY, new ArrayList<>()),
-            new ExtensionsManager(Set.of())
+            new ExtensionsManager(Set.of(), new IdentityService(Settings.EMPTY, List.of()))
         );
         actionModule.initRestHandlers(null);
         // At this point the easiest way to confirm that a handler is loaded is to try to register another one on top of it and to fail

diff --git a/server/src/test/java/org/opensearch/discovery/InitializeExtensionRequestTests.java b/server/src/test/java/org/opensearch/discovery/InitializeExtensionRequestTests.java
@@ -27,6 +27,7 @@ public class InitializeExtensionRequestTests extends OpenSearchTestCase {
     public void testInitializeExtensionRequest() throws Exception {
         String expectedUniqueId = "test uniqueid";
         Version expectedVersion = Version.fromString("2.0.0");
+        String expectedServiceAccountHeader = "test";
         ExtensionDependency expectedDependency = new ExtensionDependency(expectedUniqueId, expectedVersion);
         DiscoveryExtensionNode expectedExtensionNode = new DiscoveryExtensionNode(
             "firstExtension",
@@ -46,9 +47,14 @@ public void testInitializeExtensionRequest() throws Exception {
             Version.CURRENT
         );
 
-        InitializeExtensionRequest initializeExtensionRequest = new InitializeExtensionRequest(expectedSourceNode, expectedExtensionNode);
+        InitializeExtensionRequest initializeExtensionRequest = new InitializeExtensionRequest(
+            expectedSourceNode,
+            expectedExtensionNode,
+            expectedServiceAccountHeader
+        );
         assertEquals(expectedExtensionNode, initializeExtensionRequest.getExtension());
         assertEquals(expectedSourceNode, initializeExtensionRequest.getSourceNode());
+        assertEquals(expectedServiceAccountHeader, initializeExtensionRequest.getServiceAccountHeader());
 
         try (BytesStreamOutput out = new BytesStreamOutput()) {
             initializeExtensionRequest.writeTo(out);
@@ -58,6 +64,7 @@ public void testInitializeExtensionRequest() throws Exception {
 
                 assertEquals(expectedExtensionNode, initializeExtensionRequest.getExtension());
                 assertEquals(expectedSourceNode, initializeExtensionRequest.getSourceNode());
+                assertEquals(expectedServiceAccountHeader, initializeExtensionRequest.getServiceAccountHeader());
             }
         }
     }