Skip to content

Commit

Permalink
Provide service accounts tokens to extensions (opensearch-project#9618)
Browse files Browse the repository at this point in the history 
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>
(cherry picked from commit 994e115)
  • Loading branch information
@stephen-crawford @DarshitChanpura
stephen-crawford authored and DarshitChanpura committed Oct 2, 2023
1 parent 165721b commit a302af4
Show file tree
Hide file tree
Showing 12 changed files with 76 additions and 145 deletions.
8 changes: 4 additions & 4 deletions plugins/identity-shiro/src/main/java/org/opensearch/identity/shiro/ShiroTokenManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,16 +57,16 @@ public Optional<AuthenticationToken> translateAuthToken(org.opensearch.identity.
final BasicAuthToken basicAuthToken = (BasicAuthToken) authenticationToken;
return Optional.of(new UsernamePasswordToken(basicAuthToken.getUser(), basicAuthToken.getPassword()));
}

return Optional.empty();
}

@Override
public AuthToken issueOnBehalfOfToken(Subject subject, OnBehalfOfClaims claims) {

String password = generatePassword();
// Make a new ShiroSubject audience as name
final byte[] rawEncoded = Base64.getUrlEncoder().encode((claims.getAudience() + ":" + password).getBytes(UTF_8));
// Make a new ShiroSubject w/ audience as name

final String usernamePassword = new String(rawEncoded, UTF_8);
final String header = "Basic " + usernamePassword;
BasicAuthToken token = new BasicAuthToken(header);
Expand All @@ -79,12 +79,12 @@ public AuthToken issueOnBehalfOfToken(Subject subject, OnBehalfOfClaims claims)
public AuthToken issueServiceAccountToken(String audience) {

String password = generatePassword();
final byte[] rawEncoded = Base64.getUrlEncoder().withoutPadding().encode((audience + ":" + password).getBytes(UTF_8));
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;
}

Expand Down
12 changes: 12 additions & 0 deletions ...ins/identity-shiro/src/test/java/org/opensearch/identity/shiro/AuthTokenHandlerTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()));
}
}
14 changes: 12 additions & 2 deletions server/src/main/java/org/opensearch/discovery/InitializeExtensionRequest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand All @@ -52,6 +56,10 @@ public DiscoveryExtensionNode getExtension() {
return extension;
}

public String getServiceAccountHeader() {
return serviceAccountHeader;
}

