Fix Lease revocation on SecretLeaseContainer destroy.

Closes gh-844
spring-projects · Dec 27, 2023 · 0bc5513 · 0bc5513
1 parent 47028c9
commit 0bc5513
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 7 deletions.
diff --git a/...g-vault-core/src/main/java/org/springframework/vault/core/lease/SecretLeaseContainer.java b/...g-vault-core/src/main/java/org/springframework/vault/core/lease/SecretLeaseContainer.java
@@ -32,6 +32,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import java.util.function.BiConsumer;
 
 import org.apache.commons.logging.Log;
@@ -410,6 +411,9 @@ else if (isRotatingGenericSecret(requestedSecret, secrets)) {
 			else if (renewalScheduler.isLeaseRotateOnly(lease, requestedSecret)) {
 				scheduleLeaseRotation(requestedSecret, lease, renewalScheduler);
 			}
+			else if (lease.hasLeaseId()) {
+				renewalScheduler.associateLease(lease);
+			}
 
 			callback.accept(secrets, lease);
 		}
@@ -502,11 +506,16 @@ public void destroy() throws Exception {
 				for (Entry<RequestedSecret, LeaseRenewalScheduler> entry : this.renewals.entrySet()) {
 
 					Lease lease = entry.getValue().getLease();
+					Lease previousLease = entry.getValue().getPreviousLease();
 					entry.getValue().disableScheduleRenewal();
 
 					if (lease != null && lease.hasLeaseId()) {
 						doRevokeLease(entry.getKey(), lease);
 					}
+
+					if (previousLease != null && previousLease.hasLeaseId()) {
+						doRevokeLease(entry.getKey(), previousLease);
+					}
 				}
 				this.renewals.clear();
 
@@ -844,12 +853,19 @@ protected void doRevokeLease(RequestedSecret requestedSecret, Lease lease) {
 	 */
 	static class LeaseRenewalScheduler {
 
+		private static final AtomicReferenceFieldUpdater<LeaseRenewalScheduler, Lease> CURRENT_UPDATER = AtomicReferenceFieldUpdater
+			.newUpdater(LeaseRenewalScheduler.class, Lease.class, "currentLeaseRef");
+
 		@SuppressWarnings("FieldMayBeFinal") // allow setting via reflection.
 		private static Log logger = LogFactory.getLog(LeaseRenewalScheduler.class);
 
 		private final TaskScheduler taskScheduler;
 
-		final AtomicReference<Lease> currentLeaseRef = new AtomicReference<>();
+		@Nullable
+		volatile Lease currentLeaseRef;
+
+		@Nullable
+		volatile Lease previousLeaseRef;
 
 		final Map<Lease, ScheduledFuture<?>> schedules = new ConcurrentHashMap<>();
 
@@ -884,8 +900,8 @@ void scheduleRenewal(RequestedSecret requestedSecret, RenewLease renewLease, Lea
 				}
 			}
 
-			Lease currentLease = this.currentLeaseRef.get();
-			this.currentLeaseRef.set(lease);
+			Lease currentLease = CURRENT_UPDATER.get(this);
+			CURRENT_UPDATER.set(this, lease);
 
 			if (currentLease != null) {
 				cancelSchedule(currentLease);
@@ -898,7 +914,7 @@ public void run() {
 
 					LeaseRenewalScheduler.this.schedules.remove(lease);
 
-					if (LeaseRenewalScheduler.this.currentLeaseRef.get() != lease) {
+					if (CURRENT_UPDATER.get(LeaseRenewalScheduler.this) != lease) {
 						logger.debug("Current lease has changed. Skipping renewal");
 						return;
 					}
@@ -918,7 +934,7 @@ public void run() {
 						// Renew lease may call scheduleRenewal(…) with a different lease
 						// Id to alter set up its own renewal schedule. If it's the old
 						// lease, then renewLease() outcome controls the current LeaseId.
-						LeaseRenewalScheduler.this.currentLeaseRef.compareAndSet(lease, renewLease.renewLease(lease));
+						CURRENT_UPDATER.compareAndSet(LeaseRenewalScheduler.this, lease, renewLease.renewLease(lease));
 					}
 					catch (Exception e) {
 						logger.error(String.format("Cannot renew lease %s", lease.getLeaseId()), e);
@@ -932,6 +948,10 @@ public void run() {
 			this.schedules.put(lease, scheduledFuture);
 		}
 
+		void associateLease(Lease lease) {
+			CURRENT_UPDATER.set(this, lease);
+		}
+
 		private void cancelSchedule(Lease lease) {
 
 			ScheduledFuture<?> scheduledFuture = this.schedules.get(lease);
@@ -951,7 +971,8 @@ private void cancelSchedule(Lease lease) {
 		 */
 		void disableScheduleRenewal() {
 
-			this.currentLeaseRef.set(null);
+			// capture the previous lease to revoke it
+			this.previousLeaseRef = CURRENT_UPDATER.getAndSet(this, null);
 			Set<Lease> leases = new HashSet<>(this.schedules.keySet());
 
 			for (Lease lease : leases) {
@@ -984,7 +1005,12 @@ private boolean isLeaseRenewable(@Nullable Lease lease, RequestedSecret requeste
 
 		@Nullable
 		public Lease getLease() {
-			return this.currentLeaseRef.get();
+			return CURRENT_UPDATER.get(this);
+		}
+
+		@Nullable
+		public Lease getPreviousLease() {
+			return this.previousLeaseRef;
 		}
 
 		private boolean isLeaseRotateOnly(Lease lease, RequestedSecret requestedSecret) {

diff --git a/...ore/src/test/java/org/springframework/vault/core/lease/SecretLeaseContainerUnitTests.java b/...ore/src/test/java/org/springframework/vault/core/lease/SecretLeaseContainerUnitTests.java
@@ -630,6 +630,27 @@ void shouldNotRevokeSecretsWithoutLease() throws Exception {
 		verify(this.leaseListenerAdapter, never()).onLeaseEvent(any(AfterSecretLeaseRevocationEvent.class));
 	}
 
+	@Test
+	void shouldRevokeSecretsOnDestroy() throws Exception {
+
+		VaultResponse secrets = new VaultResponse();
+		secrets.setData(Collections.singletonMap("key", (Object) "value"));
+		secrets.setLeaseId("1234");
+		secrets.setLeaseDuration(1000);
+
+		when(this.vaultOperations.read(this.requestedSecret.getPath())).thenReturn(secrets);
+
+		this.secretLeaseContainer.addRequestedSecret(this.requestedSecret);
+		this.secretLeaseContainer.start();
+		this.secretLeaseContainer.stop();
+
+		this.secretLeaseContainer.destroy();
+
+		verify(this.leaseListenerAdapter).onLeaseEvent(any(SecretLeaseCreatedEvent.class));
+		verify(this.leaseListenerAdapter).onLeaseEvent(any(BeforeSecretLeaseRevocationEvent.class));
+		verify(this.leaseListenerAdapter).onLeaseEvent(any(AfterSecretLeaseRevocationEvent.class));
+	}
+
 	@Test
 	void shouldRequestRotatingGenericSecrets() {