Support refresh of limited use tokens (num_use > 0)

quarkiverse · Jan 23, 2025 · 90ec286 · 90ec286
1 parent da77bf6
commit 90ec286
Show file tree

Hide file tree

Showing 10 changed files with 183 additions and 74 deletions.
diff --git a/client/src/main/java/io/quarkus/vault/client/VaultClient.java b/client/src/main/java/io/quarkus/vault/client/VaultClient.java
@@ -258,7 +258,7 @@ public <T> CompletionStage<VaultResponse<T>> execute(VaultRequest<T> request) {
             if (token == null) {
                 requestBuilder.noToken();
             } else {
-                requestBuilder.token(token.getClientToken());
+                requestBuilder.token(token.getClientTokenForUsage());
             }
 
             return executor.execute(requestBuilder.rebuild());

diff --git a/client/src/main/java/io/quarkus/vault/client/auth/VaultAppRoleTokenProvider.java b/client/src/main/java/io/quarkus/vault/client/auth/VaultAppRoleTokenProvider.java
@@ -31,7 +31,7 @@ public CompletionStage<VaultToken> apply(VaultAuthRequest authRequest) {
                 .thenApply(res -> {
                     var auth = res.getAuth();
                     return VaultToken.from(auth.getClientToken(), auth.isRenewable(), auth.getLeaseDuration(),
-                            authRequest.getInstantSource());
+                            auth.getNumUses(), authRequest.getInstantSource());
                 });
     }
 }
diff --git a/client/src/main/java/io/quarkus/vault/client/auth/VaultCachingTokenProvider.java b/client/src/main/java/io/quarkus/vault/client/auth/VaultCachingTokenProvider.java
@@ -94,7 +94,7 @@ public CompletionStage<VaultToken> extend(VaultAuthRequest authRequest, String c
                 .thenApply(res -> {
                     var auth = res.getAuth();
                     var vaultToken = VaultToken.from(auth.getClientToken(), auth.isRenewable(), auth.getLeaseDuration(),
-                            authRequest.getInstantSource());
+                            auth.getNumUses(), authRequest.getInstantSource());
                     sanityCheck(vaultToken);
                     log.fine("extended login token: " + vaultToken.getConfidentialInfo(logLevel));
                     return vaultToken;

diff --git a/client/src/main/java/io/quarkus/vault/client/auth/VaultKubernetesTokenProvider.java b/client/src/main/java/io/quarkus/vault/client/auth/VaultKubernetesTokenProvider.java
@@ -44,7 +44,7 @@ public CompletionStage<VaultToken> apply(VaultAuthRequest authRequest) {
                     .thenApply(res -> {
                         var auth = res.getAuth();
                         return VaultToken.from(auth.getClientToken(), auth.isRenewable(), auth.getLeaseDuration(),
-                                authRequest.getInstantSource());
+                                auth.getNumUses(), authRequest.getInstantSource());
                     });
         });
     }

diff --git a/client/src/main/java/io/quarkus/vault/client/auth/VaultToken.java b/client/src/main/java/io/quarkus/vault/client/auth/VaultToken.java
@@ -5,42 +5,52 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.time.InstantSource;
+import java.util.concurrent.atomic.AtomicReference;
 
