Add possibility to customize JwkSource of NimbusJwtDecoder

Signed-off-by: Mark Bonnekessel <[email protected]>

Author: marbon87

oauth2/oauth2-jose/src/main/java/org/springframework/security/oauth2/jwt/NimbusJwtDecoder.java

@@ -289,6 +289,8 @@ public static final class JwkSetUriJwtDecoderBuilder {
 
 		private Consumer<ConfigurableJWTProcessor<SecurityContext>> jwtProcessorCustomizer;
 
+		private Consumer<JWKSourceBuilder<SecurityContext>> jwkSourceBuilderCustomizer;
+
 		private JwkSetUriJwtDecoderBuilder(String jwkSetUri) {
 			Assert.hasText(jwkSetUri, "jwkSetUri cannot be empty");
 			this.jwkSetUri = (rest) -> jwkSetUri;
@@ -423,6 +425,20 @@ public JwkSetUriJwtDecoderBuilder jwtProcessorCustomizer(
 			return this;
 		}
 
+		/**
+		 * Use the given {@link Consumer} to customize the {@link JWKSourceBuilder} before
+		 * passing it to the build {@link NimbusJwtDecoder}.
+		 * @param jwkSourceBuilderCustomizer the callback used to alter the builder
+		 * @return a {@link JwkSetUriJwtDecoderBuilder} for further configurations
+		 * @since 6.5
+		 */
+		public JwkSetUriJwtDecoderBuilder jwkSourceBuilderCustomizer(
+				Consumer<JWKSourceBuilder<SecurityContext>> jwkSourceBuilderCustomizer) {
+			Assert.notNull(jwkSourceBuilderCustomizer, "jwkSourceBuilderCustomizer cannot be null");
+			this.jwkSourceBuilderCustomizer = jwkSourceBuilderCustomizer;
+			return this;
+		}
+
 		JWSKeySelector<SecurityContext> jwsKeySelector(JWKSource<SecurityContext> jwkSource) {
 			if (this.signatureAlgorithms.isEmpty()) {
 				return new JWSVerificationKeySelector<>(this.defaultAlgorithms.apply(jwkSource), jwkSource);
@@ -437,11 +453,17 @@ JWSKeySelector<SecurityContext> jwsKeySelector(JWKSource<SecurityContext> jwkSou
 
 		JWKSource<SecurityContext> jwkSource() {
 			String jwkSetUri = this.jwkSetUri.apply(this.restOperations);
-			return JWKSourceBuilder.create(new SpringJWKSource<>(this.restOperations, this.cache, jwkSetUri))
+			JWKSourceBuilder<SecurityContext> jwkSourceBuilder = JWKSourceBuilder
+				.create(new SpringJWKSource<>(this.restOperations, this.cache, jwkSetUri))
 				.refreshAheadCache(false)
 				.rateLimited(false)
-				.cache(this.cache instanceof NoOpCache)
-				.build();
+				.cache(this.cache instanceof NoOpCache);
+
+			if (this.jwkSourceBuilderCustomizer != null) {
+				this.jwkSourceBuilderCustomizer.accept(jwkSourceBuilder);
+			}
+
+			return jwkSourceBuilder.build();
 		}
 
 		JWTProcessor<SecurityContext> processor() {

oauth2/oauth2-jose/src/test/java/org/springframework/security/oauth2/jwt/NimbusJwtDecoderTests.java

@@ -32,6 +32,7 @@
 import java.util.Date;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 
 import javax.crypto.SecretKey;
 
@@ -43,6 +44,7 @@
 import com.nimbusds.jose.crypto.MACSigner;
 import com.nimbusds.jose.crypto.RSASSASigner;
 import com.nimbusds.jose.jwk.source.JWKSource;
+import com.nimbusds.jose.jwk.source.JWKSourceBuilder;
 import com.nimbusds.jose.proc.BadJOSEException;
 import com.nimbusds.jose.proc.DefaultJOSEObjectTypeVerifier;
 import com.nimbusds.jose.proc.JWSKeySelector;
@@ -603,6 +605,13 @@ public void jwsKeySelectorWhenNoAlgorithmThenReturnsRS256Selector() {
 		assertThat(jwsVerificationKeySelector.isAllowed(JWSAlgorithm.RS256)).isTrue();
 	}
 
+	@Test
+	public void jwkSourceIsCustomizable() {
+		Consumer<JWKSourceBuilder<SecurityContext>> jwkSourceBuilderCustomizer = mock();
+		NimbusJwtDecoder.withJwkSetUri(JWK_SET_URI).jwkSourceBuilderCustomizer(jwkSourceBuilderCustomizer).build();
+		verify(jwkSourceBuilderCustomizer, times(1)).accept(any());
+	}
+
 	@Test
 	public void jwsKeySelectorWhenOneAlgorithmThenReturnsSingleSelector() {
 		JWKSource<SecurityContext> jwkSource = mock(JWKSource.class);