Enhance API Key Querying

docs/changelog/103192.yaml

@@ -0,0 +1,5 @@
+pr: 103192
+summary: Query api key api improvements
+area: Security
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -192,6 +192,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion INFERENCE_SERVICE_EMBEDDING_SIZE_ADDED = def(8_559_00_0);
     public static final TransportVersion ENRICH_ELASTICSEARCH_VERSION_REMOVED = def(8_560_00_0);
     public static final TransportVersion NODE_STATS_REQUEST_SIMPLIFIED = def(8_561_00_0);
+    public static final TransportVersion QUERY_API_KEY_AGGS_ADDED = def(8_562_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/search/aggregations/bucket/filter/FilterAggregationBuilder.java

@@ -41,7 +41,7 @@
 public class FilterAggregationBuilder extends AbstractAggregationBuilder<FilterAggregationBuilder> {
     public static final String NAME = "filter";
 
-    private final QueryBuilder filter;
+    private QueryBuilder filter;
 
     /**
      * @param name
@@ -59,6 +59,14 @@ public FilterAggregationBuilder(String name, QueryBuilder filter) {
         this.filter = filter;
     }
 
+    public FilterAggregationBuilder filter(QueryBuilder filter) {
+        if (filter == null) {
+            throw new IllegalArgumentException("[filter] must not be null: [" + name + "]");
+        }
+        this.filter = filter;
+        return this;
+    }
+
     protected FilterAggregationBuilder(
         FilterAggregationBuilder clone,
         AggregatorFactories.Builder factoriesBuilder,

server/src/main/java/org/elasticsearch/search/aggregations/bucket/filter/FiltersAggregator.java

@@ -61,7 +61,7 @@ public abstract class FiltersAggregator extends BucketsAggregator {
 
     public static class KeyedFilter implements Writeable, ToXContentFragment {
         private final String key;
-        private final QueryBuilder filter;
+        private QueryBuilder filter;
 
         public KeyedFilter(String key, QueryBuilder filter) {
             if (key == null) {
@@ -92,6 +92,10 @@ public String key() {
             return key;
         }
 
+        public void filter(QueryBuilder filter) {
+            this.filter = filter;
+        }
+
         public QueryBuilder filter() {
             return filter;
         }

server/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceAggregationBuilder.java

@@ -192,7 +192,6 @@ public Set<String> metricNames() {
     private String format = null;
     private Object missing = null;
     private ZoneId timeZone = null;
-    protected ValuesSourceConfig config;
 
     protected ValuesSourceAggregationBuilder(String name) {
         super(name);
@@ -209,7 +208,6 @@ protected ValuesSourceAggregationBuilder(
         this.format = clone.format;
         this.missing = clone.missing;
         this.timeZone = clone.timeZone;
-        this.config = clone.config;
         this.script = clone.script;
     }
 

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -728,6 +728,14 @@ public SearchSourceBuilder collapse(CollapseBuilder collapse) {
         return this;
     }
 
+    public SearchSourceBuilder aggregationsBuilder(AggregatorFactories.Builder aggregations) {
+        if (this.aggregations != null) {
+            throw new IllegalStateException("cannot override aggregation build [");
+        }
+        this.aggregations = aggregations;
+        return this;
+    }
+
     /**
      * Add an aggregation to perform as part of the search.
      */

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/QueryApiKeyRequest.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.searchafter.SearchAfterBuilder;
 import org.elasticsearch.search.sort.FieldSortBuilder;
 
@@ -27,6 +28,8 @@ public final class QueryApiKeyRequest extends ActionRequest {
     @Nullable
     private final QueryBuilder queryBuilder;
     @Nullable
+    private final AggregatorFactories.Builder aggsBuilder;
+    @Nullable
     private final Integer from;
     @Nullable
     private final Integer size;
@@ -42,18 +45,20 @@ public QueryApiKeyRequest() {
     }
 
     public QueryApiKeyRequest(QueryBuilder queryBuilder) {
-        this(queryBuilder, null, null, null, null, false);
+        this(queryBuilder, null, null, null, null, null, false);
     }
 
     public QueryApiKeyRequest(
         @Nullable QueryBuilder queryBuilder,
+        @Nullable AggregatorFactories.Builder aggsBuilder,
         @Nullable Integer from,
         @Nullable Integer size,
         @Nullable List<FieldSortBuilder> fieldSortBuilders,
         @Nullable SearchAfterBuilder searchAfterBuilder,
         boolean withLimitedBy
     ) {
         this.queryBuilder = queryBuilder;
+        this.aggsBuilder = aggsBuilder;
         this.from = from;
         this.size = size;
         this.fieldSortBuilders = fieldSortBuilders;
@@ -77,12 +82,21 @@ public QueryApiKeyRequest(StreamInput in) throws IOException {
         } else {
             this.withLimitedBy = false;
         }
+        if (in.getTransportVersion().onOrAfter(TransportVersions.QUERY_API_KEY_AGGS_ADDED)) {
+            this.aggsBuilder = in.readOptionalWriteable(AggregatorFactories.Builder::new);
+        } else {
+            this.aggsBuilder = null;
+        }
     }
 
     public QueryBuilder getQueryBuilder() {
         return queryBuilder;
     }
 
+    public AggregatorFactories.Builder getAggsBuilder() {
+        return aggsBuilder;
+    }
+
     public Integer getFrom() {
         return from;
     }
@@ -139,5 +153,8 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_5_0)) {
             out.writeBoolean(withLimitedBy);
         }
+        if (out.getTransportVersion().onOrAfter(TransportVersions.QUERY_API_KEY_AGGS_ADDED)) {
+            out.writeOptionalWriteable(aggsBuilder);
+        }
     }
 }

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/security/action/apikey/QueryApiKeyRequestTests.java

