Skip to content

Commit 9fceb38

Browse files
david-leifkerdavid-leifker
authored
fix(graphql): listUsers resolver incorrect access check (#15448)
1 parent ba7842d commit 9fceb38

File tree

2 files changed

+228
-48
lines changed

2 files changed

+228
-48
lines changed

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/user/ListUsersResolver.java‎

Lines changed: 40 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,7 @@
55

66
import com.linkedin.common.urn.Urn;
77
import com.linkedin.datahub.graphql.QueryContext;
8-
import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
98
import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
10-
import com.linkedin.datahub.graphql.exception.AuthorizationException;
119
import com.linkedin.datahub.graphql.generated.CorpUser;
1210
import com.linkedin.datahub.graphql.generated.ListUsersInput;
1311
import com.linkedin.datahub.graphql.generated.ListUsersResult;
@@ -45,55 +43,49 @@ public CompletableFuture<ListUsersResult> get(final DataFetchingEnvironment envi
4543

4644
final QueryContext context = environment.getContext();
4745

48-
if (AuthorizationUtils.canManageUsersAndGroups(context)) {
49-
final ListUsersInput input =
50-
bindArgument(environment.getArgument("input"), ListUsersInput.class);
51-
final Integer start = input.getStart() == null ? DEFAULT_START : input.getStart();
52-
final Integer count = input.getCount() == null ? DEFAULT_COUNT : input.getCount();
53-
final String query = input.getQuery() == null ? DEFAULT_QUERY : input.getQuery();
46+
final ListUsersInput input =
47+
bindArgument(environment.getArgument("input"), ListUsersInput.class);
48+
final Integer start = input.getStart() == null ? DEFAULT_START : input.getStart();
49+
final Integer count = input.getCount() == null ? DEFAULT_COUNT : input.getCount();
50+
final String query = input.getQuery() == null ? DEFAULT_QUERY : input.getQuery();
5451

55-
return GraphQLConcurrencyUtils.supplyAsync(
56-
() -> {
57-
try {
58-
// First, get all policy Urns.
59-
final SearchResult gmsResult =
60-
_entityClient.search(
61-
context
62-
.getOperationContext()
63-
.withSearchFlags(flags -> flags.setFulltext(true)),
64-
CORP_USER_ENTITY_NAME,
65-
query,
66-
Collections.emptyMap(),
67-
start,
68-
count);
52+
return GraphQLConcurrencyUtils.supplyAsync(
53+
() -> {
54+
try {
55+
// First, get all policy Urns.
56+
final SearchResult gmsResult =
57+
_entityClient.search(
58+
context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)),
59+
CORP_USER_ENTITY_NAME,
60+
query,
61+
Collections.emptyMap(),
62+
start,
63+
count);
6964

70-
// Then, get hydrate all users.
71-
final Map<Urn, EntityResponse> entities =
72-
_entityClient.batchGetV2(
73-
context.getOperationContext(),
74-
CORP_USER_ENTITY_NAME,
75-
new HashSet<>(
76-
gmsResult.getEntities().stream()
77-
.map(SearchEntity::getEntity)
78-
.collect(Collectors.toList())),
79-
null);
65+
// Then, get hydrate all users.
66+
final Map<Urn, EntityResponse> entities =
67+
_entityClient.batchGetV2(
68+
context.getOperationContext(),
69+
CORP_USER_ENTITY_NAME,
70+
new HashSet<>(
71+
gmsResult.getEntities().stream()
72+
.map(SearchEntity::getEntity)
73+
.collect(Collectors.toList())),
74+
null);
8075

81-
// Now that we have entities we can bind this to a result.
82-
final ListUsersResult result = new ListUsersResult();
83-
result.setStart(gmsResult.getFrom());
84-
result.setCount(gmsResult.getPageSize());
85-
result.setTotal(gmsResult.getNumEntities());
86-
result.setUsers(mapEntities(context, entities.values()));
87-
return result;
88-
} catch (Exception e) {
89-
throw new RuntimeException("Failed to list users", e);
90-
}
91-
},
92-
this.getClass().getSimpleName(),
93-
"get");
94-
}
95-
throw new AuthorizationException(
96-
"Unauthorized to perform this action. Please contact your DataHub administrator.");
76+
// Now that we have entities we can bind this to a result.
77+
final ListUsersResult result = new ListUsersResult();
78+
result.setStart(gmsResult.getFrom());
79+
result.setCount(gmsResult.getPageSize());
80+
result.setTotal(gmsResult.getNumEntities());
81+
result.setUsers(mapEntities(context, entities.values()));
82+
return result;
83+
} catch (Exception e) {
84+
throw new RuntimeException("Failed to list users", e);
85+
}
86+
},
87+
this.getClass().getSimpleName(),
88+
"get");
9789
}
9890

