Find an appropriate key based on the given KID and JWA (#21160)

* keycloak-20847 Find an appropriate key based on the given KID and JWA. Prefers matching on both inputs but will match on partials if found. Or return the first key if a match is not found.

Mark Key as fallback if it is the singular client certificate to be used for signed JWT authentication.

* Update js/apps/admin-ui/public/locales/en/clients.json

Co-authored-by: Marek Posolda <mposolda@gmail.com>

* Updating boolean variable name based on suggestions by Marek.

* Adding integration test specifically for the JWT parameters for regression #20847.

---------

Co-authored-by: Marek Posolda <mposolda@gmail.com>

core/src/main/java/org/keycloak/crypto/KeyWrapper.java

@@ -48,6 +48,7 @@ public class KeyWrapper {
     private Key privateKey;
     private X509Certificate certificate;
     private List<X509Certificate> certificateChain;
+    private boolean isDefaultClientCertificate;
 
     public String getProviderId() {
         return providerId;
@@ -167,6 +168,14 @@ public class KeyWrapper {
         this.certificateChain = certificateChain;
     }
 
+    public boolean isDefaultClientCertificate() {
+        return isDefaultClientCertificate;
+    }
+
+    public void setIsDefaultClientCertificate(boolean isDefaultClientCertificate) {
+        this.isDefaultClientCertificate = isDefaultClientCertificate;
+    }
+
     public KeyWrapper cloneKey() {
         KeyWrapper key = new KeyWrapper();
         key.providerId = this.providerId;
@@ -183,6 +192,7 @@ public class KeyWrapper {
         if (this.certificateChain != null) {
             key.certificateChain = new ArrayList<>(this.certificateChain);
         }
+        key.isDefaultClientCertificate = this.isDefaultClientCertificate;
         return key;
     }
 }

core/src/main/java/org/keycloak/crypto/PublicKeysWrapper.java

@@ -21,7 +21,9 @@ package org.keycloak.crypto;
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -46,10 +48,28 @@ public class PublicKeysWrapper {
                 .collect(Collectors.toList());
     }
 
+    /**
+     * Find an appropriate key given a KID and algorithm.
+     * Prefer matching on both parameters, but may partially match on KID only. Or if KID is not provided, the
+     * algorithm. Will use a flagged default client certificate otherwise, if a match is not found.
+     * @param kid rfc7517 KID parameter
+     * @param alg rfc7517 alg parameter
+     * @return {@link KeyWrapper} matching given parameters
+     */
     public KeyWrapper getKeyByKidAndAlg(String kid, String alg) {
-        return keys.stream()
+
-                .filter(keyWrapper -> kid == null || kid.equals(keyWrapper.getKid()))
+        Stream<KeyWrapper> potentialMatches = Stream.concat(
-                .filter(keyWrapper -> alg == null || alg.equals(keyWrapper.getAlgorithmOrDefault()) || (keyWrapper.getAlgorithm() == null && kid != null))
+            keys.stream().filter(keyWrapper -> Objects.equals(kid, keyWrapper.getKid()) && Objects.equals(alg, keyWrapper.getAlgorithm())),
-                .findFirst().orElse(null);
+            keys.stream().filter(keyWrapper -> Objects.equals(kid, keyWrapper.getKid())));
+
+        if (kid == null) {
+            potentialMatches = Stream.of(
+                    potentialMatches,
+                    keys.stream().filter(keyWrapper -> Objects.equals(alg, keyWrapper.getAlgorithmOrDefault())),
+                    keys.stream().filter(KeyWrapper::isDefaultClientCertificate)
+                ).flatMap(i -> i);
+        }
+
+        return potentialMatches.findFirst().orElse(null);
     }
 }

js/apps/admin-ui/public/locales/en/clients.json

@@ -519,7 +519,7 @@
   "browserFlow": "Browser Flow",
   "directGrant": "Direct Grant Flow",
   "jwksUrlConfig": "JWKS URL configs",
-  "keysIntro": "If \"Use JWKS URL switch\" is on, you need to fill a valid JWKS URL. After saving, admin can download keys from the JWKS URL or keys will be downloaded automatically by Keycloak server when see the stuff signed by the unknown KID",
+  "keysIntro": "If \"Use JWKS URL switch\" is on, you need to fill a valid JWKS URL. After saving, admin can download keys from the JWKS URL or keys will be downloaded automatically by Keycloak server when an unknown KID is seen during client authentication.",
   "useJwksUrl": "Use JWKS URL",
   "certificate": "Certificate",
   "jwksUrl": "JWKS URL",

model/infinispan/src/main/java/org/keycloak/keys/infinispan/InfinispanPublicKeyStorageProvider.java

@@ -124,9 +124,14 @@ public class InfinispanPublicKeyStorageProvider implements PublicKeyStorageProvi
 
     @Override
     public KeyWrapper getPublicKey(String modelKey, String kid, String algorithm, PublicKeyLoader loader) {
-        // Check if key is in cache
         PublicKeysEntry entry = keys.get(modelKey);
-        if (entry != null) {
+        int lastRequestTime = entry==null ? 0 : entry.getLastRequestTime();
+        int currentTime = Time.currentTime();
+        boolean isSendingRequestAllowed = currentTime > lastRequestTime + minTimeBetweenRequests;
+
+        // Check if key is in cache, but only if KID is provided or if the key cache has been loaded recently,
+        // in order to get a key based on partial match with alg param.
+        if (entry != null && (kid != null || !isSendingRequestAllowed)) {
             KeyWrapper publicKey = entry.getCurrentKeys().getKeyByKidAndAlg(kid, algorithm);
             if (publicKey != null) {
                 // return a copy of the key to not modify the cached one
@@ -134,12 +139,8 @@ public class InfinispanPublicKeyStorageProvider implements PublicKeyStorageProvi
             }
         }
 
-        int lastRequestTime = entry==null ? 0 : entry.getLastRequestTime();
-        int currentTime = Time.currentTime();
-
         // Check if we are allowed to send request
-        if (currentTime > lastRequestTime + minTimeBetweenRequests) {
+        if (isSendingRequestAllowed) {
-
             WrapperCallable wrapperCallable = new WrapperCallable(modelKey, loader);
             FutureTask<PublicKeysEntry> task = new FutureTask<>(wrapperCallable);
             FutureTask<PublicKeysEntry> existing = tasksInProgress.putIfAbsent(modelKey, task);

services/src/main/java/org/keycloak/keys/loader/ClientPublicKeyLoader.java

@@ -114,6 +114,7 @@ public class ClientPublicKeyLoader implements PublicKeyLoader {
             keyWrapper.setPublicKey(clientCert.getPublicKey());
             keyWrapper.setType(clientCert.getPublicKey().getAlgorithm());
             keyWrapper.setCertificate(clientCert);
+            keyWrapper.setIsDefaultClientCertificate(true);
         } else {
             PublicKey publicKey = KeycloakModelUtils.getPublicKey(encodedPublicKey);
             // Check if we have kid in DB, generate otherwise

testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/util/OAuthClient.java

@@ -2094,7 +2094,7 @@ public class OAuthClient {
         }
     }
 
-    private KeyWrapper getRealmPublicKey(String realm, String algoritm, String kid) {
+    private KeyWrapper getRealmPublicKey(String realm, String algorithm, String kid) {
         boolean loadedKeysFromServer = false;
         JSONWebKeySet jsonWebKeySet = publicKeys.get(realm);
         if (jsonWebKeySet == null) {
@@ -2103,17 +2103,17 @@ public class OAuthClient {
             loadedKeysFromServer = true;
         }
 
-        KeyWrapper key = findKey(jsonWebKeySet, algoritm, kid);
+        KeyWrapper key = findKey(jsonWebKeySet, algorithm, kid);
 
         if (key == null && !loadedKeysFromServer) {
             jsonWebKeySet = getRealmKeys(realm);
             publicKeys.put(realm, jsonWebKeySet);
 
-            key = findKey(jsonWebKeySet, algoritm, kid);
+            key = findKey(jsonWebKeySet, algorithm, kid);
         }
 
         if (key == null) {
-            throw new RuntimeException("Public key for realm:" + realm + ", algorithm: " + algoritm + " not found");
+            throw new RuntimeException("Public key for realm:" + realm + ", algorithm: " + algorithm + " not found");
         }
 
         return key;
@@ -2128,9 +2128,9 @@ public class OAuthClient {
         }
     }
 
-    private KeyWrapper findKey(JSONWebKeySet jsonWebKeySet, String algoritm, String kid) {
+    private KeyWrapper findKey(JSONWebKeySet jsonWebKeySet, String algorithm, String kid) {
         for (JWK k : jsonWebKeySet.getKeys()) {
-            if (k.getKeyId().equals(kid) && k.getAlgorithm().equals(algoritm)) {
+            if (k.getKeyId().equals(kid) && k.getAlgorithm().equals(algorithm)) {
                 PublicKey publicKey = JWKParser.create(k).toPublicKey();
 
                 KeyWrapper key = new KeyWrapper();

testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/oauth/ClientAuthSignedJWTTest.java

@@ -1040,6 +1040,35 @@ public class ClientAuthSignedJWTTest extends AbstractKeycloakTest {
         }
     }
 
+    @Test
+    public void testJWTAuthForClientCertWithOnlyAlgProvided() throws Exception {
+        ClientRepresentation clientRepresentation = app2;
+        ClientResource clientResource = getClient(testRealm.getRealm(), clientRepresentation.getId());
+        clientRepresentation = clientResource.toRepresentation();
+
+        try {
+            KeyPair keyPair = setupJwksUrl(Algorithm.ES512, clientRepresentation, clientResource);
+            PrivateKey privateKey = keyPair.getPrivate();
+            JsonWebToken assertion = createRequestToken(app2.getClientId(), getRealmInfoUrl());
+
+            SignatureSignerContext signer = oauth.createSigner(privateKey, null,  Algorithm.ES512);
+            String jws = new JWSBuilder().jsonContent(assertion).sign(signer);
+
+            List<NameValuePair> parameters = new LinkedList<>();
+            parameters.add(new BasicNameValuePair(OAuth2Constants.GRANT_TYPE, OAuth2Constants.CLIENT_CREDENTIALS));
+            parameters
+                    .add(new BasicNameValuePair(OAuth2Constants.CLIENT_ASSERTION_TYPE, OAuth2Constants.CLIENT_ASSERTION_TYPE_JWT));
+            parameters.add(new BasicNameValuePair(OAuth2Constants.CLIENT_ASSERTION, jws));
+
+            try (CloseableHttpResponse resp = sendRequest(oauth.getServiceAccountUrl(), parameters)) {
+                OAuthClient.AccessTokenResponse response = new OAuthClient.AccessTokenResponse(resp);
+                assertNotNull(response.getAccessToken());
+            }
+        } finally {
+            revertJwksUriSettings(clientRepresentation, clientResource);
+        }
+    }
+
     @Test
     public void testAssertionInvalidNotBefore() throws Exception {
         String invalidJwt = getClient1SignedJWT();
@@ -1205,7 +1234,9 @@ public class ClientAuthSignedJWTTest extends AbstractKeycloakTest {
 
     @Test
     public void testCodeToTokenRequestFailureRS256() throws Exception {
-        testCodeToTokenRequestFailure(Algorithm.RS256, OAuthErrorException.INVALID_CLIENT, "client_credentials_setup_required");
+        testCodeToTokenRequestFailure(Algorithm.RS256,
+                OAuthErrorException.INVALID_CLIENT,
+                AuthenticationFlowError.CLIENT_CREDENTIALS_SETUP_REQUIRED.toString().toLowerCase());
     }
 
     @Test
@@ -1218,7 +1249,7 @@ public class ClientAuthSignedJWTTest extends AbstractKeycloakTest {
             OIDCAdvancedConfigWrapper.fromClientRepresentation(clientRep).setTokenEndpointAuthSigningAlg(Algorithm.ES256);
             clientResource.update(clientRep);
 
-            testCodeToTokenRequestFailure(Algorithm.RS256, "invalid_client", "invalid_client_credentials");
+            testCodeToTokenRequestFailure(Algorithm.RS256, "invalid_client", Errors.INVALID_CLIENT_CREDENTIALS);
         } catch (Exception e) {
             Assert.fail();
         } finally {