Add ServerWebExchange parameter to AuthorizationRequestCustomizer

Closes gh-16306

Author: kse-music

Diff for: oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/web/server/DefaultServerOAuth2AuthorizationRequestResolver.java

@@ -22,6 +22,7 @@
 import java.util.Base64;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 
 import reactor.core.publisher.Mono;
@@ -92,7 +93,8 @@ public class DefaultServerOAuth2AuthorizationRequestResolver implements ServerOA
 
 	private final ReactiveClientRegistrationRepository clientRegistrationRepository;
 
-	private Consumer<OAuth2AuthorizationRequest.Builder> authorizationRequestCustomizer = (customizer) -> {
+	private BiConsumer<OAuth2AuthorizationRequest.Builder, ServerWebExchange> authorizationRequestCustomizer = (
+			customizer, exchange) -> {
 	};
 
 	/**
@@ -148,11 +150,29 @@ public Mono<OAuth2AuthorizationRequest> resolve(ServerWebExchange exchange, Stri
 	 * @param authorizationRequestCustomizer the {@code Consumer} to be provided the
 	 * {@link OAuth2AuthorizationRequest.Builder}
 	 * @since 5.3
+	 * @deprecated Use {@link #setAuthorizationRequestCustomizer(BiConsumer) } instead
 	 * @see OAuth2AuthorizationRequestCustomizers
 	 */
+	@Deprecated
 	public final void setAuthorizationRequestCustomizer(
 			Consumer<OAuth2AuthorizationRequest.Builder> authorizationRequestCustomizer) {
 		Assert.notNull(authorizationRequestCustomizer, "authorizationRequestCustomizer cannot be null");
+		this.authorizationRequestCustomizer = (customizer, exchange) -> authorizationRequestCustomizer
+			.accept(customizer);
+	}
+
+	/**
+	 * Sets the {@code BiConsumer} to be provided the
+	 * {@link OAuth2AuthorizationRequest.Builder} and {@link ServerWebExchange} allowing
+	 * for further customizations.
+	 * @param authorizationRequestCustomizer the {@code BiConsumer} to be provided the
+	 * {@link OAuth2AuthorizationRequest.Builder} and {@link ServerWebExchange}
+	 * @since 6.5
+	 * @see OAuth2AuthorizationRequestCustomizers
+	 */
+	public final void setAuthorizationRequestCustomizer(
+			BiConsumer<OAuth2AuthorizationRequest.Builder, ServerWebExchange> authorizationRequestCustomizer) {
+		Assert.notNull(authorizationRequestCustomizer, "authorizationRequestCustomizer cannot be null");
 		this.authorizationRequestCustomizer = authorizationRequestCustomizer;
 	}
 
@@ -175,7 +195,7 @@ private OAuth2AuthorizationRequest authorizationRequest(ServerWebExchange exchan
 				.state(DEFAULT_STATE_GENERATOR.generateKey());
 		// @formatter:on
 
-		this.authorizationRequestCustomizer.accept(builder);
+		this.authorizationRequestCustomizer.accept(builder, exchange);
 
 		return builder.build();
 	}

Diff for: oauth2/oauth2-client/src/test/java/org/springframework/security/oauth2/client/web/server/DefaultServerOAuth2AuthorizationRequestResolverTests.java

@@ -16,6 +16,8 @@
 
 package org.springframework.security.oauth2.client.web.server;
 
+import java.util.function.Consumer;
+
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -68,7 +70,8 @@ public void setup() {
 
 	@Test
 	public void setAuthorizationRequestCustomizerWhenNullThenThrowIllegalArgumentException() {
-		assertThatIllegalArgumentException().isThrownBy(() -> this.resolver.setAuthorizationRequestCustomizer(null));
+		assertThatIllegalArgumentException().isThrownBy(() -> this.resolver
+			.setAuthorizationRequestCustomizer((Consumer<OAuth2AuthorizationRequest.Builder>) null));
 	}
 
 	@Test
@@ -254,6 +257,18 @@ public void resolveWhenAuthorizationRequestCustomizerOverridesParameterThenQuery
 					+ "nonce=([a-zA-Z0-9\\-\\.\\_\\~]){43}&" + "appid=client-id");
 	}
 
+	@Test
+	public void resolveWhenAuthorizationRequestCustomizerAdditionParameterFromRequest() {
+		given(this.clientRegistrationRepository.findByRegistrationId(any()))
+			.willReturn(Mono.just(TestClientRegistrations.clientRegistration().scope(OidcScopes.OPENID).build()));
+		this.resolver.setAuthorizationRequestCustomizer((builder, exchange) -> builder.parameters((params) -> {
+			params.put("aa", exchange.getRequest().getQueryParams().getFirst("a"));
+			params.put("bb", exchange.getRequest().getQueryParams().getFirst("b"));
+		}));
+		OAuth2AuthorizationRequest authorizationRequest = resolve("/oauth2/authorization/registration-id?a=A&b=B");
+		assertThat(authorizationRequest.getAuthorizationRequestUri()).contains("aa=A&bb=B");
+	}
+
 	private OAuth2AuthorizationRequest resolve(String path) {
 		ServerWebExchange exchange = MockServerWebExchange.from(MockServerHttpRequest.get(path));
 		return this.resolver.resolve(exchange).block();