From 7604825a9e361d8600070488ac337d53ff46ba07 Mon Sep 17 00:00:00 2001 From: Pedro Silva Date: Wed, 8 Nov 2023 15:40:09 +0000 Subject: [PATCH] Add tests --- .../authorization/AuthorizerChain.java | 5 ++ .../authorization/DataHubAuthorizerTest.java | 90 ++++++++++++++++--- .../authorization/PolicyEngineTest.java | 43 +++++++++ 3 files changed, 127 insertions(+), 11 deletions(-) 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 f8eca541e1efb4..7e7a1de176f06d 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 24ecfa6fefc856..2e4152cca77648 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.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; @@ -52,6 +53,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 { @@ -61,6 +63,8 @@ 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; @@ -89,6 +93,14 @@ 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 SearchResult policySearchResult = new SearchResult(); policySearchResult.setNumEntities(3); policySearchResult.setEntities( @@ -97,7 +109,8 @@ public void setupTest() throws Exception { new SearchEntity().setEntity(activePolicyUrn), new SearchEntity().setEntity(inactivePolicyUrn), new SearchEntity().setEntity(parentDomainPolicyUrn), - new SearchEntity().setEntity(childDomainPolicyUrn) + new SearchEntity().setEntity(childDomainPolicyUrn), + new SearchEntity().setEntity(adminPolicyUrn) ) ) ); @@ -105,12 +118,13 @@ public void setupTest() throws Exception { when(_entityClient.search(eq("dataHubPolicy"), eq(""), isNull(), any(), anyInt(), anyInt(), any(), eq(new SearchFlags().setFulltext(true)))).thenReturn(policySearchResult); when(_entityClient.batchGetV2(eq(POLICY_ENTITY_NAME), - eq(ImmutableSet.of(activePolicyUrn, inactivePolicyUrn, parentDomainPolicyUrn, childDomainPolicyUrn)), eq(null), any())).thenReturn( + eq(ImmutableSet.of(activePolicyUrn, inactivePolicyUrn, parentDomainPolicyUrn, childDomainPolicyUrn, adminPolicyUrn)), eq(null), any())).thenReturn( ImmutableMap.of( activePolicyUrn, new EntityResponse().setUrn(activePolicyUrn).setAspects(activeAspectMap), inactivePolicyUrn, new EntityResponse().setUrn(inactivePolicyUrn).setAspects(inactiveAspectMap), parentDomainPolicyUrn, new EntityResponse().setUrn(parentDomainPolicyUrn).setAspects(parentDomainPolicyAspectMap), - childDomainPolicyUrn, new EntityResponse().setUrn(childDomainPolicyUrn).setAspects(childDomainPolicyAspectMap) + childDomainPolicyUrn, new EntityResponse().setUrn(childDomainPolicyUrn).setAspects(childDomainPolicyAspectMap), + adminPolicyUrn, new EntityResponse().setUrn(adminPolicyUrn).setAspects(adminPolicyAspectMap) ) ); @@ -136,6 +150,10 @@ childDomainPolicyUrn, new EntityResponse().setUrn(childDomainPolicyUrn).setAspec 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), "" @@ -270,6 +288,35 @@ 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"); @@ -310,14 +357,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")); @@ -327,6 +366,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(); @@ -397,6 +450,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 8077e5768b4fb4..2790c16ba75e62 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();