feat(roles): Add support for roles in groups in GMS (#9659)

Co-authored-by: Aseem Bansal <[email protected]>

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/CreateGroupResolver.java

@@ -43,9 +43,10 @@ public CompletableFuture<String> get(final DataFetchingEnvironment environment)
             // Create the Group key.
             final CorpGroupKey key = new CorpGroupKey();
             final String id = input.getId() != null ? input.getId() : UUID.randomUUID().toString();
+            final String description = input.getDescription() != null ? input.getDescription() : "";
             key.setName(id); // 'name' in the key really reflects nothing more than a stable "id".
             return _groupService.createNativeGroup(
-                key, input.getName(), input.getDescription(), authentication);
+                key, input.getName(), description, authentication);
           } catch (Exception e) {
             throw new RuntimeException("Failed to create group", e);
           }

metadata-models/src/main/pegasus/com/linkedin/identity/RoleMembership.pdl

@@ -3,7 +3,7 @@ namespace com.linkedin.identity
 import com.linkedin.common.Urn
 
 /**
- * Carries information about which roles a user is assigned to.
+ * Carries information about which roles a user or group is assigned to.
  */
 @Aspect = {
   "name": "roleMembership"

metadata-models/src/main/resources/entity-registry.yml

@@ -184,6 +184,7 @@ entities:
       - ownership
       - status
       - origin
+      - roleMembership
   - name: domain
     doc: A data domain within an organization.
     category: core

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

@@ -3,15 +3,19 @@
 import static com.linkedin.metadata.Constants.*;
 
 import com.datahub.authentication.Authentication;
+import com.google.common.collect.ImmutableSet;
 import com.linkedin.common.Owner;
 import com.linkedin.common.Ownership;
+import com.linkedin.common.UrnArray;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.entity.EntityResponse;
 import com.linkedin.entity.EnvelopedAspect;
 import com.linkedin.entity.EnvelopedAspectMap;
 import com.linkedin.entity.client.EntityClient;
+import com.linkedin.identity.GroupMembership;
+import com.linkedin.identity.NativeGroupMembership;
 import com.linkedin.identity.RoleMembership;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.authorization.PoliciesConfig;
@@ -26,6 +30,7 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
@@ -393,15 +398,17 @@ private Set<Urn> resolveRoles(
 
     Set<Urn> roles = new HashSet<>();
     final EnvelopedAspectMap aspectMap;
-
     try {
       Urn actorUrn = Urn.createFromString(actor);
       final EntityResponse corpUser =
           _entityClient
               .batchGetV2(
                   CORP_USER_ENTITY_NAME,
                   Collections.singleton(actorUrn),
-                  Collections.singleton(ROLE_MEMBERSHIP_ASPECT_NAME),
+                  ImmutableSet.of(
+                      ROLE_MEMBERSHIP_ASPECT_NAME,
+                      GROUP_MEMBERSHIP_ASPECT_NAME,
+                      NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME),
                   _systemAuthentication)
               .get(actorUrn);
       if (corpUser == null || !corpUser.hasAspects()) {
@@ -414,19 +421,71 @@ private Set<Urn> resolveRoles(
       return roles;
     }
 
-    if (!aspectMap.containsKey(ROLE_MEMBERSHIP_ASPECT_NAME)) {
-      return roles;
+    if (aspectMap.containsKey(ROLE_MEMBERSHIP_ASPECT_NAME)) {
+      RoleMembership roleMembership =
+          new RoleMembership(aspectMap.get(ROLE_MEMBERSHIP_ASPECT_NAME).getValue().data());
+      if (roleMembership.hasRoles()) {
+        roles.addAll(roleMembership.getRoles());
+      }
     }
 
-    RoleMembership roleMembership =
-        new RoleMembership(aspectMap.get(ROLE_MEMBERSHIP_ASPECT_NAME).getValue().data());
-    if (roleMembership.hasRoles()) {
-      roles.addAll(roleMembership.getRoles());
+    List<Urn> groups = new ArrayList<>();
+    if (aspectMap.containsKey(GROUP_MEMBERSHIP_ASPECT_NAME)) {
+      GroupMembership groupMembership =
+          new GroupMembership(aspectMap.get(GROUP_MEMBERSHIP_ASPECT_NAME).getValue().data());
+      groups.addAll(groupMembership.getGroups());
+    }
+    if (aspectMap.containsKey(NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)) {
+      NativeGroupMembership nativeGroupMembership =
+          new NativeGroupMembership(
+              aspectMap.get(NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME).getValue().data());
+      groups.addAll(nativeGroupMembership.getNativeGroups());
+    }
+    if (!groups.isEmpty()) {
+      GroupMembership memberships = new GroupMembership();
+      memberships.setGroups(new UrnArray(groups));
+      roles.addAll(getRolesFromGroups(memberships));
+    }
+
+    if (!roles.isEmpty()) {
       context.setRoles(roles);
     }
+
     return roles;
   }
 
+  private Set<Urn> getRolesFromGroups(final GroupMembership groupMembership) {
+
+    HashSet<Urn> groups = new HashSet<>(groupMembership.getGroups());
+    try {
+      Map<Urn, EntityResponse> responseMap =
+          _entityClient.batchGetV2(
+              CORP_GROUP_ENTITY_NAME,
+              groups,
+              ImmutableSet.of(ROLE_MEMBERSHIP_ASPECT_NAME),
+              _systemAuthentication);
+
+      return responseMap.keySet().stream()
+          .filter(Objects::nonNull)
+          .filter(key -> responseMap.get(key) != null)
+          .filter(key -> responseMap.get(key).hasAspects())
+          .map(key -> responseMap.get(key).getAspects())
+          .filter(aspectMap -> aspectMap.containsKey(ROLE_MEMBERSHIP_ASPECT_NAME))
+          .map(
+              aspectMap ->
+                  new RoleMembership(aspectMap.get(ROLE_MEMBERSHIP_ASPECT_NAME).getValue().data()))
+          .filter(RoleMembership::hasRoles)
+          .map(RoleMembership::getRoles)
+          .flatMap(List::stream)
+          .collect(Collectors.toSet());
+
+    } catch (Exception e) {
+      log.error(
+          String.format("Failed to fetch %s for urns %s", ROLE_MEMBERSHIP_ASPECT_NAME, groups), e);
+      return new HashSet<>();
+    }
+  }
+
   private Set<String> resolveGroups(
       ResolvedEntitySpec resolvedActorSpec, PolicyEvaluationContext context) {
     if (context.groups != null) {

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

@@ -28,6 +28,7 @@
 import com.linkedin.common.UrnArray;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.domain.DomainProperties;
 import com.linkedin.domain.Domains;
@@ -36,6 +37,7 @@
 import com.linkedin.entity.EnvelopedAspect;
 import com.linkedin.entity.EnvelopedAspectMap;
 import com.linkedin.entity.client.EntityClient;
+import com.linkedin.identity.GroupMembership;
 import com.linkedin.identity.RoleMembership;
 import com.linkedin.metadata.query.SearchFlags;
 import com.linkedin.metadata.search.ScrollResult;
@@ -254,10 +256,14 @@ public void setupTest() throws Exception {
     when(_entityClient.batchGetV2(
             any(),
             eq(Collections.singleton(USER_WITH_ADMIN_ROLE)),
-            eq(Collections.singleton(ROLE_MEMBERSHIP_ASPECT_NAME)),
+            eq(
+                ImmutableSet.of(
+                    ROLE_MEMBERSHIP_ASPECT_NAME,
+                    GROUP_MEMBERSHIP_ASPECT_NAME,
+                    NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)),
             any()))
         .thenReturn(
-            createUserRoleMembershipBatchResponse(
+            createRoleMembershipBatchResponse(
                 USER_WITH_ADMIN_ROLE, UrnUtils.getUrn("urn:li:dataHubRole:Admin")));
 
     final Authentication systemAuthentication =
@@ -460,6 +466,49 @@ public void testAuthorizationOnDomainWithoutPrivilegeIsDenied() {
     assertEquals(_dataHubAuthorizer.authorize(request).getType(), AuthorizationResult.Type.DENY);
   }
 
+  @Test
+  public void testAuthorizationGrantedBasedOnGroupRole() throws Exception {
+    final EntitySpec resourceSpec = new EntitySpec("dataset", "urn:li:dataset:custom");
+
+    final Urn userUrnWithoutPermissions = UrnUtils.getUrn("urn:li:corpuser:userWithoutRole");
+    final Urn groupWithAdminPermission = UrnUtils.getUrn("urn:li:corpGroup:groupWithRole");
+    final UrnArray groups = new UrnArray(List.of(groupWithAdminPermission));
+    final GroupMembership groupMembership = new GroupMembership();
+    groupMembership.setGroups(groups);
+
+    // User has no role associated but is part of 1 group
+    when(_entityClient.batchGetV2(
+            any(),
+            eq(Collections.singleton(userUrnWithoutPermissions)),
+            eq(
+                ImmutableSet.of(
+                    ROLE_MEMBERSHIP_ASPECT_NAME,
+                    GROUP_MEMBERSHIP_ASPECT_NAME,
+                    NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)),
+            any()))
+        .thenReturn(
+            createEntityBatchResponse(
+                userUrnWithoutPermissions, GROUP_MEMBERSHIP_ASPECT_NAME, groupMembership));
+
+    // Group has a role
+    when(_entityClient.batchGetV2(
+            any(),
+            eq(Collections.singleton(groupWithAdminPermission)),
+            eq(Collections.singleton(ROLE_MEMBERSHIP_ASPECT_NAME)),
+            any()))
+        .thenReturn(
+            createRoleMembershipBatchResponse(
+                groupWithAdminPermission, UrnUtils.getUrn("urn:li:dataHubRole:Admin")));
+
+    // This request should only be valid for actor with the admin role.
+    // Which the urn:li:corpuser:userWithoutRole does not have
+    AuthorizationRequest request =
+        new AuthorizationRequest(
+            userUrnWithoutPermissions.toString(), "EDIT_USER_PROFILE", Optional.of(resourceSpec));
+
+    assertEquals(_dataHubAuthorizer.authorize(request).getType(), AuthorizationResult.Type.ALLOW);
+  }
+
   private DataHubPolicyInfo createDataHubPolicyInfo(
       boolean active, List<String> privileges, @Nullable final Urn domain) throws Exception {
 
@@ -575,20 +624,24 @@ private Map<Urn, EntityResponse> createDomainPropertiesBatchResponse(
     return batchResponse;
   }
 
-  private Map<Urn, EntityResponse> createUserRoleMembershipBatchResponse(
-      final Urn userUrn, @Nullable final Urn roleUrn) {
-    final Map<Urn, EntityResponse> batchResponse = new HashMap<>();
-    final EntityResponse response = new EntityResponse();
-    EnvelopedAspectMap aspectMap = new EnvelopedAspectMap();
+  private Map<Urn, EntityResponse> createRoleMembershipBatchResponse(
+      final Urn actorUrn, @Nullable final Urn roleUrn) {
     final RoleMembership membership = new RoleMembership();
     if (roleUrn != null) {
       membership.setRoles(new UrnArray(roleUrn));
     }
+    return createEntityBatchResponse(actorUrn, ROLE_MEMBERSHIP_ASPECT_NAME, membership);
+  }
+
+  private Map<Urn, EntityResponse> createEntityBatchResponse(
+      final Urn actorUrn, final String aspectName, final RecordTemplate aspect) {
+    final Map<Urn, EntityResponse> batchResponse = new HashMap<>();
+    final EntityResponse response = new EntityResponse();
+    EnvelopedAspectMap aspectMap = new EnvelopedAspectMap();
     aspectMap.put(
-        ROLE_MEMBERSHIP_ASPECT_NAME,
-        new EnvelopedAspect().setValue(new com.linkedin.entity.Aspect(membership.data())));
+        aspectName, new EnvelopedAspect().setValue(new com.linkedin.entity.Aspect(aspect.data())));
     response.setAspects(aspectMap);
-    batchResponse.put(userUrn, response);
+    batchResponse.put(actorUrn, response);
     return batchResponse;
   }
 

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

@@ -83,7 +83,11 @@ public void setupTest() throws Exception {
     when(_entityClient.batchGetV2(
             eq(CORP_USER_ENTITY_NAME),
             eq(Collections.singleton(authorizedUserUrn)),
-            eq(Collections.singleton(ROLE_MEMBERSHIP_ASPECT_NAME)),
+            eq(
+                ImmutableSet.of(
+                    ROLE_MEMBERSHIP_ASPECT_NAME,
+                    GROUP_MEMBERSHIP_ASPECT_NAME,
+                    NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)),
             any()))
         .thenReturn(authorizedEntityResponseMap);
 
@@ -94,7 +98,11 @@ public void setupTest() throws Exception {
     when(_entityClient.batchGetV2(
             eq(CORP_USER_ENTITY_NAME),
             eq(Collections.singleton(unauthorizedUserUrn)),
-            eq(Collections.singleton(ROLE_MEMBERSHIP_ASPECT_NAME)),
+            eq(
+                ImmutableSet.of(
+                    ROLE_MEMBERSHIP_ASPECT_NAME,
+                    GROUP_MEMBERSHIP_ASPECT_NAME,
+                    NATIVE_GROUP_MEMBERSHIP_ASPECT_NAME)),
             any()))
         .thenReturn(unauthorizedEntityResponseMap);
 

smoke-test/tests/privileges/test_privileges.py

@@ -450,4 +450,64 @@ def test_privilege_to_create_and_manage_policies():
 
 
     # Ensure that user can't create a policy after privilege is removed by admin
-    _ensure_cant_perform_action(user_session, create_policy,"createPolicy")
+    _ensure_cant_perform_action(user_session, create_policy,"createPolicy")
+
+
+@pytest.mark.dependency(depends=["test_healthchecks"])
+def test_privilege_from_group_role_can_create_and_manage_secret():
+
+    (admin_user, admin_pass) = get_admin_credentials()
+    admin_session = login_as(admin_user, admin_pass)
+    user_session = login_as("user", "user")
+    secret_urn = "urn:li:dataHubSecret:TestSecretName"
+
+    # Verify new user can't create secrets
+    create_secret = {
+        "query": """mutation createSecret($input: CreateSecretInput!) {\n
+            createSecret(input: $input)\n}""",
+        "variables": {
+            "input":{
+                "name":"TestSecretName",
+                "value":"Test Secret Value",
+                "description":"Test Secret Description"
+            }
+        },
+    }
+    _ensure_cant_perform_action(user_session, create_secret,"createSecret")
+
+    # Create group and grant it the admin role.
+    group_urn = create_group(admin_session, "Test Group")
+
+    # Assign admin role to group
+    assign_role(admin_session,"urn:li:dataHubRole:Admin", [group_urn])
+
+    # Assign user to group
+    assign_user_to_group(admin_session, group_urn, ["urn:li:corpuser:user"])
+
+    # Verify new user with admin group can create and manage secrets
+    # Create a secret
+    _ensure_can_create_secret(user_session, create_secret, secret_urn)
+
+    # Remove a secret
+    remove_secret = {
+        "query": """mutation deleteSecret($urn: String!) {\n
+            deleteSecret(urn: $urn)\n}""",
+        "variables": {
+            "urn": secret_urn
+        },
+    }
+
+    remove_secret_response = user_session.post(f"{get_frontend_url()}/api/v2/graphql", json=remove_secret)
+    remove_secret_response.raise_for_status()
+    secret_data = remove_secret_response.json()
+
+    assert secret_data
+    assert secret_data["data"]
+    assert secret_data["data"]["deleteSecret"]
+    assert secret_data["data"]["deleteSecret"] == secret_urn
+
+    # Delete group which removes the user's admin capabilities
+    remove_group(admin_session, group_urn)
+
+    # Ensure user can't create secret after policy is removed
+    _ensure_cant_perform_action(user_session, create_secret,"createSecret")