Add support for authorizing query context params

extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQuery.java

@@ -181,6 +181,12 @@ public MaterializedViewQuery withOverriddenContext(Map<String, Object> contextOv
     return new MaterializedViewQuery(query.withOverriddenContext(contextOverride), optimizer);
   }
 
+  @Override
+  public MaterializedViewQuery withContext(Map<String, Object> context)
+  {
+    return new MaterializedViewQuery<>(query.withContext(context), optimizer);
+  }
+
   @Override
   public MaterializedViewQuery withQuerySegmentSpec(QuerySegmentSpec spec) 
   {

extensions-contrib/moving-average-query/src/main/java/org/apache/druid/query/movingaverage/MovingAverageQuery.java

@@ -254,6 +254,27 @@ public Granularity getGranularity()
     return granularity;
   }
 
+  @Override
+  public MovingAverageQuery withContext(Map<String, Object> context)
+  {
+    return new MovingAverageQuery(
+        getDataSource(),
+        getQuerySegmentSpec(),
+        dimFilter,
+        granularity,
+        dimensions,
+        aggregatorSpecs,
+        averagerSpecs,
+        postAggregatorSpecs,
+        postAveragerSpecs,
+        havingSpec,
+        limitSpec,
+        groupByQueryForLimitSpec,
+        limitFn,
+        context
+    );
+  }
+
   @JsonProperty
   public List<DimensionSpec> getDimensions()
   {

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/quantiles/sql/DoublesSketchApproxQuantileSqlAggregator.java

@@ -32,14 +32,14 @@
 import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlTypeFamily;
 import org.apache.calcite.sql.type.SqlTypeName;
-import org.apache.druid.java.util.common.Numbers;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.query.aggregation.datasketches.quantiles.DoublesSketchAggregatorFactory;
 import org.apache.druid.query.aggregation.datasketches.quantiles.DoublesSketchToQuantilePostAggregator;
 import org.apache.druid.query.aggregation.post.FieldAccessPostAggregator;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.server.QueryContext;
 import org.apache.druid.sql.calcite.aggregation.Aggregation;
 import org.apache.druid.sql.calcite.aggregation.Aggregations;
 import org.apache.druid.sql.calcite.aggregation.SqlAggregator;
@@ -50,7 +50,6 @@
 
 import javax.annotation.Nullable;
 import java.util.List;
-import java.util.Map;
 
 public class DoublesSketchApproxQuantileSqlAggregator implements SqlAggregator
 {
@@ -200,11 +199,12 @@ public Aggregation toDruidAggregation(
     );
   }
 
-  @Nullable
-  static Long getMaxStreamLengthFromQueryContext(Map<String, Object> queryContext)
+  static long getMaxStreamLengthFromQueryContext(QueryContext queryContext)
   {
-    final Object val = queryContext.get(CTX_APPROX_QUANTILE_DS_MAX_STREAM_LENGTH);
-    return val == null ? null : Numbers.parseLong(val);
+    return queryContext.getAsLong(
+        CTX_APPROX_QUANTILE_DS_MAX_STREAM_LENGTH,
+        DoublesSketchAggregatorFactory.DEFAULT_MAX_STREAM_LENGTH
+    );
   }
 
   private static class DoublesSketchApproxQuantileSqlAggFunction extends SqlAggFunction

integration-tests/docker/environment-configs/common

@@ -36,6 +36,7 @@ druid_auth_authenticator_basic_type=basic
 druid_auth_authenticatorChain=["basic"]
 druid_auth_authorizer_basic_type=basic
 druid_auth_authorizers=["basic"]
+druid_auth_authorizeQueryContextParams=true
 druid_client_https_certAlias=druid
 druid_client_https_keyManagerPassword=druid123
 druid_client_https_keyStorePassword=druid123

integration-tests/docker/environment-configs/common-ldap

@@ -46,6 +46,7 @@ druid_auth_authorizer_ldapauth_initialAdminUser=admin
 druid_auth_authorizer_ldapauth_initialAdminRole=admin
 druid_auth_authorizer_ldapauth_roleProvider_type=ldap
 druid_auth_authorizers=["ldapauth"]
+druid_auth_authorizeQueryContextParams=true
 druid_client_https_certAlias=druid
 druid_client_https_keyManagerPassword=druid123
 druid_client_https_keyStorePassword=druid123

integration-tests/docker/ldap-configs/bootstrap.ldif

@@ -71,6 +71,12 @@ cn: datasourceOnlyGroup
 description: datasourceOnlyGroup users
 uniqueMember: uid=datasourceOnlyUser,ou=Users,dc=example,dc=org
 
+dn: cn=datasourceWithContextParamsGroup,ou=Groups,dc=example,dc=org
+objectClass: groupOfUniqueNames
+cn: datasourceWithContextParamsGroup
+description: datasourceWithContextParamsGroup users
+uniqueMember: uid=datasourceWithContextParamsUser,ou=Users,dc=example,dc=org
+
 dn: uid=datasourceWithStateUser,ou=Users,dc=example,dc=org
 uid: datasourceWithStateUser
 cn: datasourceWithStateUser

integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java

@@ -105,6 +105,17 @@ public abstract class AbstractAuthConfigurationTest
       )
   );
 
