feat(auth): implement session authorization cache (#11327)

datahub-project · Sep 11, 2024 · f23765e · f23765e
1 parent dfa9775
commit f23765e
Show file tree

Hide file tree

Showing 89 changed files with 1,140 additions and 1,144 deletions.
diff --git a/datahub-frontend/app/auth/AuthModule.java b/datahub-frontend/app/auth/AuthModule.java
@@ -25,7 +25,7 @@
 import java.util.Collections;
 
 import io.datahubproject.metadata.context.ActorContext;
-import io.datahubproject.metadata.context.AuthorizerContext;
+import io.datahubproject.metadata.context.AuthorizationContext;
 import io.datahubproject.metadata.context.EntityRegistryContext;
 import io.datahubproject.metadata.context.OperationContext;
 import io.datahubproject.metadata.context.OperationContextConfig;
@@ -183,10 +183,10 @@ protected OperationContext provideOperationContext(
     return OperationContext.builder()
             .operationContextConfig(systemConfig)
             .systemActorContext(systemActorContext)
+            // Authorizer.EMPTY is fine since it doesn't actually apply to system auth
+            .authorizationContext(AuthorizationContext.builder().authorizer(Authorizer.EMPTY).build())
             .searchContext(SearchContext.EMPTY)
             .entityRegistryContext(EntityRegistryContext.builder().build(EmptyEntityRegistry.EMPTY))
-            // Authorizer.EMPTY doesn't actually apply to system auth
-            .authorizerContext(AuthorizerContext.builder().authorizer(Authorizer.EMPTY).build())
             .build(systemAuthentication);
   }
 

diff --git a/...hql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java b/...hql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java
@@ -11,7 +11,6 @@
 import com.datahub.authorization.ConjunctivePrivilegeGroup;
 import com.datahub.authorization.DisjunctivePrivilegeGroup;
 import com.datahub.authorization.EntitySpec;
-import com.datahub.plugins.auth.authorization.Authorizer;
 import com.google.common.collect.ImmutableList;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
@@ -21,7 +20,6 @@
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.List;
-import java.util.Set;
 import javax.annotation.Nonnull;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.reflect.ConstructorUtils;
@@ -40,29 +38,25 @@ public class AuthorizationUtils {
 
   public static boolean canManageUsersAndGroups(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorizedEntityType(
-        context.getActorUrn(),
-        context.getAuthorizer(),
+        context.getOperationContext(),
         MANAGE,
         List.of(CORP_USER_ENTITY_NAME, CORP_GROUP_ENTITY_NAME));
   }
 
   public static boolean canManagePolicies(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorizedEntityType(
-        context.getActorUrn(), context.getAuthorizer(), MANAGE, List.of(POLICY_ENTITY_NAME));
+        context.getOperationContext(), MANAGE, List.of(POLICY_ENTITY_NAME));
   }
 
   public static boolean canGeneratePersonalAccessToken(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorized(
-            context.getAuthorizer(),
-            context.getActorUrn(),
-            PoliciesConfig.GENERATE_PERSONAL_ACCESS_TOKENS_PRIVILEGE)
-        || AuthUtil.isAuthorized(
-            context.getAuthorizer(), context.getActorUrn(), MANAGE_ACCESS_TOKENS);
+            context.getOperationContext(), PoliciesConfig.GENERATE_PERSONAL_ACCESS_TOKENS_PRIVILEGE)
+        || AuthUtil.isAuthorized(context.getOperationContext(), MANAGE_ACCESS_TOKENS);
   }
 
   public static boolean canManageTokens(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorizedEntityType(
-        context.getActorUrn(), context.getAuthorizer(), MANAGE, List.of(ACCESS_TOKEN_ENTITY_NAME));
+        context.getOperationContext(), MANAGE, List.of(ACCESS_TOKEN_ENTITY_NAME));
   }
 
   /**
@@ -78,13 +72,12 @@ public static boolean canCreateDomains(@Nonnull QueryContext context) {
                 new ConjunctivePrivilegeGroup(
                     ImmutableList.of(PoliciesConfig.MANAGE_DOMAINS_PRIVILEGE.getType()))));
 
-    return AuthUtil.isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), orPrivilegeGroups, null);
+    return AuthUtil.isAuthorized(context.getOperationContext(), orPrivilegeGroups, null);
   }
 
   public static boolean canManageDomains(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), PoliciesConfig.MANAGE_DOMAINS_PRIVILEGE);
+        context.getOperationContext(), PoliciesConfig.MANAGE_DOMAINS_PRIVILEGE);
   }
 
   /**
@@ -100,25 +93,22 @@ public static boolean canCreateTags(@Nonnull QueryContext context) {
                 new ConjunctivePrivilegeGroup(
                     ImmutableList.of(PoliciesConfig.MANAGE_TAGS_PRIVILEGE.getType()))));
 
-    return AuthUtil.isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), orPrivilegeGroups, null);
+    return AuthUtil.isAuthorized(context.getOperationContext(), orPrivilegeGroups, null);
   }
 
   public static boolean canManageTags(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), PoliciesConfig.MANAGE_TAGS_PRIVILEGE);
+        context.getOperationContext(), PoliciesConfig.MANAGE_TAGS_PRIVILEGE);
   }
 
   public static boolean canDeleteEntity(@Nonnull Urn entityUrn, @Nonnull QueryContext context) {
     return AuthUtil.isAuthorizedEntityUrns(
-        context.getAuthorizer(), context.getActorUrn(), DELETE, List.of(entityUrn));
+        context.getOperationContext(), DELETE, List.of(entityUrn));
   }
 
   public static boolean canManageUserCredentials(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        PoliciesConfig.MANAGE_USER_CREDENTIALS_PRIVILEGE);
+        context.getOperationContext(), PoliciesConfig.MANAGE_USER_CREDENTIALS_PRIVILEGE);
   }
 
   public static boolean canEditGroupMembers(
@@ -130,12 +120,7 @@ public static boolean canEditGroupMembers(
                 new ConjunctivePrivilegeGroup(
                     ImmutableList.of(PoliciesConfig.EDIT_GROUP_MEMBERS_PRIVILEGE.getType()))));
 
-    return isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        CORP_GROUP_ENTITY_NAME,
-        groupUrnStr,
-        orPrivilegeGroups);
+    return isAuthorized(context, CORP_GROUP_ENTITY_NAME, groupUrnStr, orPrivilegeGroups);
   }
 
   public static boolean canCreateGlobalAnnouncements(@Nonnull QueryContext context) {
@@ -149,27 +134,21 @@ public static boolean canCreateGlobalAnnouncements(@Nonnull QueryContext context
                     ImmutableList.of(
                         PoliciesConfig.MANAGE_GLOBAL_ANNOUNCEMENTS_PRIVILEGE.getType()))));
 
-    return AuthUtil.isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), orPrivilegeGroups, null);
+    return AuthUtil.isAuthorized(context.getOperationContext(), orPrivilegeGroups, null);
   }
 
   public static boolean canManageGlobalAnnouncements(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        PoliciesConfig.MANAGE_GLOBAL_ANNOUNCEMENTS_PRIVILEGE);
+        context.getOperationContext(), PoliciesConfig.MANAGE_GLOBAL_ANNOUNCEMENTS_PRIVILEGE);
   }
 
   public static boolean canManageGlobalViews(@Nonnull QueryContext context) {
-    return AuthUtil.isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), PoliciesConfig.MANAGE_GLOBAL_VIEWS);
+    return AuthUtil.isAuthorized(context.getOperationContext(), PoliciesConfig.MANAGE_GLOBAL_VIEWS);
   }
 
   public static boolean canManageOwnershipTypes(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        PoliciesConfig.MANAGE_GLOBAL_OWNERSHIP_TYPES);
+        context.getOperationContext(), PoliciesConfig.MANAGE_GLOBAL_OWNERSHIP_TYPES);
   }
 
   public static boolean canEditProperties(@Nonnull Urn targetUrn, @Nonnull QueryContext context) {
@@ -183,11 +162,7 @@ public static boolean canEditProperties(@Nonnull Urn targetUrn, @Nonnull QueryCo
                     ImmutableList.of(PoliciesConfig.EDIT_ENTITY_PROPERTIES_PRIVILEGE.getType()))));
 
     return AuthorizationUtils.isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        targetUrn.getEntityType(),
-        targetUrn.toString(),
-        orPrivilegeGroups);
+        context, targetUrn.getEntityType(), targetUrn.toString(), orPrivilegeGroups);
   }
 
   public static boolean canEditEntityQueries(
@@ -202,11 +177,7 @@ public static boolean canEditEntityQueries(
         .allMatch(
             entityUrn ->
                 isAuthorized(
-                    context.getAuthorizer(),
-                    context.getActorUrn(),
-                    entityUrn.getEntityType(),
-                    entityUrn.toString(),
-                    orPrivilegeGroups));
+                    context, entityUrn.getEntityType(), entityUrn.toString(), orPrivilegeGroups));
   }
 
   public static boolean canCreateQuery(
@@ -251,28 +222,7 @@ public static boolean canView(@Nonnull OperationContext opContext, @Nonnull Urn
         && !opContext.isSystemAuth()
         && VIEW_RESTRICTED_ENTITY_TYPES.contains(urn.getEntityType())) {
 
-      return opContext
-          .getViewAuthorizationContext()
-          .map(
-              viewAuthContext -> {
-
-                // check cache
-                if (viewAuthContext.canView(Set.of(urn))) {
-                  return true;
-                }
-
-                if (!canViewEntity(
-                    opContext.getSessionAuthentication().getActor().toUrnStr(),
-                    opContext.getAuthorizerContext().getAuthorizer(),
-                    urn)) {
-                  return false;
-                }
-
-                // cache viewable urn
-                viewAuthContext.addViewableUrns(Set.of(urn));
-                return true;
-              })
-          .orElse(false);
+      return canViewEntity(opContext, urn);
     }
     return true;
   }
@@ -386,38 +336,32 @@ public static <T> T restrictEntity(@Nonnull Object entity, Class<T> clazz) {
 
   public static boolean canManageStructuredProperties(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        PoliciesConfig.MANAGE_STRUCTURED_PROPERTIES_PRIVILEGE);
+        context.getOperationContext(), PoliciesConfig.MANAGE_STRUCTURED_PROPERTIES_PRIVILEGE);
   }
 
   public static boolean canManageForms(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        PoliciesConfig.MANAGE_DOCUMENTATION_FORMS_PRIVILEGE);
+        context.getOperationContext(), PoliciesConfig.MANAGE_DOCUMENTATION_FORMS_PRIVILEGE);
   }
 
   public static boolean canManageFeatures(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), PoliciesConfig.MANAGE_FEATURES_PRIVILEGE);
+        context.getOperationContext(), PoliciesConfig.MANAGE_FEATURES_PRIVILEGE);
   }
 
   public static boolean isAuthorized(
-      @Nonnull Authorizer authorizer,
-      @Nonnull String actor,
+      @Nonnull QueryContext context,
       @Nonnull String resourceType,
       @Nonnull String resource,
       @Nonnull DisjunctivePrivilegeGroup privilegeGroup) {
     final EntitySpec resourceSpec = new EntitySpec(resourceType, resource);
-    return AuthUtil.isAuthorized(authorizer, actor, privilegeGroup, resourceSpec);
+    return AuthUtil.isAuthorized(context.getOperationContext(), privilegeGroup, resourceSpec);
   }
 
   public static boolean isViewDatasetUsageAuthorized(
       final QueryContext context, final Urn resourceUrn) {
     return AuthUtil.isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
+        context.getOperationContext(),
         PoliciesConfig.VIEW_DATASET_USAGE_PRIVILEGE,
         new EntitySpec(resourceUrn.getEntityType(), resourceUrn.toString()));
   }

diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/MeResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/MeResolver.java
@@ -108,67 +108,57 @@ public CompletableFuture<AuthenticatedUser> get(DataFetchingEnvironment environm
 
   /** Returns true if the authenticated user has privileges to view analytics. */
   private boolean canViewAnalytics(final QueryContext context) {
-    return isAuthorized(context.getActorUrn(), context.getAuthorizer(), ANALYTICS, READ);
+    return isAuthorized(context.getOperationContext(), ANALYTICS, READ);
   }
 
   /** Returns true if the authenticated user has privileges to manage policies analytics. */
   private boolean canManagePolicies(final QueryContext context) {
     return isAuthorizedEntityType(
-        context.getActorUrn(), context.getAuthorizer(), MANAGE, List.of(POLICY_ENTITY_NAME));
+        context.getOperationContext(), MANAGE, List.of(POLICY_ENTITY_NAME));
   }
 
   /** Returns true if the authenticated user has privileges to manage users & groups. */
   private boolean canManageUsersGroups(final QueryContext context) {
     return isAuthorizedEntityType(
-        context.getActorUrn(),
-        context.getAuthorizer(),
+        context.getOperationContext(),
         MANAGE,
         List.of(CORP_USER_ENTITY_NAME, CORP_GROUP_ENTITY_NAME));
   }
 
   /** Returns true if the authenticated user has privileges to generate personal access tokens */
   private boolean canGeneratePersonalAccessToken(final QueryContext context) {
     return isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        PoliciesConfig.GENERATE_PERSONAL_ACCESS_TOKENS_PRIVILEGE);
