feat(policies): Allow policies to be applied to resources based on ta…

…gs (#9684)

metadata-auth/auth-api/src/main/java/com/datahub/authorization/EntityFieldType.java

@@ -28,5 +28,7 @@ public enum EntityFieldType {
   /** Groups of which the entity (only applies to corpUser) is a member */
   GROUP_MEMBERSHIP,
   /** Data platform instance of resource */
-  DATA_PLATFORM_INSTANCE
+  DATA_PLATFORM_INSTANCE,
+  /** Tags of the entity */
+  TAG,
 }

metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultEntitySpecResolver.java

@@ -8,6 +8,7 @@
 import com.datahub.authorization.fieldresolverprovider.EntityUrnFieldResolverProvider;
 import com.datahub.authorization.fieldresolverprovider.GroupMembershipFieldResolverProvider;
 import com.datahub.authorization.fieldresolverprovider.OwnerFieldResolverProvider;
+import com.datahub.authorization.fieldresolverprovider.TagFieldResolverProvider;
 import com.google.common.collect.ImmutableList;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.util.Pair;
@@ -26,7 +27,8 @@ public DefaultEntitySpecResolver(Authentication systemAuthentication, EntityClie
             new DomainFieldResolverProvider(entityClient, systemAuthentication),
             new OwnerFieldResolverProvider(entityClient, systemAuthentication),
             new DataPlatformInstanceFieldResolverProvider(entityClient, systemAuthentication),
-            new GroupMembershipFieldResolverProvider(entityClient, systemAuthentication));
+            new GroupMembershipFieldResolverProvider(entityClient, systemAuthentication),
+            new TagFieldResolverProvider(entityClient, systemAuthentication));
   }
 
   @Override

metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/TagFieldResolverProvider.java

@@ -0,0 +1,65 @@
+package com.datahub.authorization.fieldresolverprovider;
+
+import com.datahub.authentication.Authentication;
+import com.datahub.authorization.EntityFieldType;
+import com.datahub.authorization.EntitySpec;
+import com.datahub.authorization.FieldResolver;
+import com.linkedin.common.GlobalTags;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.EnvelopedAspect;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.Constants;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+
+/** Provides field resolver for owners given entitySpec */
+@Slf4j
+@RequiredArgsConstructor
+public class TagFieldResolverProvider implements EntityFieldResolverProvider {
+
+  private final EntityClient _entityClient;
+  private final Authentication _systemAuthentication;
+
+  @Override
+  public List<EntityFieldType> getFieldTypes() {
+    return Collections.singletonList(EntityFieldType.TAG);
+  }
+
+  @Override
+  public FieldResolver getFieldResolver(EntitySpec entitySpec) {
+    return FieldResolver.getResolverFromFunction(entitySpec, this::getTags);
+  }
+
+  private FieldResolver.FieldValue getTags(EntitySpec entitySpec) {
+    Urn entityUrn = UrnUtils.getUrn(entitySpec.getEntity());
+    EnvelopedAspect globalTagsAspect;
+    try {
+      EntityResponse response =
+          _entityClient.getV2(
+              entityUrn.getEntityType(),
+              entityUrn,
+              Collections.singleton(Constants.GLOBAL_TAGS_ASPECT_NAME),
+              _systemAuthentication);
+      if (response == null
+          || !response.getAspects().containsKey(Constants.GLOBAL_TAGS_ASPECT_NAME)) {
+        return FieldResolver.emptyFieldValue();
+      }
+      globalTagsAspect = response.getAspects().get(Constants.GLOBAL_TAGS_ASPECT_NAME);
+    } catch (Exception e) {
+      log.error("Error while retrieving tags aspect for urn {}", entityUrn, e);
+      return FieldResolver.emptyFieldValue();
+    }
+    GlobalTags globalTags = new GlobalTags(globalTagsAspect.getValue().data());
+    return FieldResolver.FieldValue.builder()
+        .values(
+            globalTags.getTags().stream()
+                .map(tag -> tag.getTag().toString())
+                .collect(Collectors.toSet()))
+        .build();
+  }
+}

metadata-service/auth-impl/src/test/java/com/datahub/authorization/PolicyEngineTest.java

@@ -44,6 +44,7 @@ public class PolicyEngineTest {
   private static final String AUTHORIZED_GROUP = "urn:li:corpGroup:authorizedGroup";
   private static final String RESOURCE_URN = "urn:li:dataset:test";
   private static final String DOMAIN_URN = "urn:li:domain:domain1";
+  private static final String TAG_URN = "urn:li:tag:allowed";
   private static final String OWNERSHIP_TYPE_URN = "urn:li:ownershipType:__system__technical_owner";
   private static final String OTHER_OWNERSHIP_TYPE_URN =
       "urn:li:ownershipType:__system__data_steward";
@@ -69,7 +70,8 @@ public void setupTest() throws Exception {
             AUTHORIZED_PRINCIPAL,
             Collections.emptySet(),
             Collections.emptySet(),
-            Collections.singleton(AUTHORIZED_GROUP));
+            Collections.singleton(AUTHORIZED_GROUP),
+            Collections.emptySet());
     unauthorizedUserUrn = Urn.createFromString(UNAUTHORIZED_PRINCIPAL);
     resolvedUnauthorizedUserSpec =
         buildEntityResolvers(CORP_USER_ENTITY_NAME, UNAUTHORIZED_PRINCIPAL);
