review, tests

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java

@@ -630,14 +630,13 @@ private Optional<SqlStatementResult> getStatementStatus(
 
 
   /**
-   * This method contacts the overlord for the controller task and returns it if the requested user has the necessary
-   * permissions to do so. A user has the necessary permissions if they satisfy one of the following criterias:
-   * 1. They are the one who has submitted the query
-   * 2. They belong to a role containing the READ or WRITE permissions over the STATE resource. For endpoints like GET,
-   *   the user should have READ permission for the STATE resource, while for endpoints like POST and DELETE, the user
-   *   should have WRITE permission for the STATE resource (note: POST is only required when the user submits the query
-   *   so the currentUser should be equal to the queryUser, and we won't need to check the permissions for STATE resource
-   *   for them anyway)
+   * This method contacts the overlord for the controller task and checks if the requested user has the
+   * necessary permissions. A user has the necessary permissions if one of the following criteria is satisfied:
+   * 1. The user is the one who submitted the query
+   * 2. The user belongs to a role containing the READ or WRITE permissions over the STATE resource. For endpoints like GET,
+   *   the user should have READ permission for the STATE resource, while for endpoints like DELETE, the user should
+   *   have WRITE permission for the STATE resource. (Note: POST API does not need to check the state permissions since
+   *   the currentUser always equal to the queryUser)
    */
   private MSQControllerTask getMSQControllerTaskAndCheckPermission(
       String queryId,

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlStatementResourceTest.java

@@ -75,8 +75,13 @@
 import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.segment.column.ValueType;
 import org.apache.druid.server.mocks.MockHttpServletRequest;
+import org.apache.druid.server.security.Access;
+import org.apache.druid.server.security.Action;
 import org.apache.druid.server.security.AuthConfig;
 import org.apache.druid.server.security.AuthenticationResult;
+import org.apache.druid.server.security.Authorizer;
+import org.apache.druid.server.security.AuthorizerMapper;
+import org.apache.druid.server.security.ResourceType;
 import org.apache.druid.sql.calcite.planner.ColumnMappings;
 import org.apache.druid.sql.calcite.util.CalciteTests;
 import org.apache.druid.sql.http.ResultFormat;
@@ -339,6 +344,47 @@ public class SqlStatementResourceTest extends MSQTestBase
   private static final String FAILURE_MSG = "failure msg";
   private static SqlStatementResource resource;
 
+  private static String SUPERUSER = "superuser";
+  private static String STATE_R_USER = "stateR";
+  private static String STATE_W_USER = "stateW";
+  private static String STATE_RW_USER = "stateRW";
+
+  private AuthorizerMapper authorizerMapper = new AuthorizerMapper(null)
+  {
+    @Override
+    public Authorizer getAuthorizer(String name)
+    {
+      return (authenticationResult, resource, action) -> {
+        if (SUPERUSER.equals(authenticationResult.getIdentity())) {
+          return Access.OK;
+        }
+
+        switch (resource.getType()) {
+          case ResourceType.DATASOURCE:
+          case ResourceType.VIEW:
+          case ResourceType.QUERY_CONTEXT:
+          case ResourceType.EXTERNAL:
+            return Access.OK;
+          case ResourceType.STATE:
+            String identity = authenticationResult.getIdentity();
+            if (action == Action.READ) {
+              if (STATE_R_USER.equals(identity) || STATE_RW_USER.equals(identity)) {
+                return Access.OK;
+              }
+            } else if (action == Action.WRITE) {
+              if (STATE_W_USER.equals(identity) || STATE_RW_USER.equals(identity)) {
+                return Access.OK;
+              }
+            }
+            return Access.DENIED;
+
+          default:
+            return Access.DENIED;
+        }
+      };
+    }
+  };
+
   @Mock
   private OverlordClient overlordClient;
 
@@ -631,14 +677,24 @@ public static MockHttpServletRequest makeOkRequest()
     return makeExpectedReq(CalciteTests.REGULAR_USER_AUTH_RESULT);
   }
 
-  public static MockHttpServletRequest makeExpectedReq(AuthenticationResult authenticationResult)
+  private static MockHttpServletRequest makeExpectedReq(AuthenticationResult authenticationResult)
   {
     MockHttpServletRequest req = new MockHttpServletRequest();
     req.attributes.put(AuthConfig.DRUID_AUTHENTICATION_RESULT, authenticationResult);
     req.remoteAddr = "1.2.3.4";
     return req;
   }
 
+  private static AuthenticationResult makeAuthResultForUser(String user)
+  {
+    return new AuthenticationResult(
+        user,
+        AuthConfig.ALLOW_ALL_NAME,
+        null,
+        null
+    );
+  }
+
   @Before
   public void init() throws Exception
   {
@@ -921,7 +977,7 @@ public void testAPIBehaviourWithSuperUsers()
         Response.Status.OK.getStatusCode(),
         resource.doGetStatus(
             RUNNING_SELECT_MSQ_QUERY,
-            makeExpectedReq(CalciteTests.SUPER_USER_AUTH_RESULT)
+            makeExpectedReq(makeAuthResultForUser(SUPERUSER))
         ).getStatus()
     );
     Assert.assertEquals(
@@ -930,24 +986,22 @@ public void testAPIBehaviourWithSuperUsers()
             RUNNING_SELECT_MSQ_QUERY,
             1L,
             null,
-            makeExpectedReq(CalciteTests.SUPER_USER_AUTH_RESULT)
+            makeExpectedReq(makeAuthResultForUser(SUPERUSER))
         ).getStatus()
     );
     Assert.assertEquals(
         Response.Status.ACCEPTED.getStatusCode(),
         resource.deleteQuery(
             RUNNING_SELECT_MSQ_QUERY,
-            makeExpectedReq(CalciteTests.SUPER_USER_AUTH_RESULT)
+            makeExpectedReq(makeAuthResultForUser(SUPERUSER))
         ).getStatus()
     );
   }
 
   @Test
