Skip to content

Commit

Permalink
feat(search): search access controls unit tests (#9925)
Browse files Browse the repository at this point in the history
  • Loading branch information
david-leifker authored Feb 27, 2024
1 parent 84d053a commit a8e3c6d
Show file tree
Hide file tree
Showing 19 changed files with 1,204 additions and 36 deletions.
7 changes: 1 addition & 6 deletions datahub-frontend/app/auth/AuthModule.java
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
import com.linkedin.entity.client.SystemEntityClient;
import com.linkedin.entity.client.SystemRestliEntityClient;
import com.linkedin.metadata.restli.DefaultRestliClientFactory;
import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl;
import com.linkedin.parseq.retry.backoff.ExponentialBackoff;
import com.linkedin.util.Configuration;
import config.ConfigurationProvider;
Expand All @@ -42,15 +41,11 @@
import org.pac4j.play.store.PlayCookieSessionStore;
import org.pac4j.play.store.PlaySessionStore;
import org.pac4j.play.store.ShiroAesDataEncrypter;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.AnnotationConfigApplicationContext;
import org.springframework.context.annotation.Bean;
import play.Environment;
import play.cache.SyncCacheApi;
import utils.ConfigUtil;

import javax.annotation.Nonnull;

/** Responsible for configuring, validating, and providing authentication related components. */
@Slf4j
public class AuthModule extends AbstractModule {
Expand Down Expand Up @@ -173,7 +168,7 @@ protected OperationContext provideOperationContext(final Authentication systemAu
final ConfigurationProvider configurationProvider) {
ActorContext systemActorContext =
ActorContext.builder()
.systemAuthentication(true)
.allowSystemAuth(true)
.authentication(systemAuthentication)
.build();
OperationContextConfig systemConfig = OperationContextConfig.builder()
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,220 @@
package com.linkedin.metadata.aspect.hooks;

import static com.linkedin.metadata.Constants.DEFAULT_OWNERSHIP_TYPE_URN;
import static org.mockito.Mockito.mock;
import static org.testng.Assert.assertEquals;

import com.linkedin.common.Owner;
import com.linkedin.common.OwnerArray;
import com.linkedin.common.Ownership;
import com.linkedin.common.UrnArray;
import com.linkedin.common.UrnArrayMap;
import com.linkedin.common.urn.Urn;
import com.linkedin.common.urn.UrnUtils;
import com.linkedin.metadata.aspect.AspectRetriever;
import com.linkedin.metadata.aspect.batch.ChangeMCP;
import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig;
import com.linkedin.metadata.models.registry.EntityRegistry;
import com.linkedin.test.metadata.aspect.TestEntityRegistry;
import com.linkedin.test.metadata.aspect.batch.TestMCP;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import javax.annotation.Nullable;
import org.testng.annotations.Test;

public class OwnerTypeMapTest {
private static final AspectRetriever ASPECT_RETRIEVER = mock(AspectRetriever.class);
private static final EntityRegistry ENTITY_REGISTRY = new TestEntityRegistry();
private static final AspectPluginConfig ASPECT_PLUGIN_CONFIG =
AspectPluginConfig.builder()
.className("some class")
.enabled(true)
.supportedEntityAspectNames(
List.of(
AspectPluginConfig.EntityAspectName.builder()
.entityName("*")
.aspectName("ownership")
.build()))
.build();
private static final Urn TEST_ENTITY_URN =
UrnUtils.getUrn(
"urn:li:dataset:(urn:li:dataPlatform:bigquery,calm-pagoda-323403.jaffle_shop.orders,PROD)");
private static final Urn TEST_USER_A = UrnUtils.getUrn("urn:li:corpUser:a");
private static final Urn TEST_USER_B = UrnUtils.getUrn("urn:li:corpUser:b");
private static final Urn TEST_GROUP_A = UrnUtils.getUrn("urn:li:corpGroup:a");
private static final Urn TEST_GROUP_B = UrnUtils.getUrn("urn:li:corpGroup:b");
private static final Urn TECH_OWNER =
UrnUtils.getUrn("urn:li:ownershipType:__system__technical_owner");
private static final Urn BUS_OWNER =
UrnUtils.getUrn("urn:li:ownershipType:__system__business_owner");

@Test
public void ownershipTypeMutationNoneType() {
OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG);
Ownership ownership = buildOwnership(Map.of(TEST_USER_A, List.of(), TEST_GROUP_A, List.of()));
testHook.writeMutation(buildMCP(null, ownership), ASPECT_RETRIEVER);

assertEquals(
ownership.getOwnerTypes(),
new UrnArrayMap(
Map.of(
DEFAULT_OWNERSHIP_TYPE_URN.toString(),
new UrnArray(List.of(TEST_USER_A, TEST_GROUP_A)))),
"Expected generic owners to be grouped by `none` ownership type.");
}

@Test
public void ownershipTypeMutationNoneTypeAdd() {
OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG);
Ownership oldOwnership = buildOwnership(Map.of(TEST_USER_A, List.of()));
Ownership newOwnership =
buildOwnership(Map.of(TEST_USER_A, List.of(), TEST_GROUP_A, List.of()));
testHook.writeMutation(buildMCP(oldOwnership, newOwnership), ASPECT_RETRIEVER);

assertEquals(
newOwnership.getOwnerTypes(),
new UrnArrayMap(
Map.of(
DEFAULT_OWNERSHIP_TYPE_URN.toString(),
new UrnArray(List.of(TEST_USER_A, TEST_GROUP_A)))),
"Expected generic owners to be grouped by `none` ownership type.");
}

@Test
public void ownershipTypeMutationNoneTypeRemove() {
OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG);
Ownership oldOwnership =
buildOwnership(Map.of(TEST_USER_A, List.of(), TEST_GROUP_A, List.of()));
Ownership newOwnership = buildOwnership(Map.of(TEST_USER_A, List.of()));
testHook.writeMutation(buildMCP(oldOwnership, newOwnership), ASPECT_RETRIEVER);

assertEquals(
newOwnership.getOwnerTypes(),
new UrnArrayMap(
Map.of(DEFAULT_OWNERSHIP_TYPE_URN.toString(), new UrnArray(List.of(TEST_USER_A)))),
"Expected generic owners to be grouped by `none` ownership type.");
}

@Test
public void ownershipTypeMutationMixedType() {
OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG);
Ownership ownership =
buildOwnership(
Map.of(
TEST_USER_A,
List.of(),
TEST_GROUP_A,
List.of(),
TEST_USER_B,
List.of(BUS_OWNER),
TEST_GROUP_B,
List.of(TECH_OWNER)));
testHook.writeMutation(buildMCP(null, ownership), ASPECT_RETRIEVER);

assertEquals(
ownership.getOwnerTypes(),
new UrnArrayMap(
Map.of(
DEFAULT_OWNERSHIP_TYPE_URN.toString(),
new UrnArray(List.of(TEST_USER_A, TEST_GROUP_A)),
BUS_OWNER.toString(),
new UrnArray(List.of(TEST_USER_B)),
TECH_OWNER.toString(),
new UrnArray(List.of(TEST_GROUP_B)))),
"Expected generic owners to be grouped by `none` ownership type as well as specified types.");
}

