Add `type` parameter support, for sorting, to the Query API Key API (#104625)

This adds support for the `type` parameter, for sorting, to the Query API key API.
The type for an API Key can currently be either `rest` or `cross_cluster`.
This was overlooked in #103695 when support for the `type` parameter
was first introduced only for querying.

docs/changelog/104625.yaml

@@ -0,0 +1,6 @@
+pr: 104625
+summary: "Add support for the `type` parameter, for sorting, to the Query API Key\
+  \ API"
+area: Security
+type: enhancement
+issues: []

x-pack/plugin/security/qa/security-trial/src/javaRestTest/java/org/elasticsearch/xpack/security/apikey/ApiKeyRestIT.java

@@ -778,6 +778,58 @@ public class ApiKeyRestIT extends SecurityOnTrialLicenseRestTestCase {
         assertThat(queryResponse.evaluate("api_keys.0.name"), is("test-cross-key-query-2"));
     }
 
+    public void testSortApiKeysByType() throws IOException {
+        List<String> apiKeyIds = new ArrayList<>(2);
+        // create regular api key
+        EncodedApiKey encodedApiKey = createApiKey("test-rest-key", Map.of("tag", "rest"));
+        apiKeyIds.add(encodedApiKey.id());
+        // create cross-cluster key
+        Request createRequest = new Request("POST", "/_security/cross_cluster/api_key");
+        createRequest.setJsonEntity("""
+            {
+              "name": "test-cross-key",
+              "access": {
+                "search": [
+                  {
+                    "names": [ "whatever" ]
+                  }
+                ]
+              },
+              "metadata": { "tag": "cross" }
+            }""");
+        setUserForRequest(createRequest, MANAGE_SECURITY_USER, END_USER_PASSWORD);
+        ObjectPath createResponse = assertOKAndCreateObjectPath(client().performRequest(createRequest));
+        apiKeyIds.add(createResponse.evaluate("id"));
+
+        // desc sort all (2) keys - by type
+        Request queryRequest = new Request("GET", "/_security/_query/api_key");
+        queryRequest.addParameter("with_limited_by", String.valueOf(randomBoolean()));
+        queryRequest.setJsonEntity("""
+            {"sort":[{"type":{"order":"desc"}}]}""");
+        setUserForRequest(queryRequest, MANAGE_API_KEY_USER, END_USER_PASSWORD);
+        ObjectPath queryResponse = assertOKAndCreateObjectPath(client().performRequest(queryRequest));
+        assertThat(queryResponse.evaluate("total"), is(2));
+        assertThat(queryResponse.evaluate("count"), is(2));
+        assertThat(queryResponse.evaluate("api_keys.0.id"), is(apiKeyIds.get(0)));
+        assertThat(queryResponse.evaluate("api_keys.0.type"), is("rest"));
+        assertThat(queryResponse.evaluate("api_keys.1.id"), is(apiKeyIds.get(1)));
+        assertThat(queryResponse.evaluate("api_keys.1.type"), is("cross_cluster"));
+
+        // asc sort all (2) keys - by type
+        queryRequest = new Request("GET", "/_security/_query/api_key");
+        queryRequest.addParameter("with_limited_by", String.valueOf(randomBoolean()));
+        queryRequest.setJsonEntity("""
+            {"sort":[{"type":{"order":"asc"}}]}""");
+        setUserForRequest(queryRequest, MANAGE_API_KEY_USER, END_USER_PASSWORD);
+        queryResponse = assertOKAndCreateObjectPath(client().performRequest(queryRequest));
+        assertThat(queryResponse.evaluate("total"), is(2));
+        assertThat(queryResponse.evaluate("count"), is(2));
+        assertThat(queryResponse.evaluate("api_keys.0.id"), is(apiKeyIds.get(1)));
+        assertThat(queryResponse.evaluate("api_keys.0.type"), is("cross_cluster"));
+        assertThat(queryResponse.evaluate("api_keys.1.id"), is(apiKeyIds.get(0)));
+        assertThat(queryResponse.evaluate("api_keys.1.type"), is("rest"));
+    }
+
     public void testCreateCrossClusterApiKey() throws IOException {
         final Request createRequest = new Request("POST", "/_security/cross_cluster/api_key");
         createRequest.setJsonEntity("""

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

@@ -28,6 +28,7 @@ import org.elasticsearch.xpack.security.support.ApiKeyFieldNameTranslators;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
 
 import static org.elasticsearch.xpack.security.support.SecuritySystemIndices.SECURITY_MAIN_ALIAS;
 
@@ -81,22 +82,25 @@ public final class TransportQueryApiKeyAction extends TransportAction<QueryApiKe
         }
 
         final AtomicBoolean accessesApiKeyTypeField = new AtomicBoolean(false);
-        final ApiKeyBoolQueryBuilder apiKeyBoolQueryBuilder = ApiKeyBoolQueryBuilder.build(request.getQueryBuilder(), fieldName -> {
+        searchSourceBuilder.query(ApiKeyBoolQueryBuilder.build(request.getQueryBuilder(), fieldName -> {
             if (API_KEY_TYPE_RUNTIME_MAPPING_FIELD.equals(fieldName)) {
                 accessesApiKeyTypeField.set(true);
             }
-        }, request.isFilterForCurrentUser() ? authentication : null);
-        searchSourceBuilder.query(apiKeyBoolQueryBuilder);
+        }, request.isFilterForCurrentUser() ? authentication : null));
+
+        if (request.getFieldSortBuilders() != null) {
+            translateFieldSortBuilders(request.getFieldSortBuilders(), searchSourceBuilder, fieldName -> {
+                if (API_KEY_TYPE_RUNTIME_MAPPING_FIELD.equals(fieldName)) {
+                    accessesApiKeyTypeField.set(true);
+                }
+            });
+        }
 
         // only add the query-level runtime field to the search request if it's actually referring the "type" field
         if (accessesApiKeyTypeField.get()) {
             searchSourceBuilder.runtimeMappings(API_KEY_TYPE_RUNTIME_MAPPING);
         }
 
-        if (request.getFieldSortBuilders() != null) {
-            translateFieldSortBuilders(request.getFieldSortBuilders(), searchSourceBuilder);
-        }
-
         if (request.getSearchAfterBuilder() != null) {
             searchSourceBuilder.searchAfter(request.getSearchAfterBuilder().getSortValues());
         }
@@ -106,7 +110,11 @@ public final class TransportQueryApiKeyAction extends TransportAction<QueryApiKe
     }
 
     // package private for testing
-    static void translateFieldSortBuilders(List<FieldSortBuilder> fieldSortBuilders, SearchSourceBuilder searchSourceBuilder) {
+    static void translateFieldSortBuilders(
+        List<FieldSortBuilder> fieldSortBuilders,
+        SearchSourceBuilder searchSourceBuilder,
+        Consumer<String> fieldNameVisitor
+    ) {
         fieldSortBuilders.forEach(fieldSortBuilder -> {
             if (fieldSortBuilder.getNestedSort() != null) {
                 throw new IllegalArgumentException("nested sorting is not supported for API Key query");
@@ -115,6 +123,7 @@ public final class TransportQueryApiKeyAction extends TransportAction<QueryApiKe
                 searchSourceBuilder.sort(fieldSortBuilder);
             } else {
                 final String translatedFieldName = ApiKeyFieldNameTranslators.translate(fieldSortBuilder.getFieldName());
+                fieldNameVisitor.accept(translatedFieldName);
                 if (translatedFieldName.equals(fieldSortBuilder.getFieldName())) {
                     searchSourceBuilder.sort(fieldSortBuilder);
                 } else {

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

@@ -14,15 +14,18 @@ import org.elasticsearch.search.sort.SortMode;
 import org.elasticsearch.search.sort.SortOrder;
 import org.elasticsearch.test.ESTestCase;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
 import java.util.stream.IntStream;
 
+import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.equalTo;
 
 public class TransportQueryApiKeyActionTests extends ESTestCase {
 
     public void testTranslateFieldSortBuilders() {
+        final String metadataField = randomAlphaOfLengthBetween(3, 8);
         final List<String> fieldNames = List.of(
             "_doc",
             "username",
@@ -30,14 +33,16 @@ public class TransportQueryApiKeyActionTests extends ESTestCase {
             "name",
             "creation",
             "expiration",
+            "type",
             "invalidated",
-            "metadata." + randomAlphaOfLengthBetween(3, 8)
+            "metadata." + metadataField
         );
 
         final List<FieldSortBuilder> originals = fieldNames.stream().map(this::randomFieldSortBuilderWithName).toList();
 
+        List<String> sortFields = new ArrayList<>();
         final SearchSourceBuilder searchSourceBuilder = SearchSourceBuilder.searchSource();
-        TransportQueryApiKeyAction.translateFieldSortBuilders(originals, searchSourceBuilder);
+        TransportQueryApiKeyAction.translateFieldSortBuilders(originals, searchSourceBuilder, sortFields::add);
 
         IntStream.range(0, originals.size()).forEach(i -> {
             final FieldSortBuilder original = originals.get(i);
@@ -57,6 +62,8 @@ public class TransportQueryApiKeyActionTests extends ESTestCase {
                     assertThat(translated.getFieldName(), equalTo("api_key_invalidated"));
                 } else if (original.getFieldName().startsWith("metadata.")) {
                     assertThat(translated.getFieldName(), equalTo("metadata_flattened." + original.getFieldName().substring(9)));
+                } else if ("type".equals(original.getFieldName())) {
+                    assertThat(translated.getFieldName(), equalTo("runtime_key_type"));
                 } else {
                     fail("unrecognized field name: [" + original.getFieldName() + "]");
                 }
@@ -68,6 +75,19 @@ public class TransportQueryApiKeyActionTests extends ESTestCase {
                 assertThat(translated.sortMode(), equalTo(original.sortMode()));
             }
         });
+        assertThat(
+            sortFields,
+            containsInAnyOrder(
+                "creator.principal",
+                "creator.realm",
+                "name",
+                "creation_time",
+                "expiration_time",
+                "runtime_key_type",
+                "api_key_invalidated",
+                "metadata_flattened." + metadataField
+            )
+        );
     }
 
     public void testNestedSortingIsNotAllowed() {
@@ -75,7 +95,11 @@ public class TransportQueryApiKeyActionTests extends ESTestCase {
         fieldSortBuilder.setNestedSort(new NestedSortBuilder("name"));
         final IllegalArgumentException e = expectThrows(
             IllegalArgumentException.class,
-            () -> TransportQueryApiKeyAction.translateFieldSortBuilders(List.of(fieldSortBuilder), SearchSourceBuilder.searchSource())
+            () -> TransportQueryApiKeyAction.translateFieldSortBuilders(
+                List.of(fieldSortBuilder),
+                SearchSourceBuilder.searchSource(),
+                ignored -> {}
+            )
         );
         assertThat(e.getMessage(), equalTo("nested sorting is not supported for API Key query"));
     }