+  protected static final List<ResourceAction> DATASOURCE_QUERY_CONTEXT_PERMISSIONS = ImmutableList.of(
+      new ResourceAction(
+          new Resource("auth_test", ResourceType.DATASOURCE),
+          Action.READ
+      ),
+      new ResourceAction(
+          new Resource("auth_test_ctx", ResourceType.QUERY_CONTEXT),
+          Action.WRITE
+      )
+  );
+
   /**
    * create a ResourceAction set of permissions that can only read 'auth_test' + partial SYSTEM_TABLE, for Authorizer
    * implementations which use ResourceAction pattern matching
@@ -188,22 +199,22 @@ public abstract class AbstractAuthConfigurationTest
 
   protected HttpClient adminClient;
   protected HttpClient datasourceOnlyUserClient;
+  protected HttpClient datasourceAndContextParamsClient;
   protected HttpClient datasourceAndSysUserClient;
   protected HttpClient datasourceWithStateUserClient;
   protected HttpClient stateOnlyUserClient;
   protected HttpClient internalSystemClient;
 
 
   protected abstract void setupDatasourceOnlyUser() throws Exception;
+  protected abstract void setupDatasourceAndContextParamsUser() throws Exception;
   protected abstract void setupDatasourceAndSysTableUser() throws Exception;
   protected abstract void setupDatasourceAndSysAndStateUser() throws Exception;
   protected abstract void setupSysTableAndStateOnlyUser() throws Exception;
   protected abstract void setupTestSpecificHttpClients() throws Exception;
   protected abstract String getAuthenticatorName();
   protected abstract String getAuthorizerName();
   protected abstract String getExpectedAvaticaAuthError();
-  protected abstract Properties getAvaticaConnectionProperties();
-  protected abstract Properties getAvaticaConnectionPropertiesFailure();
 
   @Test
   public void test_systemSchemaAccess_admin() throws Exception
@@ -444,27 +455,71 @@ public void test_internalSystemUser_hasNodeAccess()
   @Test
   public void test_avaticaQuery_broker()
   {
-    testAvaticaQuery(getBrokerAvacticaUrl());
-    testAvaticaQuery(StringUtils.maybeRemoveTrailingSlash(getBrokerAvacticaUrl()));
+    final Properties properties = getAvaticaConnectionPropertiesForAdmin();
+    testAvaticaQuery(properties, getBrokerAvacticaUrl());
+    testAvaticaQuery(properties, StringUtils.maybeRemoveTrailingSlash(getBrokerAvacticaUrl()));
   }
 
   @Test
   public void test_avaticaQuery_router()
   {
-    testAvaticaQuery(getRouterAvacticaUrl());
-    testAvaticaQuery(StringUtils.maybeRemoveTrailingSlash(getRouterAvacticaUrl()));
+    final Properties properties = getAvaticaConnectionPropertiesForAdmin();
+    testAvaticaQuery(properties, getRouterAvacticaUrl());
+    testAvaticaQuery(properties, StringUtils.maybeRemoveTrailingSlash(getRouterAvacticaUrl()));
   }
 
   @Test
   public void test_avaticaQueryAuthFailure_broker() throws Exception
   {
-    testAvaticaAuthFailure(getBrokerAvacticaUrl());
+    final Properties properties = getAvaticaConnectionPropertiesForWrongAdmin();
+    testAvaticaAuthFailure(properties, getBrokerAvacticaUrl());
   }
 
   @Test
   public void test_avaticaQueryAuthFailure_router() throws Exception
   {
-    testAvaticaAuthFailure(getRouterAvacticaUrl());
+    final Properties properties = getAvaticaConnectionPropertiesForWrongAdmin();
+    testAvaticaAuthFailure(properties, getRouterAvacticaUrl());
+  }
+
+  @Test
+  public void test_avaticaQueryWithContext_datasourceOnlyUser_fail() throws Exception
+  {
+    final Properties properties = getAvaticaConnectionPropertiesForUser("datasourceOnlyUser", "helloworld");
+    properties.setProperty("auth_test_ctx", "should-be-denied");
+    testAvaticaAuthFailure(properties, getRouterAvacticaUrl());
+  }
+
+  @Test
+  public void test_avaticaQueryWithContext_datasourceAndContextParamsUser_succeed() throws Exception
+  {
+    final Properties properties = getAvaticaConnectionPropertiesForUser("datasourceAndContextParamsUser", "helloworld");
+    properties.setProperty("auth_test_ctx", "should-be-allowed");
+    testAvaticaAuthFailure(properties, getRouterAvacticaUrl());
+  }
+
+  @Test
+  public void test_sqlQueryWithContext_datasourceOnlyUser_fail() throws Exception
+  {
+    final String query = "select count(*) from auth_test";
+    StatusResponseHolder responseHolder = makeSQLQueryRequest(
+        datasourceOnlyUserClient,
+        query,
+        ImmutableMap.of("auth_test_ctx", "should-be-denied"),
+        HttpResponseStatus.FORBIDDEN
+    );
+  }
+
+  @Test
+  public void test_sqlQueryWithContext_datasourceAndContextParamsUser_succeed() throws Exception
+  {
+    final String query = "select count(*) from auth_test";
+    StatusResponseHolder responseHolder = makeSQLQueryRequest(
+        datasourceOnlyUserClient,
+        query,
+        ImmutableMap.of("auth_test_ctx", "should-be-allowed"),
+        HttpResponseStatus.OK
+    );
   }
 
   @Test
@@ -501,6 +556,7 @@ protected void setupHttpClientsAndUsers() throws Exception
   {
     setupHttpClients();
     setupDatasourceOnlyUser();
+    setupDatasourceAndContextParamsUser();
     setupDatasourceAndSysTableUser();
     setupDatasourceAndSysAndStateUser();
     setupSysTableAndStateOnlyUser();
@@ -538,11 +594,28 @@ protected void checkUnsecuredCoordinatorLoadQueuePath(HttpClient client)
     HttpUtil.makeRequest(client, HttpMethod.GET, config.getCoordinatorUrl() + "/druid/coordinator/v1/loadqueue", null);
   }
 
-  protected void testAvaticaQuery(String url)
+  private Properties getAvaticaConnectionPropertiesForAdmin()
+  {
+    return getAvaticaConnectionPropertiesForUser("admin", "priest");
+  }
+
+  private Properties getAvaticaConnectionPropertiesForWrongAdmin()
+  {
+    return getAvaticaConnectionPropertiesForUser("admin", "wrongpassword");
+  }
+
+  private Properties getAvaticaConnectionPropertiesForUser(String id, String password)
+  {
+    Properties connectionProperties = new Properties();
+    connectionProperties.setProperty("user", id);
+    connectionProperties.setProperty("password", password);
+    return connectionProperties;
+  }
+
+  protected void testAvaticaQuery(Properties connectionProperties, String url)
   {
     LOG.info("URL: " + url);
     try {
-      Properties connectionProperties = getAvaticaConnectionProperties();
       Connection connection = DriverManager.getConnection(url, connectionProperties);
       Statement statement = connection.createStatement();
       statement.setMaxRows(450);
@@ -557,11 +630,10 @@ protected void testAvaticaQuery(String url)
     }
   }
 
-  protected void testAvaticaAuthFailure(String url) throws Exception
+  protected void testAvaticaAuthFailure(Properties connectionProperties, String url) throws Exception
   {
     LOG.info("URL: " + url);
     try {
-      Properties connectionProperties = getAvaticaConnectionPropertiesFailure();
       Connection connection = DriverManager.getConnection(url, connectionProperties);
       Statement statement = connection.createStatement();
       statement.setMaxRows(450);
@@ -612,9 +684,20 @@ protected StatusResponseHolder makeSQLQueryRequest(
       String query,
       HttpResponseStatus expectedStatus
   ) throws Exception
+  {
+    return makeSQLQueryRequest(httpClient, query, ImmutableMap.of(), expectedStatus);
+  }
+
+  protected StatusResponseHolder makeSQLQueryRequest(
+      HttpClient httpClient,
+      String query,
+      Map<String, Object> context,
+      HttpResponseStatus expectedStatus
+  ) throws Exception
   {
     Map<String, Object> queryMap = ImmutableMap.of(
-        "query", query
+        "query", query,
+        "context", context
     );
     return HttpUtil.makeRequestWithExpectedStatus(
         httpClient,
@@ -758,7 +841,6 @@ protected void setupHttpClients() throws Exception
     setupTestSpecificHttpClients();
   }
 
-
   protected void setupCommonHttpClients()
   {
     adminClient = new CredentialedHttpClient(
@@ -771,6 +853,11 @@ protected void setupCommonHttpClients()
         httpClient
     );
 
+    datasourceAndContextParamsClient = new CredentialedHttpClient(
+        new BasicCredentials("datasourceAndContextParamsUser", "helloworld"),
+        httpClient
+    );
+
     datasourceAndSysUserClient = new CredentialedHttpClient(
         new BasicCredentials("datasourceAndSysUser", "helloworld"),
         httpClient

integration-tests/src/test/java/org/apache/druid/tests/security/ITBasicAuthConfigurationTest.java

@@ -36,7 +36,6 @@
 import org.testng.annotations.Test;
 
 import java.util.List;
-import java.util.Properties;
 
 @Test(groups = TestNGGroup.SECURITY)
 @Guice(moduleFactory = DruidTestModuleFactory.class)
@@ -81,6 +80,18 @@ protected void setupDatasourceOnlyUser() throws Exception
     );
   }
 
+  @Override
+  protected void setupDatasourceAndContextParamsUser() throws Exception
+  {
+    createUserAndRoleWithPermissions(
+        adminClient,
+        "datasourceAndContextParamsUser",
+        "helloworld",
+        "datasourceAndContextParamsRole",
+        DATASOURCE_QUERY_CONTEXT_PERMISSIONS
+    );
+  }
+
   @Override
   protected void setupDatasourceAndSysTableUser() throws Exception
   {
@@ -187,24 +198,6 @@ protected String getExpectedAvaticaAuthError()
     return EXPECTED_AVATICA_AUTH_ERROR;
   }
 
-  @Override
-  protected Properties getAvaticaConnectionProperties()
-  {
-    Properties connectionProperties = new Properties();
-    connectionProperties.setProperty("user", "admin");
-    connectionProperties.setProperty("password", "priest");
-    return connectionProperties;
-  }
-
-  @Override
-  protected Properties getAvaticaConnectionPropertiesFailure()
-  {
-    Properties connectionProperties = new Properties();
-    connectionProperties.setProperty("user", "admin");
-    connectionProperties.setProperty("password", "wrongpassword");
-    return connectionProperties;
-  }
-
   private void createUserAndRoleWithPermissions(
       HttpClient adminClient,
       String user,