Provide service accounts tokens to extensions

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

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();

plugins/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()));
+    }
 }

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

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;
     }

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

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;

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);

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

server/src/main/java/org/opensearch/node/Node.java

@@ -495,7 +495,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();
             }

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

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());
             }
         }
     }