Merge pull request #26437 from weixsun

* pr/26437:
  Polish "Allow direct use of exposed WebSessionIdResolver bean"
  Allow direct use of exposed WebSessionIdResolver bean

Closes gh-26437

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/reactive/WebFluxAutoConfiguration.java

@@ -17,6 +17,7 @@
 package org.springframework.boot.autoconfigure.web.reactive;
 
 import java.time.Duration;
+import java.util.function.Supplier;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -77,6 +78,7 @@ import org.springframework.web.server.i18n.FixedLocaleContextResolver;
 import org.springframework.web.server.i18n.LocaleContextResolver;
 import org.springframework.web.server.session.CookieWebSessionIdResolver;
 import org.springframework.web.server.session.DefaultWebSessionManager;
+import org.springframework.web.server.session.WebSessionIdResolver;
 import org.springframework.web.server.session.WebSessionManager;
 
 /**
@@ -307,15 +309,21 @@ public class WebFluxAutoConfiguration {
 
 		@Bean
 		@ConditionalOnMissingBean(name = WebHttpHandlerBuilder.WEB_SESSION_MANAGER_BEAN_NAME)
-		public WebSessionManager webSessionManager() {
+		public WebSessionManager webSessionManager(ObjectProvider<WebSessionIdResolver> webSessionIdResolver) {
 			DefaultWebSessionManager webSessionManager = new DefaultWebSessionManager();
-			CookieWebSessionIdResolver webSessionIdResolver = new CookieWebSessionIdResolver();
+			webSessionManager.setSessionIdResolver(webSessionIdResolver.getIfAvailable(cookieWebSessionIdResolver()));
-			webSessionIdResolver.addCookieInitializer((cookie) -> cookie
-					.sameSite(this.webFluxProperties.getSession().getCookie().getSameSite().attribute()));
-			webSessionManager.setSessionIdResolver(webSessionIdResolver);
 			return webSessionManager;
 		}
 
+		private Supplier<WebSessionIdResolver> cookieWebSessionIdResolver() {
+			return () -> {
+				CookieWebSessionIdResolver webSessionIdResolver = new CookieWebSessionIdResolver();
+				webSessionIdResolver.addCookieInitializer((cookie) -> cookie
+						.sameSite(this.webFluxProperties.getSession().getCookie().getSameSite().attribute()));
+				return webSessionIdResolver;
+			};
+		}
+
 	}
 
 	@Configuration(proxyBeanMethods = false)

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/reactive/WebFluxAutoConfigurationTests.java

@@ -28,6 +28,7 @@ import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 import javax.validation.ValidatorFactory;
 
@@ -40,8 +41,10 @@ import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.autoconfigure.validation.ValidationAutoConfiguration;
 import org.springframework.boot.autoconfigure.validation.ValidatorAdapter;
 import org.springframework.boot.autoconfigure.web.reactive.WebFluxAutoConfiguration.WebFluxConfig;
+import org.springframework.boot.test.context.runner.ContextConsumer;
 import org.springframework.boot.test.context.runner.ReactiveWebApplicationContextRunner;
 import org.springframework.boot.web.codec.CodecCustomizer;
+import org.springframework.boot.web.reactive.context.ReactiveWebApplicationContext;
 import org.springframework.boot.web.reactive.filter.OrderedHiddenHttpMethodFilter;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.Bean;
@@ -87,6 +90,8 @@ import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
 import org.springframework.web.server.i18n.AcceptHeaderLocaleContextResolver;
 import org.springframework.web.server.i18n.FixedLocaleContextResolver;
 import org.springframework.web.server.i18n.LocaleContextResolver;
+import org.springframework.web.server.session.CookieWebSessionIdResolver;
+import org.springframework.web.server.session.WebSessionIdResolver;
 import org.springframework.web.server.session.WebSessionManager;
 import org.springframework.web.util.pattern.PathPattern;
 
@@ -563,17 +568,29 @@ class WebFluxAutoConfigurationTests {
 	}
 
 	@Test
-	void customSameSteConfigurationShouldBeApplied() {
+	void customWebSessionIdResolverShouldBeApplied() {
-		this.contextRunner.withPropertyValues("spring.webflux.session.cookie.same-site:strict").run((context) -> {
+		this.contextRunner.withUserConfiguration(CustomWebSessionIdResolver.class).run(assertExchangeWithSession(
+				(exchange) -> assertThat(exchange.getResponse().getCookies().get("TEST")).isNotEmpty()));
+	}
+
+	@Test
+	void customSameSiteConfigurationShouldBeApplied() {
+		this.contextRunner.withPropertyValues("spring.webflux.session.cookie.same-site:strict").run(
+				assertExchangeWithSession((exchange) -> assertThat(exchange.getResponse().getCookies().get("SESSION"))
+						.isNotEmpty().allMatch((cookie) -> cookie.getSameSite().equals("Strict"))));
+	}
+
+	private ContextConsumer<ReactiveWebApplicationContext> assertExchangeWithSession(
+			Consumer<MockServerWebExchange> exchange) {
+		return (context) -> {
 			MockServerHttpRequest request = MockServerHttpRequest.get("/").build();
-			MockServerWebExchange exchange = MockServerWebExchange.from(request);
+			MockServerWebExchange webExchange = MockServerWebExchange.from(request);
 			WebSessionManager webSessionManager = context.getBean(WebSessionManager.class);
-			WebSession webSession = webSessionManager.getSession(exchange).block();
+			WebSession webSession = webSessionManager.getSession(webExchange).block();
 			webSession.start();
-			exchange.getResponse().setComplete().block();
+			webExchange.getResponse().setComplete().block();
-			assertThat(exchange.getResponse().getCookies().get("SESSION")).isNotEmpty()
+			exchange.accept(webExchange);
-					.allMatch((cookie) -> cookie.getSameSite().equals("Strict"));
+		};
-		});
 	}
 
 	private Map<PathPattern, Object> getHandlerMap(ApplicationContext context) {
@@ -584,6 +601,18 @@ class WebFluxAutoConfigurationTests {
 		return Collections.emptyMap();
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	static class CustomWebSessionIdResolver {
+
+		@Bean
+		WebSessionIdResolver webSessionIdResolver() {
+			CookieWebSessionIdResolver resolver = new CookieWebSessionIdResolver();
+			resolver.setCookieName("TEST");
+			return resolver;
+		}
+
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	static class CustomArgumentResolvers {