@Test
public void ownershipTypeMutationMixedTypeAdd() {
OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG);
Ownership oldOwnership =
buildOwnership(Map.of(TEST_USER_A, List.of(), TEST_USER_B, List.of(BUS_OWNER)));
Ownership newOwnership =
buildOwnership(
Map.of(
TEST_USER_A,
List.of(),
TEST_GROUP_A,
List.of(),
TEST_USER_B,
List.of(BUS_OWNER),
TEST_GROUP_B,
List.of(TECH_OWNER)));
testHook.writeMutation(buildMCP(oldOwnership, newOwnership), ASPECT_RETRIEVER);

assertEquals(
newOwnership.getOwnerTypes(),
new UrnArrayMap(
Map.of(
DEFAULT_OWNERSHIP_TYPE_URN.toString(),
new UrnArray(List.of(TEST_USER_A, TEST_GROUP_A)),
BUS_OWNER.toString(),
new UrnArray(List.of(TEST_USER_B)),
TECH_OWNER.toString(),
new UrnArray(List.of(TEST_GROUP_B)))),
"Expected generic owners to be grouped by `none` ownership type as well as specified types.");
}

@Test
public void ownershipTypeMutationMixedTypeRemove() {
OwnerTypeMap testHook = new OwnerTypeMap(ASPECT_PLUGIN_CONFIG);
Ownership oldOwnership =
buildOwnership(
Map.of(
TEST_USER_A,
List.of(),
TEST_GROUP_A,
List.of(),
TEST_USER_B,
List.of(BUS_OWNER),
TEST_GROUP_B,
List.of(TECH_OWNER)));
Ownership newOwnership =
buildOwnership(Map.of(TEST_GROUP_A, List.of(), TEST_GROUP_B, List.of(TECH_OWNER)));
testHook.writeMutation(buildMCP(oldOwnership, newOwnership), ASPECT_RETRIEVER);

assertEquals(
newOwnership.getOwnerTypes(),
new UrnArrayMap(
Map.of(
DEFAULT_OWNERSHIP_TYPE_URN.toString(),
new UrnArray(List.of(TEST_GROUP_A)),
BUS_OWNER.toString(),
new UrnArray(),
TECH_OWNER.toString(),
new UrnArray(List.of(TEST_GROUP_B)))),
"Expected generic owners to be grouped by `none` ownership type as well as specified types.");
}

