Skip to content

Commit

Permalink
Add QueryLifecycle#authorize for grpc-query-extension (#15816)
Browse files Browse the repository at this point in the history 
Proposal #13469

Original PR #14024

A new method is being added in QueryLifecycle class to authorise a query based on authentication result.
This method is required since we authenticate the query by intercepting it in the grpc extension and pass down the authentication result.
  • Loading branch information
@findingrish
findingrish authored Feb 2, 2024
1 parent 8f5b752 commit de959e5
Show file tree
Hide file tree
Showing 2 changed files with 85 additions and 20 deletions.
31 changes: 31 additions & 0 deletions server/src/main/java/org/apache/druid/server/QueryLifecycle.java
View file
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
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

0 comments on commit de959e5

Please sign in to comment.