diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/AuthorizerChain.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/AuthorizerChain.java index f8eca541e1efb..7e7a1de176f06 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/AuthorizerChain.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/AuthorizerChain.java @@ -126,11 +126,16 @@ private AuthorizedActors mergeAuthorizedActors(@Nullable AuthorizedActors origin mergedGroups = new ArrayList<>(groups); } + Set roles = new HashSet<>(original.getRoles()); + roles.addAll(other.getRoles()); + List mergedRoles = new ArrayList<>(roles); + return AuthorizedActors.builder() .allUsers(original.isAllUsers() || other.isAllUsers()) .allGroups(original.isAllGroups() || other.isAllGroups()) .users(mergedUsers) .groups(mergedGroups) + .roles(mergedRoles) .build(); } diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java index babb1c5d00ee8..b0b206001209c 100644 --- a/metadata-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java @@ -21,6 +21,7 @@ import com.linkedin.entity.EnvelopedAspect; import com.linkedin.entity.EnvelopedAspectMap; import com.linkedin.entity.client.EntityClient; +import com.linkedin.identity.RoleMembership; import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.ScrollResult; import com.linkedin.metadata.search.SearchEntity; @@ -55,6 +56,7 @@ import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; +import static org.testng.Assert.assertFalse; public class DataHubAuthorizerTest { @@ -63,6 +65,7 @@ public class DataHubAuthorizerTest { private static final Urn PARENT_DOMAIN_URN = UrnUtils.getUrn("urn:li:domain:parent"); private static final Urn CHILD_DOMAIN_URN = UrnUtils.getUrn("urn:li:domain:child"); + private static final Urn USER_WITH_ADMIN_ROLE = UrnUtils.getUrn("urn:li:corpuser:user-with-admin"); private EntityClient _entityClient; private DataHubAuthorizer _dataHubAuthorizer; @@ -92,40 +95,56 @@ public void setupTest() throws Exception { final EnvelopedAspectMap childDomainPolicyAspectMap = new EnvelopedAspectMap(); childDomainPolicyAspectMap.put(DATAHUB_POLICY_INFO_ASPECT_NAME, new EnvelopedAspect().setValue(new Aspect(childDomainPolicy.data()))); + final Urn adminPolicyUrn = Urn.createFromString("urn:li:dataHubPolicy:4"); + final DataHubActorFilter actorFilter = new DataHubActorFilter(); + actorFilter.setRoles(new UrnArray(ImmutableList.of(Urn.createFromString("urn:li:dataHubRole:Admin")))); + final DataHubPolicyInfo adminPolicy = createDataHubPolicyInfoFor(true, ImmutableList.of("EDIT_USER_PROFILE"), null, actorFilter); + final EnvelopedAspectMap adminPolicyAspectMap = new EnvelopedAspectMap(); + adminPolicyAspectMap.put(DATAHUB_POLICY_INFO_ASPECT_NAME, new EnvelopedAspect().setValue(new Aspect(adminPolicy.data()))); + final ScrollResult policySearchResult1 = new ScrollResult() .setScrollId("1") - .setNumEntities(4) + .setNumEntities(5) .setEntities( new SearchEntityArray( ImmutableList.of(new SearchEntity().setEntity(activePolicyUrn)))); final ScrollResult policySearchResult2 = new ScrollResult() .setScrollId("2") - .setNumEntities(4) + .setNumEntities(5) .setEntities( new SearchEntityArray( ImmutableList.of(new SearchEntity().setEntity(inactivePolicyUrn)))); final ScrollResult policySearchResult3 = new ScrollResult() .setScrollId("3") - .setNumEntities(4) + .setNumEntities(5) .setEntities( new SearchEntityArray( ImmutableList.of(new SearchEntity().setEntity(parentDomainPolicyUrn)))); final ScrollResult policySearchResult4 = new ScrollResult() - .setNumEntities(4) + .setScrollId("4") + .setNumEntities(5) .setEntities( new SearchEntityArray( ImmutableList.of( new SearchEntity().setEntity(childDomainPolicyUrn)))); + final ScrollResult policySearchResult5 = new ScrollResult() + .setNumEntities(5) + .setEntities( + new SearchEntityArray( + ImmutableList.of( + new SearchEntity().setEntity(adminPolicyUrn)))); + when(_entityClient.scrollAcrossEntities(eq(List.of("dataHubPolicy")), eq(""), isNull(), any(), isNull(), anyInt(), eq(new SearchFlags().setFulltext(true).setSkipAggregates(true).setSkipHighlighting(true).setSkipCache(true)), any())) .thenReturn(policySearchResult1) .thenReturn(policySearchResult2) .thenReturn(policySearchResult3) - .thenReturn(policySearchResult4); + .thenReturn(policySearchResult4) + .thenReturn(policySearchResult5); when(_entityClient.batchGetV2(eq(POLICY_ENTITY_NAME), any(), eq(null), any())).thenAnswer(args -> { Set inputUrns = args.getArgument(1); @@ -140,6 +159,8 @@ public void setupTest() throws Exception { return Map.of(parentDomainPolicyUrn, new EntityResponse().setUrn(parentDomainPolicyUrn).setAspects(parentDomainPolicyAspectMap)); case "urn:li:dataHubPolicy:3": return Map.of(childDomainPolicyUrn, new EntityResponse().setUrn(childDomainPolicyUrn).setAspects(childDomainPolicyAspectMap)); + case "urn:li:dataHubPolicy:4": + return Map.of(adminPolicyUrn, new EntityResponse().setUrn(adminPolicyUrn).setAspects(adminPolicyAspectMap)); default: throw new IllegalStateException(); } @@ -167,6 +188,10 @@ public void setupTest() throws Exception { when(_entityClient.batchGetV2(any(), eq(Collections.singleton(PARENT_DOMAIN_URN)), eq(Collections.singleton(DOMAIN_PROPERTIES_ASPECT_NAME)), any())) .thenReturn(createDomainPropertiesBatchResponse(null)); + // Mocks to reach role membership for a user urn + when(_entityClient.batchGetV2(any(), eq(Collections.singleton(USER_WITH_ADMIN_ROLE)), eq(Collections.singleton(ROLE_MEMBERSHIP_ASPECT_NAME)), any()) + ).thenReturn(createUserRoleMembershipBatchResponse(USER_WITH_ADMIN_ROLE, UrnUtils.getUrn("urn:li:dataHubRole:Admin"))); + final Authentication systemAuthentication = new Authentication( new Actor(ActorType.USER, DATAHUB_SYSTEM_CLIENT_ID), "" @@ -302,6 +327,32 @@ public void testAuthorizedActorsActivePolicy() throws Exception { )); } + @Test + public void testAuthorizedRoleActivePolicy() throws Exception { + final AuthorizedActors actors = + _dataHubAuthorizer.authorizedActors("EDIT_USER_PROFILE", // Should be inside the active policy. + Optional.of(new EntitySpec("dataset", "urn:li:dataset:1"))); + + assertFalse(actors.isAllUsers()); + assertFalse(actors.isAllGroups()); + assertEquals(new HashSet<>(actors.getUsers()), ImmutableSet.of()); + assertEquals(new HashSet<>(actors.getGroups()), ImmutableSet.of()); + assertEquals(new HashSet<>(actors.getRoles()), ImmutableSet.of(UrnUtils.getUrn("urn:li:dataHubRole:Admin"))); + } + + @Test + public void testAuthorizationBasedOnRoleIsAllowed() { + EntitySpec resourceSpec = new EntitySpec("dataset", "urn:li:dataset:test"); + + AuthorizationRequest request = new AuthorizationRequest( + USER_WITH_ADMIN_ROLE.toString(), + "EDIT_USER_PROFILE", + Optional.of(resourceSpec) + ); + + assertEquals(_dataHubAuthorizer.authorize(request).getType(), AuthorizationResult.Type.ALLOW); + } + @Test public void testAuthorizationOnDomainWithPrivilegeIsAllowed() { EntitySpec resourceSpec = new EntitySpec("dataset", "urn:li:dataset:test"); @@ -342,13 +393,6 @@ public void testAuthorizationOnDomainWithoutPrivilegeIsDenied() { } private DataHubPolicyInfo createDataHubPolicyInfo(boolean active, List privileges, @Nullable final Urn domain) throws Exception { - final DataHubPolicyInfo dataHubPolicyInfo = new DataHubPolicyInfo(); - dataHubPolicyInfo.setType(METADATA_POLICY_TYPE); - dataHubPolicyInfo.setState(active ? ACTIVE_POLICY_STATE : INACTIVE_POLICY_STATE); - dataHubPolicyInfo.setPrivileges(new StringArray(privileges)); - dataHubPolicyInfo.setDisplayName("My Test Display"); - dataHubPolicyInfo.setDescription("My test display!"); - dataHubPolicyInfo.setEditable(true); List users = ImmutableList.of(Urn.createFromString("urn:li:corpuser:user1"), Urn.createFromString("urn:li:corpuser:user2")); List groups = ImmutableList.of(Urn.createFromString("urn:li:corpGroup:group1"), Urn.createFromString("urn:li:corpGroup:group2")); @@ -359,6 +403,20 @@ private DataHubPolicyInfo createDataHubPolicyInfo(boolean active, List p actorFilter.setAllGroups(true); actorFilter.setUsers(new UrnArray(users)); actorFilter.setGroups(new UrnArray(groups)); + + return createDataHubPolicyInfoFor(active, privileges, domain, actorFilter); + } + + private DataHubPolicyInfo createDataHubPolicyInfoFor(boolean active, List privileges, + @Nullable final Urn domain, DataHubActorFilter actorFilter) throws Exception { + final DataHubPolicyInfo dataHubPolicyInfo = new DataHubPolicyInfo(); + dataHubPolicyInfo.setType(METADATA_POLICY_TYPE); + dataHubPolicyInfo.setState(active ? ACTIVE_POLICY_STATE : INACTIVE_POLICY_STATE); + dataHubPolicyInfo.setPrivileges(new StringArray(privileges)); + dataHubPolicyInfo.setDisplayName("My Test Display"); + dataHubPolicyInfo.setDescription("My test display!"); + dataHubPolicyInfo.setEditable(true); + dataHubPolicyInfo.setActors(actorFilter); final DataHubResourceFilter resourceFilter = new DataHubResourceFilter(); @@ -429,6 +487,21 @@ private Map createDomainPropertiesBatchResponse(@Nullable f return batchResponse; } + private Map createUserRoleMembershipBatchResponse(final Urn userUrn, @Nullable final Urn roleUrn) { + final Map batchResponse = new HashMap<>(); + final EntityResponse response = new EntityResponse(); + EnvelopedAspectMap aspectMap = new EnvelopedAspectMap(); + final RoleMembership membership = new RoleMembership(); + if (roleUrn != null) { + membership.setRoles(new UrnArray(roleUrn)); + } + aspectMap.put(ROLE_MEMBERSHIP_ASPECT_NAME, new EnvelopedAspect() + .setValue(new com.linkedin.entity.Aspect(membership.data()))); + response.setAspects(aspectMap); + batchResponse.put(userUrn, response); + return batchResponse; + } + private AuthorizerContext createAuthorizerContext(final Authentication systemAuthentication, final EntityClient entityClient) { return new AuthorizerContext(Collections.emptyMap(), new DefaultEntitySpecResolver(systemAuthentication, entityClient)); } 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 index 8077e5768b4fb..2790c16ba75e6 100644 --- 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 @@ -1113,11 +1113,54 @@ public void testGetMatchingActorsNoResourceMatch() throws Exception { assertFalse(actors.getAllGroups()); assertEquals(actors.getUsers(), Collections.emptyList()); assertEquals(actors.getGroups(), Collections.emptyList()); + //assertEquals(actors.getRoles(), Collections.emptyList()); // Verify no network calls verify(_entityClient, times(0)).batchGetV2(any(), any(), any(), any()); } + @Test + public void testGetMatchingActorsByRoleResourceMatch() throws Exception { + final DataHubPolicyInfo dataHubPolicyInfo = new DataHubPolicyInfo(); + dataHubPolicyInfo.setType(METADATA_POLICY_TYPE); + dataHubPolicyInfo.setState(ACTIVE_POLICY_STATE); + dataHubPolicyInfo.setPrivileges(new StringArray("EDIT_ENTITY_TAGS")); + dataHubPolicyInfo.setDisplayName("My Test Display"); + dataHubPolicyInfo.setDescription("My test display!"); + dataHubPolicyInfo.setEditable(true); + + final DataHubActorFilter actorFilter = new DataHubActorFilter(); + actorFilter.setResourceOwners(true); + actorFilter.setAllUsers(false); + actorFilter.setAllGroups(false); + actorFilter.setRoles(new UrnArray(ImmutableList.of(Urn.createFromString("urn:li:dataHubRole:Editor")))); + dataHubPolicyInfo.setActors(actorFilter); + + final DataHubResourceFilter resourceFilter = new DataHubResourceFilter(); + resourceFilter.setAllResources(false); + resourceFilter.setType("dataset"); + StringArray resourceUrns = new StringArray(); + resourceUrns.add(RESOURCE_URN); + resourceFilter.setResources(resourceUrns); + dataHubPolicyInfo.setResources(resourceFilter); + + ResolvedEntitySpec resourceSpec = buildEntityResolvers("dataset", RESOURCE_URN, ImmutableSet.of(), + Collections.emptySet(), Collections.emptySet()); + + PolicyEngine.PolicyActors actors = _policyEngine.getMatchingActors(dataHubPolicyInfo, Optional.of(resourceSpec)); + + assertFalse(actors.getAllUsers()); + assertFalse(actors.getAllGroups()); + + assertEquals(actors.getUsers(), ImmutableList.of()); + assertEquals(actors.getGroups(), ImmutableList.of()); + assertEquals(actors.getRoles(), ImmutableList.of(Urn.createFromString("urn:li:dataHubRole:Editor"))); + + // 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()); + } + private Ownership createOwnershipAspect(final Boolean addUserOwner, final Boolean addGroupOwner) throws Exception { final Ownership ownershipAspect = new Ownership(); final OwnerArray owners = new OwnerArray();