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>
@@ -535,6 +545,108 @@ 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(this.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 {