ServerCsrfTokenRepository.saveToken return Mono<CsrfToken>

Fixes gh-4800

web/src/main/java/org/springframework/security/web/server/csrf/ServerCsrfTokenRepository.java

@@ -46,7 +46,7 @@ public interface ServerCsrfTokenRepository {
 	 * @param exchange the {@link ServerWebExchange} to use
 	 * @param token the {@link CsrfToken} to save or null to delete
 	 */
-	Mono<Void> saveToken(ServerWebExchange exchange, CsrfToken token);
+	Mono<CsrfToken> saveToken(ServerWebExchange exchange, CsrfToken token);
 
 	/**
 	 * Loads the expected {@link CsrfToken} from the {@link ServerWebExchange}

web/src/main/java/org/springframework/security/web/server/csrf/WebSessionServerCsrfTokenRepository.java

@@ -50,28 +50,29 @@ public class WebSessionServerCsrfTokenRepository
 	@Override
 	public Mono<CsrfToken> generateToken(ServerWebExchange exchange) {
 		return Mono.defer(() -> Mono.just(createCsrfToken()))
-			.flatMap(token -> save(exchange, token));
+			.flatMap(token -> saveToken(exchange, token));
 	}
 
 	@Override
-	public Mono<Void> saveToken(ServerWebExchange exchange, CsrfToken token) {
-		return save(exchange, token)
-			.then();
-	}
-
-	private Mono<CsrfToken> save(ServerWebExchange exchange, CsrfToken token) {
+	public Mono<CsrfToken> saveToken(ServerWebExchange exchange, CsrfToken token) {
 		return exchange.getSession()
 			.map(WebSession::getAttributes)
 			.flatMap( attrs -> save(attrs, token));
 	}
 
 	private Mono<CsrfToken> save(Map<String,Object> attributes, CsrfToken token) {
+		return Mono.defer(() -> {
+			putToken(attributes, token);
+			return Mono.justOrEmpty(token);
+		});
+	}
+
+	private void putToken(Map<String,Object> attributes, CsrfToken token) {
 		if(token == null) {
 			attributes.remove(this.sessionAttributeName);
 		} else {
 			attributes.put(this.sessionAttributeName, token);
 		}
-		return Mono.justOrEmpty(token);
 	}
 
 	@Override

web/src/test/java/org/springframework/security/web/server/csrf/WebSessionServerCsrfTokenRepositoryTests.java

@@ -70,9 +70,10 @@ public class WebSessionServerCsrfTokenRepositoryTests {
 		CsrfToken token = new DefaultCsrfToken("h","p", "t");
 		String attrName = "ATTR";
 		this.repository.setSessionAttributeName(attrName);
-		Mono<Void> result = this.repository.saveToken(this.exchange, token);
+		Mono<CsrfToken> result = this.repository.saveToken(this.exchange, token);
 
 		StepVerifier.create(result)
+			.consumeNextWith(n -> assertThat(n).isEqualTo(token))
 			.verifyComplete();
 
 		WebSession session = this.exchange.getSession().block();
@@ -86,7 +87,7 @@ public class WebSessionServerCsrfTokenRepositoryTests {
 		CsrfToken token = new DefaultCsrfToken("h","p", "t");
 		this.repository.saveToken(this.exchange, token).block();
 
-		Mono<Void> result = this.repository.saveToken(this.exchange, null);
+		Mono<CsrfToken> result = this.repository.saveToken(this.exchange, null);
 		StepVerifier.create(result)
 			.verifyComplete();