@Override
public String toString() {
return "InitializeExtensionsRequest{" + "sourceNode=" + sourceNode + ", extension=" + extension + '}';
Expand All @@ -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
Expand Down
119 changes: 7 additions & 112 deletions server/src/main/java/org/opensearch/extensions/ExtensionsManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,15 +28,11 @@
import org.opensearch.core.transport.TransportResponse;
import org.opensearch.discovery.InitializeExtensionRequest;
import org.opensearch.discovery.InitializeExtensionResponse;
import org.opensearch.discovery.InitializeExtensionSecurityRequest;
import org.opensearch.discovery.InitializeExtensionSecurityResponse;
import org.opensearch.env.EnvironmentSettingsResponse;
import org.opensearch.extensions.ExtensionsSettings.Extension;
import org.opensearch.extensions.action.ExtensionActionRequest;
import org.opensearch.extensions.action.ExtensionActionResponse;
import org.opensearch.extensions.action.ExtensionTransportActionsHandler;
import org.opensearch.extensions.action.IssueServiceAccountRequest;
import org.opensearch.extensions.action.IssueServiceAccountResponse;
import org.opensearch.extensions.action.RegisterTransportActionsRequest;
import org.opensearch.extensions.action.RemoteExtensionActionResponse;
import org.opensearch.extensions.action.TransportActionRequestFromExtension;
Expand All @@ -60,9 +56,6 @@
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;

/**
Expand All @@ -83,10 +76,8 @@ public class ExtensionsManager {
public static final String REQUEST_EXTENSION_REGISTER_TRANSPORT_ACTIONS = "internal:discovery/registertransportactions";
public static final String REQUEST_REST_EXECUTE_ON_EXTENSION_ACTION = "internal:extensions/restexecuteonextensiontaction";
public static final String REQUEST_EXTENSION_HANDLE_TRANSPORT_ACTION = "internal:extensions/handle-transportaction";
public static final String REQUEST_EXTENSION_ISSUE_SERVICE_ACCOUNT = "internal:extensions/issue-service-account";
public static final String REQUEST_EXTENSION_HANDLE_REMOTE_TRANSPORT_ACTION = "internal:extensions/handle-remote-transportaction";
public static final String TRANSPORT_ACTION_REQUEST_FROM_EXTENSION = "internal:extensions/request-transportaction-from-extension";
public static final String REQUEST_EXTENSION_REGISTER_SECURITY_SETTINGS = "internal:discovery/registersecuritysettings";
public static final int EXTENSION_REQUEST_WAIT_TIMEOUT = 10;
private static final Logger logger = LogManager.getLogger(ExtensionsManager.class);

Expand Down Expand Up @@ -119,7 +110,7 @@ public static enum OpenSearchRequestType {
* @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>();
Expand All @@ -134,6 +125,7 @@ public ExtensionsManager(Set<Setting<?>> additionalSettings) throws IOException
}
this.client = null;
this.extensionTransportActionsHandler = null;
this.identityService = identityService;
}

/**
Expand All @@ -156,7 +148,6 @@ public void initializeServicesAndRestHandler(
NodeClient client,
IdentityService identityService
) {
this.identityService = identityService;
this.restActionsRequestHandler = new RestActionsRequestHandler(
actionModule.getRestController(),
extensionIdMap,
Expand Down Expand Up @@ -411,65 +402,13 @@ 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
);
initializeExtensionSecurity(extension);
}
});
}

private void initializeExtensionSecurity(DiscoveryExtensionNode extension) {
final CompletableFuture<InitializeExtensionSecurityResponse> inProgressFuture = new CompletableFuture<>();
final TransportResponseHandler<InitializeExtensionSecurityResponse> initializeExtensionSecurityResponseHandler =
new TransportResponseHandler<InitializeExtensionSecurityResponse>() {

@Override
public InitializeExtensionSecurityResponse read(StreamInput in) throws IOException {
return new InitializeExtensionSecurityResponse(in);
}

@Override
public void handleResponse(InitializeExtensionSecurityResponse response) {
logger.info("Registered security settings for " + response.getName());
inProgressFuture.complete(response);
}

@Override
public void handleException(TransportException exp) {
logger.error(new ParameterizedMessage("Extension initialization failed"), exp);
inProgressFuture.completeExceptionally(exp);
}

@Override
public String executor() {
return ThreadPool.Names.GENERIC;
}
};
try {
logger.info("Sending extension request type: " + REQUEST_EXTENSION_REGISTER_SECURITY_SETTINGS);
AuthToken serviceAccountToken = identityService.getTokenManager().issueServiceAccountToken(extension.getId());
transportService.sendRequest(
extension,
REQUEST_EXTENSION_REGISTER_SECURITY_SETTINGS,
new InitializeExtensionSecurityRequest(serviceAccountToken.asAuthHeaderValue()),
initializeExtensionSecurityResponseHandler
);

inProgressFuture.orTimeout(EXTENSION_REQUEST_WAIT_TIMEOUT, TimeUnit.SECONDS).join();
} catch (CompletionException | ConnectTransportException e) {
if (e.getCause() instanceof TimeoutException || e instanceof ConnectTransportException) {
logger.info("No response from extension to request.", e);
} else if (e.getCause() instanceof RuntimeException) {
throw (RuntimeException) e.getCause();
} else if (e.getCause() instanceof Error) {
throw (Error) e.getCause();
} else {
throw new RuntimeException(e.getCause());
}
}
}

/**
* Handles an {@link ExtensionRequest}.
*
Expand Down Expand Up @@ -507,52 +446,12 @@ TransportResponse handleExtensionRequest(ExtensionRequest extensionRequest) thro
}

/**
* A separate transport action handler used to issue service accounts to extensions during initialization
* A helper method called during initialization that issues a service accounts to extensions
* @param extension The extension to be issued a service account
*/
public void issueServiceAccount(Extension extension) {
DiscoveryExtensionNode discoveryExtensionNode = extensionIdMap.get(extension.getUniqueId());
AuthToken serviceAccountToken = identityService.getTokenManager().issueServiceAccountToken(extension.getUniqueId());
String authTokenAsString = serviceAccountToken.asAuthHeaderValue();
final CompletableFuture<IssueServiceAccountResponse> inProgressFuture = new CompletableFuture<>();
final TransportResponseHandler<IssueServiceAccountResponse> issueServiceAccountResponseHandler = new TransportResponseHandler<
IssueServiceAccountResponse>() {

@Override
public IssueServiceAccountResponse read(StreamInput in) throws IOException {
return new IssueServiceAccountResponse(in);
}

@Override
public void handleResponse(IssueServiceAccountResponse response) {
for (DiscoveryExtensionNode extension : extensionIdMap.values()) {
if (extension.getName().equals(response.getName())
&& (serviceAccountToken.equals(response.getServiceAccountString()))) {
logger.info("Successfully issued service account token to extension: " + extension.getName());
break;
}
}
inProgressFuture.complete(response);
}

@Override
public void handleException(TransportException exp) {
logger.error(new ParameterizedMessage("Issuance of service account token failed"), exp);
inProgressFuture.completeExceptionally(exp);
}

@Override
public String executor() {
return ThreadPool.Names.GENERIC;
}
};

transportService.sendRequest(
discoveryExtensionNode,
REQUEST_EXTENSION_ISSUE_SERVICE_ACCOUNT,
new IssueServiceAccountRequest(authTokenAsString),
issueServiceAccountResponseHandler
);
public String issueServiceAccount(DiscoveryExtensionNode extension) {
AuthToken serviceAccountToken = identityService.getTokenManager().issueServiceAccountToken(extension.getId());
return serviceAccountToken.asAuthHeaderValue();
}

static String getRequestExtensionActionName() {
Expand Down Expand Up @@ -611,10 +510,6 @@ void setCustomSettingsRequestHandler(CustomSettingsRequestHandler customSettings
this.customSettingsRequestHandler = customSettingsRequestHandler;
}

public void setIdentityService(IdentityService identityService) {
this.identityService = identityService;
}

AddSettingsUpdateConsumerRequestHandler getAddSettingsUpdateConsumerRequestHandler() {
return addSettingsUpdateConsumerRequestHandler;
}
Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/opensearch/extensions/NoopExtensionsManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,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
Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/opensearch/identity/tokens/BasicAuthToken.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/opensearch/node/Node.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
Expand Down
2 changes: 1 addition & 1 deletion server/src/test/java/org/opensearch/action/ActionModuleTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 8 additions & 1 deletion server/src/test/java/org/opensearch/discovery/InitializeExtensionRequestTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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);
Expand All @@ -58,6 +64,7 @@ public void testInitializeExtensionRequest() throws Exception {

assertEquals(expectedExtensionNode, initializeExtensionRequest.getExtension());
assertEquals(expectedSourceNode, initializeExtensionRequest.getSourceNode());
assertEquals(expectedServiceAccountHeader, initializeExtensionRequest.getServiceAccountHeader());
}
}
}
Expand Down
Loading

0 comments on commit a302af4

Please sign in to comment.