@@ -595,6 +597,7 @@ public void testEvaluatePolicyActorFilterUserResourceOwnersMatch() throws Except
             RESOURCE_URN,
             ImmutableSet.of(AUTHORIZED_PRINCIPAL),
             Collections.emptySet(),
+            Collections.emptySet(),
             Collections.emptySet());
     // Assert authorized user can edit entity tags, because he is a user owner.
     PolicyEngine.PolicyEvaluationResult result1 =
@@ -653,6 +656,7 @@ public void testEvaluatePolicyActorFilterUserResourceOwnersTypeMatch() throws Ex
             RESOURCE_URN,
             ImmutableSet.of(AUTHORIZED_PRINCIPAL),
             Collections.emptySet(),
+            Collections.emptySet(),
             Collections.emptySet());
 
     PolicyEngine.PolicyEvaluationResult result1 =
@@ -712,6 +716,7 @@ public void testEvaluatePolicyActorFilterUserResourceOwnersTypeNoMatch() throws
             RESOURCE_URN,
             ImmutableSet.of(AUTHORIZED_PRINCIPAL),
             Collections.emptySet(),
+            Collections.emptySet(),
             Collections.emptySet());
 
     PolicyEngine.PolicyEvaluationResult result1 =
@@ -767,6 +772,7 @@ public void testEvaluatePolicyActorFilterGroupResourceOwnersMatch() throws Excep
             RESOURCE_URN,
             ImmutableSet.of(AUTHORIZED_GROUP),
             Collections.emptySet(),
+            Collections.emptySet(),
             Collections.emptySet());
     // Assert authorized user can edit entity tags, because he is a user owner.
     PolicyEngine.PolicyEvaluationResult result1 =
@@ -1037,6 +1043,7 @@ public void testEvaluatePolicyResourceFilterSpecificResourceMatchDomain() throws
             RESOURCE_URN,
             Collections.emptySet(),
             Collections.singleton(DOMAIN_URN),
+            Collections.emptySet(),
             Collections.emptySet());
     PolicyEngine.PolicyEvaluationResult result =
         _policyEngine.evaluatePolicy(
@@ -1082,6 +1089,7 @@ public void testEvaluatePolicyResourceFilterSpecificResourceNoMatchDomain() thro
             RESOURCE_URN,
             Collections.emptySet(),
             Collections.singleton("urn:li:domain:domain2"),
+            Collections.emptySet(),
             Collections.emptySet()); // Domain doesn't match
     PolicyEngine.PolicyEvaluationResult result =
         _policyEngine.evaluatePolicy(
@@ -1095,6 +1103,52 @@ public void testEvaluatePolicyResourceFilterSpecificResourceNoMatchDomain() thro
     verify(_entityClient, times(0)).batchGetV2(any(), any(), any(), any());
   }
 
