From a559622c3c1f97d55c16c867015d2553d6f41bf1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jesse=20Tu=C4=9Flu?= Date: Fri, 24 Jul 2026 14:42:20 -0700 Subject: [PATCH] fix: ensure native query statusCode is 403 for blocklist --- .../EmbeddedBrokerDynamicConfigTest.java | 47 +++++++++++ .../apache/druid/server/QueryResource.java | 11 +++ .../druid/server/QueryResourceTest.java | 80 +++++++++++++++++++ 3 files changed, 138 insertions(+) diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/EmbeddedBrokerDynamicConfigTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/EmbeddedBrokerDynamicConfigTest.java index ac48e2c53466..cef70aa01e27 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/EmbeddedBrokerDynamicConfigTest.java +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/EmbeddedBrokerDynamicConfigTest.java @@ -23,6 +23,7 @@ import org.apache.druid.common.config.JacksonConfigManager; import org.apache.druid.common.utils.IdUtils; import org.apache.druid.indexing.common.task.TaskBuilder; +import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.query.QueryContext; import org.apache.druid.server.DefaultQueryBlocklistRule; import org.apache.druid.server.QueryBlocklistRule; @@ -43,6 +44,10 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; import java.util.List; import java.util.Map; import java.util.Set; @@ -127,6 +132,48 @@ public void testQueryBlocklistBlocksMatchingQueries() Assertions.assertFalse(finalResult.isBlank()); } + /** + * A blocklisted query must be rejected with 403. + */ + @Test + @Timeout(30) + public void testBlocklistedQueryReturnsForbiddenOnBothEndpoints() throws Exception + { + updateBrokerDynamicConfig( + BrokerDynamicConfig.builder() + .withQueryBlocklist(List.of( + new DefaultQueryBlocklistRule("block-test-datasource", Set.of(dataSource), null, null) + )) + .build() + ); + + try { + final String nativeQuery = StringUtils.format( + "{\"queryType\":\"timeseries\",\"dataSource\":\"%s\",\"granularity\":\"all\"," + + "\"intervals\":[\"2000/3000\"],\"aggregations\":[{\"type\":\"count\",\"name\":\"rows\"}]}", + dataSource + ); + Assertions.assertEquals(403, postToBroker("/druid/v2", nativeQuery).statusCode()); + + final String sqlQuery = StringUtils.format("{\"query\":\"SELECT COUNT(*) FROM \\\"%s\\\"\"}", dataSource); + Assertions.assertEquals(403, postToBroker("/druid/v2/sql", sqlQuery).statusCode()); + } + finally { + updateBrokerDynamicConfig(BrokerDynamicConfig.builder().build()); + } + } + + private HttpResponse postToBroker(String path, String body) throws Exception + { + final HttpRequest request = HttpRequest + .newBuilder(URI.create(getServerUrl(broker) + path)) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(body)) + .build(); + + return HttpClient.newHttpClient().send(request, HttpResponse.BodyHandlers.ofString()); + } + @Test @Timeout(30) public void testDynamicQueryContextTimeoutCausesQueryToFail() diff --git a/server/src/main/java/org/apache/druid/server/QueryResource.java b/server/src/main/java/org/apache/druid/server/QueryResource.java index 80ea2895b97b..e4be284bfb60 100644 --- a/server/src/main/java/org/apache/druid/server/QueryResource.java +++ b/server/src/main/java/org/apache/druid/server/QueryResource.java @@ -25,8 +25,10 @@ import com.fasterxml.jackson.databind.SerializerProvider; import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes; import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; import com.google.inject.Inject; +import org.apache.druid.error.DruidException; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.annotations.Json; import org.apache.druid.java.util.common.jackson.JacksonUtils; @@ -181,6 +183,15 @@ public Response doPost( try { authResult = queryLifecycle.authorize(req); } + catch (DruidException e) { + // Handled separately from QueryException so that the exception's own category (for example FORBIDDEN for a + // blocklisted query) determines the status code. Wrapping it below would erase that and report a 500. + return QueryResultPusher.handleDruidExceptionBeforeResponseStarted( + e, + MediaType.valueOf(io.getResponseWriter().getResponseType()), + ImmutableMap.of() + ); + } catch (RuntimeException e) { final QueryException qe; diff --git a/server/src/test/java/org/apache/druid/server/QueryResourceTest.java b/server/src/test/java/org/apache/druid/server/QueryResourceTest.java index d7a611bfc320..cb118f8f760a 100644 --- a/server/src/test/java/org/apache/druid/server/QueryResourceTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryResourceTest.java @@ -31,6 +31,7 @@ import com.google.common.util.concurrent.MoreExecutors; import com.google.inject.Injector; import com.google.inject.Key; +import org.apache.druid.client.BrokerViewOfBrokerConfig; import org.apache.druid.error.DruidException; import org.apache.druid.error.DruidExceptionMatcher; import org.apache.druid.error.ErrorResponse; @@ -74,6 +75,7 @@ import org.apache.druid.query.policy.NoopPolicyEnforcer; import org.apache.druid.query.policy.RowFilterPolicy; import org.apache.druid.query.timeboundary.TimeBoundaryResultValue; +import org.apache.druid.server.broker.BrokerDynamicConfig; import org.apache.druid.server.initialization.ServerConfig; import org.apache.druid.server.log.TestRequestLogger; import org.apache.druid.server.metrics.NoopServiceEmitter; @@ -103,6 +105,7 @@ import org.junit.BeforeClass; import org.junit.Test; import org.junit.internal.matchers.ThrowableMessageMatcher; +import org.mockito.Mockito; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -1835,6 +1838,83 @@ private Future eventuallyAssertAsyncResponse( }); } + @Test + public void testBlocklistedQueryReturnsForbidden() throws IOException + { + expectPermissiveHappyPathAuth(); + + final QueryResource blockingQueryResource = createQueryResourceWithBlocklist( + new DefaultQueryBlocklistRule("block-mmx", ImmutableSet.of("mmx_metrics"), null, null) + ); + + final Response response = blockingQueryResource.doPost( + new ByteArrayInputStream(SIMPLE_TIMESERIES_QUERY.getBytes(StandardCharsets.UTF_8)), + null /*pretty*/, + testServletRequest + ); + + Assert.assertNotNull(response); + Assert.assertEquals(Status.FORBIDDEN.getStatusCode(), response.getStatus()); + + MatcherAssert.assertThat( + ((ErrorResponse) response.getEntity()).getUnderlyingException(), + DruidExceptionMatcher.forbidden().expectMessageContains("blocked by rule[block-mmx]") + ); + } + + @Test + public void testNonBlocklistedQueryIsNotAffectedByBlocklist() throws IOException + { + expectPermissiveHappyPathAuth(); + + final QueryResource blockingQueryResource = createQueryResourceWithBlocklist( + new DefaultQueryBlocklistRule("block-other", ImmutableSet.of("some_other_datasource"), null, null) + ); + + final MockHttpServletResponse response = expectAsyncRequestFlow( + testServletRequest, + SIMPLE_TIMESERIES_QUERY.getBytes(StandardCharsets.UTF_8), + blockingQueryResource + ); + + Assert.assertEquals(Status.OK.getStatusCode(), response.getStatus()); + } + + private QueryResource createQueryResourceWithBlocklist(QueryBlocklistRule... rules) + { + final BrokerViewOfBrokerConfig brokerViewOfBrokerConfig = Mockito.mock(BrokerViewOfBrokerConfig.class); + Mockito.when(brokerViewOfBrokerConfig.getDynamicConfig()).thenReturn( + new BrokerDynamicConfig.Builder().withQueryBlocklist(Arrays.asList(rules)).build() + ); + + return new QueryResource( + new QueryLifecycleFactory( + CONGLOMERATE, + TEST_SEGMENT_WALKER, + new DefaultGenericQueryMetricsFactory(), + emitter, + testRequestLogger, + new AuthConfig(), + NoopPolicyEnforcer.instance(), + AuthTestUtils.TEST_AUTHORIZER_MAPPER, + new DefaultQueryConfig(Map.of()), + brokerViewOfBrokerConfig + ), + jsonMapper, + queryScheduler, + null, + new QueryResourceQueryResultPusherFactory( + jsonMapper, + ResponseContextConfig.newConfig(true), + DRUID_NODE + ), + new ResourceIOReaderWriterFactory( + jsonMapper, + smileMapper + ) + ); + } + private void expectPermissiveHappyPathAuth() { testServletRequest.setAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT, AUTHENTICATION_RESULT);