Skip to content

[Backport] Add QueryLifecycle#authorize for grpc-query-extension (#15816) #15837

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
abhishekagarwal87 merged 1 commit into from
Feb 6, 2024
Merged

[Backport] Add QueryLifecycle#authorize for grpc-query-extension (#15816) #15837

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions server/src/main/java/org/apache/druid/server/QueryLifecycle.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,6 +242,37 @@ public Access authorize(HttpServletRequest req)
);
}

/**
* Authorize the query using the authentication result.
* Will return an Access object denoting whether the query is authorized or not.
* This method is to be used by the grpc-query-extension.
*
* @param authenticationResult authentication result indicating identity of the requester
* @return authorization result of requester
*/
public Access authorize(AuthenticationResult authenticationResult)
{
transition(State.INITIALIZED, State.AUTHORIZING);
final Iterable<ResourceAction> resourcesToAuthorize = Iterables.concat(
Iterables.transform(
baseQuery.getDataSource().getTableNames(),
AuthorizationUtils.DATASOURCE_READ_RA_GENERATOR
),
Iterables.transform(
authConfig.contextKeysToAuthorize(userContextKeys),
contextParam -> new ResourceAction(new Resource(contextParam, ResourceType.QUERY_CONTEXT), Action.WRITE)
)
);
return doAuthorize(
authenticationResult,
AuthorizationUtils.authorizeAllResourceActions(
authenticationResult,
resourcesToAuthorize,
authorizerMapper
)
);
}

private void preAuthorized(final AuthenticationResult authenticationResult, final Access access)
{
// gotta transition those states, even if we are already authorized
Expand Down
74 changes: 54 additions & 20 deletions server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,15 +188,15 @@ public void testAuthorizeQueryContext_authorized()
EasyMock.expect(authenticationResult.getIdentity()).andReturn(IDENTITY).anyTimes();
EasyMock.expect(authenticationResult.getAuthorizerName()).andReturn(AUTHORIZER).anyTimes();
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource(DATASOURCE, ResourceType.DATASOURCE), Action.READ))
.andReturn(Access.OK);
.andReturn(Access.OK).times(2);
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("foo", ResourceType.QUERY_CONTEXT), Action.WRITE))
.andReturn(Access.OK);
.andReturn(Access.OK).times(2);
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("baz", ResourceType.QUERY_CONTEXT), Action.WRITE))
.andReturn(Access.OK);
.andReturn(Access.OK).times(2);

EasyMock.expect(toolChestWarehouse.getToolChest(EasyMock.anyObject()))
.andReturn(toolChest)
.once();
.times(2);

replayAll();

Expand All @@ -223,6 +223,10 @@ public void testAuthorizeQueryContext_authorized()
);

Assert.assertTrue(lifecycle.authorize(mockRequest()).isAllowed());

lifecycle = createLifecycle(authConfig);
lifecycle.initialize(query);
Assert.assertTrue(lifecycle.authorize(authenticationResult).isAllowed());
}

@Test
Expand All @@ -232,13 +236,15 @@ public void testAuthorizeQueryContext_notAuthorized()
EasyMock.expect(authenticationResult.getIdentity()).andReturn(IDENTITY).anyTimes();
EasyMock.expect(authenticationResult.getAuthorizerName()).andReturn(AUTHORIZER).anyTimes();
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource(DATASOURCE, ResourceType.DATASOURCE), Action.READ))
.andReturn(Access.OK);
.andReturn(Access.OK)
.times(2);
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("foo", ResourceType.QUERY_CONTEXT), Action.WRITE))
.andReturn(Access.DENIED);
.andReturn(Access.DENIED)
.times(2);

EasyMock.expect(toolChestWarehouse.getToolChest(EasyMock.anyObject()))
.andReturn(toolChest)
.once();
.times(2);

replayAll();

Expand All @@ -255,6 +261,10 @@ public void testAuthorizeQueryContext_notAuthorized()
QueryLifecycle lifecycle = createLifecycle(authConfig);
lifecycle.initialize(query);
Assert.assertFalse(lifecycle.authorize(mockRequest()).isAllowed());

lifecycle = createLifecycle(authConfig);
lifecycle.initialize(query);
Assert.assertFalse(lifecycle.authorize(authenticationResult).isAllowed());
}

@Test
Expand All @@ -264,11 +274,12 @@ public void testAuthorizeQueryContext_unsecuredKeys()
EasyMock.expect(authenticationResult.getIdentity()).andReturn(IDENTITY).anyTimes();
EasyMock.expect(authenticationResult.getAuthorizerName()).andReturn(AUTHORIZER).anyTimes();
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource(DATASOURCE, ResourceType.DATASOURCE), Action.READ))
.andReturn(Access.OK);
.andReturn(Access.OK)
.times(2);

