Query Blocklist

server/src/main/java/org/apache/druid/server/QueryBlocklistRule.java

@@ -0,0 +1,175 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.server;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.base.Preconditions;
+import com.google.common.base.Strings;
+import com.google.common.collect.Sets;
+import org.apache.druid.query.Query;
+
+import javax.annotation.Nullable;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+/**
+ * A rule for matching queries against blocklist criteria. A query matches this rule if ALL
+ * specified criteria match (AND logic). Null or empty criteria match everything.
+ */
+public class QueryBlocklistRule
+{
+  private final String ruleName;
+  @Nullable
+  private final Set<String> dataSources;
+  @Nullable
+  private final Set<String> queryTypes;
+  @Nullable
+  private final Map<String, String> contextMatches;
+
+  private final boolean hasDataSourceCriteria;
+  private final boolean hasQueryTypeCriteria;
+  private final boolean hasContextCriteria;
+
+  @JsonCreator
+  public QueryBlocklistRule(
+      @JsonProperty("ruleName") String ruleName,
+      @JsonProperty("dataSources") @Nullable Set<String> dataSources,
+      @JsonProperty("queryTypes") @Nullable Set<String> queryTypes,
+      @JsonProperty("contextMatches") @Nullable Map<String, String> contextMatches
+  )
+  {
+    Preconditions.checkArgument(
+        !Strings.isNullOrEmpty(ruleName),
+        "ruleName must not be null or empty"
+    );
+
+    // At least one criterion must be specified to prevent accidentally blocking all queries
+    this.hasDataSourceCriteria = dataSources != null && !dataSources.isEmpty();
+    this.hasQueryTypeCriteria = queryTypes != null && !queryTypes.isEmpty();
+    this.hasContextCriteria = contextMatches != null && !contextMatches.isEmpty();
+
+    Preconditions.checkArgument(
+        hasDataSourceCriteria || hasQueryTypeCriteria || hasContextCriteria,
+        "At least one criterion (dataSources, queryTypes, or contextMatches) must be specified. "
+        + "A rule with all null/empty criteria would block ALL queries."
+    );
+
+    this.ruleName = ruleName;
+    this.dataSources = dataSources;
+    this.queryTypes = queryTypes;
+    this.contextMatches = contextMatches;
+  }
+
+  @JsonProperty
+  public String getRuleName()
+  {
+    return ruleName;
+  }
+
+  @JsonProperty
+  @Nullable
+  public Set<String> getDataSources()
+  {
+    return dataSources;
+  }
+
+  @JsonProperty
+  @Nullable
+  public Set<String> getQueryTypes()
+  {
+    return queryTypes;
+  }
+
+  @JsonProperty
+  @Nullable
+  public Map<String, String> getContextMatches()
+  {
+    return contextMatches;
+  }
+
+  /**
+   * Returns true if the query matches ALL specified criteria (AND logic).
+   * Null or empty criteria match everything.
+   *
+   * @param query the query to check
+   * @return true if the query matches this rule, false otherwise
+   */
+  public boolean matches(Query<?> query)
+  {
+    if (hasDataSourceCriteria) {
+      Set<String> queryDatasources = query.getDataSource().getTableNames();
+      if (Sets.intersection(dataSources, queryDatasources).isEmpty()) {
+        return false;
+      }
+    }
+
+    if (hasQueryTypeCriteria) {
+      if (!queryTypes.contains(query.getType())) {
+        return false;
+      }
+    }
+
+    if (hasContextCriteria) {
+      for (Map.Entry<String, String> entry : contextMatches.entrySet()) {
+        Object contextValue = query.getContext().get(entry.getKey());
+        if (!entry.getValue().equals(String.valueOf(contextValue))) {
+          return false;
+        }
+      }
+    }
+
+    return true;
+  }
+
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    QueryBlocklistRule that = (QueryBlocklistRule) o;
+    return Objects.equals(ruleName, that.ruleName)
+           && Objects.equals(dataSources, that.dataSources)
+           && Objects.equals(queryTypes, that.queryTypes)
+           && Objects.equals(contextMatches, that.contextMatches);
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return Objects.hash(ruleName, dataSources, queryTypes, contextMatches);
+  }
+
+  @Override
+  public String toString()
+  {
+    return "QueryBlocklistRule{" +
+           "ruleName='" + ruleName + '\'' +
+           ", dataSources=" + dataSources +
+           ", queryTypes=" + queryTypes +
+           ", contextMatches=" + contextMatches +
+           '}';
+  }
+}

server/src/main/java/org/apache/druid/server/QueryLifecycle.java

@@ -23,6 +23,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.google.common.collect.Iterables;
+import org.apache.druid.client.BrokerViewOfCoordinatorConfig;
 import org.apache.druid.client.DirectDruidClient;
 import org.apache.druid.error.DruidException;
 import org.apache.druid.java.util.common.DateTimes;
@@ -49,6 +50,7 @@
 import org.apache.druid.query.QueryToolChest;
 import org.apache.druid.query.context.ResponseContext;
 import org.apache.druid.query.policy.PolicyEnforcer;
