Add OAuth2TokenCustomizer

Closes gh-199

Author: jgrandja

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

@@ -33,7 +33,9 @@
 import org.springframework.security.config.annotation.web.HttpSecurityBuilder;
 import org.springframework.security.config.annotation.web.configurers.AbstractHttpConfigurer;
 import org.springframework.security.config.annotation.web.configurers.ExceptionHandlingConfigurer;
+import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;
 import org.springframework.security.oauth2.jwt.JwtEncoder;
+import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
 import org.springframework.security.oauth2.jwt.NimbusJwsEncoder;
 import org.springframework.security.oauth2.server.authorization.InMemoryOAuth2AuthorizationService;
 import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
@@ -152,23 +154,33 @@ public void init(B builder) {
 		builder.authenticationProvider(postProcess(clientAuthenticationProvider));
 
 		JwtEncoder jwtEncoder = getJwtEncoder(builder);
+		OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer = getJwtCustomizer(builder);
 
 		OAuth2AuthorizationCodeAuthenticationProvider authorizationCodeAuthenticationProvider =
 				new OAuth2AuthorizationCodeAuthenticationProvider(
 						getAuthorizationService(builder),
 						jwtEncoder);
+		if (jwtCustomizer != null) {
+			authorizationCodeAuthenticationProvider.setJwtCustomizer(jwtCustomizer);
+		}
 		builder.authenticationProvider(postProcess(authorizationCodeAuthenticationProvider));
 
 		OAuth2RefreshTokenAuthenticationProvider refreshTokenAuthenticationProvider =
 				new OAuth2RefreshTokenAuthenticationProvider(
 						getAuthorizationService(builder),
 						jwtEncoder);
+		if (jwtCustomizer != null) {
+			refreshTokenAuthenticationProvider.setJwtCustomizer(jwtCustomizer);
+		}
 		builder.authenticationProvider(postProcess(refreshTokenAuthenticationProvider));
 
 		OAuth2ClientCredentialsAuthenticationProvider clientCredentialsAuthenticationProvider =
 				new OAuth2ClientCredentialsAuthenticationProvider(
 						getAuthorizationService(builder),
 						jwtEncoder);
+		if (jwtCustomizer != null) {
+			clientCredentialsAuthenticationProvider.setJwtCustomizer(jwtCustomizer);
+		}
 		builder.authenticationProvider(postProcess(clientCredentialsAuthenticationProvider));
 
 		OAuth2TokenRevocationAuthenticationProvider tokenRevocationAuthenticationProvider =
@@ -314,6 +326,19 @@ private static <B extends HttpSecurityBuilder<B>> JWKSource<SecurityContext> get
 		return jwkSource;
 	}
 
+	@SuppressWarnings("unchecked")
+	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);
+			jwtCustomizer = getOptionalBean(builder, type);
+			if (jwtCustomizer != null) {
+				builder.setSharedObject(OAuth2TokenCustomizer.class, jwtCustomizer);
+			}
+		}
+		return jwtCustomizer;
+	}
+
 	private static <B extends HttpSecurityBuilder<B>> ProviderSettings getProviderSettings(B builder) {
 		ProviderSettings providerSettings = builder.getSharedObject(ProviderSettings.class);
 		if (providerSettings == null) {
@@ -353,4 +378,14 @@ private static <B extends HttpSecurityBuilder<B>, T> T getOptionalBean(B builder
 		}
 		return (!beansMap.isEmpty() ? beansMap.values().iterator().next() : null);
 	}
+
+	@SuppressWarnings("unchecked")
+	private static <B extends HttpSecurityBuilder<B>, T> T getOptionalBean(B builder, ResolvableType type) {
+		ApplicationContext context = builder.getSharedObject(ApplicationContext.class);
+		String[] names = context.getBeanNamesForType(type);
+		if (names.length > 1) {
+			throw new NoUniqueBeanDefinitionException(type, names);
+		}
+		return names.length == 1 ? (T) context.getBean(names[0]) : null;
+	}
 }

oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/core/context/Context.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2020-2021 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.core.context;
+
+import java.util.Map;
+
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+
+/**
+ * A facility for holding information associated to a specific context.
+ *
+ * @author Joe Grandja
+ * @since 0.1.0
+ */
+public interface Context {
+
+	@Nullable
+	<V> V get(Object key);
+
+	@Nullable
+	default <V> V get(Class<V> key) {
+		Assert.notNull(key, "key cannot be null");
+		V value = get((Object) key);
+		return key.isInstance(value) ? value : null;
+	}
+
+	boolean hasKey(Object key);
+
+	static Context of(Map<Object, Object> context) {
+		return new DefaultContext(context);
+	}
+
+}

oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/core/context/DefaultContext.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2020-2021 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.core.context;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+
+/**
+ * @author Joe Grandja
+ * @since 0.1.0
+ */
+final class DefaultContext implements Context {
+	private final Map<Object, Object> context;
+
+	DefaultContext(Map<Object, Object> context) {
+		Assert.notNull(context, "context cannot be null");
+		this.context = Collections.unmodifiableMap(new HashMap<>(context));
+	}
+
+	@SuppressWarnings("unchecked")
+	@Override
+	@Nullable
+	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);
+	}
+
+}

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

@@ -23,8 +23,6 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 
 import com.nimbusds.jose.JOSEException;
@@ -46,7 +44,6 @@
 import com.nimbusds.jwt.SignedJWT;
 
 import org.springframework.core.convert.converter.Converter;
