Add builder to create NimbusJwtDecoder with JwkSource

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

Author: marbon87

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

@@ -261,6 +261,16 @@ public static SecretKeyJwtDecoderBuilder withSecretKey(SecretKey secretKey) {
 		return new SecretKeyJwtDecoderBuilder(secretKey);
 	}
 
+	/**
+	 * Use the given {@code JWKSource} to create a JwkSourceJwtDecoderBuilder.
+	 * @param jwkSource the JWK Source to use
+	 * @return a {@link JwkSetUriJwtDecoderBuilder} for further configurations
+	 * @since 7.0
+	 */
+	public static JwkSourceJwtDecoderBuilder withJwkSource(JWKSource<SecurityContext> jwkSource) {
+		return new JwkSourceJwtDecoderBuilder(jwkSource);
+	}
+
 	/**
 	 * A builder for creating {@link NimbusJwtDecoder} instances based on a
 	 * <a target="_blank" href="https://tools.ietf.org/html/rfc7517#section-5">JWK Set</a>
@@ -274,7 +284,7 @@ public static final class JwkSetUriJwtDecoderBuilder {
 		private static final JOSEObjectTypeVerifier<SecurityContext> NO_TYPE_VERIFIER = (header, context) -> {
 		};
 
-		private final Function<RestOperations, String> jwkSetUri;
+		private Function<RestOperations, String> jwkSetUri;
 
 		private Function<JWKSource<SecurityContext>, Set<JWSAlgorithm>> defaultAlgorithms = (source) -> Set
 			.of(JWSAlgorithm.RS256);
@@ -289,6 +299,8 @@ public static final class JwkSetUriJwtDecoderBuilder {
 
 		private Consumer<ConfigurableJWTProcessor<SecurityContext>> jwtProcessorCustomizer;
 
+		private JWKSource<SecurityContext> jwkSource;
+
 		private JwkSetUriJwtDecoderBuilder(String jwkSetUri) {
 			Assert.hasText(jwkSetUri, "jwkSetUri cannot be empty");
 			this.jwkSetUri = (rest) -> jwkSetUri;
@@ -306,6 +318,13 @@ private JwkSetUriJwtDecoderBuilder(Function<RestOperations, String> jwkSetUri,
 			};
 		}
 
+		private JwkSetUriJwtDecoderBuilder(JWKSource<SecurityContext> jwkSource) {
+			Assert.notNull(jwkSource, "jwkSource cannot be null");
+			this.jwkSource = jwkSource;
+			this.jwtProcessorCustomizer = (processor) -> {
+			};
+		}
+
 		/**
 		 * Whether to use Nimbus's typ header verification. This is {@code true} by
 		 * default, however it may change to {@code false} in a future major release.
@@ -436,6 +455,9 @@ JWSKeySelector<SecurityContext> jwsKeySelector(JWKSource<SecurityContext> jwkSou
 		}
 
 		JWKSource<SecurityContext> jwkSource() {
+			if (this.jwkSource != null) {
+				return this.jwkSource;
+			}
 			String jwkSetUri = this.jwkSetUri.apply(this.restOperations);
 			return JWKSourceBuilder.create(new SpringJWKSource<>(this.restOperations, this.cache, jwkSetUri))
 				.refreshAheadCache(false)
@@ -535,6 +557,106 @@ public void close() {
 
 	}
 
+	/**
+	 * A builder for creating {@link NimbusJwtDecoder} instances based on a {@code JWKSource}.
+	 */
+	public static final class JwkSourceJwtDecoderBuilder  {
+
+		private static final JOSEObjectTypeVerifier<SecurityContext> NO_TYPE_VERIFIER = (header, context) -> {
+		};
+
+		private final Function<JWKSource<SecurityContext>, Set<JWSAlgorithm>> defaultAlgorithms = (source) -> Set
+				.of(JWSAlgorithm.RS256);
+
+		private final JOSEObjectTypeVerifier<SecurityContext> typeVerifier = NO_TYPE_VERIFIER;
+
+		private final Set<SignatureAlgorithm> signatureAlgorithms = new HashSet<>();
+
+		private Consumer<ConfigurableJWTProcessor<SecurityContext>> jwtProcessorCustomizer;
+
+		private final JWKSource<SecurityContext> jwkSource;
+
+		private JwkSourceJwtDecoderBuilder (JWKSource<SecurityContext> jwkSource) {
+			Assert.notNull(jwkSource, "jwkSource cannot be null");
+			this.jwkSource = jwkSource;
+			this.jwtProcessorCustomizer = (processor) -> {
+			};
+		}
+
+		/**
+		 * Append the given signing
+		 * <a href="https://tools.ietf.org/html/rfc7515#section-4.1.1" target=
+		 * "_blank">algorithm</a> to the set of algorithms to use.
+		 * @param signatureAlgorithm the algorithm to use
+		 * @return a {@link JwkSourceJwtDecoderBuilder } for further configurations
+		 */
+		public JwkSourceJwtDecoderBuilder  jwsAlgorithm(SignatureAlgorithm signatureAlgorithm) {
+			Assert.notNull(signatureAlgorithm, "signatureAlgorithm cannot be null");
+			this.signatureAlgorithms.add(signatureAlgorithm);
+			return this;
+		}
+
+		/**
+		 * Configure the list of
+		 * <a href="https://tools.ietf.org/html/rfc7515#section-4.1.1" target=
+		 * "_blank">algorithms</a> to use with the given {@link Consumer}.
+		 * @param signatureAlgorithmsConsumer a {@link Consumer} for further configuring
+		 * the algorithm list
+		 * @return a {@link JwkSourceJwtDecoderBuilder } for further configurations
+		 */
+		public JwkSourceJwtDecoderBuilder  jwsAlgorithms(Consumer<Set<SignatureAlgorithm>> signatureAlgorithmsConsumer) {
+			Assert.notNull(signatureAlgorithmsConsumer, "signatureAlgorithmsConsumer cannot be null");
+			signatureAlgorithmsConsumer.accept(this.signatureAlgorithms);
+			return this;
+		}
+
+		/**
+		 * Use the given {@link Consumer} to customize the {@link JWTProcessor
+		 * ConfigurableJWTProcessor} before passing it to the build
+		 * {@link NimbusJwtDecoder}.
+		 * @param jwtProcessorCustomizer the callback used to alter the processor
+		 * @return a {@link JwkSourceJwtDecoderBuilder } for further configurations
+		 * @since 5.4
+		 */
+		public JwkSourceJwtDecoderBuilder  jwtProcessorCustomizer(
+				Consumer<ConfigurableJWTProcessor<SecurityContext>> jwtProcessorCustomizer) {
+			Assert.notNull(jwtProcessorCustomizer, "jwtProcessorCustomizer cannot be null");
+			this.jwtProcessorCustomizer = jwtProcessorCustomizer;
+			return this;
+		}
+
+		JWSKeySelector<SecurityContext> jwsKeySelector(JWKSource<SecurityContext> jwkSource) {
+			if (this.signatureAlgorithms.isEmpty()) {
+				return new JWSVerificationKeySelector<>(this.defaultAlgorithms.apply(jwkSource), jwkSource);
+			}
+			Set<JWSAlgorithm> jwsAlgorithms = new HashSet<>();
+			for (SignatureAlgorithm signatureAlgorithm : this.signatureAlgorithms) {
+				JWSAlgorithm jwsAlgorithm = JWSAlgorithm.parse(signatureAlgorithm.getName());
+				jwsAlgorithms.add(jwsAlgorithm);
+			}
+			return new JWSVerificationKeySelector<>(jwsAlgorithms, jwkSource);
+		}
+
+		JWTProcessor<SecurityContext> processor() {
+			ConfigurableJWTProcessor<SecurityContext> jwtProcessor = new DefaultJWTProcessor<>();
+			jwtProcessor.setJWSTypeVerifier(this.typeVerifier);
+			jwtProcessor.setJWSKeySelector(jwsKeySelector(jwkSource));
+			// Spring Security validates the claim set independent from Nimbus
+			jwtProcessor.setJWTClaimsSetVerifier((claims, context) -> {
+			});
+			this.jwtProcessorCustomizer.accept(jwtProcessor);
+			return jwtProcessor;
+		}
+
+		/**
+		 * Build the configured {@link NimbusJwtDecoder}.
+		 * @return the configured {@link NimbusJwtDecoder}
+		 */
+		public NimbusJwtDecoder build() {
+			return new NimbusJwtDecoder(processor());
+		}
+	}
+
 	/**
 	 * A builder for creating {@link NimbusJwtDecoder} instances based on a public key.
 	 */

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

