Customize OAuth2AuthorizationConsent prior to saving

Closes spring-projectsgh-436

Author: Steve Riesenberg

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

@@ -29,6 +29,7 @@
 
 import org.springframework.security.authentication.AnonymousAuthenticationToken;
 import org.springframework.security.authentication.AuthenticationProvider;
+import org.springframework.security.config.Customizer;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.AuthenticationException;
 import org.springframework.security.crypto.keygen.Base64StringKeyGenerator;
@@ -82,6 +83,7 @@ public final class OAuth2AuthorizationCodeRequestAuthenticationProvider implemen
 	private final OAuth2AuthorizationConsentService authorizationConsentService;
 	private Supplier<String> authorizationCodeGenerator = DEFAULT_AUTHORIZATION_CODE_GENERATOR::generateKey;
 	private Function<String, OAuth2AuthenticationValidator> authenticationValidatorResolver = DEFAULT_AUTHENTICATION_VALIDATOR_RESOLVER;
+	private Customizer<OAuth2AuthenticationContext> authorizationConsentCustomizer;
 
 	/**
 	 * Constructs an {@code OAuth2AuthorizationCodeRequestAuthenticationProvider} using the provided parameters.
@@ -145,6 +147,31 @@ public void setAuthenticationValidatorResolver(Function<String, OAuth2Authentica
 		this.authenticationValidatorResolver = authenticationValidatorResolver;
 	}
 
+	/**
+	 * Sets the {@link Customizer} providing access to the {@link OAuth2AuthenticationContext} containing an
+	 * {@link OAuth2AuthorizationConsent.Builder}.
+	 *
+	 * <p>
+	 * The {@link OAuth2AuthenticationContext} gives the mapper access to the
+	 * {@link OAuth2AuthorizationCodeRequestAuthenticationToken}.
+	 * In addition, the following context attributes are supported:
+	 * <ul>
+	 * <li>{@code OAuth2AuthorizationConsent.Builder.class} - The {@link OAuth2AuthorizationConsent.Builder} used to build the
+	 * authorization consent prior to {@link OAuth2AuthorizationConsentService#save(OAuth2AuthorizationConsent)}</li>
+	 * <li>{@code OAuth2Authorization.class} - The {@link OAuth2Authorization} associated with the state token presented
+	 * in the authorization consent request.</li>
+	 * <li>{@code OAuth2AuthorizationRequest.class} - The {@link OAuth2AuthorizationRequest} requiring the resource
+	 * owner's consent.</li>
+	 * </ul>
+	 *
+	 * @param authorizationConsentCustomizer the {@link Customizer} providing access to the
+	 * {@link OAuth2AuthenticationContext} containing an {@link OAuth2AuthorizationConsent.Builder}
+	 */
+	public void setAuthorizationConsentCustomizer(Customizer<OAuth2AuthenticationContext> authorizationConsentCustomizer) {
+		Assert.notNull(authorizationConsentCustomizer, "authorizationConsentCustomizer cannot be null");
+		this.authorizationConsentCustomizer = authorizationConsentCustomizer;
+	}
+
 	private Authentication authenticateAuthorizationRequest(Authentication authentication) throws AuthenticationException {
 		OAuth2AuthorizationCodeRequestAuthenticationToken authorizationCodeRequestAuthentication =
 				(OAuth2AuthorizationCodeRequestAuthenticationToken) authentication;
@@ -330,6 +357,17 @@ private Authentication authenticateAuthorizationConsent(Authentication authentic
 						authorization.getRegisteredClientId(), authorization.getPrincipalName());
 			}
 			authorizedScopes.forEach(authorizationConsentBuilder::scope);
+
+			if (this.authorizationConsentCustomizer != null) {
+				Map<Object, Object> context = new HashMap<>();
+				context.put(OAuth2AuthorizationConsent.Builder.class, authorizationConsentBuilder);
+				context.put(OAuth2Authorization.class, authorization);
+				context.put(OAuth2AuthorizationRequest.class, authorizationRequest);
+				OAuth2AuthenticationContext authenticationContext = new OAuth2AuthenticationContext(
+						authorizationCodeRequestAuthentication, context);
+				this.authorizationConsentCustomizer.customize(authenticationContext);
+			}
+
 			OAuth2AuthorizationConsent authorizationConsent = authorizationConsentBuilder.build();
 			this.authorizationConsentService.save(authorizationConsent);
 		}

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

@@ -29,19 +29,21 @@
 import org.mockito.ArgumentCaptor;
 
 import org.springframework.security.authentication.TestingAuthenticationToken;
+import org.springframework.security.config.Customizer;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.oauth2.core.AuthorizationGrantType;
+import org.springframework.security.oauth2.core.OAuth2AuthorizationCode;
 import org.springframework.security.oauth2.core.OAuth2Error;
 import org.springframework.security.oauth2.core.OAuth2ErrorCodes;
 import org.springframework.security.oauth2.core.OAuth2TokenType;
