From 65ce1454e003c4f77ee527040e2251eefe0bd7ad Mon Sep 17 00:00:00 2001
From: weixsun <weixsun@aliyun.com>
Date: Wed, 12 May 2021 15:56:06 +0800
Subject: [PATCH] Allow direct use of exposed WebSessionIdResolver bean

See gh-26437
---
 .../reactive/WebFluxAutoConfiguration.java    | 13 ++++++---
 .../WebFluxAutoConfigurationTests.java        | 27 +++++++++++++++++++
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/reactive/WebFluxAutoConfiguration.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/reactive/WebFluxAutoConfiguration.java
index b1e6969ee7..fc24a76274 100644
--- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/reactive/WebFluxAutoConfiguration.java
+++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/reactive/WebFluxAutoConfiguration.java
@@ -77,6 +77,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,12 +308,16 @@ public class WebFluxAutoConfiguration {
 
 		@Bean
 		@ConditionalOnMissingBean(name = WebHttpHandlerBuilder.WEB_SESSION_MANAGER_BEAN_NAME)
-		public WebSessionManager webSessionManager() {
+		public WebSessionManager webSessionManager(ObjectProvider<WebSessionIdResolver> webSessionIdResolvers) {
 			DefaultWebSessionManager webSessionManager = new DefaultWebSessionManager();
-			CookieWebSessionIdResolver webSessionIdResolver = new CookieWebSessionIdResolver();
-			webSessionIdResolver.addCookieInitializer((cookie) -> cookie
+			if (webSessionIdResolvers.getIfAvailable() != null) {
+				webSessionManager.setSessionIdResolver(webSessionIdResolvers.getIfAvailable());
+				return webSessionManager;
+			}
+			CookieWebSessionIdResolver cookieWebSessionIdResolver = new CookieWebSessionIdResolver();
+			cookieWebSessionIdResolver.addCookieInitializer((cookie) -> cookie
 					.sameSite(this.webFluxProperties.getSession().getCookie().getSameSite().attribute()));
-			webSessionManager.setSessionIdResolver(webSessionIdResolver);
+			webSessionManager.setSessionIdResolver(cookieWebSessionIdResolver);
 			return webSessionManager;
 		}
 
diff --git a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/reactive/WebFluxAutoConfigurationTests.java b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/reactive/WebFluxAutoConfigurationTests.java
index d63e701e7d..306c1cc1e4 100644
--- a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/reactive/WebFluxAutoConfigurationTests.java
+++ b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/reactive/WebFluxAutoConfigurationTests.java
@@ -87,6 +87,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;
 
@@ -562,6 +564,19 @@ class WebFluxAutoConfigurationTests {
 								HighPrecedenceConfigurer.class, WebFluxConfig.class, LowPrecedenceConfigurer.class));
 	}
 
+	@Test
+	void customWebSessionIdResolverShouldBeApplied() {
+		this.contextRunner.withUserConfiguration(CustomWebSessionIdResolvers.class).run((context) -> {
+			MockServerHttpRequest request = MockServerHttpRequest.get("/").build();
+			MockServerWebExchange exchange = MockServerWebExchange.from(request);
+			WebSessionManager webSessionManager = context.getBean(WebSessionManager.class);
+			WebSession webSession = webSessionManager.getSession(exchange).block();
+			webSession.start();
+			exchange.getResponse().setComplete().block();
+			assertThat(exchange.getResponse().getCookies().get("JSESSIONID")).isNotEmpty();
+		});
+	}
+
 	@Test
 	void customSameSteConfigurationShouldBeApplied() {
 		this.contextRunner.withPropertyValues("spring.webflux.session.cookie.same-site:strict").run((context) -> {
@@ -584,6 +599,18 @@ class WebFluxAutoConfigurationTests {
 		return Collections.emptyMap();
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	static class CustomWebSessionIdResolvers {
+
+		@Bean
+		WebSessionIdResolver webSessionIdResolver() {
+			CookieWebSessionIdResolver resolver = new CookieWebSessionIdResolver();
+			resolver.setCookieName("JSESSIONID");
+			return resolver;
+		}
+
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	static class CustomArgumentResolvers {