private static Ownership buildOwnership(Map<Urn, List<Urn>> ownershipTypes) {
Ownership ownership = new Ownership();
ownership.setOwners(
ownershipTypes.entrySet().stream()
.flatMap(
entry -> {
if (entry.getValue().isEmpty()) {
Owner owner = new Owner();
owner.setOwner(entry.getKey());
return Stream.of(owner);
} else {
return entry.getValue().stream()
.map(
typeUrn -> {
Owner owner = new Owner();
owner.setOwner(entry.getKey());
owner.setTypeUrn(typeUrn);
return owner;
});
}
})
.collect(Collectors.toCollection(OwnerArray::new)));
return ownership;
}

private static Set<ChangeMCP> buildMCP(@Nullable Ownership oldOwnership, Ownership newOwnership) {
return TestMCP.ofOneMCP(TEST_ENTITY_URN, oldOwnership, newOwnership, ENTITY_REGISTRY);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import com.datahub.plugins.Plugin;
import com.linkedin.common.urn.Urn;
import com.linkedin.policy.DataHubPolicyInfo;
import java.util.Collection;
import java.util.Collections;
import java.util.Map;
import java.util.Optional;
Expand Down Expand Up @@ -60,7 +61,7 @@ default Set<DataHubPolicyInfo> getActorPolicies(@Nonnull Urn actorUrn) {
}

/** Given the actor's urn retrieve the actor's groups */
default Set<Urn> getActorGroups(@Nonnull Urn actorUrn) {
return Collections.emptySet();
default Collection<Urn> getActorGroups(@Nonnull Urn actorUrn) {
return Collections.emptyList();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -66,10 +66,10 @@ public static Optional<QueryBuilder> buildAccessControlFilters(
If search authorization is enabled AND we're also not the system performing the query
*/
if (opContext.getOperationContextConfig().getSearchAuthorizationConfiguration().isEnabled()
&& !opContext.getSessionActorContext().isSystemAuthentication()
&& !opContext.getSessionActorContext().isAllowSystemAuth()
&& !opContext.getSearchContext().isRestrictedSearch()) {

BoolQueryBuilder builder = QueryBuilders.boolQuery();
builder.minimumShouldMatch(1);

// Apply access policies
streamViewQueries(opContext).distinct().forEach(builder::should);
Expand All @@ -78,7 +78,8 @@ public static Optional<QueryBuilder> buildAccessControlFilters(
// default deny if no filters
return Optional.of(builder.mustNot(MATCH_ALL));
} else if (!builder.should().contains(MATCH_ALL)) {
// if MATCH_ALL is not present, apply filters
// if MATCH_ALL is not present, apply filters requiring at least 1
builder.minimumShouldMatch(1);
response = Optional.of(builder);
}
}
Expand Down Expand Up @@ -181,24 +182,22 @@ private static QueryBuilder buildActorQuery(
OperationContext opContext, DataHubPolicyInfo policy) {
DataHubActorFilter actorFilter = policy.getActors();

if (!policy.hasActors() || !policy.getActors().isResourceOwners()) {
if (!policy.hasActors()
|| !(actorFilter.isResourceOwners() || actorFilter.hasResourceOwnersTypes())) {
// no owner restriction
return MATCH_ALL;
}

ActorContext actorContext = opContext.getSessionActorContext();

// policy might apply to the actor via user or group
Set<String> actorAndGroupUrns =
List<String> actorAndGroupUrns =
Stream.concat(
Stream.of(actorContext.getAuthentication().getActor().toUrnStr()),
actorContext.getGroupMembership().stream()
.filter(
groupUrn ->
actorFilter.getGroups() != null
&& actorFilter.getGroups().contains(groupUrn))
.map(Urn::toString))
.collect(Collectors.toSet());
actorContext.getGroupMembership().stream().map(Urn::toString))
.map(String::toLowerCase)
.distinct()
.collect(Collectors.toList());

if (!actorFilter.hasResourceOwnersTypes()) {
// owners without owner type restrictions
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ public abstract class SearchGraphServiceTestBase extends GraphServiceTestBase {
@Nonnull
protected abstract ESIndexBuilder getIndexBuilder();

private final IndexConvention _indexConvention = new IndexConventionImpl(null);
private final IndexConvention _indexConvention = IndexConventionImpl.NO_PREFIX;
private final String _indexName = _indexConvention.getIndexName(INDEX_NAME);
private ElasticSearchGraphService _client;
private boolean _enableMultiPathSearch =
Expand Down
Loading

0 comments on commit a8e3c6d

Please sign in to comment.