9991
private static List<CorpUser> mapEntities(

‎datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/user/ListUsersResolverTest.java‎

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
package com.linkedin.datahub.graphql.resolvers.user;
2+
3+
import static com.linkedin.datahub.graphql.TestUtils.*;
4+
import static com.linkedin.metadata.Constants.CORP_USER_ENTITY_NAME;
5+
import static org.mockito.ArgumentMatchers.any;
6+
import static org.testng.Assert.*;
7+
8+
import com.google.common.collect.ImmutableSet;
9+
import com.linkedin.common.urn.Urn;
10+
import com.linkedin.datahub.graphql.QueryContext;
11+
import com.linkedin.datahub.graphql.generated.ListUsersInput;
12+
import com.linkedin.entity.Aspect;
13+
import com.linkedin.entity.EntityResponse;
14+
import com.linkedin.entity.EnvelopedAspect;
15+
import com.linkedin.entity.EnvelopedAspectMap;
16+
import com.linkedin.entity.client.EntityClient;
17+
import com.linkedin.identity.CorpUserInfo;
18+
import com.linkedin.metadata.Constants;
19+
import com.linkedin.metadata.key.CorpUserKey;
20+
import com.linkedin.metadata.search.SearchEntity;
21+
import com.linkedin.metadata.search.SearchEntityArray;
22+
import com.linkedin.metadata.search.SearchResult;
23+
import com.linkedin.r2.RemoteInvocationException;
24+
import graphql.schema.DataFetchingEnvironment;
25+
import java.util.Collections;
26+
import java.util.HashMap;
27+
import java.util.Map;
28+
import java.util.concurrent.CompletionException;
29+
import org.mockito.Mockito;
30+
import org.testng.annotations.Test;
31+
32+
public class ListUsersResolverTest {
33+
34+
private static final Urn TEST_USER_URN = Urn.createFromTuple("corpuser", "test");
35+
36+
private static final ListUsersInput TEST_INPUT = new ListUsersInput();
37+
38+
@Test
39+
public void testGetSuccess() throws Exception {
40+
// Create resolver
41+
EntityClient mockClient = Mockito.mock(EntityClient.class);
42+
43+
// Mock search result
44+
SearchResult searchResult =
45+
new SearchResult()
46+
.setFrom(0)
47+
.setPageSize(1)
48+
.setNumEntities(1)
49+
.setEntities(
50+
new SearchEntityArray(
51+
ImmutableSet.of(new SearchEntity().setEntity(TEST_USER_URN))));
52+
53+
Mockito.when(
54+
mockClient.search(
55+
any(),
56+
Mockito.eq(CORP_USER_ENTITY_NAME),
57+
Mockito.eq(""),
58+
Mockito.eq(Collections.emptyMap()),
59+
Mockito.eq(0),
60+
Mockito.eq(20)))
61+
.thenReturn(searchResult);
62+
63+
// Mock batchGetV2 result
64+
EntityResponse entityResponse = createMockEntityResponse(TEST_USER_URN);
65+
Map<Urn, EntityResponse> entityMap = new HashMap<>();
66+
entityMap.put(TEST_USER_URN, entityResponse);
67+
68+
Mockito.when(
69+
mockClient.batchGetV2(
70+
any(), Mockito.eq(CORP_USER_ENTITY_NAME), any(), Mockito.isNull()))
71+
.thenReturn(entityMap);
72+
73+
ListUsersResolver resolver = new ListUsersResolver(mockClient);
74+
75+
// Execute resolver
76+
QueryContext mockContext = getMockAllowContext();
77+
DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class);
78+
Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT);
79+
Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
80+
81+
// Data Assertions
82+
var result = resolver.get(mockEnv).get();
83+
assertEquals((int) result.getStart(), 0);
84+
assertEquals((int) result.getCount(), 1);
85+
assertEquals((int) result.getTotal(), 1);
86+
assertNotNull(result.getUsers());
87+
assertEquals(result.getUsers().size(), 1);
88+
assertEquals(result.getUsers().get(0).getUrn(), TEST_USER_URN.toString());
89+
}
90+
91+
@Test
92+
public void testGetWithCustomInput() throws Exception {
93+
// Create resolver
94+
EntityClient mockClient = Mockito.mock(EntityClient.class);
95+
96+
// Create custom input
97+
ListUsersInput customInput = new ListUsersInput();
98+
customInput.setStart(10);
99+
customInput.setCount(5);
100+
customInput.setQuery("test query");
101+
102+
// Mock search result
103+
SearchResult searchResult =
104+
new SearchResult()
105+
.setFrom(10)
106+
.setPageSize(5)
107+
.setNumEntities(0)
108+
.setEntities(new SearchEntityArray(ImmutableSet.of()));
109+
110+
Mockito.when(
111+
mockClient.search(
112+
any(),
113+
Mockito.eq(CORP_USER_ENTITY_NAME),
114+
Mockito.eq("test query"),
115+
Mockito.eq(Collections.emptyMap()),
116+
Mockito.eq(10),
117+
Mockito.eq(5)))
118+
.thenReturn(searchResult);
119+
120+
// Mock batchGetV2 result (empty)
121+
Mockito.when(
122+
mockClient.batchGetV2(
123+
any(), Mockito.eq(CORP_USER_ENTITY_NAME), any(), Mockito.isNull()))
124+
.thenReturn(Collections.emptyMap());
125+
126+
ListUsersResolver resolver = new ListUsersResolver(mockClient);
127+
128+
// Execute resolver
129+
QueryContext mockContext = getMockAllowContext();
130+
DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class);
131+
Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(customInput);
132+
Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
133+
134+
// Data Assertions
135+
var result = resolver.get(mockEnv).get();
136+
assertEquals((int) result.getStart(), 10);
137+
assertEquals((int) result.getCount(), 5);
138+
assertEquals((int) result.getTotal(), 0);
139+
assertNotNull(result.getUsers());
140+
assertEquals(result.getUsers().size(), 0);
141+
}
142+
143+
@Test
144+
public void testGetEntityClientException() throws Exception {
145+
// Create resolver
146+
EntityClient mockClient = Mockito.mock(EntityClient.class);
147+
Mockito.doThrow(RemoteInvocationException.class)
148+
.when(mockClient)
149+
.search(any(), any(), Mockito.eq(""), Mockito.anyMap(), Mockito.anyInt(), Mockito.anyInt());
150+
ListUsersResolver resolver = new ListUsersResolver(mockClient);
151+
152+
// Execute resolver
153+
DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class);
154+
QueryContext mockContext = getMockAllowContext();
155+
Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT);
156+
Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
157+
158+
assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join());
159+
}
160+
161+
private EntityResponse createMockEntityResponse(Urn userUrn) {
162+
EntityResponse entityResponse = new EntityResponse();
163+
entityResponse.setUrn(userUrn);
164+
165+
// Create CorpUserKey aspect
166+
CorpUserKey corpUserKey = new CorpUserKey();
167+
corpUserKey.setUsername(userUrn.getId());
168+
169+
EnvelopedAspect keyAspect = new EnvelopedAspect();
170+
keyAspect.setValue(new Aspect(corpUserKey.data()));
171+
172+
// Create CorpUserInfo aspect
173+
CorpUserInfo corpUserInfo = new CorpUserInfo();
174+
corpUserInfo.setActive(true);
175+
corpUserInfo.setFullName("Test User");
176+
corpUserInfo.setEmail("[email protected]");
177+
178+
EnvelopedAspect infoAspect = new EnvelopedAspect();
179+
infoAspect.setValue(new Aspect(corpUserInfo.data()));
180+
181+
Map<String, EnvelopedAspect> aspects = new HashMap<>();
182+
aspects.put(Constants.CORP_USER_KEY_ASPECT_NAME, keyAspect);
183+
aspects.put(Constants.CORP_USER_INFO_ASPECT_NAME, infoAspect);
184+
185+
entityResponse.setAspects(new EnvelopedAspectMap(aspects));
186+
return entityResponse;
187+
}
188+
}

0 commit comments

Comments
 (0)