Support JWK Selection Strategy

Closes gh-16170

Signed-off-by: douxiaofeng99 <[email protected]>

Author: douxiaofeng99

Diff for: oauth2/oauth2-jose/src/main/java/org/springframework/security/oauth2/jwt/NimbusJwtEncoder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2025 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -46,6 +46,7 @@
 import com.nimbusds.jwt.JWTClaimsSet;
 import com.nimbusds.jwt.SignedJWT;
 
+import org.springframework.core.convert.converter.Converter;
 import org.springframework.security.oauth2.jose.jws.SignatureAlgorithm;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
@@ -86,6 +87,19 @@ public final class NimbusJwtEncoder implements JwtEncoder {
 
 	private final JWKSource<SecurityContext> jwkSource;
 
+	private Converter<List<JWK>, JWK> jwkSelector= (jwks)->{
+		if (jwks.size() > 1) {
+			throw new JwtEncodingException(String.format(
+					"Failed to select a key since there are multiple for the signing algorithm [%s]; " +
+							"please specify a selector in NimbusJwsEncoder#setJwkSelector",jwks.get(0).getAlgorithm()));
+		}
+		if (jwks.isEmpty()) {
+			throw new JwtEncodingException(
+					String.format(ENCODING_ERROR_MESSAGE_TEMPLATE, "Failed to select a JWK signing key"));
+		}
+		return jwks.get(0);
+	};
+
 	/**
 	 * Constructs a {@code NimbusJwtEncoder} using the provided parameters.
 	 * @param jwkSource the {@code com.nimbusds.jose.jwk.source.JWKSource}
@@ -94,6 +108,18 @@ public NimbusJwtEncoder(JWKSource<SecurityContext> jwkSource) {
 		Assert.notNull(jwkSource, "jwkSource cannot be null");
 		this.jwkSource = jwkSource;
 	}
+	/**
+	 * Use this strategy to reduce the list of matching JWKs down to a since one.
+	 * <p> For example, you can call {@code setJwkSelector(List::getFirst)} in order
+	 * to have this encoder select the first match.
+	 *
+	 * <p> By default, the class with throw an exception if there is more than one result.
+	 * @since 6.5
+	 */
+	public void setJwkSelector(Converter<List<JWK>, JWK> jwkSelector) {
+		if(null!=jwkSelector)
+			this.jwkSelector = jwkSelector;
+	}
 
 	@Override
 	public Jwt encode(JwtEncoderParameters parameters) throws JwtEncodingException {
@@ -123,18 +149,7 @@ private JWK selectJwk(JwsHeader headers) {
 			throw new JwtEncodingException(String.format(ENCODING_ERROR_MESSAGE_TEMPLATE,
 					"Failed to select a JWK signing key -> " + ex.getMessage()), ex);
 		}
-
-		if (jwks.size() > 1) {
-			throw new JwtEncodingException(String.format(ENCODING_ERROR_MESSAGE_TEMPLATE,
-					"Found multiple JWK signing keys for algorithm '" + headers.getAlgorithm().getName() + "'"));
-		}
-
-		if (jwks.isEmpty()) {
-			throw new JwtEncodingException(
-					String.format(ENCODING_ERROR_MESSAGE_TEMPLATE, "Failed to select a JWK signing key"));
-		}
-
-		return jwks.get(0);
+		return this.jwkSelector.convert(jwks);
 	}
 
 	private String serialize(JwsHeader headers, JwtClaimsSet claims, JWK jwk) {