Add authorizationCustomizer pattern for OAuth2Authorization customization

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

@@ -23,6 +23,7 @@
 
 import org.springframework.lang.Nullable;
 import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
+import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
 import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationConsent;
 import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
 import org.springframework.util.Assert;
@@ -89,6 +90,15 @@ public OAuth2AuthorizationConsent getAuthorizationConsent() {
 		return get(OAuth2AuthorizationConsent.class);
 	}
 
+	/**
+	 * Returns the {@link OAuth2Authorization.Builder authorization builder}.
+	 * @return the {@link OAuth2Authorization.Builder}
+	 */
+	@Nullable
+	public OAuth2Authorization.Builder getAuthorizationBuilder() {
+		return get(OAuth2Authorization.Builder.class);
+	}
+
 	/**
 	 * Constructs a new {@link Builder} with the provided
 	 * {@link OAuth2AuthorizationCodeRequestAuthenticationToken}.
@@ -138,6 +148,15 @@ public Builder authorizationConsent(OAuth2AuthorizationConsent authorizationCons
 			return put(OAuth2AuthorizationConsent.class, authorizationConsent);
 		}
 
+		/**
+		 * Sets the {@link OAuth2Authorization.Builder authorization builder}.
+		 * @param authorizationBuilder the {@link OAuth2Authorization.Builder}
+		 * @return the {@link Builder} for further configuration
+		 */
+		public Builder authorizationBuilder(OAuth2Authorization.Builder authorizationBuilder) {
+			return put(OAuth2Authorization.Builder.class, authorizationBuilder);
+		}
+
 		/**
 		 * Builds a new {@link OAuth2AuthorizationCodeRequestAuthenticationContext}.
 		 * @return the {@link OAuth2AuthorizationCodeRequestAuthenticationContext}

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

@@ -100,6 +100,9 @@ public final class OAuth2AuthorizationCodeRequestAuthenticationProvider implemen
 
 	private Predicate<OAuth2AuthorizationCodeRequestAuthenticationContext> authorizationConsentRequired = OAuth2AuthorizationCodeRequestAuthenticationProvider::isAuthorizationConsentRequired;
 
+	private Consumer<OAuth2AuthorizationCodeRequestAuthenticationContext> authorizationCustomizer = (context) -> {
+	};
+
 	/**
 	 * Constructs an {@code OAuth2AuthorizationCodeRequestAuthenticationProvider} using
 	 * the provided parameters.
@@ -257,9 +260,13 @@ public Authentication authenticate(Authentication authentication) throws Authent
 			}
 
 			String state = DEFAULT_STATE_GENERATOR.generateKey();
-			OAuth2Authorization authorization = authorizationBuilder(registeredClient, principal, authorizationRequest)
-				.attribute(OAuth2ParameterNames.STATE, state)
-				.build();
+			OAuth2Authorization.Builder authorizationBuilder = authorizationBuilder(registeredClient, principal, authorizationRequest)
+				.attribute(OAuth2ParameterNames.STATE, state);
+
+			authenticationContextBuilder.authorizationBuilder(authorizationBuilder);
+			this.authorizationCustomizer.accept(authenticationContextBuilder.build());
+
+			OAuth2Authorization authorization = authorizationBuilder.build();
 
 			if (this.logger.isTraceEnabled()) {
 				this.logger.trace("Generated authorization consent state");
@@ -299,10 +306,14 @@ public Authentication authenticate(Authentication authentication) throws Authent
 			this.logger.trace("Generated authorization code");
 		}
 
-		OAuth2Authorization authorization = authorizationBuilder(registeredClient, principal, authorizationRequest)
+		OAuth2Authorization.Builder authorizationBuilder = authorizationBuilder(registeredClient, principal, authorizationRequest)
 			.authorizedScopes(authorizationRequest.getScopes())
-			.token(authorizationCode)
-			.build();
+			.token(authorizationCode);
+
+		authenticationContextBuilder.authorizationBuilder(authorizationBuilder);
+		this.authorizationCustomizer.accept(authenticationContextBuilder.build());
+
+		OAuth2Authorization authorization = authorizationBuilder.build();
 		this.authorizationService.save(authorization);
 
 		if (this.logger.isTraceEnabled()) {
@@ -395,6 +406,29 @@ public void setAuthorizationConsentRequired(
 		this.authorizationConsentRequired = authorizationConsentRequired;
 	}
 
+	/**
+	 * Sets the {@code Consumer} providing access to the
+	 * {@link OAuth2AuthorizationCodeRequestAuthenticationContext} and is responsible for
+	 * customizing the {@link OAuth2Authorization.Builder} prior to building.
+	 * <p>
+	 * The following context attributes are available:
+	 * <ul>
+	 * <li>The {@link RegisteredClient} associated with the authorization request.</li>
+	 * <li>The {@link OAuth2AuthorizationRequest} containing the authorization request
+	 * parameters.</li>
+	 * <li>The {@link OAuth2Authorization.Builder} to be customized.</li>
+	 * </ul>
+	 * @param authorizationCustomizer the {@code Consumer} providing access to the
+	 * {@link OAuth2AuthorizationCodeRequestAuthenticationContext} and is responsible for
+	 * customizing the {@link OAuth2Authorization.Builder} prior to building
+	 * @since 1.4
+	 */
+	public void setAuthorizationCustomizer(
+			Consumer<OAuth2AuthorizationCodeRequestAuthenticationContext> authorizationCustomizer) {
+		Assert.notNull(authorizationCustomizer, "authorizationCustomizer cannot be null");
+		this.authorizationCustomizer = authorizationCustomizer;
+	}
+
 	private static boolean isAuthorizationConsentRequired(
 			OAuth2AuthorizationCodeRequestAuthenticationContext authenticationContext) {
 		if (!authenticationContext.getRegisteredClient().getClientSettings().isRequireAuthorizationConsent()) {

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

@@ -86,6 +86,15 @@ public OAuth2Authorization getAuthorization() {
 		return get(OAuth2Authorization.class);
 	}
 
+	/**
+	 * Returns the {@link OAuth2Authorization.Builder authorization builder}.
+	 * @return the {@link OAuth2Authorization.Builder}
+	 */
+	@Nullable
+	public OAuth2Authorization.Builder getAuthorizationBuilder() {
+		return get(OAuth2Authorization.Builder.class);
+	}
+
 	/**
 	 * Returns the {@link OAuth2AuthorizationRequest authorization request}.
 	 * @return the {@link OAuth2AuthorizationRequest}
@@ -142,6 +151,15 @@ public Builder authorization(OAuth2Authorization authorization) {
 			return put(OAuth2Authorization.class, authorization);
 		}
 
+		/**
+		 * Sets the {@link OAuth2Authorization.Builder authorization builder}.
+		 * @param authorizationBuilder the {@link OAuth2Authorization.Builder}
+		 * @return the {@link Builder} for further configuration
+		 */
+		public Builder authorizationBuilder(OAuth2Authorization.Builder authorizationBuilder) {
+			return put(OAuth2Authorization.Builder.class, authorizationBuilder);
+		}
+
 		/**
 		 * Sets the {@link OAuth2AuthorizationRequest authorization request}.
 		 * @param authorizationRequest the {@link OAuth2AuthorizationRequest}

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

@@ -80,6 +80,9 @@ public final class OAuth2AuthorizationConsentAuthenticationProvider implements A
 
 	private Consumer<OAuth2AuthorizationConsentAuthenticationContext> authorizationConsentCustomizer;
 
+	private Consumer<OAuth2AuthorizationConsentAuthenticationContext> authorizationCustomizer = (context) -> {
+	};
+
 	/**
 	 * Constructs an {@code OAuth2AuthorizationConsentAuthenticationProvider} using the
 	 * provided parameters.
@@ -239,11 +242,22 @@ public Authentication authenticate(Authentication authentication) throws Authent
 			this.logger.trace("Generated authorization code");
 		}
 
-		OAuth2Authorization updatedAuthorization = OAuth2Authorization.from(authorization)
+		OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.from(authorization)
 			.authorizedScopes(authorizedScopes)
 			.token(authorizationCode)
-			.attributes((attrs) -> attrs.remove(OAuth2ParameterNames.STATE))
-			.build();
+			.attributes((attrs) -> attrs.remove(OAuth2ParameterNames.STATE));
+
+		OAuth2AuthorizationConsentAuthenticationContext authorizationConsentAuthenticationContext =
+				OAuth2AuthorizationConsentAuthenticationContext.with(authorizationConsentAuthentication)
+						.authorizationConsent(authorizationConsentBuilder)
+						.registeredClient(registeredClient)
+						.authorization(authorization)
+						.authorizationBuilder(authorizationBuilder)
+						.authorizationRequest(authorizationRequest)
+						.build();
+		this.authorizationCustomizer.accept(authorizationConsentAuthenticationContext);
+
+		OAuth2Authorization updatedAuthorization = authorizationBuilder.build();
 		this.authorizationService.save(updatedAuthorization);
 
 		if (this.logger.isTraceEnabled()) {
@@ -310,6 +324,28 @@ public void setAuthorizationConsentCustomizer(
 		this.authorizationConsentCustomizer = authorizationConsentCustomizer;
 	}
 
+	/**
+	 * Sets the {@code Consumer} providing access to the
+	 * {@link OAuth2AuthorizationConsentAuthenticationContext} and is responsible for
+	 * customizing the {@link OAuth2Authorization.Builder} prior to building.
+	 * <p>
+	 * The following context attributes are available:
+	 * <ul>
+	 * <li>The {@link RegisteredClient} associated with the authorization request.</li>
+	 * <li>The {@link OAuth2Authorization} associated with the state token.</li>
+	 * <li>The {@link OAuth2Authorization.Builder} to be customized.</li>
+	 * </ul>
+	 * @param authorizationCustomizer the {@code Consumer} providing access to the
+	 * {@link OAuth2AuthorizationConsentAuthenticationContext} and is responsible for
+	 * customizing the {@link OAuth2Authorization.Builder} prior to building
+	 * @since 1.4
+	 */
+	public void setAuthorizationCustomizer(
+			Consumer<OAuth2AuthorizationConsentAuthenticationContext> authorizationCustomizer) {
+		Assert.notNull(authorizationCustomizer, "authorizationCustomizer cannot be null");
+		this.authorizationCustomizer = authorizationCustomizer;
+	}
+
 	private static OAuth2TokenContext createAuthorizationCodeTokenContext(
 			OAuth2AuthorizationConsentAuthenticationToken authorizationConsentAuthentication,
 			RegisteredClient registeredClient, OAuth2Authorization authorization, Set<String> authorizedScopes) {

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

@@ -21,6 +21,7 @@
 import java.util.function.Consumer;
 
 import org.springframework.lang.Nullable;
+import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
 import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
 import org.springframework.util.Assert;
 
@@ -64,6 +65,15 @@ public RegisteredClient getRegisteredClient() {
 		return get(RegisteredClient.class);
 	}
 
+	/**
+	 * Returns the {@link OAuth2Authorization.Builder authorization builder}.
+	 * @return the {@link OAuth2Authorization.Builder}
+	 */
+	@Nullable
+	public OAuth2Authorization.Builder getAuthorizationBuilder() {
+		return get(OAuth2Authorization.Builder.class);
+	}
+
 	/**
 	 * Constructs a new {@link Builder} with the provided
 	 * {@link OAuth2ClientCredentialsAuthenticationToken}.
@@ -92,6 +102,15 @@ public Builder registeredClient(RegisteredClient registeredClient) {
 			return put(RegisteredClient.class, registeredClient);
 		}
 
+		/**
+		 * Sets the {@link OAuth2Authorization.Builder authorization builder}.
+		 * @param authorizationBuilder the {@link OAuth2Authorization.Builder}
+		 * @return the {@link Builder} for further configuration
+		 */
+		public Builder authorizationBuilder(OAuth2Authorization.Builder authorizationBuilder) {
+			return put(OAuth2Authorization.Builder.class, authorizationBuilder);
+		}
+
 		/**
 		 * Builds a new {@link OAuth2ClientCredentialsAuthenticationContext}.
 		 * @return the {@link OAuth2ClientCredentialsAuthenticationContext}

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

@@ -73,6 +73,9 @@ public final class OAuth2ClientCredentialsAuthenticationProvider implements Auth
 
 	private Consumer<OAuth2ClientCredentialsAuthenticationContext> authenticationValidator = new OAuth2ClientCredentialsAuthenticationValidator();
 
+	private Consumer<OAuth2ClientCredentialsAuthenticationContext> authorizationCustomizer = (context) -> {
+	};
+
 	/**
 	 * Constructs an {@code OAuth2ClientCredentialsAuthenticationProvider} using the
 	 * provided parameters.
@@ -160,6 +163,13 @@ public Authentication authenticate(Authentication authentication) throws Authent
 		OAuth2AccessToken accessToken = OAuth2AuthenticationProviderUtils.accessToken(authorizationBuilder,
 				generatedAccessToken, tokenContext);
 
+		authenticationContext = OAuth2ClientCredentialsAuthenticationContext
+				.with(clientCredentialsAuthentication)
+				.registeredClient(registeredClient)
+				.authorizationBuilder(authorizationBuilder)
+				.build();
+		this.authorizationCustomizer.accept(authenticationContext);
+
 		OAuth2Authorization authorization = authorizationBuilder.build();
 
 		this.authorizationService.save(authorization);
@@ -199,4 +209,25 @@ public void setAuthenticationValidator(
 		this.authenticationValidator = authenticationValidator;
 	}
 
+	/**
+	 * Sets the {@code Consumer} providing access to the
+	 * {@link OAuth2ClientCredentialsAuthenticationContext} and is responsible for
+	 * customizing the {@link OAuth2Authorization.Builder} prior to building.
+	 * <p>
+	 * The following context attributes are available:
+	 * <ul>
+	 * <li>The {@link RegisteredClient} associated with the authentication.</li>
+	 * <li>The {@link OAuth2Authorization.Builder} to be customized.</li>
+	 * </ul>
+	 * @param authorizationCustomizer the {@code Consumer} providing access to the
+	 * {@link OAuth2ClientCredentialsAuthenticationContext} and is responsible for
+	 * customizing the {@link OAuth2Authorization.Builder} prior to building
+	 * @since 1.5
+	 */
+	public void setAuthorizationCustomizer(
+			Consumer<OAuth2ClientCredentialsAuthenticationContext> authorizationCustomizer) {
+		Assert.notNull(authorizationCustomizer, "authorizationCustomizer cannot be null");
+		this.authorizationCustomizer = authorizationCustomizer;
+	}
+
 }

oauth2-authorization-server/src/test/java/org/springframework/security/oauth2/server/authorization/authentication/OAuth2AuthorizationCodeRequestAuthenticationProviderTests.java

@@ -814,4 +814,43 @@ private static void assertAuthenticationException(
 		assertThat(authorizationCodeRequestAuthentication.getRedirectUri()).isEqualTo(redirectUri);
 	}
 
+	@Test
+	public void setAuthorizationCustomizerWhenNullThenThrowIllegalArgumentException() {
+		assertThatThrownBy(() -> this.authenticationProvider.setAuthorizationCustomizer(null))
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessage("authorizationCustomizer cannot be null");
+	}
+
+	@Test
+	public void authenticateWhenCustomAuthorizationCustomizerThenUsed() {
+		RegisteredClient registeredClient = TestRegisteredClients.registeredClient().scope("scope1", "scope2").build();
+		given(this.registeredClientRepository.findByClientId(eq(registeredClient.getClientId())))
+			.willReturn(registeredClient);
+
+		Set<String> requestedScopes = registeredClient.getScopes();
+		OAuth2AuthorizationCodeRequestAuthenticationToken authentication = OAuth2AuthorizationCodeRequestAuthenticationToken
+			.with(AUTHORIZATION_URI, registeredClient.getClientId(), this.principal, requestedScopes, STATE)
+			.build();
+
+		Set<String> filteredScopes = Collections.singleton("scope1");
+		this.authenticationProvider.setAuthorizationCustomizer((context) -> {
+			OAuth2Authorization.Builder builder = context.getAuthorizationBuilder();
+			if (builder != null) {
+				builder.authorizedScopes(filteredScopes);
+			}
+		});
+
+		OAuth2AuthorizationCode authorizationCode = createAuthorizationCode();
+		given(this.authorizationService.save(any())).willAnswer((invocation) -> invocation.getArgument(0));
+
+		OAuth2AuthorizationCodeRequestAuthenticationToken authenticationResult = (OAuth2AuthorizationCodeRequestAuthenticationToken) this.authenticationProvider
+			.authenticate(authentication);
+
+		ArgumentCaptor<OAuth2Authorization> authorizationCaptor = ArgumentCaptor.forClass(OAuth2Authorization.class);
+		verify(this.authorizationService).save(authorizationCaptor.capture());
+		OAuth2Authorization authorization = authorizationCaptor.getValue();
+		assertThat(authorization.getAuthorizedScopes()).isEqualTo(filteredScopes);
+		assertThat(authenticationResult.getScopes()).isEqualTo(filteredScopes);
+	}
+
 }