Skip to content

Commit

Permalink
Add tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@pedro93
pedro93 committed Nov 8, 2023
1 parent dd159c3 commit 6651afe
Show file tree
Hide file tree
Showing 3 changed files with 133 additions and 12 deletions.
5 changes: 5 additions & 0 deletions metadata-service/auth-impl/src/main/java/com/datahub/authorization/AuthorizerChain.java
View file
Expand Up @@ -126,11 +126,16 @@ private AuthorizedActors mergeAuthorizedActors(@Nullable AuthorizedActors origin
mergedGroups = new ArrayList<>(groups);
}

Set<Urn> roles = new HashSet<>(original.getRoles());
roles.addAll(other.getRoles());
List<Urn> mergedRoles = new ArrayList<>(roles);

return AuthorizedActors.builder()
.allUsers(original.isAllUsers() || other.isAllUsers())
.allGroups(original.isAllGroups() || other.isAllGroups())
.users(mergedUsers)
.groups(mergedGroups)
.roles(mergedRoles)
.build();
}

Expand Down
97 changes: 85 additions & 12 deletions ...data-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java
View file
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {
Expand All @@ -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;
Expand Down Expand Up @@ -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<Urn> inputUrns = args.getArgument(1);
Expand All @@ -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();
}
Expand Down Expand Up @@ -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),
""
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -342,13 +393,6 @@ public void testAuthorizationOnDomainWithoutPrivilegeIsDenied() {
}

private DataHubPolicyInfo createDataHubPolicyInfo(boolean active, List<String> 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<Urn> users = ImmutableList.of(Urn.createFromString("urn:li:corpuser:user1"), Urn.createFromString("urn:li:corpuser:user2"));
List<Urn> groups = ImmutableList.of(Urn.createFromString("urn:li:corpGroup:group1"), Urn.createFromString("urn:li:corpGroup:group2"));
Expand All @@ -359,6 +403,20 @@ private DataHubPolicyInfo createDataHubPolicyInfo(boolean active, List<String> 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<String> 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();
Expand Down Expand Up @@ -429,6 +487,21 @@ private Map<Urn, EntityResponse> createDomainPropertiesBatchResponse(@Nullable f
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();
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));
}
Expand Down
43 changes: 43 additions & 0 deletions metadata-service/auth-impl/src/test/java/com/datahub/authorization/PolicyEngineTest.java
View file
Expand Up @@ -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();
Expand Down

0 comments on commit 6651afe

Please sign in to comment.