From de959e513d120348dfbfcf3e13d265b200a7db6c Mon Sep 17 00:00:00 2001 From: Rishabh Singh <6513075+findingrish@users.noreply.github.com> Date: Fri, 2 Feb 2024 21:49:57 +0530 Subject: [PATCH] Add QueryLifecycle#authorize for grpc-query-extension (#15816) 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. --- .../apache/druid/server/QueryLifecycle.java | 31 ++++++++ .../druid/server/QueryLifecycleTest.java | 74 ++++++++++++++----- 2 files changed, 85 insertions(+), 20 deletions(-) diff --git a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java index 0f46e5da4d9..e0bb9875240 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java @@ -242,6 +242,37 @@ public class QueryLifecycle ); } + /** + * 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 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 diff --git a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java index 578661ea768..a2691297d0b 100644 --- a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java @@ -188,15 +188,15 @@ public class QueryLifecycleTest 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(); @@ -223,6 +223,10 @@ public class QueryLifecycleTest ); Assert.assertTrue(lifecycle.authorize(mockRequest()).isAllowed()); + + lifecycle = createLifecycle(authConfig); + lifecycle.initialize(query); + Assert.assertTrue(lifecycle.authorize(authenticationResult).isAllowed()); } @Test @@ -232,13 +236,15 @@ public class QueryLifecycleTest 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(); @@ -255,6 +261,10 @@ public class QueryLifecycleTest 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 @@ -264,11 +274,12 @@ public class QueryLifecycleTest 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(); @@ -296,6 +307,10 @@ public class QueryLifecycleTest ); Assert.assertTrue(lifecycle.authorize(mockRequest()).isAllowed()); + + lifecycle = createLifecycle(authConfig); + lifecycle.initialize(query); + Assert.assertTrue(lifecycle.authorize(authenticationResult).isAllowed()); } @Test @@ -305,11 +320,12 @@ public class QueryLifecycleTest 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(); @@ -338,6 +354,10 @@ public class QueryLifecycleTest ); Assert.assertTrue(lifecycle.authorize(mockRequest()).isAllowed()); + + lifecycle = createLifecycle(authConfig); + lifecycle.initialize(query); + Assert.assertTrue(lifecycle.authorize(authenticationResult).isAllowed()); } @Test @@ -347,13 +367,15 @@ public class QueryLifecycleTest 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(); @@ -373,6 +395,10 @@ public class QueryLifecycleTest 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 @@ -382,22 +408,26 @@ public class QueryLifecycleTest 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); @@ -408,6 +438,10 @@ public class QueryLifecycleTest 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()