+import io.quarkus.vault.client.VaultException;
 import io.quarkus.vault.client.logging.LogConfidentialityLevel;
 
 public class VaultToken extends VaultTimeLimited {
 
     private final String clientToken;
     private final boolean fromCache;
+    private final AtomicReference<Integer> allowedUsesRemaining;
 
-    public static VaultToken from(String clientToken, boolean renewable, Duration leaseDuration, InstantSource instantSource) {
-        return new VaultToken(clientToken, renewable, leaseDuration, false, instantSource);
+    public static VaultToken from(String clientToken, boolean renewable, Duration leaseDuration, Integer allowedUsesRemaining,
+            InstantSource instantSource) {
+        return new VaultToken(clientToken, renewable, leaseDuration, allowedUsesRemaining, false, instantSource);
     }
 
-    public static VaultToken renewable(String clientToken, Duration leaseDuration, InstantSource instantSource) {
-        return from(clientToken, true, leaseDuration, instantSource);
+    public static VaultToken renewable(String clientToken, Duration leaseDuration, Integer allowedUsesRemaining,
+            InstantSource instantSource) {
+        return from(clientToken, true, leaseDuration, allowedUsesRemaining, instantSource);
     }
 
     public static VaultToken neverExpires(String clientToken, InstantSource instantSource) {
-        return from(clientToken, false, Duration.ofSeconds(Long.MAX_VALUE), instantSource);
+        return from(clientToken, false, Duration.ofSeconds(Long.MAX_VALUE), null, instantSource);
     }
 
-    public VaultToken(String clientToken, boolean renewable, Duration leaseDuration, boolean fromCache,
+    public VaultToken(String clientToken, boolean renewable, Duration leaseDuration, Integer allowedUsesRemaining,
+            boolean fromCache,
             InstantSource instantSource) {
         super(renewable, leaseDuration, instantSource);
         this.clientToken = clientToken;
         this.fromCache = fromCache;
+        this.allowedUsesRemaining = new AtomicReference<>(
+                allowedUsesRemaining != null && allowedUsesRemaining > 0 ? allowedUsesRemaining : Integer.MAX_VALUE);
     }
 
     public VaultToken(String clientToken, boolean renewable, Duration leaseDuration, Instant created, boolean fromCache,
-            InstantSource instantSource) {
+            AtomicReference<Integer> allowedUsesRemaining, InstantSource instantSource) {
         super(renewable, leaseDuration, created, instantSource);
         this.clientToken = clientToken;
         this.fromCache = fromCache;
+        this.allowedUsesRemaining = allowedUsesRemaining;
     }
 
     public VaultToken cached() {
-        return new VaultToken(clientToken, isRenewable(), getLeaseDuration(), getCreated(), true, getInstantSource());
+        return new VaultToken(clientToken, isRenewable(), getLeaseDuration(), getCreated(), true, allowedUsesRemaining,
+                getInstantSource());
     }
 
     public String getClientToken() {
@@ -51,6 +61,28 @@ public boolean isFromCache() {
         return fromCache;
     }
 
+    public String getClientTokenForUsage() {
+        if (isNotUsable()) {
+            throw new VaultException("Token has exhausted its allowed uses of " + allowedUsesRemaining.get());
+        }
+        allowedUsesRemaining.updateAndGet(remaining -> remaining > 0 ? remaining - 1 : 0);
+        return clientToken;
+    }
+
+    public boolean isNotUsable() {
+        return allowedUsesRemaining.get() <= 0;
+    }
+
+    @Override
+    public boolean isExpired() {
+        return isNotUsable() || super.isExpired();
+    }
+
+    @Override
+    public boolean isExpiringWithin(Duration gracePeriod) {
+        return isNotUsable() || super.isExpiringWithin(gracePeriod);
+    }
+
     public String getConfidentialInfo(LogConfidentialityLevel level) {
         return "{clientToken: " + level.maskWithTolerance(clientToken, LOW) + ", " + super.info() + "}";
     }

diff --git a/client/src/main/java/io/quarkus/vault/client/auth/VaultUserPassTokenProvider.java b/client/src/main/java/io/quarkus/vault/client/auth/VaultUserPassTokenProvider.java
@@ -31,7 +31,7 @@ public CompletionStage<VaultToken> apply(VaultAuthRequest authRequest) {
                     .thenApply(res -> {
                         var auth = res.getAuth();
                         return VaultToken.from(auth.getClientToken(), auth.isRenewable(), auth.getLeaseDuration(),
-                                authRequest.getInstantSource());
+                                auth.getNumUses(), authRequest.getInstantSource());
                     });
         });
     }