Skip to content

Commit

Permalink
Resume lease renewal after restart.
Browse files Browse the repository at this point in the history
Secret renewals are resumed after restarting secrets by re-registering LeaseRenewalScheduler.

Also, guard concurrent renewals that might have been in progress while the renewal scheduler has been in progress.

Closes gh-867
  • Loading branch information
mp911de committed Jun 6, 2024
1 parent 3314fef commit 62acae0
Show file tree
Hide file tree
Showing 2 changed files with 91 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -549,6 +549,7 @@ void restartSecrets() {
Lease previousLease = getPreviousLease(previousLeases, requestedSecret);

try {
this.renewals.put(requestedSecret, renewalScheduler);
doStart(requestedSecret, renewalScheduler, (secrets, lease) -> {
onSecretsRotated(requestedSecret, previousLease, lease, secrets.getRequiredData());
}, () -> {
Expand Down Expand Up @@ -807,7 +808,19 @@ private Lease doRenew(Lease lease) {
protected void onLeaseExpired(RequestedSecret requestedSecret, Lease lease) {

if (requestedSecret.getMode() == Mode.ROTATE) {
doStart(requestedSecret, this.renewals.get(requestedSecret), (secrets, currentLease) -> {
LeaseRenewalScheduler renewalScheduler = this.renewals.get(requestedSecret);
// prevent races for concurrent renewals of the same secret using different
// leases
if (renewalScheduler == null || !renewalScheduler.leaseEquals(lease)) {
logger.debug(String.format(
"Skipping rotation after renewal expiry for secret %s with lease %s as no LeaseRenewalScheduler is found. This can happen if leases have been restarted while concurrent expiry processing.",
requestedSecret.getPath(), lease.getLeaseId()));

super.onLeaseExpired(requestedSecret, lease);
return;
}

doStart(requestedSecret, renewalScheduler, (secrets, currentLease) -> {
onSecretsRotated(requestedSecret, lease, currentLease, secrets.getRequiredData());
}, () -> super.onLeaseExpired(requestedSecret, lease));
}
Expand Down Expand Up @@ -1023,6 +1036,10 @@ private boolean isLeaseRotateOnly(Lease lease, RequestedSecret requestedSecret)
&& requestedSecret.getMode() == Mode.ROTATE;
}

boolean leaseEquals(Lease lease) {
return getLease() == lease;
}

}

private class LeaseAuthenticationEventListener implements AuthenticationListener, AuthenticationErrorListener {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
import org.springframework.http.HttpStatus;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.Trigger;
import org.springframework.test.util.ReflectionTestUtils;
import org.springframework.vault.VaultException;
import org.springframework.vault.core.RestOperationsCallback;
import org.springframework.vault.core.VaultOperations;
Expand Down Expand Up @@ -479,6 +480,78 @@ void shouldNotRotateExpiringLease() {
assertThat(((SecretLeaseCreatedEvent) events.get(1)).getSecrets()).containsOnlyKeys("foo");
}

@SuppressWarnings("unchecked")
@Test
void secretShouldContinueOnRestartSecrets() {

when(this.taskScheduler.schedule(any(Runnable.class), any(Trigger.class))).thenReturn(this.scheduledFuture);

VaultResponse first = createSecrets();
VaultResponse second = createSecrets();
second.setData(Collections.singletonMap("foo", (Object) "bar"));

when(this.vaultOperations.read(this.requestedSecret.getPath())).thenReturn(first, second);
when(this.vaultOperations.doWithSession(any(RestOperationsCallback.class)))
.thenReturn(Lease.of("after_restart", Duration.ofSeconds(1), true));

RequestedSecret secret = RequestedSecret.rotating("my-secret");

this.secretLeaseContainer.setExpiryThreshold(Duration.ofSeconds(1));
this.secretLeaseContainer.addRequestedSecret(secret);
this.secretLeaseContainer.start();

ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
this.secretLeaseContainer.restartSecrets();

verify(this.taskScheduler, times(2)).schedule(captor.capture(), any(Trigger.class));
captor.getAllValues().get(0).run(); // old lease run
captor.getAllValues().get(1).run(); // new lease run

// one from restartSecrets, the second from the scheduler detecting expiry
verify(this.leaseListenerAdapter, times(2)).onLeaseEvent(any(SecretLeaseRotatedEvent.class));
verify(this.leaseListenerAdapter, never()).onLeaseEvent(any(SecretLeaseExpiredEvent.class));
verify(this.leaseListenerAdapter, never()).onLeaseEvent(any(SecretLeaseErrorEvent.class));
}

@SuppressWarnings("unchecked")
@Test
void concurrentRenewalAfterRestartSecretShouldContinueOnRestartSecrets() {

when(this.taskScheduler.schedule(any(Runnable.class), any(Trigger.class))).thenReturn(this.scheduledFuture);

VaultResponse first = createSecrets();
VaultResponse second = createSecrets();
second.setData(Collections.singletonMap("foo", (Object) "bar"));

when(this.vaultOperations.read(this.requestedSecret.getPath())).thenReturn(first, second);
when(this.vaultOperations.doWithSession(any(RestOperationsCallback.class)))
.thenReturn(Lease.of("after_restart", Duration.ofSeconds(1), true));

RequestedSecret secret = RequestedSecret.rotating("my-secret");

this.secretLeaseContainer.setExpiryThreshold(Duration.ofSeconds(1));
this.secretLeaseContainer.addRequestedSecret(secret);
this.secretLeaseContainer.start();

Map<RequestedSecret, SecretLeaseContainer.LeaseRenewalScheduler> renewals = (Map<RequestedSecret, SecretLeaseContainer.LeaseRenewalScheduler>) ReflectionTestUtils.getField(this.secretLeaseContainer, "renewals");

SecretLeaseContainer.LeaseRenewalScheduler scheduler = renewals.get(secret);
Lease lease = scheduler.getLease();
ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
this.secretLeaseContainer.restartSecrets();

scheduler.associateLease(lease);

verify(this.taskScheduler, times(2)).schedule(captor.capture(), any(Trigger.class));
captor.getAllValues().get(1).run(); // new lease run
captor.getAllValues().get(0).run(); // old lease run

// one from restartSecrets, the second from the scheduler detecting expiry
verify(this.leaseListenerAdapter, times(2)).onLeaseEvent(any(SecretLeaseRotatedEvent.class));
verify(this.leaseListenerAdapter).onLeaseEvent(any(SecretLeaseExpiredEvent.class));
verify(this.leaseListenerAdapter, never()).onLeaseEvent(any(SecretLeaseErrorEvent.class));
}

@Test
void scheduleRenewalShouldApplyExpiryThreshold() {

Expand Down

0 comments on commit 62acae0

Please sign in to comment.