diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index 232cf5b22a22..1d20f5ddf1a6 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -268,8 +268,7 @@ public Response cancelQuery( lifecycles.forEach(SqlLifecycle::cancel); return Response.status(Status.ACCEPTED).build(); } else { - // Return 404 for authorization failures as well - return Response.status(Status.NOT_FOUND).build(); + return Response.status(Status.FORBIDDEN).build(); } } } diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index e1d68ef2ecce..870342afb5c1 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -61,6 +61,7 @@ import org.apache.druid.server.scheduling.HiLoQueryLaningStrategy; import org.apache.druid.server.scheduling.ManualQueryPrioritizationStrategy; import org.apache.druid.server.security.AuthConfig; +import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.security.ForbiddenException; import org.apache.druid.sql.SqlLifecycle; import org.apache.druid.sql.SqlLifecycleFactory; @@ -986,7 +987,7 @@ public void testTooManyRequests() throws Exception ImmutableMap.of("priority", -5, "sqlQueryId", sqlQueryId), null ), - makeExpectedReq() + makeRegularUserReq() ); } catch (Exception e) { @@ -1053,7 +1054,7 @@ public void testCancelBetweenValidateAndPlan() throws Exception Future future = executorService.submit( () -> resource.doPost( createSimpleQueryWithId(sqlQueryId, "SELECT DISTINCT dim1 FROM foo"), - makeExpectedReq() + makeRegularUserReq() ) ); Assert.assertTrue(validateAndAuthorizeLatch.await(1, TimeUnit.SECONDS)); @@ -1084,7 +1085,7 @@ public void testCancelBetweenPlanAndExecute() throws Exception Future future = executorService.submit( () -> resource.doPost( createSimpleQueryWithId(sqlQueryId, "SELECT DISTINCT dim1 FROM foo"), - makeExpectedReq() + makeRegularUserReq() ) ); Assert.assertTrue(planLatch.await(1, TimeUnit.SECONDS)); @@ -1114,7 +1115,7 @@ public void testCancelInvalidQuery() throws Exception Future future = executorService.submit( () -> resource.doPost( createSimpleQueryWithId(sqlQueryId, "SELECT DISTINCT dim1 FROM foo"), - makeExpectedReq() + makeRegularUserReq() ) ); Assert.assertTrue(planLatch.await(1, TimeUnit.SECONDS)); @@ -1128,6 +1129,31 @@ public void testCancelInvalidQuery() throws Exception Assert.assertEquals(Status.OK.getStatusCode(), response.getStatus()); } + @Test + public void testCancelForbidden() throws Exception + { + final String sqlQueryId = "toCancel"; + CountDownLatch planLatch = new CountDownLatch(1); + planLatchSupplier.set(new NonnullPair<>(planLatch, true)); + CountDownLatch execLatch = new CountDownLatch(1); + executeLatchSupplier.set(new NonnullPair<>(execLatch, false)); + Future future = executorService.submit( + () -> resource.doPost( + createSimpleQueryWithId(sqlQueryId, "SELECT DISTINCT dim1 FROM forbiddenDatasource"), + makeSuperUserReq() + ) + ); + Assert.assertTrue(planLatch.await(1, TimeUnit.SECONDS)); + Response response = resource.cancelQuery(sqlQueryId, mockRequestForCancel()); + Assert.assertEquals(Status.FORBIDDEN.getStatusCode(), response.getStatus()); + + Assert.assertFalse(lifecycleManager.getAll(sqlQueryId).isEmpty()); + + execLatch.countDown(); + response = future.get(); + Assert.assertEquals(Status.OK.getStatusCode(), response.getStatus()); + } + @SuppressWarnings("unchecked") private void checkSqlRequestLog(boolean success) { @@ -1222,24 +1248,34 @@ private Pair doPostRaw(final SqlQuery query, final HttpS } } - private HttpServletRequest makeExpectedReq() + private HttpServletRequest makeSuperUserReq() + { + return makeExpectedReq(CalciteTests.SUPER_USER_AUTH_RESULT); + } + + private HttpServletRequest makeRegularUserReq() + { + return makeExpectedReq(CalciteTests.REGULAR_USER_AUTH_RESULT); + } + + private HttpServletRequest makeExpectedReq(AuthenticationResult authenticationResult) { HttpServletRequest req = EasyMock.createStrictMock(HttpServletRequest.class); EasyMock.expect(req.getRemoteAddr()).andReturn(null).once(); EasyMock.expect(req.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)) - .andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT) + .andReturn(authenticationResult) .anyTimes(); EasyMock.expect(req.getAttribute(AuthConfig.DRUID_ALLOW_UNSECURED_PATH)).andReturn(null).anyTimes(); EasyMock.expect(req.getAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED)) .andReturn(null) .anyTimes(); EasyMock.expect(req.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)) - .andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT) + .andReturn(authenticationResult) .anyTimes(); req.setAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED, true); EasyMock.expectLastCall().anyTimes(); EasyMock.expect(req.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)) - .andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT) + .andReturn(authenticationResult) .anyTimes(); EasyMock.replay(req); return req;