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.

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<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

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);
+            .andReturn(Access.OK)
-    EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("baz", ResourceType.QUERY_CONTEXT), Action.WRITE)).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)
+                                      .setAuthorizeQueryContextParams(true)
-        .build();
+                                      .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()