Introduce OAuth2TokenGenerator

oauth2-authorization-server/src/main/java/org/springframework/security/config/annotation/web/configurers/oauth2/server/authorization/OAuth2AuthorizationServerConfigurer.java

@@ -37,8 +37,10 @@
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.Transient;
 import org.springframework.security.core.context.SecurityContext;
+import org.springframework.security.oauth2.core.OAuth2Token;
 import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationConsentService;
 import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
+import org.springframework.security.oauth2.server.authorization.OAuth2TokenGenerator;
 import org.springframework.security.oauth2.server.authorization.authentication.OAuth2TokenIntrospectionAuthenticationProvider;
 import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
 import org.springframework.security.oauth2.server.authorization.config.ProviderSettings;
@@ -146,6 +148,19 @@ public OAuth2AuthorizationServerConfigurer<B> providerSettings(ProviderSettings
 		return this;
 	}
 
+	/**
+	 * Sets the token generator.
+	 *
+	 * @param tokenGenerator the token generator
+	 * @return the {@link OAuth2AuthorizationServerConfigurer} for further configuration
+	 * @since 0.2.3
+	 */
+	public OAuth2AuthorizationServerConfigurer<B> tokenGenerator(OAuth2TokenGenerator<? extends OAuth2Token> tokenGenerator) {
+		Assert.notNull(tokenGenerator, "tokenGenerator cannot be null");
+		getBuilder().setSharedObject(OAuth2TokenGenerator.class, tokenGenerator);
+		return this;
+	}
+
 	/**
 	 * Configures OAuth 2.0 Client Authentication.
 	 *

oauth2-authorization-server/src/main/java/org/springframework/security/config/annotation/web/configurers/oauth2/server/authorization/OAuth2ConfigurerUtils.java

@@ -26,14 +26,17 @@
 import org.springframework.context.ApplicationContext;
 import org.springframework.core.ResolvableType;
 import org.springframework.security.config.annotation.web.HttpSecurityBuilder;
+import org.springframework.security.oauth2.core.OAuth2Token;
 import org.springframework.security.oauth2.jwt.JwtEncoder;
 import org.springframework.security.oauth2.jwt.NimbusJwsEncoder;
 import org.springframework.security.oauth2.server.authorization.InMemoryOAuth2AuthorizationConsentService;
 import org.springframework.security.oauth2.server.authorization.InMemoryOAuth2AuthorizationService;
 import org.springframework.security.oauth2.server.authorization.JwtEncodingContext;
+import org.springframework.security.oauth2.server.authorization.JwtGenerator;
 import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationConsentService;
 import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
 import org.springframework.security.oauth2.server.authorization.OAuth2TokenCustomizer;
+import org.springframework.security.oauth2.server.authorization.OAuth2TokenGenerator;
 import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
 import org.springframework.security.oauth2.server.authorization.config.ProviderSettings;
 import org.springframework.util.StringUtils;
@@ -82,7 +85,25 @@ static <B extends HttpSecurityBuilder<B>> OAuth2AuthorizationConsentService getA
 		return authorizationConsentService;
 	}
 
-	static <B extends HttpSecurityBuilder<B>> JwtEncoder getJwtEncoder(B builder) {
+	@SuppressWarnings("unchecked")
+	static <B extends HttpSecurityBuilder<B>> OAuth2TokenGenerator<? extends OAuth2Token> getTokenGenerator(B builder) {
+		OAuth2TokenGenerator<? extends OAuth2Token> tokenGenerator = builder.getSharedObject(OAuth2TokenGenerator.class);
+		if (tokenGenerator == null) {
+			tokenGenerator = getOptionalBean(builder, OAuth2TokenGenerator.class);
+			if (tokenGenerator == null) {
+				JwtGenerator jwtGenerator = new JwtGenerator(getJwtEncoder(builder));
+				OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer = getJwtCustomizer(builder);
+				if (jwtCustomizer != null) {
+					jwtGenerator.setJwtCustomizer(jwtCustomizer);
+				}
+				tokenGenerator = jwtGenerator;
+			}
+			builder.setSharedObject(OAuth2TokenGenerator.class, tokenGenerator);
+		}
+		return tokenGenerator;
+	}
+
+	private static <B extends HttpSecurityBuilder<B>> JwtEncoder getJwtEncoder(B builder) {
 		JwtEncoder jwtEncoder = builder.getSharedObject(JwtEncoder.class);
 		if (jwtEncoder == null) {
 			jwtEncoder = getOptionalBean(builder, JwtEncoder.class);
@@ -107,7 +128,7 @@ static <B extends HttpSecurityBuilder<B>> JWKSource<SecurityContext> getJwkSourc
 	}
 
 	@SuppressWarnings("unchecked")
-	static <B extends HttpSecurityBuilder<B>> OAuth2TokenCustomizer<JwtEncodingContext> getJwtCustomizer(B builder) {
+	private static <B extends HttpSecurityBuilder<B>> OAuth2TokenCustomizer<JwtEncodingContext> getJwtCustomizer(B builder) {
 		OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer = builder.getSharedObject(OAuth2TokenCustomizer.class);
 		if (jwtCustomizer == null) {
 			ResolvableType type = ResolvableType.forClassWithGenerics(OAuth2TokenCustomizer.class, JwtEncodingContext.class);

oauth2-authorization-server/src/main/java/org/springframework/security/config/annotation/web/configurers/oauth2/server/authorization/OAuth2TokenEndpointConfigurer.java

@@ -28,10 +28,10 @@
 import org.springframework.security.config.annotation.web.HttpSecurityBuilder;
 import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
 import org.springframework.security.oauth2.core.OAuth2Error;
+import org.springframework.security.oauth2.core.OAuth2Token;
 import org.springframework.security.oauth2.core.endpoint.OAuth2AccessTokenResponse;
-import org.springframework.security.oauth2.jwt.JwtEncoder;
-import org.springframework.security.oauth2.server.authorization.JwtEncodingContext;
-import org.springframework.security.oauth2.server.authorization.OAuth2TokenCustomizer;
+import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
+import org.springframework.security.oauth2.server.authorization.OAuth2TokenGenerator;
 import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
 import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AuthorizationCodeAuthenticationProvider;
 import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AuthorizationGrantAuthenticationToken;
@@ -160,34 +160,19 @@ RequestMatcher getRequestMatcher() {
 	private <B extends HttpSecurityBuilder<B>> List<AuthenticationProvider> createDefaultAuthenticationProviders(B builder) {
 		List<AuthenticationProvider> authenticationProviders = new ArrayList<>();
 
-		JwtEncoder jwtEncoder = OAuth2ConfigurerUtils.getJwtEncoder(builder);
-		OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer = OAuth2ConfigurerUtils.getJwtCustomizer(builder);
+		OAuth2AuthorizationService authorizationService = OAuth2ConfigurerUtils.getAuthorizationService(builder);
+		OAuth2TokenGenerator<? extends OAuth2Token> tokenGenerator = OAuth2ConfigurerUtils.getTokenGenerator(builder);
 
 		OAuth2AuthorizationCodeAuthenticationProvider authorizationCodeAuthenticationProvider =
-				new OAuth2AuthorizationCodeAuthenticationProvider(
-						OAuth2ConfigurerUtils.getAuthorizationService(builder),
-						jwtEncoder);
-		if (jwtCustomizer != null) {
-			authorizationCodeAuthenticationProvider.setJwtCustomizer(jwtCustomizer);
-		}
+				new OAuth2AuthorizationCodeAuthenticationProvider(authorizationService, tokenGenerator);
 		authenticationProviders.add(authorizationCodeAuthenticationProvider);
 
 		OAuth2RefreshTokenAuthenticationProvider refreshTokenAuthenticationProvider =
-				new OAuth2RefreshTokenAuthenticationProvider(
-						OAuth2ConfigurerUtils.getAuthorizationService(builder),
-						jwtEncoder);
-		if (jwtCustomizer != null) {
-			refreshTokenAuthenticationProvider.setJwtCustomizer(jwtCustomizer);
-		}
+				new OAuth2RefreshTokenAuthenticationProvider(authorizationService, tokenGenerator);
 		authenticationProviders.add(refreshTokenAuthenticationProvider);
 
 		OAuth2ClientCredentialsAuthenticationProvider clientCredentialsAuthenticationProvider =
-				new OAuth2ClientCredentialsAuthenticationProvider(
-						OAuth2ConfigurerUtils.getAuthorizationService(builder),
-						jwtEncoder);
-		if (jwtCustomizer != null) {
-			clientCredentialsAuthenticationProvider.setJwtCustomizer(jwtCustomizer);
-		}
+				new OAuth2ClientCredentialsAuthenticationProvider(authorizationService, tokenGenerator);
 		authenticationProviders.add(clientCredentialsAuthenticationProvider);
 
 		return authenticationProviders;

oauth2-authorization-server/src/main/java/org/springframework/security/config/annotation/web/configurers/oauth2/server/authorization/OidcClientRegistrationEndpointConfigurer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2021 the original author or authors.
+ * Copyright 2020-2022 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.
@@ -57,7 +57,7 @@ <B extends HttpSecurityBuilder<B>> void init(B builder) {
 				new OidcClientRegistrationAuthenticationProvider(
 						OAuth2ConfigurerUtils.getRegisteredClientRepository(builder),
 						OAuth2ConfigurerUtils.getAuthorizationService(builder),
-						OAuth2ConfigurerUtils.getJwtEncoder(builder));
+						OAuth2ConfigurerUtils.getTokenGenerator(builder));
 		builder.authenticationProvider(postProcess(oidcClientRegistrationAuthenticationProvider));
 	}
 

oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/DefaultOAuth2TokenContext.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2020-2022 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.security.oauth2.server.authorization;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+
+/**
+ * Default implementation of {@link OAuth2TokenContext}.
+ *
+ * @author Joe Grandja
+ * @since 0.2.3
+ * @see OAuth2TokenContext
+ */
+public final class DefaultOAuth2TokenContext implements OAuth2TokenContext {
+	private final Map<Object, Object> context;
+
+	private DefaultOAuth2TokenContext(Map<Object, Object> context) {
+		this.context = Collections.unmodifiableMap(new HashMap<>(context));
+	}
+
+	@SuppressWarnings("unchecked")
+	@Nullable
+	@Override
+	public <V> V get(Object key) {
+		return hasKey(key) ? (V) this.context.get(key) : null;
+	}
+
+	@Override
+	public boolean hasKey(Object key) {
+		Assert.notNull(key, "key cannot be null");
+		return this.context.containsKey(key);
+	}
+
+	/**
+	 * Returns a new {@link Builder}.
+	 *
+	 * @return the {@link Builder}
+	 */
+	public static Builder builder() {
+		return new Builder();
+	}
+
+	/**
+	 * A builder for {@link DefaultOAuth2TokenContext}.
+	 */
+	public static final class Builder extends AbstractBuilder<DefaultOAuth2TokenContext, Builder> {
+
+		private Builder() {
+		}
+
+		/**
+		 * Builds a new {@link DefaultOAuth2TokenContext}.
+		 *
+		 * @return the {@link DefaultOAuth2TokenContext}
+		 */
+		public DefaultOAuth2TokenContext build() {
+			return new DefaultOAuth2TokenContext(getContext());
+		}
+
+	}
+
+}