From f6e96dc0d3b280a693530ab3ae545fbde1943146 Mon Sep 17 00:00:00 2001 From: Aseem Bansal Date: Thu, 8 Feb 2024 20:51:12 +0530 Subject: [PATCH] feat(token): helper to debug owner of raw token (#9793) --- .../datahub/graphql/GmsGraphQLEngine.java | 8 +- .../auth/GetAccessTokenMetadataResolver.java | 59 ++++++++++ .../src/main/resources/auth.graphql | 6 + .../token/StatefulTokenService.java | 16 ++- .../tokens/revokable_access_token_test.py | 108 +++++++++++------- 5 files changed, 142 insertions(+), 55 deletions(-) create mode 100644 datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/GetAccessTokenMetadataResolver.java diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java index 28b3a982c7..3cb1d7ab9e 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java @@ -109,10 +109,7 @@ import com.linkedin.datahub.graphql.resolvers.MeResolver; import com.linkedin.datahub.graphql.resolvers.assertion.AssertionRunEventResolver; import com.linkedin.datahub.graphql.resolvers.assertion.DeleteAssertionResolver; import com.linkedin.datahub.graphql.resolvers.assertion.EntityAssertionsResolver; -import com.linkedin.datahub.graphql.resolvers.auth.CreateAccessTokenResolver; -import com.linkedin.datahub.graphql.resolvers.auth.GetAccessTokenResolver; -import com.linkedin.datahub.graphql.resolvers.auth.ListAccessTokensResolver; -import com.linkedin.datahub.graphql.resolvers.auth.RevokeAccessTokenResolver; +import com.linkedin.datahub.graphql.resolvers.auth.*; import com.linkedin.datahub.graphql.resolvers.browse.BrowsePathsResolver; import com.linkedin.datahub.graphql.resolvers.browse.BrowseResolver; import com.linkedin.datahub.graphql.resolvers.browse.EntityBrowsePathsResolver; @@ -931,6 +928,9 @@ public class GmsGraphQLEngine { .dataFetcher("getEntityCounts", new EntityCountsResolver(this.entityClient)) .dataFetcher("getAccessToken", new GetAccessTokenResolver(statefulTokenService)) .dataFetcher("listAccessTokens", new ListAccessTokensResolver(this.entityClient)) + .dataFetcher( + "getAccessTokenMetadata", + new GetAccessTokenMetadataResolver(statefulTokenService, this.entityClient)) .dataFetcher("container", getResolver(containerType)) .dataFetcher("listDomains", new ListDomainsResolver(this.entityClient)) .dataFetcher("listSecrets", new ListSecretsResolver(this.entityClient)) diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/GetAccessTokenMetadataResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/GetAccessTokenMetadataResolver.java new file mode 100644 index 0000000000..c3e14565e0 --- /dev/null +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/GetAccessTokenMetadataResolver.java @@ -0,0 +1,59 @@ +package com.linkedin.datahub.graphql.resolvers.auth; + +import com.datahub.authentication.token.StatefulTokenService; +import com.google.common.collect.ImmutableList; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.authorization.AuthorizationUtils; +import com.linkedin.datahub.graphql.exception.AuthorizationException; +import com.linkedin.datahub.graphql.generated.AccessTokenMetadata; +import com.linkedin.datahub.graphql.types.auth.AccessTokenMetadataType; +import com.linkedin.entity.client.EntityClient; +import graphql.execution.DataFetcherResult; +import graphql.schema.DataFetcher; +import graphql.schema.DataFetchingEnvironment; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +public class GetAccessTokenMetadataResolver + implements DataFetcher> { + + private final StatefulTokenService _tokenService; + private final EntityClient _entityClient; + + public GetAccessTokenMetadataResolver( + final StatefulTokenService tokenService, EntityClient entityClient) { + _tokenService = tokenService; + _entityClient = entityClient; + } + + @Override + public CompletableFuture get(final DataFetchingEnvironment environment) + throws Exception { + return CompletableFuture.supplyAsync( + () -> { + final QueryContext context = environment.getContext(); + final String token = environment.getArgument("token"); + log.info("User {} requesting access token metadata information.", context.getActorUrn()); + if (!AuthorizationUtils.canManageTokens(context)) { + throw new AuthorizationException( + "Unauthorized to perform this action. Please contact your DataHub administrator."); + } + + AccessTokenMetadataType metadataType = new AccessTokenMetadataType(_entityClient); + final String tokenHash = _tokenService.hash(token); + final String tokenUrn = _tokenService.tokenUrnFromKey(tokenHash).toString(); + try { + List> batchLoad = + metadataType.batchLoad(ImmutableList.of(tokenUrn), context); + if (batchLoad.isEmpty()) { + return null; + } + return batchLoad.get(0).getData(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } +} diff --git a/datahub-graphql-core/src/main/resources/auth.graphql b/datahub-graphql-core/src/main/resources/auth.graphql index b76aa132c2..c7dc6be137 100644 --- a/datahub-graphql-core/src/main/resources/auth.graphql +++ b/datahub-graphql-core/src/main/resources/auth.graphql @@ -11,6 +11,12 @@ extend type Query { List access tokens stored in DataHub. """ listAccessTokens(input: ListAccessTokenInput!): ListAccessTokenResult! + + """ + Fetches the metadata of an access token. + This is useful to debug when you have a raw token but don't know the actor. + """ + getAccessTokenMetadata(token: String!): AccessTokenMetadata! } extend type Mutation { diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java b/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java index 0d1da4a768..50e357331b 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java @@ -16,12 +16,7 @@ import com.linkedin.metadata.key.DataHubAccessTokenKey; import com.linkedin.metadata.utils.AuditStampUtils; import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeProposal; -import java.util.Base64; -import java.util.Date; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Objects; +import java.util.*; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import javax.annotation.Nonnull; @@ -58,8 +53,7 @@ public class StatefulTokenService extends StatelessTokenService { new CacheLoader() { @Override public Boolean load(final String key) { - final Urn accessUrn = - Urn.createFromTuple(Constants.ACCESS_TOKEN_ENTITY_NAME, key); + final Urn accessUrn = tokenUrnFromKey(key); return !_entityService.exists(accessUrn, true); } }); @@ -173,10 +167,14 @@ public class StatefulTokenService extends StatelessTokenService { } } + public Urn tokenUrnFromKey(String tokenHash) { + return Urn.createFromTuple(Constants.ACCESS_TOKEN_ENTITY_NAME, tokenHash); + } + public void revokeAccessToken(@Nonnull String hashedToken) throws TokenException { try { if (!_revokedTokenCache.get(hashedToken)) { - final Urn tokenUrn = Urn.createFromTuple(Constants.ACCESS_TOKEN_ENTITY_NAME, hashedToken); + final Urn tokenUrn = tokenUrnFromKey(hashedToken); _entityService.deleteUrn(tokenUrn); _revokedTokenCache.put(hashedToken, true); return; diff --git a/smoke-test/tests/tokens/revokable_access_token_test.py b/smoke-test/tests/tokens/revokable_access_token_test.py index 10332b32b9..6e8deb41f1 100644 --- a/smoke-test/tests/tokens/revokable_access_token_test.py +++ b/smoke-test/tests/tokens/revokable_access_token_test.py @@ -40,10 +40,10 @@ def custom_user_setup(): # Test getting the invite token get_invite_token_json = { - "query": """query getInviteToken($input: GetInviteTokenInput!) {\n - getInviteToken(input: $input){\n - inviteToken\n - }\n + "query": """query getInviteToken($input: GetInviteTokenInput!) { + getInviteToken(input: $input){ + inviteToken + } }""", "variables": {"input": {}}, } @@ -131,6 +131,7 @@ def access_token_setup(): @pytest.mark.dependency(depends=["test_healthchecks"]) def test_admin_can_create_list_and_revoke_tokens(wait_for_healthchecks): admin_session = login_as(admin_user, admin_pass) + admin_user_urn = f"urn:li:corpuser:{admin_user}" # Using a super account, there should be no tokens res_data = listAccessTokens(admin_session) @@ -140,19 +141,26 @@ def test_admin_can_create_list_and_revoke_tokens(wait_for_healthchecks): assert len(res_data["data"]["listAccessTokens"]["tokens"]) == 0 # Using a super account, generate a token for itself. - res_data = generateAccessToken_v2(admin_session, f"urn:li:corpuser:{admin_user}") + res_data = generateAccessToken_v2(admin_session, admin_user_urn) assert res_data assert res_data["data"] assert res_data["data"]["createAccessToken"] assert res_data["data"]["createAccessToken"]["accessToken"] assert ( - res_data["data"]["createAccessToken"]["metadata"]["actorUrn"] - == f"urn:li:corpuser:{admin_user}" + res_data["data"]["createAccessToken"]["metadata"]["actorUrn"] == admin_user_urn ) + access_token = res_data["data"]["createAccessToken"]["accessToken"] admin_tokenId = res_data["data"]["createAccessToken"]["metadata"]["id"] # Sleep for eventual consistency wait_for_writes_to_sync() + res_data = getAccessTokenMetadata(admin_session, access_token) + assert res_data + assert res_data["data"] + assert res_data["data"]["getAccessTokenMetadata"] + assert res_data["data"]["getAccessTokenMetadata"]["ownerUrn"] == admin_user_urn + assert res_data["data"]["getAccessTokenMetadata"]["actorUrn"] == admin_user_urn + # Using a super account, list the previously created token. res_data = listAccessTokens(admin_session) assert res_data @@ -160,12 +168,10 @@ def test_admin_can_create_list_and_revoke_tokens(wait_for_healthchecks): assert res_data["data"]["listAccessTokens"]["total"] is not None assert len(res_data["data"]["listAccessTokens"]["tokens"]) == 1 assert ( - res_data["data"]["listAccessTokens"]["tokens"][0]["actorUrn"] - == f"urn:li:corpuser:{admin_user}" + res_data["data"]["listAccessTokens"]["tokens"][0]["actorUrn"] == admin_user_urn ) assert ( - res_data["data"]["listAccessTokens"]["tokens"][0]["ownerUrn"] - == f"urn:li:corpuser:{admin_user}" + res_data["data"]["listAccessTokens"]["tokens"][0]["ownerUrn"] == admin_user_urn ) # Check that the super account can revoke tokens that it created @@ -396,17 +402,17 @@ def test_non_admin_can_not_generate_tokens_for_others(wait_for_healthchecks): def generateAccessToken_v2(session, actorUrn): # Create new token json = { - "query": """mutation createAccessToken($input: CreateAccessTokenInput!) {\n - createAccessToken(input: $input) {\n - accessToken\n - metadata {\n - id\n - actorUrn\n - ownerUrn\n - name\n - description\n + "query": """mutation createAccessToken($input: CreateAccessTokenInput!) { + createAccessToken(input: $input) { + accessToken + metadata { + id + actorUrn + ownerUrn + name + description } - }\n + } }""", "variables": { "input": { @@ -434,18 +440,18 @@ def listAccessTokens(session, filters=[]): input["filters"] = filters json = { - "query": """query listAccessTokens($input: ListAccessTokenInput!) {\n - listAccessTokens(input: $input) {\n - start\n - count\n - total\n - tokens {\n - urn\n - id\n - actorUrn\n - ownerUrn\n - }\n - }\n + "query": """query listAccessTokens($input: ListAccessTokenInput!) { + listAccessTokens(input: $input) { + start + count + total + tokens { + urn + id + actorUrn + ownerUrn + } + } }""", "variables": {"input": input}, } @@ -458,7 +464,7 @@ def listAccessTokens(session, filters=[]): def revokeAccessToken(session, tokenId): # Revoke token json = { - "query": """mutation revokeAccessToken($tokenId: String!) {\n + "query": """mutation revokeAccessToken($tokenId: String!) { revokeAccessToken(tokenId: $tokenId) }""", "variables": {"tokenId": tokenId}, @@ -470,10 +476,28 @@ def revokeAccessToken(session, tokenId): return response.json() +def getAccessTokenMetadata(session, token): + json = { + "query": """ + query getAccessTokenMetadata($token: String!) { + getAccessTokenMetadata(token: $token) { + id + ownerUrn + actorUrn + } + }""", + "variables": {"token": token}, + } + response = session.post(f"{get_frontend_url()}/api/v2/graphql", json=json) + response.raise_for_status() + + return response.json() + + def removeUser(session, urn): # Remove user json = { - "query": """mutation removeUser($urn: String!) {\n + "query": """mutation removeUser($urn: String!) { removeUser(urn: $urn) }""", "variables": {"urn": urn}, @@ -493,13 +517,13 @@ def listUsers(session): # list users json = { - "query": """query listUsers($input: ListUsersInput!) {\n - listUsers(input: $input) {\n - start\n - count\n - total\n - users {\n - username\n + "query": """query listUsers($input: ListUsersInput!) { + listUsers(input: $input) { + start + count + total + users { + username } } }""",