+import org.apache.druid.server.coordinator.CoordinatorDynamicConfig;
 import org.apache.druid.server.log.RequestLogger;
 import org.apache.druid.server.security.Action;
 import org.apache.druid.server.security.AuthConfig;
@@ -99,6 +101,7 @@ public class QueryLifecycle
   private final DefaultQueryConfig defaultQueryConfig;
   private final AuthConfig authConfig;
   private final PolicyEnforcer policyEnforcer;
+  private final BrokerViewOfCoordinatorConfig brokerViewOfCoordinatorConfig;
   private final long startMs;
   private final long startNs;
 
@@ -121,6 +124,7 @@ public QueryLifecycle(
       final DefaultQueryConfig defaultQueryConfig,
       final AuthConfig authConfig,
       final PolicyEnforcer policyEnforcer,
+      @Nullable final BrokerViewOfCoordinatorConfig brokerViewOfCoordinatorConfig,
       final long startMs,
       final long startNs
   )
@@ -134,6 +138,7 @@ public QueryLifecycle(
     this.defaultQueryConfig = defaultQueryConfig;
     this.authConfig = authConfig;
     this.policyEnforcer = policyEnforcer;
+    this.brokerViewOfCoordinatorConfig = brokerViewOfCoordinatorConfig;
     this.startMs = startMs;
     this.startNs = startNs;
   }
@@ -159,6 +164,8 @@ public <T> QueryResponse<T> runSimple(
   {
     initialize(query);
 
+    checkQueryBlocklist();
+
     final Sequence<T> results;
 
     final QueryResponse<T> queryResponse;
@@ -310,6 +317,36 @@ private void preAuthorized(
     }
   }
 
+  /**
+   * Checks if the query matches any blocklist rules. If a rule matches, throws a DruidException.
+   * Rules are evaluated in order, and the first match wins.
+   *
+   * @throws DruidException if the query is blocklisted
+   */
+  private void checkQueryBlocklist()
+  {
+    if (brokerViewOfCoordinatorConfig == null) {
+      return; // Not running on broker, skip blocklist check
+    }
+
+    CoordinatorDynamicConfig config = brokerViewOfCoordinatorConfig.getDynamicConfig();
+    if (config == null) {
+      return; // Config not loaded yet, allow query (best effort)
+    }
+
+    for (QueryBlocklistRule rule : config.getQueryBlocklist()) {
+      if (rule.matches(this.baseQuery)) {
+        throw DruidException.forPersona(DruidException.Persona.USER)
+                            .ofCategory(DruidException.Category.FORBIDDEN)
+                            .build(
+                                "Query[%s] blocked by rule[%s]",
+                                this.baseQuery.getId(),
+                                rule.getRuleName()
+                            );
+      }
+    }
+  }
+
   private AuthorizationResult doAuthorize(
       final AuthenticationResult authenticationResult,
       final AuthorizationResult authorizationResult

server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java

@@ -21,6 +21,7 @@
 
 import com.google.common.base.Supplier;
 import com.google.inject.Inject;
+import org.apache.druid.client.BrokerViewOfCoordinatorConfig;
 import org.apache.druid.guice.LazySingleton;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
 import org.apache.druid.query.DefaultQueryConfig;
@@ -32,6 +33,8 @@
 import org.apache.druid.server.security.AuthConfig;
 import org.apache.druid.server.security.AuthorizerMapper;
 
+import javax.annotation.Nullable;
+
 @LazySingleton
 public class QueryLifecycleFactory
 {
@@ -44,6 +47,7 @@ public class QueryLifecycleFactory
   private final DefaultQueryConfig defaultQueryConfig;
   private final AuthConfig authConfig;
   private final PolicyEnforcer policyEnforcer;
+  private final BrokerViewOfCoordinatorConfig brokerViewOfCoordinatorConfig;
 
   @Inject
   public QueryLifecycleFactory(
@@ -55,7 +59,8 @@ public QueryLifecycleFactory(
       final AuthConfig authConfig,
       final PolicyEnforcer policyEnforcer,
       final AuthorizerMapper authorizerMapper,
-      final Supplier<DefaultQueryConfig> queryConfigSupplier
+      final Supplier<DefaultQueryConfig> queryConfigSupplier,
+      @Nullable final BrokerViewOfCoordinatorConfig brokerViewOfCoordinatorConfig
   )
   {
     this.conglomerate = conglomerate;
@@ -67,6 +72,7 @@ public QueryLifecycleFactory(
     this.defaultQueryConfig = queryConfigSupplier.get();
     this.authConfig = authConfig;
     this.policyEnforcer = policyEnforcer;
+    this.brokerViewOfCoordinatorConfig = brokerViewOfCoordinatorConfig;
   }
 
   public QueryLifecycle factorize()
@@ -81,6 +87,7 @@ public QueryLifecycle factorize()
         defaultQueryConfig,
         authConfig,
         policyEnforcer,
+        brokerViewOfCoordinatorConfig,
         System.currentTimeMillis(),
         System.nanoTime()
     );