Add expiry Predicate to SecretLeaseContainer to determine whether a Lease is expired.

Closes gh-809

Author: mp911de

spring-vault-core/src/main/java/org/springframework/vault/core/lease/SecretLeaseContainer.java

@@ -34,6 +34,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import java.util.function.BiConsumer;
+import java.util.function.Predicate;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -127,6 +128,16 @@ public class SecretLeaseContainer extends SecretLeaseEventPublisher
 	private static final AtomicIntegerFieldUpdater<SecretLeaseContainer> UPDATER = AtomicIntegerFieldUpdater
 		.newUpdater(SecretLeaseContainer.class, "status");
 
+	/**
+	 * {@link Predicate} to test whether a {@link Lease} has no lease identifier.
+	 */
+	public static Predicate<Lease> NO_LEASE_ID = Predicate.not(Lease::hasLeaseId);
+
+	/**
+	 * {@link Predicate} to test whether a {@link Lease} has no lease identifier.
+	 */
+	public static Predicate<Lease> NO_LEASE_DURATION = forDuration(Duration::isZero);
+
 	private static final AtomicInteger poolId = new AtomicInteger();
 
 	private static final int STATUS_INITIAL = 0;
@@ -154,6 +165,10 @@ public class SecretLeaseContainer extends SecretLeaseEventPublisher
 
 	private Duration minRenewal = Duration.ofSeconds(10);
 
+	private @Nullable Predicate<Lease> isExpired;
+
+	private Predicate<Lease> isExpiredFallback = createIsExpiredPredicate(this.minRenewal);
+
 	private Duration expiryThreshold = Duration.ofSeconds(60);
 
 	private LeaseStrategy leaseStrategy = LeaseStrategy.dropOnError();
@@ -241,6 +256,20 @@ public void setMinRenewal(Duration minRenewal) {
 		Assert.isTrue(!minRenewal.isNegative(), "Minimal renewal time must not be negative");
 
 		this.minRenewal = minRenewal;
+		this.isExpiredFallback = createIsExpiredPredicate(this.minRenewal);
+	}
+
+	/**
+	 * Sets the {@link Predicate} to determine whether a {@link Lease} is expired.
+	 * Defaults to comparing whether a lease {@link Lease#hasLeaseId() has no identifier},
+	 * its remaining TTL is zero or less or equal to {@code minRenewal}.
+	 * @since 3.2
+	 */
+	public void setExpiryPredicate(Predicate<Lease> isExpired) {
+
+		Assert.notNull(isExpired, "Expiry predicate must not be null");
+
+		this.isExpired = isExpired;
 	}
 
 	/**
@@ -737,8 +766,7 @@ protected Lease doRenewLease(RequestedSecret requestedSecret, Lease lease) {
 		try {
 			Lease renewed = lease.hasLeaseId() ? doRenew(lease) : lease;
 
-			if (!renewed.hasLeaseId() || renewed.getLeaseDuration().isZero()
-					|| renewed.getLeaseDuration().getSeconds() < this.minRenewal.getSeconds()) {
+			if (isExpired(renewed)) {
 
 				onLeaseExpired(requestedSecret, lease);
 				return Lease.none();
@@ -778,6 +806,10 @@ protected Lease doRenewLease(RequestedSecret requestedSecret, Lease lease) {
 		}
 	}
 
+	boolean isExpired(Lease lease) {
+		return isExpired == null ? isExpiredFallback.test(lease) : isExpired.test(lease);
+	}
+
 	@Nullable
 	private HttpStatusCodeException potentiallyUnwrapHttpStatusCodeException(RuntimeException e) {
 
@@ -857,6 +889,18 @@ protected void doRevokeLease(RequestedSecret requestedSecret, Lease lease) {
 		}
 	}
 
+	private Predicate<Lease> createIsExpiredPredicate(Duration minRenewal) {
+		return NO_LEASE_ID.or(NO_LEASE_DURATION).or(forDuration(isLessOrEqual(minRenewal)));
+	}
+
+	private static <T extends Comparable<T>> Predicate<T> isLessOrEqual(T other) {
+		return it -> it.compareTo(other) <= 0;
+	}
+
+	private static Predicate<Lease> forDuration(Predicate<Duration> predicate) {
+		return lease -> predicate.test(lease.getLeaseDuration());
+	}
+
 	/**
 	 * Abstracts scheduled lease renewal. A {@link LeaseRenewalScheduler} can be accessed
 	 * concurrently to schedule lease renewal. Each renewal run checks if the previously

spring-vault-core/src/test/java/org/springframework/vault/core/lease/SecretLeaseContainerUnitTests.java

@@ -25,6 +25,7 @@
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Predicate;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -112,6 +113,36 @@ void shouldSetProperties() {
 		assertThat(this.secretLeaseContainer.getExpiryThresholdSeconds()).isEqualTo(180);
 	}
 
+	@Test
+	void defaultIsExpiredShouldCalculateCorrectResult() {
+
+		assertThat((Predicate<Lease>) secretLeaseContainer::isExpired).accepts(Lease.none())
+			.accepts(Lease.of("foo", Duration.ZERO, false))
+			.accepts(Lease.fromTimeToLive(Duration.ofHours(1)))
+			.accepts(Lease.of("foo", Duration.ofSeconds(9), false))
+			.accepts(Lease.of("foo", Duration.ofSeconds(10), false))
+			.rejects(Lease.of("foo", Duration.ofSeconds(11), false));
+
+		this.secretLeaseContainer.setMinRenewal(Duration.ofMinutes(2));
+
+		assertThat((Predicate<Lease>) secretLeaseContainer::isExpired).accepts(Lease.none())
+			.accepts(Lease.of("foo", Duration.ZERO, false))
+			.accepts(Lease.fromTimeToLive(Duration.ofHours(1)))
+			.accepts(Lease.of("foo", Duration.ofSeconds(9), false))
+			.accepts(Lease.of("foo", Duration.ofSeconds(120), false))
+			.rejects(Lease.of("foo", Duration.ofSeconds(121), false));
+	}
+
+	@Test
+	void configuredExpiryPredicateShouldBeEvaluated() {
+
+		secretLeaseContainer.setExpiryPredicate(lease -> "expired".equals(lease.getLeaseId()));
+
+		assertThat((Predicate<Lease>) secretLeaseContainer::isExpired).rejects(Lease.none())
+			.rejects(Lease.of("not-expired", Duration.ZERO, false))
+			.accepts(Lease.of("expired", Duration.ZERO, false));
+	}
+
 	@Test
 	void shouldWorkIfNoSecretsRequested() {