fix(tokens): Fix stale cache problem, reduce cache timeout for access…

… tokens + fix listing owner tokens (#5140)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolver.java

@@ -101,6 +101,6 @@ public CompletableFuture<ListAccessTokenResult> get(DataFetchingEnvironment envi
    */
   private boolean isListingSelfTokens(final List<FacetFilterInput> filters, final QueryContext context) {
     return AuthorizationUtils.canGeneratePersonalAccessToken(context) && filters.stream()
-        .anyMatch(filter -> filter.getField().equals("actorUrn") && filter.getValue().equals(context.getActorUrn()));
+        .anyMatch(filter -> filter.getField().equals("ownerUrn") && filter.getValue().equals(context.getActorUrn()));
   }
 }

metadata-service/auth-api/src/main/java/com/datahub/authentication/AuthenticationConstants.java

@@ -30,6 +30,8 @@ public class AuthenticationConstants {
   public static final String SYSTEM_CLIENT_SECRET_CONFIG = "systemClientSecret";
 
   public static final String ENTITY_SERVICE = "entityService";
+  public static final String TOKEN_SERVICE = "tokenService";
+
 
   private AuthenticationConstants() { }
 }

metadata-service/auth-filter/src/main/java/com/datahub/authentication/filter/AuthenticationFilter.java

@@ -11,6 +11,7 @@
 import com.datahub.authentication.authenticator.AuthenticatorChain;
 import com.datahub.authentication.authenticator.DataHubSystemAuthenticator;
 import com.datahub.authentication.authenticator.NoOpAuthenticator;
+import com.datahub.authentication.token.StatefulTokenService;
 import com.google.common.collect.ImmutableMap;
 import com.linkedin.gms.factory.config.ConfigurationProvider;
 import com.linkedin.metadata.entity.EntityService;
@@ -49,6 +50,10 @@ public class AuthenticationFilter implements Filter {
   @Named("entityService")
   private EntityService _entityService;
 
+  @Inject
+  @Named("dataHubTokenService")
+  private StatefulTokenService _tokenService;
+
   private AuthenticatorChain authenticatorChain;
 
   @Override
@@ -109,8 +114,12 @@ private void buildAuthenticatorChain() {
     boolean isAuthEnabled = this.configurationProvider.getAuthentication().isEnabled();
 
     // Create authentication context object to pass to authenticator instances. They can use it as needed.
-    final AuthenticatorContext authenticatorContext = new AuthenticatorContext(ImmutableMap.of(ENTITY_SERVICE,
-        this._entityService));
+    final AuthenticatorContext authenticatorContext = new AuthenticatorContext(ImmutableMap.of(
+        ENTITY_SERVICE,
+        this._entityService,
+        TOKEN_SERVICE,
+        this._tokenService
+    ));
 
     if (isAuthEnabled) {
       log.info("Auth is enabled. Building authenticator chain...");

metadata-service/auth-impl/src/main/java/com/datahub/authentication/authenticator/DataHubTokenAuthenticator.java

@@ -62,9 +62,7 @@ public void init(@Nonnull final Map<String, Object> config, final AuthenticatorC
           "Unable to initialize DataHubTokenAuthenticator, entity service reference is not of type: "
               + "EntityService.class, found: " + entityService.getClass());
     }
-    this._statefulTokenService =
-        new StatefulTokenService(signingKey, signingAlgorithm, DEFAULT_ISSUER, (EntityService) entityService,
-            salt);
+    this._statefulTokenService = (StatefulTokenService) Objects.requireNonNull(context.data().get(TOKEN_SERVICE));
   }
 
   @Override

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

@@ -56,7 +56,7 @@ public StatefulTokenService(@Nonnull final String signingKey, @Nonnull final Str
     this._entityService = entityService;
     this._revokedTokenCache = CacheBuilder.newBuilder()
         .maximumSize(10000)
-        .expireAfterWrite(6, TimeUnit.HOURS)
+        .expireAfterWrite(5, TimeUnit.MINUTES)
         .build(new CacheLoader<String, Boolean>() {
           @Override
           public Boolean load(final String key) {

metadata-service/auth-impl/src/test/java/com/datahub/authentication/authenticator/DataHubTokenAuthenticatorTest.java

@@ -6,6 +6,7 @@
 import com.datahub.authentication.AuthenticationException;
 import com.datahub.authentication.AuthenticationRequest;
 import com.datahub.authentication.AuthenticatorContext;
+import com.datahub.authentication.token.StatefulTokenService;
 import com.datahub.authentication.token.TokenType;
 import com.google.common.collect.ImmutableMap;
 import com.linkedin.common.urn.Urn;
@@ -20,8 +21,7 @@
 import java.util.Collections;
 import java.util.Map;
 
-import static com.datahub.authentication.AuthenticationConstants.AUTHORIZATION_HEADER_NAME;
-import static com.datahub.authentication.AuthenticationConstants.ENTITY_SERVICE;
+import static com.datahub.authentication.AuthenticationConstants.*;
 import static com.datahub.authentication.authenticator.DataHubTokenAuthenticator.SALT_CONFIG_NAME;
 import static com.datahub.authentication.authenticator.DataHubTokenAuthenticator.SIGNING_ALG_CONFIG_NAME;
 import static com.datahub.authentication.authenticator.DataHubTokenAuthenticator.SIGNING_KEY_CONFIG_NAME;
@@ -40,12 +40,13 @@ public class DataHubTokenAuthenticatorTest {
     private static final String TEST_SALT = "WnEdIeTG/VVCLQqGwC/BAkqyY0k+H8NEAtWGejrBI93=";
 
     final EntityService mockService = Mockito.mock(EntityService.class);
+    final StatefulTokenService statefulTokenService = new StatefulTokenService(TEST_SIGNING_KEY, "HS256", null, mockService, TEST_SALT);
 
     @Test
     public void testInit() {
         final DataHubTokenAuthenticator authenticator = new DataHubTokenAuthenticator();
         AuthenticatorContext authenticatorContext =
-            new AuthenticatorContext(ImmutableMap.of(ENTITY_SERVICE, mockService));
+            new AuthenticatorContext(ImmutableMap.of(ENTITY_SERVICE, mockService, TOKEN_SERVICE, statefulTokenService));
         assertThrows(() -> authenticator.init(null, authenticatorContext));
         assertThrows(() -> authenticator.init(Collections.emptyMap(), authenticatorContext));
         assertThrows(() -> authenticator.init(ImmutableMap.of(SIGNING_KEY_CONFIG_NAME, TEST_SIGNING_KEY,
@@ -64,7 +65,7 @@ public void testAuthenticateFailureMissingAuthorizationHeader() {
 
         authenticator.init(ImmutableMap.of(SIGNING_KEY_CONFIG_NAME, TEST_SIGNING_KEY, SALT_CONFIG_NAME,
                         TEST_SALT, SIGNING_ALG_CONFIG_NAME, "HS256"),
-            new AuthenticatorContext(ImmutableMap.of(ENTITY_SERVICE, mockService)));
+            new AuthenticatorContext(ImmutableMap.of(ENTITY_SERVICE, mockService, TOKEN_SERVICE, statefulTokenService)));
 
         final AuthenticationRequest context = new AuthenticationRequest(Collections.emptyMap());
         assertThrows(AuthenticationException.class, () -> authenticator.authenticate(context));
@@ -73,10 +74,9 @@ public void testAuthenticateFailureMissingAuthorizationHeader() {
     @Test
     public void testAuthenticateFailureMissingBearerCredentials() {
         final DataHubTokenAuthenticator authenticator = new DataHubTokenAuthenticator();
-
         authenticator.init(ImmutableMap.of(SIGNING_KEY_CONFIG_NAME, TEST_SIGNING_KEY, SALT_CONFIG_NAME,
                         TEST_SALT, SIGNING_ALG_CONFIG_NAME, "HS256"),
-            new AuthenticatorContext(ImmutableMap.of(ENTITY_SERVICE, mockService)));
+            new AuthenticatorContext(ImmutableMap.of(ENTITY_SERVICE, mockService, TOKEN_SERVICE, statefulTokenService)));
 
         final AuthenticationRequest context = new AuthenticationRequest(
                 ImmutableMap.of(AUTHORIZATION_HEADER_NAME, "Basic username:password")
@@ -90,7 +90,7 @@ public void testAuthenticateFailureInvalidToken() {
 
         authenticator.init(ImmutableMap.of(SIGNING_KEY_CONFIG_NAME, TEST_SIGNING_KEY, SALT_CONFIG_NAME,
                         TEST_SALT, SIGNING_ALG_CONFIG_NAME, "HS256"),
-            new AuthenticatorContext(ImmutableMap.of(ENTITY_SERVICE, mockService)));
+            new AuthenticatorContext(ImmutableMap.of(ENTITY_SERVICE, mockService, TOKEN_SERVICE, statefulTokenService)));
 
         final AuthenticationRequest context = new AuthenticationRequest(
                 ImmutableMap.of(AUTHORIZATION_HEADER_NAME, "Bearer someRandomToken")
@@ -111,7 +111,7 @@ public void testAuthenticateSuccess() throws Exception {
         final DataHubTokenAuthenticator authenticator = new DataHubTokenAuthenticator();
         authenticator.init(ImmutableMap.of(SIGNING_KEY_CONFIG_NAME, TEST_SIGNING_KEY, SALT_CONFIG_NAME,
                         TEST_SALT, SIGNING_ALG_CONFIG_NAME, "HS256"),
-            new AuthenticatorContext(ImmutableMap.of(ENTITY_SERVICE, mockService)));
+            new AuthenticatorContext(ImmutableMap.of(ENTITY_SERVICE, mockService, TOKEN_SERVICE, statefulTokenService)));
 
         final Actor datahub = new Actor(ActorType.USER, "datahub");
         final String validToken = authenticator._statefulTokenService.generateAccessToken(

smoke-test/tests/tokens/revokable_access_token_test.py

@@ -131,7 +131,7 @@ def test_non_admin_can_create_list_revoke_tokens():
     user_tokenId = res_data["data"]["createAccessToken"]["metadata"]["id"]
 
     # User should be able to list his own token
-    res_data = listAccessTokens(user_session, [{"field": "actorUrn","value": "urn:li:corpuser:user"}])
+    res_data = listAccessTokens(user_session, [{"field": "ownerUrn","value": "urn:li:corpuser:user"}])
     assert res_data
     assert res_data["data"]
     assert res_data["data"]["listAccessTokens"]["total"] is not None
@@ -148,7 +148,7 @@ def test_non_admin_can_create_list_revoke_tokens():
     assert res_data["data"]["revokeAccessToken"] == True
 
     # Using a normal account, check that all its tokens where removed.
-    res_data = listAccessTokens(user_session, [{"field": "actorUrn","value": "urn:li:corpuser:user"}])
+    res_data = listAccessTokens(user_session, [{"field": "ownerUrn","value": "urn:li:corpuser:user"}])
     assert res_data
     assert res_data["data"]
     assert res_data["data"]["listAccessTokens"]["total"] is not None