WebSession supports changeSessionId

Issue: SPR-15571

Author: rstoyanchev

spring-web/src/main/java/org/springframework/web/server/WebSession.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2016 the original author or authors.
+ * Copyright 2002-2017 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -101,6 +101,14 @@ default <T> T getAttributeOrDefault(String name, T defaultValue) {
 	 */
 	boolean isStarted();
 
+	/**
+	 * Generate a new id for the session and update the underlying session
+	 * storage to reflect the new id. After a successful call {@link #getId()}
+	 * reflects the new session id.
+	 * @return completion notification (success or error)
+	 */
+	Mono<Void> changeSessionId();
+
 	/**
 	 * Save the session persisting attributes (e.g. if stored remotely) and also
 	 * sending the session id to the client if the session is new.

spring-web/src/main/java/org/springframework/web/server/session/DefaultWebSession.java

@@ -21,11 +21,13 @@
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 
 import reactor.core.publisher.Mono;
 
 import org.springframework.util.Assert;
+import org.springframework.util.IdGenerator;
 import org.springframework.web.server.WebSession;
 
 /**
@@ -36,12 +38,16 @@
  */
 class DefaultWebSession implements WebSession {
 
-	private final String id;
+	private final AtomicReference<String> id;
+
+	private final IdGenerator idGenerator;
 
 	private final Map<String, Object> attributes;
 
 	private final Clock clock;
 
+	private final BiFunction<String, WebSession, Mono<Void>> changeIdOperation;
+
 	private final Function<WebSession, Mono<Void>> saveOperation;
 
 	private final Instant creationTime;
@@ -55,14 +61,22 @@ class DefaultWebSession implements WebSession {
 
 	/**
 	 * Constructor for creating a brand, new session.
-	 * @param id the session id
+	 * @param idGenerator the session id generator
 	 * @param clock for access to current time
 	 */
-	DefaultWebSession(String id, Clock clock, Function<WebSession, Mono<Void>> saveOperation) {
-		Assert.notNull(id, "'id' is required.");
+	DefaultWebSession(IdGenerator idGenerator, Clock clock,
+			BiFunction<String, WebSession, Mono<Void>> changeIdOperation,
+			Function<WebSession, Mono<Void>> saveOperation) {
+
+		Assert.notNull(idGenerator, "'idGenerator' is required.");
 		Assert.notNull(clock, "'clock' is required.");
-		this.id = id;
+		Assert.notNull(changeIdOperation, "'changeIdOperation' is required.");
+		Assert.notNull(saveOperation, "'saveOperation' is required.");
+
+		this.id = new AtomicReference<>(String.valueOf(idGenerator.generateId()));
+		this.idGenerator = idGenerator;
 		this.clock = clock;
+		this.changeIdOperation = changeIdOperation;
 		this.saveOperation = saveOperation;
 		this.attributes = new ConcurrentHashMap<>();
 		this.creationTime = Instant.now(clock);
@@ -81,12 +95,14 @@ class DefaultWebSession implements WebSession {
 			Function<WebSession, Mono<Void>> saveOperation) {
 
 		this.id = existingSession.id;
+		this.idGenerator = existingSession.idGenerator;
 		this.attributes = existingSession.attributes;
 		this.clock = existingSession.clock;
+		this.changeIdOperation = existingSession.changeIdOperation;
+		this.saveOperation = saveOperation;
 		this.creationTime = existingSession.creationTime;
 		this.lastAccessTime = lastAccessTime;
 		this.maxIdleTime = existingSession.maxIdleTime;
-		this.saveOperation = saveOperation;
 		this.state = existingSession.state;
 	}
 
@@ -95,19 +111,21 @@ class DefaultWebSession implements WebSession {
 	 */
 	DefaultWebSession(DefaultWebSession existingSession, Instant lastAccessTime) {
 		this.id = existingSession.id;
+		this.idGenerator = existingSession.idGenerator;
 		this.attributes = existingSession.attributes;
 		this.clock = existingSession.clock;
+		this.changeIdOperation = existingSession.changeIdOperation;
+		this.saveOperation = existingSession.saveOperation;
 		this.creationTime = existingSession.creationTime;
 		this.lastAccessTime = lastAccessTime;
 		this.maxIdleTime = existingSession.maxIdleTime;
-		this.saveOperation = existingSession.saveOperation;
 		this.state = existingSession.state;
 	}
 
 
 	@Override
 	public String getId() {
-		return this.id;
+		return this.id.get();
 	}
 
 	@Override
@@ -151,6 +169,14 @@ public boolean isStarted() {
 		return (State.STARTED.equals(value) || (State.NEW.equals(value) && !getAttributes().isEmpty()));
 	}
 
+	@Override
+	public Mono<Void> changeSessionId() {
+		String oldId = this.id.get();
+		String newId = String.valueOf(this.idGenerator.generateId());
+		this.id.set(newId);
+		return this.changeIdOperation.apply(oldId, this).doOnError(ex -> this.id.set(oldId));
+	}
+
 	@Override
 	public Mono<Void> save() {
 		return this.saveOperation.apply(this);

spring-web/src/main/java/org/springframework/web/server/session/DefaultWebSessionManager.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2016 the original author or authors.
+ * Copyright 2002-2017 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,12 +19,13 @@
 import java.time.Instant;
 import java.time.ZoneId;
 import java.util.List;
-import java.util.UUID;
 
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import org.springframework.util.Assert;
+import org.springframework.util.IdGenerator;
+import org.springframework.util.JdkIdGenerator;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebSession;
 
@@ -39,6 +40,9 @@
  */
 public class DefaultWebSessionManager implements WebSessionManager {
 
+	private static final IdGenerator idGenerator = new JdkIdGenerator();
+
+
 	private WebSessionIdResolver sessionIdResolver = new CookieWebSessionIdResolver();
 
 	private WebSessionStore sessionStore = new InMemoryWebSessionStore();
@@ -159,10 +163,10 @@ private boolean hasNewSessionId(ServerWebExchange exchange, WebSession session)
 	}
 
 	private Mono<DefaultWebSession> createSession(ServerWebExchange exchange) {
-		return Mono.fromSupplier(() -> {
-			String id = UUID.randomUUID().toString();
-			return new DefaultWebSession(id, getClock(), sess -> saveSession(exchange, sess));
-		});
+		return Mono.fromSupplier(() ->
+				new DefaultWebSession(idGenerator, getClock(),
+						(oldId, session) -> this.sessionStore.changeSessionId(oldId, session),
+						session -> saveSession(exchange, session)));
 	}
 
 }

spring-web/src/main/java/org/springframework/web/server/session/InMemoryWebSessionStore.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2016 the original author or authors.
+ * Copyright 2002-2017 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -44,6 +44,13 @@ public Mono<WebSession> retrieveSession(String id) {
 		return (this.sessions.containsKey(id) ? Mono.just(this.sessions.get(id)) : Mono.empty());
 	}
 
+	@Override
+	public Mono<Void> changeSessionId(String oldId, WebSession session) {
+		this.sessions.remove(oldId);
+		this.sessions.put(session.getId(), session);
+		return Mono.empty();
+	}
+
 	@Override
 	public Mono<Void> removeSession(String id) {
 		this.sessions.remove(id);

spring-web/src/main/java/org/springframework/web/server/session/WebSessionStore.java

@@ -41,6 +41,18 @@ public interface WebSessionStore {
 	 */
 	Mono<WebSession> retrieveSession(String sessionId);
 
+	/**
+	 * Update WebSession data storage to reflect a change in session id.
+	 * <p>Note that the same can be achieved via a combination of
+	 * {@link #removeSession} + {@link #storeSession}. The purpose of this method
+	 * is to allow a more efficient replacement of the session id mapping
+	 * without replacing and storing the session with all of its data.
+	 * @param oldId the previous session id
+	 * @param session the session reflecting the changed session id
+	 * @return completion notification (success or error)
+	 */
+	Mono<Void> changeSessionId(String oldId, WebSession session);
+
 	/**
 	 * Remove the WebSession for the specified id.
 	 * @param sessionId the id of the session to remove

spring-web/src/test/java/org/springframework/web/server/session/DefaultWebSessionManagerTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2016 the original author or authors.
+ * Copyright 2002-2017 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
+import java.time.ZoneId;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -31,6 +32,8 @@
 import org.springframework.lang.Nullable;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpRequest;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpResponse;
+import org.springframework.util.IdGenerator;
+import org.springframework.util.JdkIdGenerator;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebSession;
 import org.springframework.web.server.adapter.DefaultServerWebExchange;
@@ -44,10 +47,16 @@
 import static org.junit.Assert.assertSame;
 
 /**
+ * Unit tests for {@link DefaultWebSessionManager}.
  * @author Rossen Stoyanchev
  */
 public class DefaultWebSessionManagerTests {
 
+	private static final Clock CLOCK = Clock.system(ZoneId.of("GMT"));
+
+	private static final IdGenerator idGenerator = new JdkIdGenerator();
+
+
 	private DefaultWebSessionManager manager;
 
 	private TestWebSessionIdResolver idResolver;
@@ -105,9 +114,10 @@ public void startSessionImplicitly() throws Exception {
 
 	@Test
 	public void existingSession() throws Exception {
-		DefaultWebSession existing = new DefaultWebSession("1", Clock.systemDefaultZone(), s -> Mono.empty());
+		DefaultWebSession existing = createDefaultWebSession();
+		String id = existing.getId();
 		this.manager.getSessionStore().storeSession(existing);
-		this.idResolver.setIdsToResolve(Collections.singletonList("1"));
+		this.idResolver.setIdsToResolve(Collections.singletonList(id));
 
 		WebSession actual = this.manager.getSession(this.exchange).block();
 		assertNotNull(actual);
@@ -116,10 +126,9 @@ public void existingSession() throws Exception {
 
 	@Test
 	public void existingSessionIsExpired() throws Exception {
-		Clock clock = Clock.systemDefaultZone();
-		DefaultWebSession existing = new DefaultWebSession("1", clock, s -> Mono.empty());
+		DefaultWebSession existing = createDefaultWebSession();
 		existing.start();
-		Instant lastAccessTime = Instant.now(clock).minus(Duration.ofMinutes(31));
+		Instant lastAccessTime = Instant.now(CLOCK).minus(Duration.ofMinutes(31));
 		existing = new DefaultWebSession(existing, lastAccessTime, s -> Mono.empty());
 		this.manager.getSessionStore().storeSession(existing);
 		this.idResolver.setIdsToResolve(Collections.singletonList("1"));
@@ -129,16 +138,21 @@ public void existingSessionIsExpired() throws Exception {
 	}
 
 	@Test
-	public void multipleSessions() throws Exception {
-		DefaultWebSession existing = new DefaultWebSession("3", Clock.systemDefaultZone(), s -> Mono.empty());
+	public void multipleSessionIds() throws Exception {
+		DefaultWebSession existing = createDefaultWebSession();
+		String id = existing.getId();
 		this.manager.getSessionStore().storeSession(existing);
-		this.idResolver.setIdsToResolve(Arrays.asList("1", "2", "3"));
+		this.idResolver.setIdsToResolve(Arrays.asList("neither-this", "nor-that", id));
 
 		WebSession actual = this.manager.getSession(this.exchange).block();
 		assertNotNull(actual);
 		assertEquals(existing.getId(), actual.getId());
 	}
 
+	private DefaultWebSession createDefaultWebSession() {
+		return new DefaultWebSession(idGenerator, CLOCK, (s, session) -> Mono.empty(), s -> Mono.empty());
+	}
+
 
 	private static class TestWebSessionIdResolver implements WebSessionIdResolver {