@@ -42,6 +42,7 @@
 import com.nimbusds.jose.JWSSigner;
 import com.nimbusds.jose.crypto.MACSigner;
 import com.nimbusds.jose.crypto.RSASSASigner;
+import com.nimbusds.jose.jwk.JWKSet;
 import com.nimbusds.jose.jwk.source.JWKSource;
 import com.nimbusds.jose.proc.BadJOSEException;
 import com.nimbusds.jose.proc.DefaultJOSEObjectTypeVerifier;
@@ -557,6 +558,15 @@ public void decodeWhenUsingSecretKeyWithKidThenStillUsesKey() throws Exception {
 		// @formatter:on
 	}
 
+	@Test
+	public void withJwkSourceWhenDefaultsThenUsesProvidedJwkSource() throws Exception {
+		JWKSource<SecurityContext> source = mock(JWKSource.class);
+		given(source.get(any(), any())).willReturn(JWKSet.parse(JWK_SET).getKeys());
+		NimbusJwtDecoder decoder = NimbusJwtDecoder.withJwkSource(source).build();
+		Jwt jwt = decoder.decode(SIGNED_JWT);
+		assertThat(jwt.getClaimAsString("sub")).isEqualTo("test-subject");
+	}
+
 	// gh-8730
 	@Test
 	public void withSecretKeyWhenUsingCustomTypeHeaderThenSuccessfullyDecodes() throws Exception {