+        context.getOperationContext(), PoliciesConfig.GENERATE_PERSONAL_ACCESS_TOKENS_PRIVILEGE);
   }
 
   /** Returns true if the authenticated user has privileges to view tests. */
   private boolean canViewTests(final QueryContext context) {
-    return isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), PoliciesConfig.VIEW_TESTS_PRIVILEGE);
+    return isAuthorized(context.getOperationContext(), PoliciesConfig.VIEW_TESTS_PRIVILEGE);
   }
 
   /** Returns true if the authenticated user has privileges to manage (add or remove) tests. */
   private boolean canManageTests(final QueryContext context) {
-    return isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), PoliciesConfig.MANAGE_TESTS_PRIVILEGE);
+    return isAuthorized(context.getOperationContext(), PoliciesConfig.MANAGE_TESTS_PRIVILEGE);
   }
 
   /** Returns true if the authenticated user has privileges to manage domains */
   private boolean canManageDomains(final QueryContext context) {
-    return isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), PoliciesConfig.MANAGE_DOMAINS_PRIVILEGE);
+    return isAuthorized(context.getOperationContext(), PoliciesConfig.MANAGE_DOMAINS_PRIVILEGE);
   }
 
   /** Returns true if the authenticated user has privileges to manage access tokens */
   private boolean canManageTokens(final QueryContext context) {
-    return isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), PoliciesConfig.MANAGE_ACCESS_TOKENS);
+    return isAuthorized(context.getOperationContext(), PoliciesConfig.MANAGE_ACCESS_TOKENS);
   }
 
   /** Returns true if the authenticated user has privileges to manage glossaries */
   private boolean canManageGlossaries(final QueryContext context) {
-    return isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), PoliciesConfig.MANAGE_GLOSSARIES_PRIVILEGE);
+    return isAuthorized(context.getOperationContext(), PoliciesConfig.MANAGE_GLOSSARIES_PRIVILEGE);
   }
 
   /** Returns true if the authenticated user has privileges to manage user credentials */
   private boolean canManageUserCredentials(@Nonnull QueryContext context) {
     return isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        PoliciesConfig.MANAGE_USER_CREDENTIALS_PRIVILEGE);
