Merge pull request #104 from scalecube/enh1

Added async support for vault-token-supplier
scalecube · Jan 27, 2025 · 60c42f9 · 60c42f9
2 parents c1aa627 + 4339aeb
commit 60c42f9
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 19 deletions.
diff --git a/vault/src/main/java/io/scalecube/security/vault/VaultServiceRolesInstaller.java b/vault/src/main/java/io/scalecube/security/vault/VaultServiceRolesInstaller.java
@@ -19,6 +19,8 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.StringJoiner;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import org.slf4j.Logger;
@@ -37,14 +39,16 @@ public class VaultServiceRolesInstaller {
       new ObjectMapper(new YAMLFactory()).setVisibility(PropertyAccessor.FIELD, Visibility.ANY);
 
   private final String vaultAddress;
-  private final Supplier<String> vaultTokenSupplier;
+  private final CompletableFuture<String> vaultTokenSupplier;
   private final Supplier<String> keyNameSupplier;
   private final Function<String, String> roleNameBuilder;
   private final List<Supplier<ServiceRoles>> serviceRolesSources;
   private final String keyAlgorithm;
   private final String keyRotationPeriod;
   private final String keyVerificationTtl;
   private final String roleTtl;
+  private final long timeout;
+  private final TimeUnit timeUnit;
 
   private VaultServiceRolesInstaller(Builder builder) {
     this.vaultAddress = builder.vaultAddress;
@@ -56,6 +60,8 @@ private VaultServiceRolesInstaller(Builder builder) {
     this.keyRotationPeriod = builder.keyRotationPeriod;
     this.keyVerificationTtl = builder.keyVerificationTtl;
     this.roleTtl = builder.roleTtl;
+    this.timeout = builder.timeout;
+    this.timeUnit = builder.timeUnit;
   }
 
   /**
@@ -74,19 +80,30 @@ public void install() {
       return;
     }
 
-    final String token = vaultTokenSupplier.get();
-    final Rest rest = new Rest().header(VAULT_TOKEN_HEADER, token);
-
-    final String keyName = keyNameSupplier.get();
-    createVaultIdentityKey(rest.url(buildVaultIdentityKeyUri(keyName)), keyName);
-
-    for (Role role : serviceRoles.roles) {
-      String roleName = roleNameBuilder.apply(role.role);
-      createVaultIdentityRole(
-          rest.url(buildVaultIdentityRoleUri(roleName)), keyName, roleName, role.permissions);
+    try {
+      vaultTokenSupplier
+          .thenAcceptAsync(
+              token -> {
+                final var rest = new Rest().header(VAULT_TOKEN_HEADER, token);
+                final var keyName = keyNameSupplier.get();
+
+                createVaultIdentityKey(rest.url(buildVaultIdentityKeyUri(keyName)), keyName);
+
+                for (var role : serviceRoles.roles) {
+                  String roleName = roleNameBuilder.apply(role.role);
+                  createVaultIdentityRole(
+                      rest.url(buildVaultIdentityRoleUri(roleName)),
+                      keyName,
+                      roleName,
+                      role.permissions);
+                }
+
+                LOGGER.debug("Installed serviceRoles ({})", serviceRoles);
+              })
+          .get(timeout, timeUnit);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
     }
-
-    LOGGER.debug("Installed serviceRoles ({})", serviceRoles);
   }
 
   private ServiceRoles loadServiceRoles() {
@@ -338,14 +355,16 @@ public String toString() {
   public static class Builder {
 
     private String vaultAddress;
-    private Supplier<String> vaultTokenSupplier;
+    private CompletableFuture<String> vaultTokenSupplier;
     private Supplier<String> keyNameSupplier;
     private Function<String, String> roleNameBuilder;
     private List<Supplier<ServiceRoles>> serviceRolesSources = DEFAULT_SERVICE_ROLES_SOURCES;
     private String keyAlgorithm = "RS256";
     private String keyRotationPeriod = "1h";
     private String keyVerificationTtl = "1h";
     private String roleTtl = "1m";
+    private long timeout = 10;
+    private TimeUnit timeUnit = TimeUnit.SECONDS;
 
     public Builder() {}
 
@@ -354,7 +373,7 @@ public Builder vaultAddress(String vaultAddress) {
       return this;
     }
 
-    public Builder vaultTokenSupplier(Supplier<String> vaultTokenSupplier) {
+    public Builder vaultTokenSupplier(CompletableFuture<String> vaultTokenSupplier) {
       this.vaultTokenSupplier = vaultTokenSupplier;
       return this;
     }
@@ -399,6 +418,12 @@ public Builder roleTtl(String roleTtl) {
       return this;
     }
 
+    public Builder timeout(long timeout, TimeUnit timeUnit) {
+      this.timeout = timeout;
+      this.timeUnit = timeUnit;
+      return this;
+    }
+
     public VaultServiceRolesInstaller build() {
       return new VaultServiceRolesInstaller(this);
     }

diff --git a/vault/src/main/java/io/scalecube/security/vault/VaultServiceTokenSupplier.java b/vault/src/main/java/io/scalecube/security/vault/VaultServiceTokenSupplier.java
@@ -7,8 +7,8 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.StringJoiner;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.BiFunction;
-import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -20,7 +20,7 @@ public class VaultServiceTokenSupplier {
 
   private final String vaultAddress;
   private final String serviceRole;
-  private final Supplier<String> vaultTokenSupplier;
+  private final CompletableFuture<String> vaultTokenSupplier;
   private final BiFunction<String, Map<String, String>, String> serviceTokenNameBuilder;
 
   private VaultServiceTokenSupplier(Builder builder) {
@@ -93,7 +93,7 @@ public static class Builder {
 
     private String vaultAddress;
     private String serviceRole;
-    private Supplier<String> vaultTokenSupplier;
+    private CompletableFuture<String> vaultTokenSupplier;
     private BiFunction<String, Map<String, String>, String> serviceTokenNameBuilder;
 
     public Builder() {}
@@ -108,7 +108,7 @@ public Builder serviceRole(String serviceRole) {
       return this;
     }
 
-    public Builder vaultTokenSupplier(Supplier<String> vaultTokenSupplier) {
+    public Builder vaultTokenSupplier(CompletableFuture<String> vaultTokenSupplier) {
       this.vaultTokenSupplier = vaultTokenSupplier;
       return this;
     }