From 3b43fdac40a431925fe8dc32bfc8133435df80ce Mon Sep 17 00:00:00 2001
From: Mark Paluch <mpaluch@pivotal.io>
Date: Fri, 24 Jan 2020 15:14:48 +0100
Subject: [PATCH] Pass-thru VaultToken obtained from ReactiveSessionManager

We now pass-thru the token obtained from ReactiveSessionManager instead of caching it.
Previously, we used a method reference and the method reference cached the SessionToken Mono instead of requesting a fresh token each time from the session manager.

Closes gh-382.
---
 .../VaultReactiveBootstrapConfiguration.java  |  9 +++++-
 ...ltReactiveBootstrapConfigurationTests.java | 31 ++++++++++++++++++-
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/spring-cloud-vault-config/src/main/java/org/springframework/cloud/vault/config/VaultReactiveBootstrapConfiguration.java b/spring-cloud-vault-config/src/main/java/org/springframework/cloud/vault/config/VaultReactiveBootstrapConfiguration.java
index 5390d9c3..670fa34c 100644
--- a/spring-cloud-vault-config/src/main/java/org/springframework/cloud/vault/config/VaultReactiveBootstrapConfiguration.java
+++ b/spring-cloud-vault-config/src/main/java/org/springframework/cloud/vault/config/VaultReactiveBootstrapConfiguration.java
@@ -55,6 +55,7 @@ import org.springframework.vault.core.ReactiveVaultOperations;
 import org.springframework.vault.core.ReactiveVaultTemplate;
 import org.springframework.vault.support.ClientOptions;
 import org.springframework.vault.support.SslConfiguration;
+import org.springframework.vault.support.VaultToken;
 import org.springframework.web.reactive.function.client.WebClient;
 
 /**
@@ -156,7 +157,13 @@ public class VaultReactiveBootstrapConfiguration {
 	@Bean
 	@ConditionalOnMissingBean
 	public SessionManager vaultSessionManager(ReactiveSessionManager sessionManager) {
-		return sessionManager.getSessionToken()::block;
+		return () -> {
+
+			VaultToken token = sessionManager.getSessionToken().block();
+			Assert.state(token != null,
+					"ReactiveSessionManager returned a null VaultToken");
+			return token;
+		};
 	}
 
 	/**
diff --git a/spring-cloud-vault-config/src/test/java/org/springframework/cloud/vault/config/VaultReactiveBootstrapConfigurationTests.java b/spring-cloud-vault-config/src/test/java/org/springframework/cloud/vault/config/VaultReactiveBootstrapConfigurationTests.java
index cc56fd06..5736c3fd 100644
--- a/spring-cloud-vault-config/src/test/java/org/springframework/cloud/vault/config/VaultReactiveBootstrapConfigurationTests.java
+++ b/spring-cloud-vault-config/src/test/java/org/springframework/cloud/vault/config/VaultReactiveBootstrapConfigurationTests.java
@@ -16,6 +16,8 @@
 
 package org.springframework.cloud.vault.config;
 
+import java.util.concurrent.atomic.AtomicLong;
+
 import org.junit.Test;
 import reactor.core.publisher.Mono;
 
@@ -27,6 +29,7 @@ import org.springframework.context.annotation.Configuration;
 import org.springframework.vault.authentication.AuthenticationSteps;
 import org.springframework.vault.authentication.AuthenticationStepsFactory;
 import org.springframework.vault.authentication.LifecycleAwareSessionManager;
+import org.springframework.vault.authentication.ReactiveSessionManager;
 import org.springframework.vault.authentication.SessionManager;
 import org.springframework.vault.authentication.SimpleSessionManager;
 import org.springframework.vault.authentication.VaultTokenSupplier;
@@ -109,6 +112,20 @@ public class VaultReactiveBootstrapConfigurationTests {
 				});
 	}
 
+	@Test
+	public void sessionManagerBridgeShouldNotCacheTokens() {
+		this.contextRunner.withUserConfiguration(TokenSupplierConfiguration.class,
+				CustomSessionManager.class).run(context -> {
+
+					SessionManager sessionManager = context.getBean(SessionManager.class);
+
+					assertThat(sessionManager.getSessionToken().getToken())
+							.isEqualTo("token-1");
+					assertThat(sessionManager.getSessionToken().getToken())
+							.isEqualTo("token-2");
+				});
+	}
+
 	@Configuration
 	static class AuthenticationFactoryConfiguration {
 
@@ -124,7 +141,19 @@ public class VaultReactiveBootstrapConfigurationTests {
 
 		@Bean
 		VaultTokenSupplier vaultTokenSupplier() {
-			return () -> Mono.just(VaultToken.of("foo"));
+			AtomicLong counter = new AtomicLong();
+			return () -> Mono.just(VaultToken.of("token-" + counter.incrementAndGet()));
+		}
+
+	}
+
+	@Configuration
+	static class CustomSessionManager {
+
+		@Bean
+		ReactiveSessionManager reactiveVaultSessionManager(
+				VaultTokenSupplier tokenSupplier) {
+			return tokenSupplier::getVaultToken;
 		}
 
 	}