yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-12-15 20:16:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Aseem Bansal 2024-02-08 20:51:12 +05:30 committed by GitHub
parent a5e473812e
commit f6e96dc0d3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 142 additions and 55 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
View File

@ -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))

59
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/GetAccessTokenMetadataResolver.java Normal file
View File

@ -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);
}
});
}
}

6
datahub-graphql-core/src/main/resources/auth.graphql
View File

@ -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 {

16
metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java
View File

@ -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<String, Boolean>() {
@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;

108
smoke-test/tests/tokens/revokable_access_token_test.py
View File

@ -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
}
}
}""",