feat(token): helper to debug owner of raw token (#9793)

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

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<CompletableFuture<AccessTokenMetadata>> {
+
+  private final StatefulTokenService _tokenService;
+  private final EntityClient _entityClient;
+
+  public GetAccessTokenMetadataResolver(
+      final StatefulTokenService tokenService, EntityClient entityClient) {
+    _tokenService = tokenService;
+    _entityClient = entityClient;
+  }
+
+  @Override
+  public CompletableFuture<AccessTokenMetadata> 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<DataFetcherResult<AccessTokenMetadata>> batchLoad =
+                metadataType.batchLoad(ImmutableList.of(tokenUrn), context);
+            if (batchLoad.isEmpty()) {
+              return null;
+            }
+            return batchLoad.get(0).getData();
+          } catch (Exception e) {
+            throw new RuntimeException(e);
+          }
+        });
+  }
+}

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 {

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.*;
-import java.util.Date;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
 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<String, Boolean>() {
                   @Override
                   public Boolean load(final String key) {
-                    final Urn accessUrn =
+                    final Urn accessUrn = tokenUrnFromKey(key);
-                        Urn.createFromTuple(Constants.ACCESS_TOKEN_ENTITY_NAME, 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;

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
+        "query": """query getInviteToken($input: GetInviteTokenInput!) {
-            getInviteToken(input: $input){\n
+            getInviteToken(input: $input){
-              inviteToken\n
+              inviteToken
-            }\n
+            }
         }""",
         "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"]
+        res_data["data"]["createAccessToken"]["metadata"]["actorUrn"] == admin_user_urn
-        == f"urn:li:corpuser:{admin_user}"
     )
+    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"]
+        res_data["data"]["listAccessTokens"]["tokens"][0]["actorUrn"] == admin_user_urn
-        == f"urn:li:corpuser:{admin_user}"
     )
     assert (
-        res_data["data"]["listAccessTokens"]["tokens"][0]["ownerUrn"]
+        res_data["data"]["listAccessTokens"]["tokens"][0]["ownerUrn"] == admin_user_urn
-        == f"urn:li:corpuser:{admin_user}"
     )
 
     # 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
+        "query": """mutation createAccessToken($input: CreateAccessTokenInput!) {
-            createAccessToken(input: $input) {\n
+            createAccessToken(input: $input) {
-              accessToken\n
+              accessToken
-              metadata {\n
+              metadata {
-                id\n
+                id
-                actorUrn\n
+                actorUrn
-                ownerUrn\n
+                ownerUrn
-                name\n
+                name
-                description\n
+                description
               }
-            }\n
+            }
         }""",
         "variables": {
             "input": {
@@ -434,18 +440,18 @@ def listAccessTokens(session, filters=[]):
         input["filters"] = filters
 
     json = {
-        "query": """query listAccessTokens($input: ListAccessTokenInput!) {\n
+        "query": """query listAccessTokens($input: ListAccessTokenInput!) {
-            listAccessTokens(input: $input) {\n
+            listAccessTokens(input: $input) {
-              start\n
+              start
-              count\n
+              count
-              total\n
+              total
-              tokens {\n
+              tokens {
-                urn\n
+                urn
-                id\n
+                id
-                actorUrn\n
+                actorUrn
-                ownerUrn\n
+                ownerUrn
-              }\n
+              }
-            }\n
+            }
         }""",
         "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
+        "query": """query listUsers($input: ListUsersInput!) {
-            listUsers(input: $input) {\n
+            listUsers(input: $input) {
-              start\n
+              start
-              count\n
+              count
-              total\n
+              total
-              users {\n
+              users {
-                username\n
+                username
               }
             }
         }""",