EasyMock.expect(toolChestWarehouse.getToolChest(EasyMock.anyObject()))
.andReturn(toolChest)
.once();
.times(2);

replayAll();

Expand Down Expand Up @@ -296,6 +307,10 @@ public void testAuthorizeQueryContext_unsecuredKeys()
);

Assert.assertTrue(lifecycle.authorize(mockRequest()).isAllowed());

lifecycle = createLifecycle(authConfig);
lifecycle.initialize(query);
Assert.assertTrue(lifecycle.authorize(authenticationResult).isAllowed());
}

@Test
Expand All @@ -305,11 +320,12 @@ public void testAuthorizeQueryContext_securedKeys()
EasyMock.expect(authenticationResult.getIdentity()).andReturn(IDENTITY).anyTimes();
EasyMock.expect(authenticationResult.getAuthorizerName()).andReturn(AUTHORIZER).anyTimes();
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource(DATASOURCE, ResourceType.DATASOURCE), Action.READ))
.andReturn(Access.OK);
.andReturn(Access.OK)
.times(2);

EasyMock.expect(toolChestWarehouse.getToolChest(EasyMock.anyObject()))
.andReturn(toolChest)
.once();
.times(2);

replayAll();

Expand Down Expand Up @@ -338,6 +354,10 @@ public void testAuthorizeQueryContext_securedKeys()
);

Assert.assertTrue(lifecycle.authorize(mockRequest()).isAllowed());

lifecycle = createLifecycle(authConfig);
lifecycle.initialize(query);
Assert.assertTrue(lifecycle.authorize(authenticationResult).isAllowed());
}

@Test
Expand All @@ -347,13 +367,15 @@ public void testAuthorizeQueryContext_securedKeysNotAuthorized()
EasyMock.expect(authenticationResult.getIdentity()).andReturn(IDENTITY).anyTimes();
EasyMock.expect(authenticationResult.getAuthorizerName()).andReturn(AUTHORIZER).anyTimes();
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource(DATASOURCE, ResourceType.DATASOURCE), Action.READ))
.andReturn(Access.OK);
.andReturn(Access.OK)
.times(2);
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("foo", ResourceType.QUERY_CONTEXT), Action.WRITE))
.andReturn(Access.DENIED);
.andReturn(Access.DENIED)
.times(2);

EasyMock.expect(toolChestWarehouse.getToolChest(EasyMock.anyObject()))
.andReturn(toolChest)
.once();
.times(2);

replayAll();

Expand All @@ -373,6 +395,10 @@ public void testAuthorizeQueryContext_securedKeysNotAuthorized()
QueryLifecycle lifecycle = createLifecycle(authConfig);
lifecycle.initialize(query);
Assert.assertFalse(lifecycle.authorize(mockRequest()).isAllowed());

lifecycle = createLifecycle(authConfig);
lifecycle.initialize(query);
Assert.assertFalse(lifecycle.authorize(authenticationResult).isAllowed());
}

@Test
Expand All @@ -382,22 +408,26 @@ public void testAuthorizeLegacyQueryContext_authorized()
EasyMock.expect(authenticationResult.getIdentity()).andReturn(IDENTITY).anyTimes();
EasyMock.expect(authenticationResult.getAuthorizerName()).andReturn(AUTHORIZER).anyTimes();
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("fake", ResourceType.DATASOURCE), Action.READ))
.andReturn(Access.OK);
.andReturn(Access.OK)
.times(2);
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("foo", ResourceType.QUERY_CONTEXT), Action.WRITE))
.andReturn(Access.OK);
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("baz", ResourceType.QUERY_CONTEXT), Action.WRITE)).andReturn(Access.OK);
.andReturn(Access.OK)
.times(2);
EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("baz", ResourceType.QUERY_CONTEXT), Action.WRITE))
.andReturn(Access.OK)
.times(2);

EasyMock.expect(toolChestWarehouse.getToolChest(EasyMock.anyObject()))
.andReturn(toolChest)
.once();
.times(2);

replayAll();

final QueryContextTest.LegacyContextQuery query = new QueryContextTest.LegacyContextQuery(ImmutableMap.of("foo", "bar", "baz", "qux"));

AuthConfig authConfig = AuthConfig.newBuilder()
.setAuthorizeQueryContextParams(true)
.build();
.setAuthorizeQueryContextParams(true)
.build();
QueryLifecycle lifecycle = createLifecycle(authConfig);
lifecycle.initialize(query);

Expand All @@ -408,6 +438,10 @@ public void testAuthorizeLegacyQueryContext_authorized()
Assert.assertTrue(revisedContext.containsKey("queryId"));

Assert.assertTrue(lifecycle.authorize(mockRequest()).isAllowed());

lifecycle = createLifecycle(authConfig);
lifecycle.initialize(query);
Assert.assertTrue(lifecycle.authorize(mockRequest()).isAllowed());
}

private HttpServletRequest mockRequest()
Expand Down