From 1752928d962970949de2ce7a02b639d3b4dd1570 Mon Sep 17 00:00:00 2001
From: Greg Turnquist <gturnquist@pivotal.io>
Date: Thu, 7 Sep 2017 16:54:58 -0500
Subject: [PATCH] Configure WebSessionManager's WebSessionIdResolver by bean
 definition

Allow a WebSessionIdResolver registered as a Spring bean to be wired into the WebSessionManager.
---
 .../SpringWebSessionConfiguration.java        | 13 ++++++
 .../SpringWebSessionConfigurationTests.java   | 40 +++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/spring-session-core/src/main/java/org/springframework/session/SpringWebSessionConfiguration.java b/spring-session-core/src/main/java/org/springframework/session/SpringWebSessionConfiguration.java
index e8e7210e..123ff39f 100644
--- a/spring-session-core/src/main/java/org/springframework/session/SpringWebSessionConfiguration.java
+++ b/spring-session-core/src/main/java/org/springframework/session/SpringWebSessionConfiguration.java
@@ -15,11 +15,13 @@
  */
 package org.springframework.session;
 
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.session.web.server.session.SpringSessionWebSessionStore;
 import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
 import org.springframework.web.server.session.DefaultWebSessionManager;
+import org.springframework.web.server.session.WebSessionIdResolver;
 import org.springframework.web.server.session.WebSessionManager;
 
 /**
@@ -34,6 +36,12 @@ import org.springframework.web.server.session.WebSessionManager;
 @Configuration
 public class SpringWebSessionConfiguration {
 
+	/**
+	 * Optional override of default {@link WebSessionIdResolver}.
+	 */
+	@Autowired(required = false)
+	WebSessionIdResolver webSessionIdResolver;
+
 	/**
 	 * Configure a {@link WebSessionManager} using a provided {@link ReactorSessionRepository}.
 	 *
@@ -45,6 +53,11 @@ public class SpringWebSessionConfiguration {
 		SpringSessionWebSessionStore<? extends Session> sessionStore = new SpringSessionWebSessionStore<>(repository);
 		DefaultWebSessionManager manager = new DefaultWebSessionManager();
 		manager.setSessionStore(sessionStore);
+
+		if (this.webSessionIdResolver != null) {
+			manager.setSessionIdResolver(this.webSessionIdResolver);
+		}
+		
 		return manager;
 	}
 }
diff --git a/spring-session-core/src/test/java/org/springframework/session/SpringWebSessionConfigurationTests.java b/spring-session-core/src/test/java/org/springframework/session/SpringWebSessionConfigurationTests.java
index e3da0b55..77f4d969 100644
--- a/spring-session-core/src/test/java/org/springframework/session/SpringWebSessionConfigurationTests.java
+++ b/spring-session-core/src/test/java/org/springframework/session/SpringWebSessionConfigurationTests.java
@@ -21,6 +21,10 @@ import org.springframework.beans.factory.UnsatisfiedDependencyException;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
+import org.springframework.web.server.session.CookieWebSessionIdResolver;
+import org.springframework.web.server.session.DefaultWebSessionManager;
+import org.springframework.web.server.session.HeaderWebSessionIdResolver;
+import org.springframework.web.server.session.WebSessionIdResolver;
 import org.springframework.web.server.session.WebSessionManager;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -62,6 +66,28 @@ public class SpringWebSessionConfigurationTests {
 				.withMessageContaining("No qualifying bean of type '" + ReactorSessionRepository.class.getCanonicalName());
 	}
 
+	@Test
+	public void defaultSessionIdResolverShouldBeCookieBased() {
+
+		AnnotationConfigApplicationContext ctx = new AnnotationConfigApplicationContext();
+		ctx.register(GoodConfig.class);
+		ctx.refresh();
+
+		DefaultWebSessionManager manager = ctx.getBean(DefaultWebSessionManager.class);
+		assertThat(manager.getSessionIdResolver().getClass()).isAssignableFrom(CookieWebSessionIdResolver.class);
+	}
+
+	@Test
+	public void providedSessionIdResolverShouldBePickedUpAutomatically() {
+
+		AnnotationConfigApplicationContext ctx = new AnnotationConfigApplicationContext();
+		ctx.register(OverrideSessionIdResolver.class);
+		ctx.refresh();
+
+		DefaultWebSessionManager manager = ctx.getBean(DefaultWebSessionManager.class);
+		assertThat(manager.getSessionIdResolver().getClass()).isAssignableFrom(HeaderWebSessionIdResolver.class);
+	}
+
 	/**
 	 * A configuration with all the right parts.
 	 */
@@ -84,4 +110,18 @@ public class SpringWebSessionConfigurationTests {
 	static class BadConfig {
 
 	}
+
+	@EnableSpringWebSession
+	static class OverrideSessionIdResolver {
+
+		@Bean
+		ReactorSessionRepository<?> reactorSessionRepository() {
+			return new MapReactorSessionRepository();
+		}
+
+		@Bean
+		WebSessionIdResolver alternateWebSessionIdResolver() {
+			return new HeaderWebSessionIdResolver();
+		}
+	}
 }