-  public void testAPIBehaviourWithForbiddenUser()
+  public void testAPIBehaviourWithDifferentUserAndNoStatePermission()
   {
-    AuthenticationResult differentUserAuthResult = new AuthenticationResult(
-        "differentUser", AuthConfig.ALLOW_ALL_NAME, null, null
-    );
+    AuthenticationResult differentUserAuthResult = makeAuthResultForUser("differentUser");
     Assert.assertEquals(
         Response.Status.FORBIDDEN.getStatusCode(),
         resource.doGetStatus(
@@ -973,6 +1027,93 @@ public void testAPIBehaviourWithForbiddenUser()
     );
   }
 
+  @Test
+  public void testAPIBehaviourWithDifferentUserAndStateRPermission()
+  {
+    AuthenticationResult differentUserAuthResult = makeAuthResultForUser(STATE_R_USER);
+    Assert.assertEquals(
+        Response.Status.OK.getStatusCode(),
+        resource.doGetStatus(
+            RUNNING_SELECT_MSQ_QUERY,
+            makeExpectedReq(differentUserAuthResult)
+        ).getStatus()
+    );
+    Assert.assertEquals(
+        Response.Status.BAD_REQUEST.getStatusCode(),
+        resource.doGetResults(
+            RUNNING_SELECT_MSQ_QUERY,
+            1L,
+            null,
+            makeExpectedReq(differentUserAuthResult)
+        ).getStatus()
+    );
+    Assert.assertEquals(
+        Response.Status.FORBIDDEN.getStatusCode(),
+        resource.deleteQuery(
+            RUNNING_SELECT_MSQ_QUERY,
+            makeExpectedReq(differentUserAuthResult)
+        ).getStatus()
+    );
+  }
+
+  @Test
+  public void testAPIBehaviourWithDifferentUserAndStateWPermission()
+  {
+    AuthenticationResult differentUserAuthResult = makeAuthResultForUser(STATE_W_USER);
+    Assert.assertEquals(
+        Response.Status.FORBIDDEN.getStatusCode(),
+        resource.doGetStatus(
+            RUNNING_SELECT_MSQ_QUERY,
+            makeExpectedReq(differentUserAuthResult)
+        ).getStatus()
+    );
+    Assert.assertEquals(
+        Response.Status.FORBIDDEN.getStatusCode(),
+        resource.doGetResults(
+            RUNNING_SELECT_MSQ_QUERY,
+            1L,
+            null,
+            makeExpectedReq(differentUserAuthResult)
+        ).getStatus()
+    );
+    Assert.assertEquals(
+        Response.Status.ACCEPTED.getStatusCode(),
+        resource.deleteQuery(
+            RUNNING_SELECT_MSQ_QUERY,
+            makeExpectedReq(differentUserAuthResult)
+        ).getStatus()
+    );
+  }
+
+  @Test
+  public void testAPIBehaviourWithDifferentUserAndStateRWPermission()
+  {
+    AuthenticationResult differentUserAuthResult = makeAuthResultForUser(STATE_RW_USER);
+    Assert.assertEquals(
+        Response.Status.OK.getStatusCode(),
+        resource.doGetStatus(
+            RUNNING_SELECT_MSQ_QUERY,
+            makeExpectedReq(differentUserAuthResult)
+        ).getStatus()
+    );
+    Assert.assertEquals(
+        Response.Status.BAD_REQUEST.getStatusCode(),
+        resource.doGetResults(
+            RUNNING_SELECT_MSQ_QUERY,
+            1L,
+            null,
+            makeExpectedReq(differentUserAuthResult)
+        ).getStatus()
+    );
+    Assert.assertEquals(
+        Response.Status.ACCEPTED.getStatusCode(),
+        resource.deleteQuery(
+            RUNNING_SELECT_MSQ_QUERY,
+            makeExpectedReq(differentUserAuthResult)
+        ).getStatus()
+    );
+  }
+
   @Test
   public void testTaskIdNotFound()
   {

sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java

@@ -209,13 +209,15 @@ public AuthenticationResult createEscalatedAuthenticationResult()
   public static final AuthenticationResult REGULAR_USER_AUTH_RESULT = new AuthenticationResult(
       AuthConfig.ALLOW_ALL_NAME,
       AuthConfig.ALLOW_ALL_NAME,
-      null, null
+      null,
+      null
   );
 
   public static final AuthenticationResult SUPER_USER_AUTH_RESULT = new AuthenticationResult(
       TEST_SUPERUSER_NAME,
       AuthConfig.ALLOW_ALL_NAME,
-      null, null
+      null,
+      null
   );
 
   public static final Injector INJECTOR = new CalciteTestInjectorBuilder()