-import org.springframework.security.oauth2.jose.jws.JwsAlgorithm;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.StringUtils;
@@ -88,9 +85,6 @@ public final class NimbusJwsEncoder implements JwtEncoder {
 
 	private final JWKSource<SecurityContext> jwkSource;
 
-	private BiConsumer<JoseHeader.Builder, JwtClaimsSet.Builder> jwtCustomizer = (headers, claims) -> {
-	};
-
 	/**
 	 * Constructs a {@code NimbusJwsEncoder} using the provided parameters.
 	 * @param jwkSource the {@code com.nimbusds.jose.jwk.source.JWKSource}
@@ -100,32 +94,12 @@ public NimbusJwsEncoder(JWKSource<SecurityContext> jwkSource) {
 		this.jwkSource = jwkSource;
 	}
 
-	/**
-	 * Sets the {@link Jwt} customizer to be provided the {@link JoseHeader.Builder} and
-	 * {@link JwtClaimsSet.Builder} allowing for further customizations.
-	 * @param jwtCustomizer the {@link Jwt} customizer to be provided the
-	 * {@link JoseHeader.Builder} and {@link JwtClaimsSet.Builder}
-	 */
-	public void setJwtCustomizer(BiConsumer<JoseHeader.Builder, JwtClaimsSet.Builder> jwtCustomizer) {
-		Assert.notNull(jwtCustomizer, "jwtCustomizer cannot be null");
-		this.jwtCustomizer = jwtCustomizer;
-	}
-
 	@Override
 	public Jwt encode(JoseHeader headers, JwtClaimsSet claims) throws JwtEncodingException {
 		Assert.notNull(headers, "headers cannot be null");
 		Assert.notNull(claims, "claims cannot be null");
 
-		// @formatter:off
-		JoseHeader.Builder headersBuilder = JoseHeader.from(headers)
-				.type(JOSEObjectType.JWT.getType());
-		JwtClaimsSet.Builder claimsBuilder = JwtClaimsSet.from(claims)
-				.id(UUID.randomUUID().toString());
-		// @formatter:on
-
-		this.jwtCustomizer.accept(headersBuilder, claimsBuilder);
-
-		JWK jwk = selectJwk(headersBuilder);
+		JWK jwk = selectJwk(headers);
 		if (jwk == null) {
 			throw new JwtEncodingException(
 					String.format(ENCODING_ERROR_MESSAGE_TEMPLATE, "Failed to select a JWK signing key"));
@@ -135,8 +109,15 @@ else if (!StringUtils.hasText(jwk.getKeyID())) {
 					"The \"kid\" (key ID) from the selected JWK cannot be empty"));
 		}
 
-		headers = headersBuilder.keyId(jwk.getKeyID()).build();
-		claims = claimsBuilder.build();
+		// @formatter:off
+		headers = JoseHeader.from(headers)
+				.type(JOSEObjectType.JWT.getType())
+				.keyId(jwk.getKeyID())
+				.build();
+		claims = JwtClaimsSet.from(claims)
+				.id(UUID.randomUUID().toString())
+				.build();
+		// @formatter:on
 
 		JWSHeader jwsHeader = JWS_HEADER_CONVERTER.convert(headers);
 		JWTClaimsSet jwtClaimsSet = JWT_CLAIMS_SET_CONVERTER.convert(claims);
@@ -164,13 +145,9 @@ else if (!StringUtils.hasText(jwk.getKeyID())) {
 		return new Jwt(jws, claims.getIssuedAt(), claims.getExpiresAt(), headers.getHeaders(), claims.getClaims());
 	}
 
-	private JWK selectJwk(JoseHeader.Builder headersBuilder) {
-		final AtomicReference<JWSAlgorithm> jwsAlgorithm = new AtomicReference<>();
-		headersBuilder.headers((h) -> {
-			JwsAlgorithm jwsAlg = (JwsAlgorithm) h.get(JoseHeaderNames.ALG);
-			jwsAlgorithm.set(JWSAlgorithm.parse(jwsAlg.getName()));
-		});
-		JWSHeader jwsHeader = new JWSHeader(jwsAlgorithm.get());
+	private JWK selectJwk(JoseHeader headers) {
+		JWSAlgorithm jwsAlgorithm = JWSAlgorithm.parse(headers.getJwsAlgorithm().getName());
+		JWSHeader jwsHeader = new JWSHeader(jwsAlgorithm);
 		JWKSelector jwkSelector = new JWKSelector(JWKMatcher.forJWSHeader(jwsHeader));
 
 		List<JWK> jwks;
@@ -184,7 +161,7 @@ private JWK selectJwk(JoseHeader.Builder headersBuilder) {
 
 		if (jwks.size() > 1) {
 			throw new JwtEncodingException(String.format(ENCODING_ERROR_MESSAGE_TEMPLATE,
-					"Found multiple JWK signing keys for algorithm '" + jwsAlgorithm.get().getName() + "'"));
+					"Found multiple JWK signing keys for algorithm '" + jwsAlgorithm.getName() + "'"));
 		}
 
 		return !jwks.isEmpty() ? jwks.get(0) : null;

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2020 the original author or authors.
+ * Copyright 2020-2021 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.
@@ -56,4 +56,9 @@ public interface OAuth2AuthorizationAttributeNames {
 	 */
 	String ACCESS_TOKEN_ATTRIBUTES = OAuth2Authorization.class.getName().concat(".ACCESS_TOKEN_ATTRIBUTES");
 
+	/**
+	 * The name of the attribute used for the resource owner {@code Principal}.
+	 */
+	String PRINCIPAL = OAuth2Authorization.class.getName().concat(".PRINCIPAL");
+
 }