+        context.getOperationContext(), PoliciesConfig.MANAGE_USER_CREDENTIALS_PRIVILEGE);
   }
 }
diff --git a/...l-core/src/main/java/com/linkedin/datahub/graphql/resolvers/assertion/AssertionUtils.java b/...l-core/src/main/java/com/linkedin/datahub/graphql/resolvers/assertion/AssertionUtils.java
@@ -18,10 +18,6 @@ public static boolean isAuthorizedToEditAssertionFromAssertee(
                 new ConjunctivePrivilegeGroup(
                     ImmutableList.of(PoliciesConfig.EDIT_ENTITY_ASSERTIONS_PRIVILEGE.getType()))));
     return AuthorizationUtils.isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        asserteeUrn.getEntityType(),
-        asserteeUrn.toString(),
-        orPrivilegeGroups);
+        context, asserteeUrn.getEntityType(), asserteeUrn.toString(), orPrivilegeGroups);
   }
 }
diff --git a/...c/main/java/com/linkedin/datahub/graphql/resolvers/assertion/DeleteAssertionResolver.java b/...c/main/java/com/linkedin/datahub/graphql/resolvers/assertion/DeleteAssertionResolver.java
@@ -113,11 +113,7 @@ private boolean isAuthorizedToDeleteAssertionFromAssertee(
                 new ConjunctivePrivilegeGroup(
                     ImmutableList.of(PoliciesConfig.EDIT_ENTITY_ASSERTIONS_PRIVILEGE.getType()))));
     return AuthorizationUtils.isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        asserteeUrn.getEntityType(),
