Skip to content

Commit

Permalink
feat(auth): Add roles to policy engine validation logic (#9178)
Browse files Browse the repository at this point in the history
  • Loading branch information
@pedro93
pedro93 authored Nov 8, 2023
1 parent f87983d commit f38c808
Show file tree
Hide file tree
Showing 7 changed files with 162 additions and 48 deletions.
1 change: 1 addition & 0 deletions metadata-auth/auth-api/src/main/java/com/datahub/authorization/AuthorizedActors.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ public class AuthorizedActors {
String privilege;
List<Urn> users;
List<Urn> groups;
List<Urn> roles;
boolean allUsers;
boolean allGroups;
}
5 changes: 5 additions & 0 deletions metadata-service/auth-impl/src/main/java/com/datahub/authorization/AuthorizerChain.java
View file
Original file line number Diff line number Diff line change
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
8 changes: 5 additions & 3 deletions metadata-service/auth-impl/src/main/java/com/datahub/authorization/DataHubAuthorizer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,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;

Expand All @@ -153,16 +154,17 @@ public AuthorizedActors authorizedActors(
// Step 3: For each matching policy, add actors that are authorized.
authorizedUsers.addAll(matchingActors.getUsers());
authorizedGroups.addAll(matchingActors.getGroups());
if (matchingActors.allUsers()) {
authorizedRoles.addAll(matchingActors.getRoles());
if (matchingActors.getAllUsers()) {
allUsers = true;
}
if (matchingActors.allGroups()) {
if (matchingActors.getAllGroups()) {
allGroups = true;
}
}

// 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);
}

/**
Expand Down
43 changes: 15 additions & 28 deletions metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyEngine.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,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.*;
Expand Down Expand Up @@ -75,6 +78,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)) {
Expand All @@ -96,6 +100,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()) {
Expand All @@ -104,7 +111,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(
Expand Down Expand Up @@ -438,34 +445,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) {
Expand Down
97 changes: 85 additions & 12 deletions ...data-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java
View file
Original file line number Diff line number Diff line change
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
Loading

0 comments on commit f38c808

Please sign in to comment.