Customize JWKSourceBuilder in NimbusJwtDecoder

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

@@ -87,6 +87,7 @@
  * @author Mykyta Bezverkhyi
  * @author Daeho Kwon
  * @author Andrey Litvitski
+ * @author Ngoc Nhan
  * @since 5.2
  */
 public final class NimbusJwtDecoder implements JwtDecoder {
@@ -306,6 +307,8 @@ public static final class JwkSetUriJwtDecoderBuilder {
 
 		private Consumer<ConfigurableJWTProcessor<SecurityContext>> jwtProcessorCustomizer;
 
+		private @Nullable Consumer<JWKSourceBuilder<SecurityContext>> jwkSourceBuilderCustomizer;
+
 		private OAuth2TokenValidator<Jwt> validator = JwtValidators.createDefault();
 
 		private JwkSetUriJwtDecoderBuilder(String jwkSetUri) {
@@ -448,6 +451,31 @@ public JwkSetUriJwtDecoderBuilder jwtProcessorCustomizer(
 			return this;
 		}
 
+		/**
+		 * Use the given {@link Consumer} to customize the {@link JWKSourceBuilder} before
+		 * passing it to the {@link NimbusJwtDecoder} build process.
+		 *
+		 * <p>
+		 * By default, the {@link JWKSourceBuilder} is configured as follows:
+		 *
+		 * <pre>
+		 * jwkSourceBuilder
+		 *     .refreshAheadCache(false)
+		 *     .rateLimited(false)
+		 *     .cache(this.cache instanceof NoOpCache);
+		 * </pre>
+		 * @param jwkSourceBuilderCustomizer the callback used to customize the
+		 * {@link JWKSourceBuilder}
+		 * @return a {@link JwkSetUriJwtDecoderBuilder} for further configuration
+		 * @since 7.1
+		 */
+		public JwkSetUriJwtDecoderBuilder jwkSourceBuilderCustomizer(
+				Consumer<JWKSourceBuilder<SecurityContext>> jwkSourceBuilderCustomizer) {
+			Assert.notNull(jwkSourceBuilderCustomizer, "jwkSourceBuilderCustomizer cannot be null");
+			this.jwkSourceBuilderCustomizer = jwkSourceBuilderCustomizer;
+			return this;
+		}
+
 		JwkSetUriJwtDecoderBuilder validator(OAuth2TokenValidator<Jwt> validator) {
 			Assert.notNull(validator, "validator cannot be null");
 			this.validator = validator;
@@ -468,11 +496,15 @@ 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

@@ -43,7 +43,10 @@
 import com.nimbusds.jose.crypto.MACSigner;
 import com.nimbusds.jose.crypto.RSASSASigner;
 import com.nimbusds.jose.jwk.JWKSet;
+import com.nimbusds.jose.jwk.source.CachingJWKSetSource;
 import com.nimbusds.jose.jwk.source.JWKSource;
+import com.nimbusds.jose.jwk.source.RefreshAheadCachingJWKSetSource;
+import com.nimbusds.jose.jwk.source.RetryingJWKSetSource;
 import com.nimbusds.jose.proc.BadJOSEException;
 import com.nimbusds.jose.proc.DefaultJOSEObjectTypeVerifier;
 import com.nimbusds.jose.proc.JWSKeySelector;
@@ -99,6 +102,7 @@
  * @author Joe Grandja
  * @author Mykyta Bezverkhyi
  * @author Andrey Litvitski
+ * @author Ngoc Nhan
  */
 public class NimbusJwtDecoderTests {
 
@@ -929,6 +933,46 @@ public void decodeWhenSecretKeyValidateTypeFalseThenSkipsNimbusTypeValidation()
 		jwtDecoder.decode(jwt.serialize());
 	}
 
+	@Test
+	public void shouldUseCachingJwkSetSource() {
+
+		DefaultJWTProcessor<SecurityContext> jwtProcessor = (DefaultJWTProcessor<SecurityContext>) NimbusJwtDecoder
+			.withJwkSetUri(JWK_SET_URI)
+			.processor();
+		Object jwkSource = ReflectionTestUtils.getField(jwtProcessor.getJWSKeySelector(),
+				JWSVerificationKeySelector.class, "jwkSource");
+		assertThat(jwkSource).isNotNull();
+		assertThat(ReflectionTestUtils.getField(jwkSource, "source")).isInstanceOf(CachingJWKSetSource.class);
+	}
+
+	@Test
+	public void shouldUseRefreshAheadCachingJWKSetSource() {
+
+		DefaultJWTProcessor<SecurityContext> jwtProcessor = (DefaultJWTProcessor<SecurityContext>) NimbusJwtDecoder
+			.withJwkSetUri(JWK_SET_URI)
+			.jwkSourceBuilderCustomizer(
+					(jwkSourceBuilder) -> jwkSourceBuilder.refreshAheadCache(true).rateLimited(true))
+			.processor();
+		Object jwkSource = ReflectionTestUtils.getField(jwtProcessor.getJWSKeySelector(),
+				JWSVerificationKeySelector.class, "jwkSource");
+		assertThat(jwkSource).isNotNull();
+		assertThat(ReflectionTestUtils.getField(jwkSource, "source"))
+			.isInstanceOf(RefreshAheadCachingJWKSetSource.class);
+	}
+
+	@Test
+	public void shouldRetryingJWKSetSource() {
+
+		DefaultJWTProcessor<SecurityContext> jwtProcessor = (DefaultJWTProcessor<SecurityContext>) NimbusJwtDecoder
+			.withJwkSetUri(JWK_SET_URI)
+			.jwkSourceBuilderCustomizer((jwkSourceBuilder) -> jwkSourceBuilder.cache(false).retrying(true))
+			.processor();
+		Object jwkSource = ReflectionTestUtils.getField(jwtProcessor.getJWSKeySelector(),
+				JWSVerificationKeySelector.class, "jwkSource");
+		assertThat(jwkSource).isNotNull();
+		assertThat(ReflectionTestUtils.getField(jwkSource, "source")).isInstanceOf(RetryingJWKSetSource.class);
+	}
+
 	private RSAPublicKey key() throws InvalidKeySpecException {
 		byte[] decoded = Base64.getDecoder().decode(VERIFY_KEY.getBytes());
 		EncodedKeySpec spec = new X509EncodedKeySpec(decoded);