+  @Test
+  public void testEvaluatePolicyResourceFilterSpecificResourceMatchTag() throws Exception {
+    final DataHubPolicyInfo dataHubPolicyInfo = new DataHubPolicyInfo();
+    dataHubPolicyInfo.setType(METADATA_POLICY_TYPE);
+    dataHubPolicyInfo.setState(ACTIVE_POLICY_STATE);
+    dataHubPolicyInfo.setPrivileges(new StringArray("VIEW_ENTITY_PAGE"));
+    dataHubPolicyInfo.setDisplayName("Tag-based policy");
+    dataHubPolicyInfo.setDescription("Allow viewing entity pages based on tags");
+    dataHubPolicyInfo.setEditable(true);
+
+    final DataHubActorFilter actorFilter = new DataHubActorFilter();
+    actorFilter.setResourceOwners(true);
+    actorFilter.setAllUsers(true);
+    actorFilter.setAllGroups(true);
+    dataHubPolicyInfo.setActors(actorFilter);
+
+    final DataHubResourceFilter resourceFilter = new DataHubResourceFilter();
+    resourceFilter.setFilter(
+        FilterUtils.newFilter(
+            ImmutableMap.of(
+                EntityFieldType.TYPE,
+                Collections.singletonList("dataset"),
+                EntityFieldType.TAG,
+                Collections.singletonList(TAG_URN))));
+    dataHubPolicyInfo.setResources(resourceFilter);
+
+    ResolvedEntitySpec resourceSpec =
+        buildEntityResolvers(
+            "dataset",
+            RESOURCE_URN,
+            Collections.emptySet(),
+            Collections.emptySet(),
+            Collections.emptySet(),
+            Collections.singleton(TAG_URN));
+    PolicyEngine.PolicyEvaluationResult result =
+        _policyEngine.evaluatePolicy(
+            dataHubPolicyInfo,
+            resolvedAuthorizedUserSpec,
+            "VIEW_ENTITY_PAGE",
+            Optional.of(resourceSpec));
+    assertTrue(result.isGranted());
+
+    // Verify no network calls
+    verify(_entityClient, times(0)).batchGetV2(any(), any(), any(), any());
+  }
+
   @Test
   public void testGetGrantedPrivileges() throws Exception {
     // Policy 1, match dataset type and domain
@@ -1180,6 +1234,7 @@ public void testGetGrantedPrivileges() throws Exception {
             RESOURCE_URN,
             Collections.emptySet(),
             Collections.singleton(DOMAIN_URN),
+            Collections.emptySet(),
             Collections.emptySet()); // Everything matches
     assertEquals(
         _policyEngine.getGrantedPrivileges(
@@ -1192,6 +1247,7 @@ public void testGetGrantedPrivileges() throws Exception {
             RESOURCE_URN,
             Collections.emptySet(),
             Collections.singleton("urn:li:domain:domain2"),
+            Collections.emptySet(),
             Collections.emptySet()); // Domain doesn't match
     assertEquals(
         _policyEngine.getGrantedPrivileges(
@@ -1204,6 +1260,7 @@ public void testGetGrantedPrivileges() throws Exception {
             "urn:li:dataset:random",
             Collections.emptySet(),
             Collections.singleton(DOMAIN_URN),
+            Collections.emptySet(),
             Collections.emptySet()); // Resource doesn't match
     assertEquals(
         _policyEngine.getGrantedPrivileges(
@@ -1228,6 +1285,7 @@ public void testGetGrantedPrivileges() throws Exception {
             RESOURCE_URN,
             Collections.singleton(AUTHORIZED_PRINCIPAL),
             Collections.singleton(DOMAIN_URN),
+            Collections.emptySet(),
             Collections.emptySet()); // Is owner
     assertEquals(
         _policyEngine.getGrantedPrivileges(
@@ -1240,6 +1298,7 @@ public void testGetGrantedPrivileges() throws Exception {
             RESOURCE_URN,
             Collections.singleton(AUTHORIZED_PRINCIPAL),
             Collections.singleton(DOMAIN_URN),
+            Collections.emptySet(),
             Collections.emptySet()); // Resource type doesn't match
     assertEquals(
         _policyEngine.getGrantedPrivileges(
@@ -1289,6 +1348,7 @@ public void testGetMatchingActorsResourceMatch() throws Exception {
             RESOURCE_URN,
             ImmutableSet.of(AUTHORIZED_PRINCIPAL, AUTHORIZED_GROUP),
             Collections.emptySet(),
+            Collections.emptySet(),
             Collections.emptySet());
     PolicyEngine.PolicyActors actors =
         _policyEngine.getMatchingActors(dataHubPolicyInfo, Optional.of(resourceSpec));
@@ -1406,6 +1466,7 @@ public void testGetMatchingActorsByRoleResourceMatch() throws Exception {
             RESOURCE_URN,
             ImmutableSet.of(),
             Collections.emptySet(),
+            Collections.emptySet(),
             Collections.emptySet());
 
     PolicyEngine.PolicyActors actors =
@@ -1506,6 +1567,7 @@ public static ResolvedEntitySpec buildEntityResolvers(String entityType, String
         entityUrn,
         Collections.emptySet(),
         Collections.emptySet(),
+        Collections.emptySet(),
         Collections.emptySet());
   }
 
@@ -1514,7 +1576,8 @@ public static ResolvedEntitySpec buildEntityResolvers(
       String entityUrn,
       Set<String> owners,
       Set<String> domains,
-      Set<String> groups) {
+      Set<String> groups,
+      Set<String> tags) {
     return new ResolvedEntitySpec(
         new EntitySpec(entityType, entityUrn),
         ImmutableMap.of(
@@ -1527,6 +1590,8 @@ public static ResolvedEntitySpec buildEntityResolvers(
             EntityFieldType.DOMAIN,
             FieldResolver.getResolverFromValues(domains),
             EntityFieldType.GROUP_MEMBERSHIP,
-            FieldResolver.getResolverFromValues(groups)));
+            FieldResolver.getResolverFromValues(groups),
+            EntityFieldType.TAG,
+            FieldResolver.getResolverFromValues(tags)));
   }
 }