@@ -65,6 +65,7 @@ public void testReadWrite() throws IOException {
 
         final QueryApiKeyRequest request3 = new QueryApiKeyRequest(
             QueryBuilders.matchAllQuery(),
+            null,
             42,
             20,
             List.of(
@@ -91,6 +92,7 @@ public void testReadWrite() throws IOException {
 
     public void testValidate() {
         final QueryApiKeyRequest request1 = new QueryApiKeyRequest(
+            null,
             null,
             randomIntBetween(0, Integer.MAX_VALUE),
             randomIntBetween(0, Integer.MAX_VALUE),
@@ -101,6 +103,7 @@ public void testValidate() {
         assertThat(request1.validate(), nullValue());
 
         final QueryApiKeyRequest request2 = new QueryApiKeyRequest(
+            null,
             null,
             randomIntBetween(Integer.MIN_VALUE, -1),
             randomIntBetween(0, Integer.MAX_VALUE),
@@ -111,6 +114,7 @@ public void testValidate() {
         assertThat(request2.validate().getMessage(), containsString("[from] parameter cannot be negative"));
 
         final QueryApiKeyRequest request3 = new QueryApiKeyRequest(
+            null,
             null,
             randomIntBetween(0, Integer.MAX_VALUE),
             randomIntBetween(Integer.MIN_VALUE, -1),

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/security/authz/privilege/ManageOwnApiKeyClusterPrivilegeTests.java

@@ -259,7 +259,7 @@ public void testCheckQueryApiKeyRequest() {
         final ClusterPermission clusterPermission = ManageOwnApiKeyClusterPrivilege.INSTANCE.buildPermission(ClusterPermission.builder())
             .build();
 
-        final QueryApiKeyRequest queryApiKeyRequest = new QueryApiKeyRequest(null, null, null, null, null, randomBoolean());
+        final QueryApiKeyRequest queryApiKeyRequest = new QueryApiKeyRequest(null, null, null, null, null, null, randomBoolean());
         if (randomBoolean()) {
             queryApiKeyRequest.setFilterForCurrentUser();
         }
@@ -278,7 +278,7 @@ public void testAuthenticationWithApiKeyAllowsDeniesQueryApiKeyWithLimitedBy() {
             .build();
 
         final boolean withLimitedBy = randomBoolean();
-        final QueryApiKeyRequest queryApiKeyRequest = new QueryApiKeyRequest(null, null, null, null, null, withLimitedBy);
+        final QueryApiKeyRequest queryApiKeyRequest = new QueryApiKeyRequest(null, null, null, null, null, null, withLimitedBy);
         queryApiKeyRequest.setFilterForCurrentUser();
         assertThat(
             clusterPermission.check(QueryApiKeyAction.NAME, queryApiKeyRequest, AuthenticationTestHelper.builder().apiKey().build(false)),

x-pack/plugin/security/build.gradle

@@ -27,6 +27,9 @@ dependencies {
   testImplementation project(path: xpackModule('wildcard'))
   testImplementation project(path: ':modules:legacy-geo')
   testImplementation project(path: ':modules:percolator')
+//  Conflict found for the following module:
+//  - org.ow2.asm:asm between versions 8.0.1 and 7.2
+//  testImplementation project(path: ':modules:lang-painless')
   testImplementation project(path: xpackModule('sql:sql-action'))
   testImplementation project(path: ':modules:analysis-common')
   testImplementation project(path: ':modules:reindex')

x-pack/plugin/security/src/internalClusterTest/java/org/elasticsearch/xpack/security/authc/ApiKeyIntegTests.java

@@ -2852,7 +2852,7 @@ private ApiKey getApiKeyInfo(Client client, String apiKeyId, boolean withLimited
             final PlainActionFuture<QueryApiKeyResponse> future = new PlainActionFuture<>();
             client.execute(
                 QueryApiKeyAction.INSTANCE,
-                new QueryApiKeyRequest(QueryBuilders.idsQuery().addIds(apiKeyId), null, null, null, null, withLimitedBy),
+                new QueryApiKeyRequest(QueryBuilders.idsQuery().addIds(apiKeyId), null, null, null, null, null, withLimitedBy),
                 future
             );
             final QueryApiKeyResponse queryApiKeyResponse = future.actionGet();
@@ -2871,7 +2871,7 @@ private ApiKey[] getAllApiKeyInfo(Client client, boolean withLimitedBy) {
             final PlainActionFuture<QueryApiKeyResponse> future = new PlainActionFuture<>();
             client.execute(
                 QueryApiKeyAction.INSTANCE,
-                new QueryApiKeyRequest(QueryBuilders.matchAllQuery(), null, 1000, null, null, withLimitedBy),
+                new QueryApiKeyRequest(QueryBuilders.matchAllQuery(), null, null, 1000, null, null, withLimitedBy),
                 future
             );
             final QueryApiKeyResponse queryApiKeyResponse = future.actionGet();

x-pack/plugin/security/src/internalClusterTest/java/org/elasticsearch/xpack/security/authc/apikey/ApiKeySingleNodeTests.java

@@ -625,6 +625,7 @@ public void testCreateCrossClusterApiKey() throws IOException {
             null,
             null,
             null,
+            null,
             randomBoolean()
         );
         final QueryApiKeyResponse queryApiKeyResponse = client().execute(QueryApiKeyAction.INSTANCE, queryApiKeyRequest).actionGet();
@@ -722,6 +723,7 @@ public void testUpdateCrossClusterApiKey() throws IOException {
             null,
             null,
             null,
+            null,
             randomBoolean()
         );
         final QueryApiKeyResponse queryApiKeyResponse = client().execute(QueryApiKeyAction.INSTANCE, queryApiKeyRequest).actionGet();

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/action/apikey/TransportQueryApiKeyAction.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.action.support.HandledTransportAction;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
+import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.sort.FieldSortBuilder;
 import org.elasticsearch.tasks.Task;
@@ -23,15 +24,23 @@
 import org.elasticsearch.xpack.core.security.action.apikey.QueryApiKeyResponse;
 import org.elasticsearch.xpack.core.security.authc.Authentication;
 import org.elasticsearch.xpack.security.authc.ApiKeyService;
+import org.elasticsearch.xpack.security.support.ApiKeyAggregationsBuilder;
 import org.elasticsearch.xpack.security.support.ApiKeyBoolQueryBuilder;
 import org.elasticsearch.xpack.security.support.ApiKeyFieldNameTranslators;
 
 import java.util.List;
+import java.util.Map;
 
 import static org.elasticsearch.xpack.security.support.SecuritySystemIndices.SECURITY_MAIN_ALIAS;
 
 public final class TransportQueryApiKeyAction extends HandledTransportAction<QueryApiKeyRequest, QueryApiKeyResponse> {
 
+    public static final String API_KEY_TYPE_RUNTIME_MAPPING_FIELD = "runtime_key_type";
+    private static final Map<String, Object> API_KEY_TYPE_RUNTIME_MAPPING = Map.of(
+        API_KEY_TYPE_RUNTIME_MAPPING_FIELD,
+        Map.of("type", "keyword", "script", Map.of("source", "emit(doc['type'].value ?: \"rest\");"))
+    );
+
     private final ApiKeyService apiKeyService;
     private final SecurityContext securityContext;
 
@@ -49,10 +58,13 @@ public TransportQueryApiKeyAction(
 
     @Override
     protected void doExecute(Task task, QueryApiKeyRequest request, ActionListener<QueryApiKeyResponse> listener) {
-        final Authentication authentication = securityContext.getAuthentication();
-        if (authentication == null) {
+        Authentication filteringAuthentication = securityContext.getAuthentication();
+        if (filteringAuthentication == null) {
             listener.onFailure(new IllegalStateException("authentication is required"));
         }
+        if (request.isFilterForCurrentUser() == false) {
+            filteringAuthentication = null;
+        }
 
         final SearchSourceBuilder searchSourceBuilder = SearchSourceBuilder.searchSource()
             .version(false)
@@ -68,10 +80,18 @@ protected void doExecute(Task task, QueryApiKeyRequest request, ActionListener<Q
 
         final ApiKeyBoolQueryBuilder apiKeyBoolQueryBuilder = ApiKeyBoolQueryBuilder.build(
             request.getQueryBuilder(),
-            request.isFilterForCurrentUser() ? authentication : null
+            filteringAuthentication
         );
         searchSourceBuilder.query(apiKeyBoolQueryBuilder);
 
+        AggregatorFactories.Builder aggsBuilder = request.getAggsBuilder();
+        // this modifies the aggsBuilder in-place
+        ApiKeyAggregationsBuilder.verifyRequested(aggsBuilder, filteringAuthentication);
+        searchSourceBuilder.aggregationsBuilder(aggsBuilder);
+
+        // doesn't work because of unmet test dependency on the lang-painless module
+        // searchSourceBuilder.runtimeMappings(API_KEY_TYPE_RUNTIME_MAPPING);
+
         if (request.getFieldSortBuilders() != null) {
             translateFieldSortBuilders(request.getFieldSortBuilders(), searchSourceBuilder);
         }

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestQueryApiKeyAction.java

@@ -17,6 +17,7 @@
 import org.elasticsearch.rest.Scope;
 import org.elasticsearch.rest.ServerlessScope;
 import org.elasticsearch.rest.action.RestToXContentListener;
+import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.searchafter.SearchAfterBuilder;
 import org.elasticsearch.search.sort.FieldSortBuilder;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
@@ -32,6 +33,7 @@
 import static org.elasticsearch.index.query.AbstractQueryBuilder.parseTopLevelQuery;
 import static org.elasticsearch.rest.RestRequest.Method.GET;
 import static org.elasticsearch.rest.RestRequest.Method.POST;
+import static org.elasticsearch.search.aggregations.AggregatorFactories.parseAggregators;
 import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg;
 
 /**
@@ -43,11 +45,19 @@ public final class RestQueryApiKeyAction extends ApiKeyBaseRestHandler {
     @SuppressWarnings("unchecked")
     private static final ConstructingObjectParser<Payload, Void> PARSER = new ConstructingObjectParser<>(
         "query_api_key_request_payload",
-        a -> new Payload((QueryBuilder) a[0], (Integer) a[1], (Integer) a[2], (List<FieldSortBuilder>) a[3], (SearchAfterBuilder) a[4])
+        a -> new Payload(
+            (QueryBuilder) a[0],
+            (AggregatorFactories.Builder) a[1],
+            (Integer) a[2],
+            (Integer) a[3],
+            (List<FieldSortBuilder>) a[4],
+            (SearchAfterBuilder) a[5]
+        )
     );
 
     static {
         PARSER.declareObject(optionalConstructorArg(), (p, c) -> parseTopLevelQuery(p), new ParseField("query"));
+        PARSER.declareObject(optionalConstructorArg(), (p, c) -> parseAggregators(p), new ParseField("aggs"));
         PARSER.declareInt(optionalConstructorArg(), new ParseField("from"));
         PARSER.declareInt(optionalConstructorArg(), new ParseField("size"));
         PARSER.declareObjectArray(optionalConstructorArg(), (p, c) -> {
@@ -98,20 +108,22 @@ protected RestChannelConsumer innerPrepareRequest(final RestRequest request, fin
             final Payload payload = PARSER.parse(request.contentOrSourceParamParser(), null);
             queryApiKeyRequest = new QueryApiKeyRequest(
                 payload.queryBuilder,
+                payload.aggsBuilder,
                 payload.from,
                 payload.size,
                 payload.fieldSortBuilders,
                 payload.searchAfterBuilder,
                 withLimitedBy
             );
         } else {
-            queryApiKeyRequest = new QueryApiKeyRequest(null, null, null, null, null, withLimitedBy);
+            queryApiKeyRequest = new QueryApiKeyRequest(null, null, null, null, null, null, withLimitedBy);
         }
         return channel -> client.execute(QueryApiKeyAction.INSTANCE, queryApiKeyRequest, new RestToXContentListener<>(channel));
     }
 
     private record Payload(
         @Nullable QueryBuilder queryBuilder,
+        @Nullable AggregatorFactories.Builder aggsBuilder,
         @Nullable Integer from,
         @Nullable Integer size,
         @Nullable List<FieldSortBuilder> fieldSortBuilders,