From 8fc9c755ebd2d2ad570d9c1eaaa5f51030850cf4 Mon Sep 17 00:00:00 2001 From: Derek Ho Date: Fri, 28 Feb 2025 13:55:29 -0500 Subject: [PATCH 1/3] Change to action listeners and limit to 100 tokens Signed-off-by: Derek Ho --- .../action/apitokens/ApiTokenAction.java | 191 +++++++++++------- .../apitokens/ApiTokenIndexHandler.java | 78 +++---- .../action/apitokens/ApiTokenRepository.java | 18 +- 3 files changed, 167 insertions(+), 120 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 eddafd79ee..36d38fe82d 100644 --- a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenAction.java +++ b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenAction.java @@ -146,30 +146,32 @@ RestChannelConsumer doPrepareRequest(RestRequest request, NodeClient client) { private RestChannelConsumer handleGet(RestRequest request, NodeClient client) { return channel -> { - final XContentBuilder builder = channel.newBuilder(); - BytesRestResponse response; - try { - Map tokens = apiTokenRepository.getApiTokens(); - - builder.startArray(); - for (ApiToken token : tokens.values()) { - builder.startObject(); - builder.field(NAME_FIELD, token.getName()); - builder.field(CREATION_TIME_FIELD, token.getCreationTime().toEpochMilli()); - builder.field(EXPIRATION_FIELD, token.getExpiration()); - builder.field(CLUSTER_PERMISSIONS_FIELD, token.getClusterPermissions()); - builder.field(INDEX_PERMISSIONS_FIELD, token.getIndexPermissions()); - builder.endObject(); + apiTokenRepository.getApiTokens(ActionListener.wrap(tokens -> { + try { + XContentBuilder builder = channel.newBuilder(); + builder.startArray(); + for (ApiToken token : tokens.values()) { + builder.startObject(); + builder.field(NAME_FIELD, token.getName()); + builder.field(CREATION_TIME_FIELD, token.getCreationTime().toEpochMilli()); + builder.field(EXPIRATION_FIELD, token.getExpiration()); + builder.field(CLUSTER_PERMISSIONS_FIELD, token.getClusterPermissions()); + builder.field(INDEX_PERMISSIONS_FIELD, token.getIndexPermissions()); + builder.endObject(); + } + builder.endArray(); + + BytesRestResponse response = new BytesRestResponse(RestStatus.OK, builder); + builder.close(); + channel.sendResponse(response); + } catch (final Exception exception) { + sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, exception.getMessage()); } - builder.endArray(); + }, exception -> { + sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, exception.getMessage()); + + })); - response = new BytesRestResponse(RestStatus.OK, builder); - } catch (final Exception exception) { - builder.startObject().field("error", exception.getMessage()).endObject(); - response = new BytesRestResponse(RestStatus.INTERNAL_SERVER_ERROR, builder); - } - builder.close(); - channel.sendResponse(response); }; } @@ -181,39 +183,66 @@ private RestChannelConsumer handlePost(RestRequest request, NodeClient client) { List clusterPermissions = extractClusterPermissions(requestBody); List indexPermissions = extractIndexPermissions(requestBody); - - String token = apiTokenRepository.createApiToken( - (String) requestBody.get(NAME_FIELD), - clusterPermissions, - indexPermissions, - (Long) requestBody.getOrDefault(EXPIRATION_FIELD, Instant.now().toEpochMilli() + TimeUnit.DAYS.toMillis(30)) + String name = (String) requestBody.get(NAME_FIELD); + long expiration = (Long) requestBody.getOrDefault( + EXPIRATION_FIELD, + Instant.now().toEpochMilli() + TimeUnit.DAYS.toMillis(30) ); - // Then trigger the update action - ApiTokenUpdateRequest updateRequest = new ApiTokenUpdateRequest(); - client.execute(ApiTokenUpdateAction.INSTANCE, updateRequest, new ActionListener() { - @Override - public void onResponse(ApiTokenUpdateResponse updateResponse) { - try { - XContentBuilder builder = channel.newBuilder(); - builder.startObject(); - builder.field("Api Token: ", token); - builder.endObject(); - - BytesRestResponse response = new BytesRestResponse(RestStatus.OK, builder); - channel.sendResponse(response); - } catch (IOException e) { - sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, "Failed to send response after token creation"); - } + // First check token count + apiTokenRepository.getTokenCount(ActionListener.wrap(tokenCount -> { + if (tokenCount >= 100) { + sendErrorResponse( + channel, + RestStatus.TOO_MANY_REQUESTS, + "Maximum limit of 100 API tokens reached. Please delete existing tokens before creating new ones." + ); + return; } - @Override - public void onFailure(Exception e) { - sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, "Failed to propagate token creation"); - } - }); - } catch (final Exception exception) { - sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, exception.getMessage()); + // If count is ok, create the token + apiTokenRepository.createApiToken(name, clusterPermissions, indexPermissions, expiration, ActionListener.wrap(token -> { + // After successful creation, trigger the update action + ApiTokenUpdateRequest updateRequest = new ApiTokenUpdateRequest(); + client.execute(ApiTokenUpdateAction.INSTANCE, updateRequest, ActionListener.wrap(updateResponse -> { + try { + XContentBuilder builder = channel.newBuilder(); + builder.startObject(); + builder.field("Api Token: ", token); + builder.endObject(); + channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder)); + builder.close(); + } catch (IOException e) { + sendErrorResponse( + channel, + RestStatus.INTERNAL_SERVER_ERROR, + "Failed to send response after token creation" + ); + } + }, + updateException -> sendErrorResponse( + channel, + RestStatus.INTERNAL_SERVER_ERROR, + "Failed to propagate token creation: " + updateException.getMessage() + ) + )); + }, + createException -> sendErrorResponse( + channel, + RestStatus.INTERNAL_SERVER_ERROR, + "Failed to create token: " + createException.getMessage() + ) + )); + }, + countException -> sendErrorResponse( + channel, + RestStatus.INTERNAL_SERVER_ERROR, + "Failed to get token count: " + countException.getMessage() + ) + )); + + } catch (Exception e) { + sendErrorResponse(channel, RestStatus.BAD_REQUEST, "Invalid request: " + e.getMessage()); } }; } @@ -303,32 +332,44 @@ private RestChannelConsumer handleDelete(RestRequest request, NodeClient client) final Map requestBody = request.contentOrSourceParamParser().map(); validateRequestParameters(requestBody); - apiTokenRepository.deleteApiToken((String) requestBody.get(NAME_FIELD)); - - ApiTokenUpdateRequest updateRequest = new ApiTokenUpdateRequest(); - client.execute(ApiTokenUpdateAction.INSTANCE, updateRequest, new ActionListener() { - @Override - public void onResponse(ApiTokenUpdateResponse updateResponse) { - try { - XContentBuilder builder = channel.newBuilder(); - builder.startObject(); - builder.field("message", "token " + requestBody.get(NAME_FIELD) + " deleted successfully."); - builder.endObject(); - - BytesRestResponse response = new BytesRestResponse(RestStatus.OK, builder); - channel.sendResponse(response); - } catch (Exception e) { - sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, "Failed to send response after token update"); - } + apiTokenRepository.deleteApiToken((String) requestBody.get(NAME_FIELD), ActionListener.wrap(ignored -> { + try { + ApiTokenUpdateRequest updateRequest = new ApiTokenUpdateRequest(); + client.execute(ApiTokenUpdateAction.INSTANCE, updateRequest, new ActionListener() { + @Override + public void onResponse(ApiTokenUpdateResponse updateResponse) { + try { + XContentBuilder builder = channel.newBuilder(); + builder.startObject(); + builder.field("message", "Token " + requestBody.get(NAME_FIELD) + " deleted successfully."); + builder.endObject(); + + BytesRestResponse response = new BytesRestResponse(RestStatus.OK, builder); + channel.sendResponse(response); + } catch (Exception e) { + sendErrorResponse( + channel, + RestStatus.INTERNAL_SERVER_ERROR, + "Failed to send response after token update" + ); + } + } + + @Override + public void onFailure(Exception e) { + sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, "Failed to propagate token deletion"); + } + }); + } catch (IOException e) { + sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, "Failed to build success response: " + e.getMessage()); } - - @Override - public void onFailure(Exception e) { - sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, "Failed to propagate token deletion"); + }, exception -> { + RestStatus status = RestStatus.INTERNAL_SERVER_ERROR; + if (exception instanceof ApiTokenException) { + status = RestStatus.NOT_FOUND; } - }); - } catch (final ApiTokenException exception) { - sendErrorResponse(channel, RestStatus.NOT_FOUND, exception.getMessage()); + sendErrorResponse(channel, status, exception.getMessage()); + })); } catch (final Exception exception) { sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, exception.getMessage()); } 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 9145ee4bb1..f9a42cce42 100644 --- a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandler.java +++ b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandler.java @@ -21,9 +21,7 @@ import org.opensearch.action.admin.indices.create.CreateIndexRequest; import org.opensearch.action.index.IndexRequest; -import org.opensearch.action.index.IndexResponse; import org.opensearch.action.search.SearchRequest; -import org.opensearch.action.search.SearchResponse; import org.opensearch.client.Client; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.xcontent.XContentFactory; @@ -35,7 +33,6 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.index.query.QueryBuilders; -import org.opensearch.index.reindex.BulkByScrollResponse; import org.opensearch.index.reindex.DeleteByQueryAction; import org.opensearch.index.reindex.DeleteByQueryRequest; import org.opensearch.search.SearchHit; @@ -55,66 +52,69 @@ public ApiTokenIndexHandler(Client client, ClusterService clusterService) { this.clusterService = clusterService; } - public void indexTokenMetadata(ApiToken token) { + public void indexTokenMetadata(ApiToken token, ActionListener listener) { 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 -> { + client.index(request, ActionListener.wrap(indexResponse -> { LOGGER.info("Created {} entry.", ConfigConstants.OPENSEARCH_API_TOKENS_INDEX); - }, (failResponse) -> { - LOGGER.error(failResponse.getMessage()); + listener.onResponse(null); + }, exception -> { + LOGGER.error(exception.getMessage()); LOGGER.info("Failed to create {} entry.", ConfigConstants.OPENSEARCH_API_TOKENS_INDEX); - }); - client.index(request, irListener); + listener.onFailure(exception); + })); } catch (IOException e) { throw new RuntimeException(e); } - } - public void deleteToken(String name) throws ApiTokenException { + public void deleteToken(String name, ActionListener listener) { 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(); - - long deletedDocs = response.getDeleted(); - - if (deletedDocs == 0) { - throw new ApiTokenException("No token found with name " + name); - } + client.execute(DeleteByQueryAction.INSTANCE, request, ActionListener.wrap(response -> { + long deletedDocs = response.getDeleted(); + if (deletedDocs == 0) { + listener.onFailure(new ApiTokenException("No token found with name " + name)); + } else { + listener.onResponse(null); + } + }, exception -> listener.onFailure(exception))); } - public Map getTokenMetadatas() { + public void getTokenMetadatas(ActionListener> listener) { 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); + client.search(searchRequest, ActionListener.wrap(response -> { + try { + 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); + } + } + listener.onResponse(tokens); + } catch (IOException e) { + listener.onFailure(e); } - } - return tokens; - } catch (IOException e) { - throw new RuntimeException(e); + }, listener::onFailure)); + } catch (Exception e) { + listener.onFailure(e); } } diff --git a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenRepository.java b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenRepository.java index 850fad0b66..cb696da005 100644 --- a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenRepository.java +++ b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenRepository.java @@ -21,6 +21,7 @@ import org.opensearch.client.Client; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.core.action.ActionListener; import org.opensearch.index.IndexNotFoundException; import org.opensearch.security.authtoken.jwt.ExpiringBearerAuthToken; import org.opensearch.security.identity.SecurityTokenManager; @@ -86,21 +87,26 @@ public String createApiToken( String name, List clusterPermissions, List indexPermissions, - Long expiration + Long expiration, + ActionListener listener ) { apiTokenIndexHandler.createApiTokenIndexIfAbsent(); ExpiringBearerAuthToken token = securityTokenManager.issueApiToken(name, expiration); ApiToken apiToken = new ApiToken(name, clusterPermissions, indexPermissions, expiration); - apiTokenIndexHandler.indexTokenMetadata(apiToken); + apiTokenIndexHandler.indexTokenMetadata(apiToken, listener); return token.getCompleteToken(); } - public void deleteApiToken(String name) throws ApiTokenException, IndexNotFoundException { - apiTokenIndexHandler.deleteToken(name); + public void deleteApiToken(String name, ActionListener listener) throws ApiTokenException, IndexNotFoundException { + apiTokenIndexHandler.deleteToken(name, listener); } - public Map getApiTokens() throws IndexNotFoundException { - return apiTokenIndexHandler.getTokenMetadatas(); + public void getApiTokens(ActionListener> listener) { + apiTokenIndexHandler.getTokenMetadatas(listener); + } + + public void getTokenCount(ActionListener listener) { + getApiTokens(ActionListener.wrap(tokens -> listener.onResponse((long) tokens.size()), listener::onFailure)); } } From b5f8a814fa3f7a40f45cecf30f9437531a9a10f7 Mon Sep 17 00:00:00 2001 From: Derek Ho Date: Thu, 20 Mar 2025 14:17:24 -0400 Subject: [PATCH 2/3] Cleanup and fix CI, adjust tests to be actionlistener model Signed-off-by: Derek Ho --- .../action/apitokens/ApiTokenAction.java | 2 +- .../apitokens/ApiTokenIndexHandler.java | 2 +- .../action/apitokens/ApiTokenRepository.java | 38 +++-- .../apitokens/ApiTokenIndexHandlerTest.java | 122 +++++++------- .../apitokens/ApiTokenRepositoryTest.java | 149 +++++++++++++----- .../action/apitokens/ApiTokenTest.java | 4 +- .../security/util/ActionListenerUtils.java | 71 +++++++++ 7 files changed, 274 insertions(+), 114 deletions(-) create mode 100644 src/test/java/org/opensearch/security/util/ActionListenerUtils.java 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 36d38fe82d..ceeed23ada 100644 --- a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenAction.java +++ b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenAction.java @@ -24,7 +24,6 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.opensearch.client.node.NodeClient; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.Settings; @@ -48,6 +47,7 @@ import org.opensearch.security.ssl.transport.PrincipalExtractor; import org.opensearch.security.support.ConfigConstants; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.client.node.NodeClient; import static org.opensearch.rest.RestRequest.Method.DELETE; import static org.opensearch.rest.RestRequest.Method.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 f9a42cce42..d34368b34a 100644 --- a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandler.java +++ b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandler.java @@ -22,7 +22,6 @@ import org.opensearch.action.admin.indices.create.CreateIndexRequest; import org.opensearch.action.index.IndexRequest; import org.opensearch.action.search.SearchRequest; -import org.opensearch.client.Client; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.common.xcontent.XContentType; @@ -38,6 +37,7 @@ import org.opensearch.search.SearchHit; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.security.support.ConfigConstants; +import org.opensearch.transport.client.Client; import static org.opensearch.security.action.apitokens.ApiToken.NAME_FIELD; diff --git a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenRepository.java b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenRepository.java index cb696da005..817bdb23c7 100644 --- a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenRepository.java +++ b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenRepository.java @@ -13,19 +13,20 @@ import java.util.List; import java.util.Map; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import com.google.common.annotations.VisibleForTesting; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.opensearch.client.Client; import org.opensearch.cluster.service.ClusterService; import org.opensearch.core.action.ActionListener; import org.opensearch.index.IndexNotFoundException; import org.opensearch.security.authtoken.jwt.ExpiringBearerAuthToken; import org.opensearch.security.identity.SecurityTokenManager; import org.opensearch.security.user.User; +import org.opensearch.transport.client.Client; import static org.opensearch.security.http.ApiTokenAuthenticator.API_TOKEN_USER_PREFIX; @@ -37,11 +38,26 @@ public class ApiTokenRepository { private final Map jtis = new ConcurrentHashMap<>(); void reloadApiTokensFromIndex() { - Map tokensFromIndex = apiTokenIndexHandler.getTokenMetadatas(); - jtis.keySet().removeIf(key -> !tokensFromIndex.containsKey(key)); - tokensFromIndex.forEach( - (key, apiToken) -> jtis.put(key, new Permissions(apiToken.getClusterPermissions(), apiToken.getIndexPermissions())) - ); + CompletableFuture> future = new CompletableFuture<>(); + + apiTokenIndexHandler.getTokenMetadatas(new ActionListener>() { + @Override + public void onResponse(Map tokensFromIndex) { + future.complete(tokensFromIndex); + } + + @Override + public void onFailure(Exception e) { + future.completeExceptionally(e); + } + }); + + future.thenAccept(tokensFromIndex -> { + jtis.keySet().removeIf(key -> !tokensFromIndex.containsKey(key)); + tokensFromIndex.forEach( + (key, apiToken) -> jtis.put(key, new Permissions(apiToken.getClusterPermissions(), apiToken.getIndexPermissions())) + ); + }); } public Permissions getApiTokenPermissionsForUser(User user) { @@ -83,18 +99,20 @@ static ApiTokenRepository forTest(ApiTokenIndexHandler apiTokenIndexHandler, Sec return new ApiTokenRepository(apiTokenIndexHandler, securityTokenManager); } - public String createApiToken( + public void createApiToken( String name, List clusterPermissions, List indexPermissions, Long expiration, - ActionListener listener + ActionListener listener ) { apiTokenIndexHandler.createApiTokenIndexIfAbsent(); ExpiringBearerAuthToken token = securityTokenManager.issueApiToken(name, expiration); ApiToken apiToken = new ApiToken(name, clusterPermissions, indexPermissions, expiration); - apiTokenIndexHandler.indexTokenMetadata(apiToken, listener); - return token.getCompleteToken(); + apiTokenIndexHandler.indexTokenMetadata( + apiToken, + ActionListener.wrap(unused -> listener.onResponse(token.getCompleteToken()), exception -> listener.onFailure(exception)) + ); } public void deleteApiToken(String name, ActionListener listener) throws ApiTokenException, IndexNotFoundException { diff --git a/src/test/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandlerTest.java b/src/test/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandlerTest.java index 9b3b8638e2..b74e7073fc 100644 --- a/src/test/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandlerTest.java +++ b/src/test/java/org/opensearch/security/action/apitokens/ApiTokenIndexHandlerTest.java @@ -26,9 +26,6 @@ import org.opensearch.action.index.IndexResponse; import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchResponse; -import org.opensearch.action.support.PlainActionFuture; -import org.opensearch.client.Client; -import org.opensearch.client.IndicesAdminClient; import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.Settings; @@ -36,7 +33,6 @@ import org.opensearch.common.xcontent.XContentType; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.bytes.BytesReference; -import org.opensearch.core.index.shard.ShardId; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.index.query.MatchQueryBuilder; @@ -46,6 +42,9 @@ import org.opensearch.search.SearchHit; import org.opensearch.search.SearchHits; import org.opensearch.security.support.ConfigConstants; +import org.opensearch.security.util.ActionListenerUtils.TestActionListener; +import org.opensearch.transport.client.Client; +import org.opensearch.transport.client.IndicesAdminClient; import org.mockito.ArgumentCaptor; import org.mockito.Mock; @@ -56,7 +55,6 @@ import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.is; import static org.opensearch.security.action.apitokens.ApiToken.NAME_FIELD; -import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.RETURNS_DEEP_STUBS; import static org.mockito.Mockito.doAnswer; @@ -66,6 +64,7 @@ import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; +@SuppressWarnings("unchecked") public class ApiTokenIndexHandlerTest { @Mock @@ -107,7 +106,6 @@ public void testCreateApiTokenIndexWhenIndexNotExist() { indexHandler.createApiTokenIndexIfAbsent(); ArgumentCaptor captor = ArgumentCaptor.forClass(CreateIndexRequest.class); - verify(indicesAdminClient).create(captor.capture()); assertThat(captor.getValue().index(), equalTo(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX)); } @@ -125,16 +123,27 @@ public void testCreateApiTokenIndexWhenIndexExists() { public void testDeleteApiTokeCallsDeleteByQueryWithSuppliedName() { when(metadata.hasConcreteIndex(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX)).thenReturn(true); String tokenName = "token"; - try { - indexHandler.deleteToken(tokenName); - } catch (Exception e) { - // Ignore - } + + TestActionListener listener = new TestActionListener<>(); + + doAnswer(invocation -> { + DeleteByQueryRequest request = invocation.getArgument(1); + ActionListener parentListener = invocation.getArgument(2); + + BulkByScrollResponse response = mock(BulkByScrollResponse.class); + when(response.getDeleted()).thenReturn(1L); + + parentListener.onResponse(response); + return null; + }).when(client).execute(eq(DeleteByQueryAction.INSTANCE), any(DeleteByQueryRequest.class), any(ActionListener.class)); + + indexHandler.deleteToken(tokenName, listener); ArgumentCaptor captor = ArgumentCaptor.forClass(DeleteByQueryRequest.class); - verify(client).execute(eq(DeleteByQueryAction.INSTANCE), captor.capture()); + verify(client).execute(eq(DeleteByQueryAction.INSTANCE), captor.capture(), any(ActionListener.class)); + + listener.assertSuccess(); - // Captured request has the correct name DeleteByQueryRequest capturedRequest = captor.getValue(); MatchQueryBuilder query = (MatchQueryBuilder) capturedRequest.getSearchRequest().source().query(); assertThat(query.fieldName(), equalTo(NAME_FIELD)); @@ -145,38 +154,39 @@ public void testDeleteApiTokeCallsDeleteByQueryWithSuppliedName() { public void testDeleteTokenThrowsExceptionWhenNoDocumentsDeleted() { when(metadata.hasConcreteIndex(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX)).thenReturn(true); - PlainActionFuture future = new PlainActionFuture<>(); - BulkByScrollResponse response = mock(BulkByScrollResponse.class); - when(response.getDeleted()).thenReturn(0L); - future.onResponse(response); - when(client.execute(eq(DeleteByQueryAction.INSTANCE), any(DeleteByQueryRequest.class))).thenReturn(future); + doAnswer(invocation -> { + ActionListener listener = invocation.getArgument(2); + BulkByScrollResponse response = mock(BulkByScrollResponse.class); + when(response.getDeleted()).thenReturn(0L); + listener.onResponse(response); + return null; + }).when(client).execute(eq(DeleteByQueryAction.INSTANCE), any(DeleteByQueryRequest.class), any(ActionListener.class)); String tokenName = "nonexistent-token"; - try { - indexHandler.deleteToken(tokenName); - fail("Expected ApiTokenException to be thrown"); - } catch (ApiTokenException e) { - assertThat(e.getMessage(), equalTo("No token found with name " + tokenName)); - } + TestActionListener listener = new TestActionListener<>(); + indexHandler.deleteToken(tokenName, listener); + + Exception e = listener.assertException(ApiTokenException.class); + assertThat(e.getMessage(), containsString("No token found with name " + tokenName)); } @Test public void testDeleteTokenSucceedsWhenDocumentIsDeleted() { when(metadata.hasConcreteIndex(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX)).thenReturn(true); - // 1 deleted document - PlainActionFuture future = new PlainActionFuture<>(); - BulkByScrollResponse response = mock(BulkByScrollResponse.class); - when(response.getDeleted()).thenReturn(1L); - future.onResponse(response); - when(client.execute(eq(DeleteByQueryAction.INSTANCE), any(DeleteByQueryRequest.class))).thenReturn(future); + doAnswer(invocation -> { + ActionListener listener = invocation.getArgument(2); + BulkByScrollResponse response = mock(BulkByScrollResponse.class); + when(response.getDeleted()).thenReturn(1L); + listener.onResponse(response); + return null; + }).when(client).execute(eq(DeleteByQueryAction.INSTANCE), any(DeleteByQueryRequest.class), any(ActionListener.class)); String tokenName = "existing-token"; - try { - indexHandler.deleteToken(tokenName); - } catch (ApiTokenException e) { - fail("Should not have thrown exception"); - } + TestActionListener listener = new TestActionListener<>(); + indexHandler.deleteToken(tokenName, listener); + + listener.assertSuccess(); } @Test @@ -198,35 +208,24 @@ public void testIndexTokenStoresTokenPayload() { Long.MAX_VALUE ); - // Mock the index method with ActionListener - @SuppressWarnings("unchecked") - ArgumentCaptor> listenerCaptor = - ArgumentCaptor.forClass((Class>) (Class) ActionListener.class); - + // Mock the index response doAnswer(invocation -> { - ActionListener listener = listenerCaptor.getValue(); - listener.onResponse(new IndexResponse( - new ShardId(".opensearch_security_api_tokens", "_na_", 1), - "1", - 0, - 1, - 1, - true - )); + ActionListener listener = invocation.getArgument(1); + listener.onResponse(mock(IndexResponse.class)); return null; - }).when(client).index(any(IndexRequest.class), listenerCaptor.capture()); + }).when(client).index(any(IndexRequest.class), any(ActionListener.class)); + TestActionListener listener = new TestActionListener<>(); + indexHandler.indexTokenMetadata(token, listener); - indexHandler.indexTokenMetadata(token); + listener.assertSuccess(); - // Verify the index request ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(IndexRequest.class); - verify(client).index(requestCaptor.capture(), listenerCaptor.capture()); + verify(client).index(requestCaptor.capture(), any(ActionListener.class)); IndexRequest capturedRequest = requestCaptor.getValue(); assertThat(capturedRequest.index(), equalTo(ConfigConstants.OPENSEARCH_API_TOKENS_INDEX)); - // verify contents String source = capturedRequest.source().utf8ToString(); assertThat(source, containsString("test-token-description")); assertThat(source, containsString("cluster:admin/something")); @@ -280,14 +279,16 @@ public void testGetTokenPayloads() throws IOException { SearchHits searchHits = new SearchHits(hits, new TotalHits(2, TotalHits.Relation.EQUAL_TO), 1.0f); when(searchResponse.getHits()).thenReturn(searchHits); - // Mock client search call - PlainActionFuture future = new PlainActionFuture<>(); - future.onResponse(searchResponse); - when(client.search(any(SearchRequest.class))).thenReturn(future); + doAnswer(invocation -> { + ActionListener listener = invocation.getArgument(1); + listener.onResponse(searchResponse); + return null; + }).when(client).search(any(SearchRequest.class), any(ActionListener.class)); - // Get tokens and verify - Map resultTokens = indexHandler.getTokenMetadatas(); + TestActionListener> listener = new TestActionListener<>(); + indexHandler.getTokenMetadatas(listener); + Map resultTokens = listener.assertSuccess(); assertThat(resultTokens.size(), equalTo(2)); assertThat(resultTokens.containsKey("token1-description"), is(true)); assertThat(resultTokens.containsKey("token2-description"), is(true)); @@ -298,5 +299,4 @@ public void testGetTokenPayloads() throws IOException { ApiToken resultToken2 = resultTokens.get("token2-description"); assertThat(resultToken2.getClusterPermissions(), contains("cluster:admin/other")); } - } diff --git a/src/test/java/org/opensearch/security/action/apitokens/ApiTokenRepositoryTest.java b/src/test/java/org/opensearch/security/action/apitokens/ApiTokenRepositoryTest.java index 89f8b950cd..a7a43cb862 100644 --- a/src/test/java/org/opensearch/security/action/apitokens/ApiTokenRepositoryTest.java +++ b/src/test/java/org/opensearch/security/action/apitokens/ApiTokenRepositoryTest.java @@ -13,35 +13,43 @@ import java.io.IOException; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.TimeUnit; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; +import org.opensearch.core.action.ActionListener; import org.opensearch.index.IndexNotFoundException; import org.opensearch.security.authtoken.jwt.ExpiringBearerAuthToken; import org.opensearch.security.identity.SecurityTokenManager; import org.opensearch.security.user.User; +import org.opensearch.security.util.ActionListenerUtils.TestActionListener; import org.mockito.Mock; import org.mockito.junit.MockitoJUnitRunner; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; +import static org.awaitility.Awaitility.await; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.any; import static org.mockito.Mockito.argThat; -import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +@SuppressWarnings("unchecked") @RunWith(MockitoJUnitRunner.class) public class ApiTokenRepositoryTest { @Mock @@ -61,9 +69,17 @@ public void setUp() { public void testDeleteApiToken() throws ApiTokenException { String tokenName = "test-token"; - repository.deleteApiToken(tokenName); + doAnswer(invocation -> { + ActionListener listener = invocation.getArgument(1); + listener.onResponse(null); + return null; + }).when(apiTokenIndexHandler).deleteToken(eq(tokenName), any(ActionListener.class)); - verify(apiTokenIndexHandler).deleteToken(tokenName); + TestActionListener listener = new TestActionListener<>(); + repository.deleteApiToken(tokenName, listener); + + listener.assertSuccess(); + verify(apiTokenIndexHandler).deleteToken(eq(tokenName), any(ActionListener.class)); } @Test @@ -86,19 +102,25 @@ public void testGetApiTokenPermissionsForUser() throws ApiTokenException { assertEquals(List.of("cluster_all"), permissionsForApiTokenExists.getClusterPerm()); assertEquals(List.of("*"), permissionsForApiTokenExists.getIndexPermission().getFirst().getAllowedActions()); assertEquals(List.of("*"), permissionsForApiTokenExists.getIndexPermission().getFirst().getIndexPatterns()); - } @Test public void testGetApiTokens() throws IndexNotFoundException { Map expectedTokens = new HashMap<>(); expectedTokens.put("token1", new ApiToken("token1", Arrays.asList("perm1"), Arrays.asList(), Long.MAX_VALUE)); - when(apiTokenIndexHandler.getTokenMetadatas()).thenReturn(expectedTokens); - Map result = repository.getApiTokens(); + doAnswer(invocation -> { + ActionListener> listener = invocation.getArgument(0); + listener.onResponse(expectedTokens); + return null; + }).when(apiTokenIndexHandler).getTokenMetadatas(any(ActionListener.class)); + TestActionListener> listener = new TestActionListener<>(); + repository.getApiTokens(listener); + + Map result = listener.assertSuccess(); assertThat(result, equalTo(expectedTokens)); - verify(apiTokenIndexHandler).getTokenMetadatas(); + verify(apiTokenIndexHandler).getTokenMetadatas(any(ActionListener.class)); } @Test @@ -115,35 +137,68 @@ public void testCreateApiToken() { when(bearerToken.getCompleteToken()).thenReturn(completeToken); when(securityTokenManager.issueApiToken(any(), any())).thenReturn(bearerToken); - String result = repository.createApiToken(tokenName, clusterPermissions, indexPermissions, expiration); - - verify(apiTokenIndexHandler).createApiTokenIndexIfAbsent(); - verify(securityTokenManager).issueApiToken(any(), any()); - verify(apiTokenIndexHandler).indexTokenMetadata( - argThat( - token -> token.getName().equals(tokenName) - && token.getClusterPermissions().equals(clusterPermissions) - && token.getIndexPermissions().equals(indexPermissions) - && token.getExpiration().equals(expiration) - ) - ); - assertThat(result, equalTo(completeToken)); + doAnswer(invocation -> { + ActionListener listener = invocation.getArgument(1); + listener.onResponse(null); + return null; + }).when(apiTokenIndexHandler).indexTokenMetadata(any(ApiToken.class), any(ActionListener.class)); + + TestActionListener listener = new TestActionListener() { + @Override + public void onResponse(String result) { + try { + assertThat(result, equalTo(completeToken)); + verify(apiTokenIndexHandler).createApiTokenIndexIfAbsent(); + verify(securityTokenManager).issueApiToken(any(), any()); + verify(apiTokenIndexHandler).indexTokenMetadata( + argThat( + token -> token.getName().equals(tokenName) + && token.getClusterPermissions().equals(clusterPermissions) + && token.getIndexPermissions().equals(indexPermissions) + && token.getExpiration().equals(expiration) + ), + any(ActionListener.class) + ); + } finally { + super.onResponse(result); + } + } + }; + + repository.createApiToken(tokenName, clusterPermissions, indexPermissions, expiration, listener); + listener.assertSuccess(); } - @Test(expected = IndexNotFoundException.class) - public void testGetApiTokensThrowsIndexNotFoundException() throws IndexNotFoundException { - when(apiTokenIndexHandler.getTokenMetadatas()).thenThrow(new IndexNotFoundException("test-index")); - - repository.getApiTokens(); - + @Test + public void testGetApiTokensThrowsIndexNotFoundException() { + doAnswer(invocation -> { + ActionListener> listener = invocation.getArgument(0); + listener.onFailure(new IndexNotFoundException("test-index")); + return null; + }).when(apiTokenIndexHandler).getTokenMetadatas(any(ActionListener.class)); + + TestActionListener> listener = new TestActionListener<>(); + repository.getApiTokens(listener); + + Exception e = listener.assertException(IndexNotFoundException.class); + assertThat(e.getMessage(), containsString("test-index")); } - @Test(expected = ApiTokenException.class) - public void testDeleteApiTokenThrowsApiTokenException() throws ApiTokenException { + @Test + public void testDeleteApiTokenThrowsApiTokenException() { String tokenName = "test-token"; - doThrow(new ApiTokenException("Token not found")).when(apiTokenIndexHandler).deleteToken(tokenName); - repository.deleteApiToken(tokenName); + doAnswer(invocation -> { + ActionListener listener = invocation.getArgument(1); + listener.onFailure(new ApiTokenException("Token not found")); + return null; + }).when(apiTokenIndexHandler).deleteToken(eq(tokenName), any(ActionListener.class)); + + TestActionListener listener = new TestActionListener<>(); + repository.deleteApiToken(tokenName, listener); + + Exception e = listener.assertException(ApiTokenException.class); + assertThat(e.getMessage(), containsString("Token not found")); } @Test @@ -161,24 +216,40 @@ public void testJtisOperations() { @Test public void testClearJtis() { repository.getJtis().put("testJti", new Permissions(List.of("read"), List.of(new ApiToken.IndexPermission(List.of(), List.of())))); + + doAnswer(invocation -> { + ActionListener> listener = invocation.getArgument(0); + listener.onResponse(Collections.emptyMap()); + return null; + }).when(apiTokenIndexHandler).getTokenMetadatas(any(ActionListener.class)); + repository.reloadApiTokensFromIndex(); - assertTrue("Jtis should be empty after clear", repository.getJtis().isEmpty()); + await().atMost(5, TimeUnit.SECONDS) + .untilAsserted(() -> assertTrue("Jtis should be empty after clear", repository.getJtis().isEmpty())); } @Test public void testReloadApiTokensFromIndexAndParse() throws IOException { - when(apiTokenIndexHandler.getTokenMetadatas()).thenReturn(Map.of("test", new ApiToken("test", List.of("cluster:monitor"), List.of(), Long.MAX_VALUE))); + // Setup mock response + Map expectedTokens = Map.of("test", new ApiToken("test", List.of("cluster:monitor"), List.of(), Long.MAX_VALUE)); + + doAnswer(invocation -> { + ActionListener> listener = invocation.getArgument(0); + listener.onResponse(expectedTokens); + return null; + }).when(apiTokenIndexHandler).getTokenMetadatas(any(ActionListener.class)); // Execute the reload repository.reloadApiTokensFromIndex(); - // Verify the cache was updated - assertFalse("Jtis should not be empty after reload", repository.getJtis().isEmpty()); - assertEquals("Should have one JTI entry", 1, repository.getJtis().size()); - assertTrue("Should contain testJti", repository.getJtis().containsKey("test")); - // Verify extraction works - assertEquals("Should have one cluster action", List.of("cluster:monitor"), repository.getJtis().get("test").getClusterPerm()); - assertEquals("Should have no index actions", List.of(), repository.getJtis().get("test").getIndexPermission()); + // Wait for and verify the async updates + await().atMost(5, TimeUnit.SECONDS).untilAsserted(() -> { + assertFalse("Jtis should not be empty after reload", repository.getJtis().isEmpty()); + assertEquals("Should have one JTI entry", 1, repository.getJtis().size()); + assertTrue("Should contain testJti", repository.getJtis().containsKey("test")); + assertEquals("Should have one cluster action", List.of("cluster:monitor"), repository.getJtis().get("test").getClusterPerm()); + assertEquals("Should have no index actions", List.of(), repository.getJtis().get("test").getIndexPermission()); + }); } } diff --git a/src/test/java/org/opensearch/security/action/apitokens/ApiTokenTest.java b/src/test/java/org/opensearch/security/action/apitokens/ApiTokenTest.java index 4951507359..922bfaff1e 100644 --- a/src/test/java/org/opensearch/security/action/apitokens/ApiTokenTest.java +++ b/src/test/java/org/opensearch/security/action/apitokens/ApiTokenTest.java @@ -18,8 +18,6 @@ import org.junit.Before; import org.junit.Test; -import org.opensearch.client.Client; -import org.opensearch.client.IndicesAdminClient; import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.Settings; @@ -30,6 +28,8 @@ import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.transport.client.Client; +import org.opensearch.transport.client.IndicesAdminClient; import org.mockito.Mock; diff --git a/src/test/java/org/opensearch/security/util/ActionListenerUtils.java b/src/test/java/org/opensearch/security/util/ActionListenerUtils.java new file mode 100644 index 0000000000..8ec37649b1 --- /dev/null +++ b/src/test/java/org/opensearch/security/util/ActionListenerUtils.java @@ -0,0 +1,71 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.security.util; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import org.opensearch.core.action.ActionListener; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.instanceOf; +import static org.junit.Assert.fail; + +public class ActionListenerUtils { + public static class TestActionListener implements ActionListener { + private final CountDownLatch latch = new CountDownLatch(1); + private final AtomicReference response = new AtomicReference<>(); + private final AtomicReference exception = new AtomicReference<>(); + + @Override + public void onResponse(T result) { + response.set(result); + latch.countDown(); + } + + @Override + public void onFailure(Exception e) { + exception.set(e); + latch.countDown(); + } + + public T assertSuccess() { + waitForCompletion(); + if (exception.get() != null) { + fail("Expected success but got exception: " + exception.get()); + } + return response.get(); + } + + public Exception assertException(Class expectedExceptionClass) { + waitForCompletion(); + Exception actualException = exception.get(); + if (actualException == null) { + fail("Expected exception of type " + expectedExceptionClass.getSimpleName() + " but operation succeeded"); + } + assertThat("Exception type mismatch", actualException, instanceOf(expectedExceptionClass)); + return actualException; + } + + void waitForCompletion() { + try { + if (!latch.await(5, TimeUnit.SECONDS)) { + fail("Test timed out waiting for response"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + fail("Test interrupted: " + e.getMessage()); + } + } + } +} From cdfc2aaf99e3aa79e6848554a4370a7eb932e317 Mon Sep 17 00:00:00 2001 From: Derek Ho Date: Fri, 21 Mar 2025 11:07:59 -0400 Subject: [PATCH 3/3] PR comments Signed-off-by: Derek Ho --- .../security/action/apitokens/ApiToken.java | 6 +- .../action/apitokens/ApiTokenAction.java | 129 ++++++++---------- 2 files changed, 61 insertions(+), 74 deletions(-) diff --git a/src/main/java/org/opensearch/security/action/apitokens/ApiToken.java b/src/main/java/org/opensearch/security/action/apitokens/ApiToken.java index 6a81ad9f4d..b790f0d38f 100644 --- a/src/main/java/org/opensearch/security/action/apitokens/ApiToken.java +++ b/src/main/java/org/opensearch/security/action/apitokens/ApiToken.java @@ -22,7 +22,7 @@ public class ApiToken implements ToXContent { public static final String NAME_FIELD = "name"; - public static final String CREATION_TIME_FIELD = "creation_time"; + public static final String ISSUED_AT_FIELD = "iat"; public static final String CLUSTER_PERMISSIONS_FIELD = "cluster_permissions"; public static final String INDEX_PERMISSIONS_FIELD = "index_permissions"; public static final String INDEX_PATTERN_FIELD = "index_pattern"; @@ -149,7 +149,7 @@ public static ApiToken fromXContent(XContentParser parser) throws IOException { case NAME_FIELD: name = parser.text(); break; - case CREATION_TIME_FIELD: + case ISSUED_AT_FIELD: creationTime = Instant.ofEpochMilli(parser.longValue()); break; case EXPIRATION_FIELD: @@ -227,7 +227,7 @@ public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params xContentBuilder.field(NAME_FIELD, name); xContentBuilder.field(CLUSTER_PERMISSIONS_FIELD, clusterPermissions); xContentBuilder.field(INDEX_PERMISSIONS_FIELD, indexPermissions); - xContentBuilder.field(CREATION_TIME_FIELD, creationTime.toEpochMilli()); + xContentBuilder.field(ISSUED_AT_FIELD, creationTime.toEpochMilli()); xContentBuilder.endObject(); return xContentBuilder; } 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 ceeed23ada..448c5fecc1 100644 --- a/src/main/java/org/opensearch/security/action/apitokens/ApiTokenAction.java +++ b/src/main/java/org/opensearch/security/action/apitokens/ApiTokenAction.java @@ -54,10 +54,10 @@ import static org.opensearch.rest.RestRequest.Method.POST; import static org.opensearch.security.action.apitokens.ApiToken.ALLOWED_ACTIONS_FIELD; import static org.opensearch.security.action.apitokens.ApiToken.CLUSTER_PERMISSIONS_FIELD; -import static org.opensearch.security.action.apitokens.ApiToken.CREATION_TIME_FIELD; import static org.opensearch.security.action.apitokens.ApiToken.EXPIRATION_FIELD; import static org.opensearch.security.action.apitokens.ApiToken.INDEX_PATTERN_FIELD; import static org.opensearch.security.action.apitokens.ApiToken.INDEX_PERMISSIONS_FIELD; +import static org.opensearch.security.action.apitokens.ApiToken.ISSUED_AT_FIELD; import static org.opensearch.security.action.apitokens.ApiToken.NAME_FIELD; import static org.opensearch.security.dlic.rest.support.Utils.addRoutesPrefix; import static org.opensearch.security.support.ConfigConstants.SECURITY_RESTAPI_ADMIN_ENABLED; @@ -153,7 +153,7 @@ private RestChannelConsumer handleGet(RestRequest request, NodeClient client) { for (ApiToken token : tokens.values()) { builder.startObject(); builder.field(NAME_FIELD, token.getName()); - builder.field(CREATION_TIME_FIELD, token.getCreationTime().toEpochMilli()); + builder.field(ISSUED_AT_FIELD, token.getCreationTime().toEpochMilli()); builder.field(EXPIRATION_FIELD, token.getExpiration()); builder.field(CLUSTER_PERMISSIONS_FIELD, token.getClusterPermissions()); builder.field(INDEX_PERMISSIONS_FIELD, token.getIndexPermissions()); @@ -200,39 +200,27 @@ private RestChannelConsumer handlePost(RestRequest request, NodeClient client) { return; } - // If count is ok, create the token - apiTokenRepository.createApiToken(name, clusterPermissions, indexPermissions, expiration, ActionListener.wrap(token -> { - // After successful creation, trigger the update action - ApiTokenUpdateRequest updateRequest = new ApiTokenUpdateRequest(); - client.execute(ApiTokenUpdateAction.INSTANCE, updateRequest, ActionListener.wrap(updateResponse -> { - try { - XContentBuilder builder = channel.newBuilder(); - builder.startObject(); - builder.field("Api Token: ", token); - builder.endObject(); - channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder)); - builder.close(); - } catch (IOException e) { - sendErrorResponse( - channel, - RestStatus.INTERNAL_SERVER_ERROR, - "Failed to send response after token creation" - ); - } + apiTokenRepository.createApiToken( + name, + clusterPermissions, + indexPermissions, + expiration, + wrapWithCacheRefresh(ActionListener.wrap(token -> { + XContentBuilder builder = channel.newBuilder(); + builder.startObject(); + builder.field("token", token); + builder.endObject(); + channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder)); + builder.close(); + }, - updateException -> sendErrorResponse( + createException -> sendErrorResponse( channel, RestStatus.INTERNAL_SERVER_ERROR, - "Failed to propagate token creation: " + updateException.getMessage() + "Failed to create token: " + createException.getMessage() ) - )); - }, - createException -> sendErrorResponse( - channel, - RestStatus.INTERNAL_SERVER_ERROR, - "Failed to create token: " + createException.getMessage() - ) - )); + ), client) + ); }, countException -> sendErrorResponse( channel, @@ -247,6 +235,24 @@ private RestChannelConsumer handlePost(RestRequest request, NodeClient client) { }; } + private ActionListener wrapWithCacheRefresh(ActionListener listener, NodeClient client) { + return ActionListener.wrap(response -> { + try { + ApiTokenUpdateRequest updateRequest = new ApiTokenUpdateRequest(); + client.execute( + ApiTokenUpdateAction.INSTANCE, + updateRequest, + ActionListener.wrap( + updateResponse -> listener.onResponse(response), + exception -> listener.onFailure(new ApiTokenException("Failed to refresh cache", exception)) + ) + ); + } catch (Exception e) { + listener.onFailure(new ApiTokenException("Failed to refresh cache after operation", e)); + } + }, listener::onFailure); + } + /** * Extracts cluster permissions from the request body */ @@ -332,49 +338,30 @@ private RestChannelConsumer handleDelete(RestRequest request, NodeClient client) final Map requestBody = request.contentOrSourceParamParser().map(); validateRequestParameters(requestBody); - apiTokenRepository.deleteApiToken((String) requestBody.get(NAME_FIELD), ActionListener.wrap(ignored -> { - try { - ApiTokenUpdateRequest updateRequest = new ApiTokenUpdateRequest(); - client.execute(ApiTokenUpdateAction.INSTANCE, updateRequest, new ActionListener() { - @Override - public void onResponse(ApiTokenUpdateResponse updateResponse) { - try { - XContentBuilder builder = channel.newBuilder(); - builder.startObject(); - builder.field("message", "Token " + requestBody.get(NAME_FIELD) + " deleted successfully."); - builder.endObject(); - - BytesRestResponse response = new BytesRestResponse(RestStatus.OK, builder); - channel.sendResponse(response); - } catch (Exception e) { - sendErrorResponse( - channel, - RestStatus.INTERNAL_SERVER_ERROR, - "Failed to send response after token update" - ); - } - } - - @Override - public void onFailure(Exception e) { - sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, "Failed to propagate token deletion"); - } - }); - } catch (IOException e) { - sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, "Failed to build success response: " + e.getMessage()); - } - }, exception -> { - RestStatus status = RestStatus.INTERNAL_SERVER_ERROR; - if (exception instanceof ApiTokenException) { - status = RestStatus.NOT_FOUND; - } - sendErrorResponse(channel, status, exception.getMessage()); - })); + apiTokenRepository.deleteApiToken( + (String) requestBody.get(NAME_FIELD), + wrapWithCacheRefresh(ActionListener.wrap(ignored -> { + XContentBuilder builder = channel.newBuilder(); + builder.startObject(); + builder.field("message", "Token " + requestBody.get(NAME_FIELD) + " deleted successfully."); + builder.endObject(); + channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder)); + }, + deleteException -> sendErrorResponse( + channel, + RestStatus.INTERNAL_SERVER_ERROR, + "Failed to delete token: " + deleteException.getMessage() + ) + ), client) + ); } catch (final Exception exception) { - sendErrorResponse(channel, RestStatus.INTERNAL_SERVER_ERROR, exception.getMessage()); + RestStatus status = RestStatus.INTERNAL_SERVER_ERROR; + if (exception instanceof ApiTokenException) { + status = RestStatus.NOT_FOUND; + } + sendErrorResponse(channel, status, exception.getMessage()); } }; - } private void sendErrorResponse(RestChannel channel, RestStatus status, String errorMessage) {