From e13c055e063ae3e5ee40386a53155e5c65893b58 Mon Sep 17 00:00:00 2001 From: Derek Ho Date: Thu, 9 Jan 2025 16:17:01 -0500 Subject: [PATCH] Fix test and clean up code Signed-off-by: Derek Ho --- .../action/apitokens/ApiTokenAction.java | 33 +++- .../apitokens/ApiTokenIndexHandler.java | 146 +++++++----------- .../api/RestApiAdminPrivilegesEvaluator.java | 2 +- 3 files changed, 79 insertions(+), 102 deletions(-) diff --git a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenAction.java b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenAction.java index 5da43dd47b..b612656307 100644 --- a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenAction.java +++ b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenAction.java @@ -20,6 +20,7 @@ import java.util.Map; import java.util.Set; import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; import java.util.stream.Collectors; import com.google.common.collect.ImmutableList; @@ -32,6 +33,7 @@ import org.opensearch.client.node.NodeClient; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.transport.TransportAddress; import org.opensearch.core.rest.RestStatus; @@ -47,6 +49,7 @@ import org.opensearch.security.dlic.rest.api.RestApiAdminPrivilegesEvaluator; import org.opensearch.security.dlic.rest.api.RestApiPrivilegesEvaluator; import org.opensearch.security.dlic.rest.api.SecurityApiDependencies; +import org.opensearch.security.dlic.rest.support.Utils; import org.opensearch.security.identity.SecurityTokenManager; import org.opensearch.security.privileges.PrivilegesEvaluator; import org.opensearch.security.securityconf.DynamicConfigFactory; @@ -138,12 +141,18 @@ protected RestChannelConsumer prepareRequest(final RestRequest request, final No } RestChannelConsumer doPrepareRequest(RestRequest request, NodeClient client) { - return switch (request.method()) { - case POST -> handlePost(request, client); - case DELETE -> handleDelete(request, client); - case GET -> handleGet(request, client); - default -> throw new IllegalArgumentException(request.method() + " not supported"); - }; + final var originalUserAndRemoteAddress = Utils.userAndRemoteAddressFrom(client.threadPool().getThreadContext()); + try (final ThreadContext.StoredContext ctx = client.threadPool().getThreadContext().stashContext()) { + client.threadPool() + .getThreadContext() + .putTransient(ConfigConstants.OPENDISTRO_SECURITY_USER, originalUserAndRemoteAddress.getLeft()); + return switch (request.method()) { + case POST -> handlePost(request, client); + case DELETE -> handleDelete(request, client); + case GET -> handleGet(request, client); + default -> throw new IllegalArgumentException(request.method() + " not supported"); + }; + } } private RestChannelConsumer handleGet(RestRequest request, NodeClient client) { @@ -177,7 +186,6 @@ private RestChannelConsumer handleGet(RestRequest request, NodeClient client) { private RestChannelConsumer handlePost(RestRequest request, NodeClient client) { return channel -> { - final XContentBuilder builder = channel.newBuilder(); try { final Map requestBody = request.contentOrSourceParamParser().map(); validateRequestParameters(requestBody); @@ -305,7 +313,6 @@ void validateIndexPermissionsList(List> indexPermsList) { private RestChannelConsumer handleDelete(RestRequest request, NodeClient client) { return channel -> { - final XContentBuilder builder = channel.newBuilder(); try { final Map requestBody = request.contentOrSourceParamParser().map(); @@ -449,4 +456,14 @@ protected void authorizeSecurityAccess(RestRequest request) throws IOException { throw new SecurityException("User does not have required security API access"); } } + + private T withSecurityContext(NodeClient client, Supplier operation) { + final var originalUserAndRemoteAddress = Utils.userAndRemoteAddressFrom(client.threadPool().getThreadContext()); + try (final ThreadContext.StoredContext ctx = client.threadPool().getThreadContext().stashContext()) { + client.threadPool() + .getThreadContext() + .putTransient(ConfigConstants.OPENDISTRO_SECURITY_USER, originalUserAndRemoteAddress.getLeft()); + return operation.get(); + } + } } diff --git a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandler.java b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandler.java index d34366733e..9145ee4bb1 100644 --- a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandler.java +++ b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandler.java @@ -14,7 +14,6 @@ import java.io.IOException; import java.util.HashMap; import java.util.Map; -import java.util.function.Supplier; import com.google.common.collect.ImmutableMap; import org.apache.logging.log4j.LogManager; @@ -27,7 +26,6 @@ import org.opensearch.action.search.SearchResponse; import org.opensearch.client.Client; import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.common.xcontent.XContentType; import org.opensearch.core.action.ActionListener; @@ -42,7 +40,6 @@ import org.opensearch.index.reindex.DeleteByQueryRequest; import org.opensearch.search.SearchHit; import org.opensearch.search.builder.SearchSourceBuilder; -import org.opensearch.security.dlic.rest.support.Utils; import org.opensearch.security.support.ConfigConstants; import static org.opensearch.security.action.apitokens.ApiToken.NAME_FIELD; @@ -59,116 +56,79 @@ public ApiTokenIndexHandler(Client client, ClusterService clusterService) { } public void indexTokenMetadata(ApiToken token) { - withSecurityContext(() -> { - try { - - XContentBuilder builder = XContentFactory.jsonBuilder(); - String jsonString = token.toXContent(builder, ToXContent.EMPTY_PARAMS).toString(); - - IndexRequest request = new IndexRequest(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX).source(jsonString, XContentType.JSON); - - ActionListener irListener = ActionListener.wrap(idxResponse -> { - LOGGER.info("Created {} entry.", ConfigConstants.OPENSEARCH_API_TOKENS_INDEX); - }, (failResponse) -> { - LOGGER.error(failResponse.getMessage()); - LOGGER.info("Failed to create {} entry.", ConfigConstants.OPENSEARCH_API_TOKENS_INDEX); - }); - client.index(request, irListener); - } catch (IOException e) { - throw new RuntimeException(e); - } - }); + try { + + XContentBuilder builder = XContentFactory.jsonBuilder(); + String jsonString = token.toXContent(builder, ToXContent.EMPTY_PARAMS).toString(); + + IndexRequest request = new IndexRequest(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX).source(jsonString, XContentType.JSON); + + ActionListener irListener = ActionListener.wrap(idxResponse -> { + LOGGER.info("Created {} entry.", ConfigConstants.OPENSEARCH_API_TOKENS_INDEX); + }, (failResponse) -> { + LOGGER.error(failResponse.getMessage()); + LOGGER.info("Failed to create {} entry.", ConfigConstants.OPENSEARCH_API_TOKENS_INDEX); + }); + client.index(request, irListener); + } catch (IOException e) { + throw new RuntimeException(e); + } } public void deleteToken(String name) throws ApiTokenException { - withSecurityContext(() -> { - DeleteByQueryRequest request = new DeleteByQueryRequest(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX).setQuery( - QueryBuilders.matchQuery(NAME_FIELD, name) - ).setRefresh(true); + DeleteByQueryRequest request = new DeleteByQueryRequest(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX).setQuery( + QueryBuilders.matchQuery(NAME_FIELD, name) + ).setRefresh(true); - BulkByScrollResponse response = client.execute(DeleteByQueryAction.INSTANCE, request).actionGet(); + BulkByScrollResponse response = client.execute(DeleteByQueryAction.INSTANCE, request).actionGet(); - long deletedDocs = response.getDeleted(); + long deletedDocs = response.getDeleted(); - if (deletedDocs == 0) { - throw new ApiTokenException("No token found with name " + name); - } - }); + if (deletedDocs == 0) { + throw new ApiTokenException("No token found with name " + name); + } } public Map getTokenMetadatas() { - return withSecurityContext(() -> { - try { - SearchRequest searchRequest = new SearchRequest(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX); - searchRequest.source(new SearchSourceBuilder()); - - SearchResponse response = client.search(searchRequest).actionGet(); - - Map tokens = new HashMap<>(); - for (SearchHit hit : response.getHits().getHits()) { - try ( - XContentParser parser = XContentType.JSON.xContent() - .createParser( - NamedXContentRegistry.EMPTY, - DeprecationHandler.THROW_UNSUPPORTED_OPERATION, - hit.getSourceRef().streamInput() - ) - ) { - - ApiToken token = ApiToken.fromXContent(parser); - tokens.put(token.getName(), token); - } + try { + SearchRequest searchRequest = new SearchRequest(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX); + searchRequest.source(new SearchSourceBuilder()); + + SearchResponse response = client.search(searchRequest).actionGet(); + + Map tokens = new HashMap<>(); + for (SearchHit hit : response.getHits().getHits()) { + try ( + XContentParser parser = XContentType.JSON.xContent() + .createParser( + NamedXContentRegistry.EMPTY, + DeprecationHandler.THROW_UNSUPPORTED_OPERATION, + hit.getSourceRef().streamInput() + ) + ) { + + ApiToken token = ApiToken.fromXContent(parser); + tokens.put(token.getName(), token); } - return tokens; - } catch (IOException e) { - throw new RuntimeException(e); } - }); + return tokens; + } catch (IOException e) { + throw new RuntimeException(e); + } } public Boolean apiTokenIndexExists() { return clusterService.state().metadata().hasConcreteIndex(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX); } - private T withSecurityContext(Supplier operation) { - final var originalUserAndRemoteAddress = Utils.userAndRemoteAddressFrom(client.threadPool().getThreadContext()); - try (final ThreadContext.StoredContext ctx = client.threadPool().getThreadContext().stashContext()) { - client.threadPool() - .getThreadContext() - .putTransient(ConfigConstants.OPENDISTRO_SECURITY_USER, originalUserAndRemoteAddress.getLeft()); - return operation.get(); - } - } - - private void withSecurityContext(Runnable operation) { - final var originalUserAndRemoteAddress = Utils.userAndRemoteAddressFrom(client.threadPool().getThreadContext()); - try (final ThreadContext.StoredContext ctx = client.threadPool().getThreadContext().stashContext()) { - client.threadPool() - .getThreadContext() - .putTransient(ConfigConstants.OPENDISTRO_SECURITY_USER, originalUserAndRemoteAddress.getLeft()); - operation.run(); - } - } - public void createApiTokenIndexIfAbsent() { if (!apiTokenIndexExists()) { - final var originalUserAndRemoteAddress = Utils.userAndRemoteAddressFrom(client.threadPool().getThreadContext()); - try (final ThreadContext.StoredContext ctx = client.threadPool().getThreadContext().stashContext()) { - client.threadPool() - .getThreadContext() - .putTransient(ConfigConstants.OPENDISTRO_SECURITY_USER, originalUserAndRemoteAddress.getLeft()); - final Map indexSettings = ImmutableMap.of( - "index.number_of_shards", - 1, - "index.auto_expand_replicas", - "0-all" - ); - final CreateIndexRequest createIndexRequest = new CreateIndexRequest(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX).settings( - indexSettings - ); - client.admin().indices().create(createIndexRequest); - } + final Map indexSettings = ImmutableMap.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); + final CreateIndexRequest createIndexRequest = new CreateIndexRequest(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX).settings( + indexSettings + ); + client.admin().indices().create(createIndexRequest); } } diff --git a/src/main/java/org/opensearch/security/dlic/rest/api/RestApiAdminPrivilegesEvaluator.java b/src/main/java/org/opensearch/security/dlic/rest/api/RestApiAdminPrivilegesEvaluator.java index 3714ac4c3d..768f9d2f70 100644 --- a/src/main/java/org/opensearch/security/dlic/rest/api/RestApiAdminPrivilegesEvaluator.java +++ b/src/main/java/org/opensearch/security/dlic/rest/api/RestApiAdminPrivilegesEvaluator.java @@ -70,8 +70,8 @@ default String build() { .put(Endpoint.ROLES, action -> buildEndpointPermission(Endpoint.ROLES)) .put(Endpoint.ROLESMAPPING, action -> buildEndpointPermission(Endpoint.ROLESMAPPING)) .put(Endpoint.TENANTS, action -> buildEndpointPermission(Endpoint.TENANTS)) - .put(Endpoint.SSL, action -> buildEndpointActionPermission(Endpoint.SSL, action)) .put(Endpoint.APITOKENS, action -> buildEndpointPermission(Endpoint.APITOKENS)) + .put(Endpoint.SSL, action -> buildEndpointActionPermission(Endpoint.SSL, action)) .build(); private final ThreadContext threadContext;