feat(auth): Add roles to policy engine validation logic

datahub-project · Nov 5, 2023 · 27710ae · 27710ae
1 parent 81daae8
commit 27710ae
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 30 deletions.
diff --git a/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthorizedActors.java b/metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthorizedActors.java
@@ -15,6 +15,7 @@ public class AuthorizedActors {
   String privilege;
   List<Urn> users;
   List<Urn> groups;
+  List<Urn> roles;
   boolean allUsers;
   boolean allGroups;
 }
diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DataHubAuthorizer.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/DataHubAuthorizer.java
@@ -131,6 +131,7 @@ public AuthorizedActors authorizedActors(
 
     final List<Urn> authorizedUsers = new ArrayList<>();
     final List<Urn> authorizedGroups = new ArrayList<>();
+    final List<Urn> authorizedRoles = new ArrayList<>();
     boolean allUsers = false;
     boolean allGroups = false;
 
@@ -151,6 +152,7 @@ public AuthorizedActors authorizedActors(
       // Step 3: For each matching policy, add actors that are authorized.
       authorizedUsers.addAll(matchingActors.getUsers());
       authorizedGroups.addAll(matchingActors.getGroups());
+      authorizedRoles.add(matchingActors.getRoles());
       if (matchingActors.allUsers()) {
         allUsers = true;
       }
@@ -160,7 +162,7 @@ public AuthorizedActors authorizedActors(
     }
 
     // Step 4: Return all authorized users and groups.
-    return new AuthorizedActors(privilege, authorizedUsers, authorizedGroups, allUsers, allGroups);
+    return new AuthorizedActors(privilege, authorizedUsers, authorizedGroups, authorizedRoles, allUsers, allGroups);
   }
 
   /**

diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyEngine.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyEngine.java
@@ -13,6 +13,7 @@
 import com.linkedin.identity.RoleMembership;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.authorization.PoliciesConfig;
+import com.linkedin.pegasus2avro.common.AccessLevel;
 import com.linkedin.policy.DataHubActorFilter;
 import com.linkedin.policy.DataHubPolicyInfo;
 import com.linkedin.policy.DataHubResourceFilter;
@@ -32,7 +33,10 @@
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
 
+import lombok.AccessLevel;
+import lombok.AllArgsConstructor;
 import lombok.RequiredArgsConstructor;
+import lombok.Value;
 import lombok.extern.slf4j.Slf4j;
 
 import static com.linkedin.metadata.Constants.*;
@@ -75,6 +79,7 @@ public PolicyActors getMatchingActors(
       final Optional<ResolvedEntitySpec> resource) {
     final List<Urn> users = new ArrayList<>();
     final List<Urn> groups = new ArrayList<>();
+    final List<Urn> roles = new ArrayList<>();
     boolean allUsers = false;
     boolean allGroups = false;
     if (policyMatchesResource(policy, resource)) {
@@ -96,6 +101,9 @@ public PolicyActors getMatchingActors(
       if (actorFilter.getGroups() != null) {
         groups.addAll(actorFilter.getGroups());
       }
+      if (actorFilter.getRoles() != null) {
+        roles.addAll(actorFilter.getRoles());
+      }
 
       // 2. Fetch Actors based on resource ownership.
       if (actorFilter.isResourceOwners() && resource.isPresent()) {
@@ -104,7 +112,7 @@ public PolicyActors getMatchingActors(
         groups.addAll(groupOwners(owners));
       }
     }
-    return new PolicyActors(users, groups, allUsers, allGroups);
+    return new PolicyActors(users, groups, roles, allUsers, allGroups);
   }
 
   private boolean isPolicyApplicable(
@@ -438,34 +446,14 @@ public boolean isGranted() {
   /**
    * Class used to represent all valid users of a policy.
    */
+  @Value
+  @AllArgsConstructor(access = AccessLevel.PUBLIC)
   public static class PolicyActors {
-    final List<Urn> _users;
-    final List<Urn> _groups;
-    final Boolean _allUsers;
-    final Boolean _allGroups;
-
-    public PolicyActors(final List<Urn> users, final List<Urn> groups, final Boolean allUsers, final Boolean allGroups) {
-      _users = users;
-      _groups = groups;
-      _allUsers = allUsers;
-      _allGroups = allGroups;
-    }
-
-    public List<Urn> getUsers() {
-      return _users;
-    }
-
-    public List<Urn> getGroups() {
-      return _groups;
-    }
-
-    public Boolean allUsers() {
-      return _allUsers;
-    }
-
-    public Boolean allGroups() {
-      return _allGroups;
-    }
+    List<Urn> users;
+    List<Urn> groups;
+    List<Urn> roles;
+    Boolean allUsers;
+    Boolean allGroups;
   }
 
   private List<Urn> userOwners(final Set<String> owners) {

diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authorization/PolicyEngineTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/PolicyEngineTest.java
@@ -1,6 +1,7 @@
 package com.datahub.authorization;
 
 import com.datahub.authentication.Authentication;
+import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -1041,6 +1042,7 @@ public void testGetMatchingActorsResourceMatch() throws Exception {
         Urn.createFromString("urn:li:corpuser:user2"))));
     actorFilter.setGroups(new UrnArray(ImmutableList.of(Urn.createFromString("urn:li:corpGroup:group1"),
         Urn.createFromString("urn:li:corpGroup:group2"))));
+    actorFilter.setRoles(new UrnArray(ImmutableList.of(Urn.createFromString("urn:li:role:Admin"))));
     dataHubPolicyInfo.setActors(actorFilter);
 
     final DataHubResourceFilter resourceFilter = new DataHubResourceFilter();
@@ -1068,6 +1070,8 @@ public void testGetMatchingActorsResourceMatch() throws Exception {
         Urn.createFromString("urn:li:corpGroup:group2"), Urn.createFromString(AUTHORIZED_GROUP) // Resource Owner
     ));
 
+    assertEquals(actors.getRoles(), ImmutableList.of(Urn.createFromString("urn:li:role:Admin")));
+
     // Verify aspect client called, entity client not called.
     verify(_entityClient, times(0)).batchGetV2(eq(CORP_USER_ENTITY_NAME), eq(Collections.singleton(authorizedUserUrn)),
         eq(null), any());

diff --git a/...n/src/test/sample-test-plugins/src/main/java/com/datahub/plugins/test/TestAuthorizer.java b/...n/src/test/sample-test-plugins/src/main/java/com/datahub/plugins/test/TestAuthorizer.java
@@ -75,7 +75,7 @@ public AuthorizationResult authorize(@Nonnull AuthorizationRequest request) {
 
   @Override
   public AuthorizedActors authorizedActors(String privilege, Optional<EntitySpec> resourceSpec) {
-    return new AuthorizedActors("ALL", null, null, true, true);
+    return new AuthorizedActors("ALL", null, null, null, true, true);
   }
 }