-        asserteeUrn.toString(),
-        orPrivilegeGroups);
+        context, asserteeUrn.getEntityType(), asserteeUrn.toString(), orPrivilegeGroups);
   }
 
   private Urn getAsserteeUrnFromInfo(final AssertionInfo info) {

diff --git a/...edin/datahub/graphql/resolvers/businessattribute/BusinessAttributeAuthorizationUtils.java b/...edin/datahub/graphql/resolvers/businessattribute/BusinessAttributeAuthorizationUtils.java
@@ -20,8 +20,7 @@ public static boolean canCreateBusinessAttribute(@Nonnull QueryContext context)
                 new ConjunctivePrivilegeGroup(
                     ImmutableList.of(
                         PoliciesConfig.MANAGE_BUSINESS_ATTRIBUTE_PRIVILEGE.getType()))));
-    return AuthUtil.isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), orPrivilegeGroups, null);
+    return AuthUtil.isAuthorized(context.getOperationContext(), orPrivilegeGroups, null);
   }
 
   public static boolean canManageBusinessAttribute(@Nonnull QueryContext context) {
@@ -31,7 +30,6 @@ public static boolean canManageBusinessAttribute(@Nonnull QueryContext context)
                 new ConjunctivePrivilegeGroup(
                     ImmutableList.of(
                         PoliciesConfig.MANAGE_BUSINESS_ATTRIBUTE_PRIVILEGE.getType()))));
-    return AuthUtil.isAuthorized(
-        context.getAuthorizer(), context.getActorUrn(), orPrivilegeGroups, null);
+    return AuthUtil.isAuthorized(context.getOperationContext(), orPrivilegeGroups, null);
   }
 }
diff --git a/...core/src/main/java/com/linkedin/datahub/graphql/resolvers/connection/ConnectionUtils.java b/...core/src/main/java/com/linkedin/datahub/graphql/resolvers/connection/ConnectionUtils.java
@@ -14,9 +14,7 @@ public class ConnectionUtils {
    */
   public static boolean canManageConnections(@Nonnull QueryContext context) {
     return AuthUtil.isAuthorized(
-        context.getAuthorizer(),
-        context.getActorUrn(),
-        PoliciesConfig.MANAGE_CONNECTIONS_PRIVILEGE);
+        context.getOperationContext(), PoliciesConfig.MANAGE_CONNECTIONS_PRIVILEGE);
   }
 
   private ConnectionUtils() {}