+import org.springframework.security.oauth2.core.authentication.OAuth2AuthenticationContext;
 import org.springframework.security.oauth2.core.authentication.OAuth2AuthenticationValidator;
 import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
 import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationResponseType;
 import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
 import org.springframework.security.oauth2.core.endpoint.PkceParameterNames;
 import org.springframework.security.oauth2.core.oidc.OidcScopes;
 import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
-import org.springframework.security.oauth2.core.OAuth2AuthorizationCode;
 import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationConsent;
 import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationConsentService;
 import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
@@ -129,6 +131,13 @@ public void setAuthenticationValidatorResolverWhenNullThenThrowIllegalArgumentEx
 				.hasMessage("authenticationValidatorResolver cannot be null");
 	}
 
+	@Test
+	public void setAuthorizationConsentCustomizerWhenNullThenThrowIllegalArgumentException() {
+		assertThatThrownBy(() -> this.authenticationProvider.setAuthorizationConsentCustomizer(null))
+				.isInstanceOf(IllegalArgumentException.class)
+				.hasMessage("authorizationConsentCustomizer cannot be null");
+	}
+
 	@Test
 	public void authenticateWhenInvalidClientIdThenThrowOAuth2AuthorizationCodeRequestAuthenticationException() {
 		RegisteredClient registeredClient = TestRegisteredClients.registeredClient().build();
@@ -773,6 +782,53 @@ public void authenticateWhenConsentRequestApproveAllThenReturnAuthorizationCode(
 		OAuth2AuthorizationCodeRequestAuthenticationToken authenticationResult =
 				(OAuth2AuthorizationCodeRequestAuthenticationToken) this.authenticationProvider.authenticate(authentication);
 
+		assertAuthorizationConsentRequestWithAuthorizationCodeResult(registeredClient, authorization, authenticationResult);
+	}
+
+	@Test
+	public void authenticateWhenCustomAuthorizationConsentCustomizerThenUsed() {
+		RegisteredClient registeredClient = TestRegisteredClients.registeredClient()
+				.build();
+		when(this.registeredClientRepository.findByClientId(eq(registeredClient.getClientId())))
+				.thenReturn(registeredClient);
+		OAuth2Authorization authorization = TestOAuth2Authorizations.authorization(registeredClient)
+				.principalName(this.principal.getName())
+				.build();
+		OAuth2AuthorizationRequest authorizationRequest = authorization.getAttribute(OAuth2AuthorizationRequest.class.getName());
+		Set<String> authorizedScopes = authorizationRequest.getScopes();
+		OAuth2AuthorizationCodeRequestAuthenticationToken authentication =
+				authorizationConsentRequestAuthentication(registeredClient, this.principal)
+						.scopes(authorizedScopes)		// Approve all scopes
+						.build();
+		when(this.authorizationService.findByToken(eq(authentication.getState()), eq(STATE_TOKEN_TYPE)))
+				.thenReturn(authorization);
+
+		@SuppressWarnings("unchecked")
+		Customizer<OAuth2AuthenticationContext> authorizationConsentCustomizer = mock(Customizer.class);
+		this.authenticationProvider.setAuthorizationConsentCustomizer(authorizationConsentCustomizer);
+
+		OAuth2AuthorizationCodeRequestAuthenticationToken authenticationResult =
+				(OAuth2AuthorizationCodeRequestAuthenticationToken) this.authenticationProvider.authenticate(authentication);
+
+		assertAuthorizationConsentRequestWithAuthorizationCodeResult(registeredClient, authorization, authenticationResult);
+
+		ArgumentCaptor<OAuth2AuthenticationContext> contextCaptor = ArgumentCaptor.forClass(OAuth2AuthenticationContext.class);
+		verify(authorizationConsentCustomizer).customize(contextCaptor.capture());
+
+		OAuth2AuthenticationContext context = contextCaptor.getValue();
+		assertThat((Authentication) context.getAuthentication()).isEqualTo(authentication);
+		assertThat(context.get(OAuth2AuthorizationConsent.Builder.class)).isInstanceOf(OAuth2AuthorizationConsent.Builder.class);
+		assertThat(context.get(OAuth2Authorization.class)).isInstanceOf(OAuth2Authorization.class);
+		assertThat(context.get(OAuth2AuthorizationRequest.class)).isInstanceOf(OAuth2AuthorizationRequest.class);
+	}
+
+	private void assertAuthorizationConsentRequestWithAuthorizationCodeResult(
+			RegisteredClient registeredClient,
+			OAuth2Authorization authorization,
+			OAuth2AuthorizationCodeRequestAuthenticationToken authenticationResult) {
+		OAuth2AuthorizationRequest authorizationRequest = authorization.getAttribute(OAuth2AuthorizationRequest.class.getName());
+		Set<String> authorizedScopes = authorizationRequest.getScopes();
+
 		ArgumentCaptor<OAuth2AuthorizationConsent> authorizationConsentCaptor = ArgumentCaptor.forClass(OAuth2AuthorizationConsent.class);
 		verify(this.authorizationConsentService).save(authorizationConsentCaptor.capture());
 		OAuth2AuthorizationConsent authorizationConsent = authorizationConsentCaptor.getValue();