fix: ensure native query statusCode is 403 for blocklist

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<String> 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()

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;
 

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<Boolean> 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);