Support List of proxy principals in AthenzIdentityProvider

4 files changed, 56 insertions, 20 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
index 1a921c3b3ad..101d40be491 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
@@ -141,16 +141,21 @@ public class DefaultZtsClient extends ClientBase implements ZtsClient {
     }
 
     @Override
-    public AthenzAccessToken getAccessToken(AthenzDomain domain,  List<AthenzIdentity> proxyPrincipals) {
+    public AthenzAccessToken getAccessToken(AthenzDomain domain, List<AthenzIdentity> proxyPrincipals) {
         return this.getAccessTokenImpl(List.of(new AthenzResourceName(domain, "domain")), proxyPrincipals);
     }
 
     @Override
     public AthenzAccessToken getAccessToken(List<AthenzRole> athenzRole) {
+        return getAccessToken(athenzRole, List.of());
+    }
+
+    @Override
+    public AthenzAccessToken getAccessToken(List<AthenzRole> athenzRole, List<AthenzIdentity> proxyPrincipals) {
         List<AthenzResourceName> athenzResourceNames = athenzRole.stream()
                 .map(AthenzRole::toResourceName)
                 .toList();
-        return this.getAccessTokenImpl(athenzResourceNames, List.of());
+        return this.getAccessTokenImpl(athenzResourceNames, proxyPrincipals);
     }
 
     private AthenzAccessToken getAccessTokenImpl(List<AthenzResourceName> resources, List<AthenzIdentity> proxyPrincipals) {
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/ZtsClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/ZtsClient.java
index 67a53715cec..0528042da16 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/ZtsClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/ZtsClient.java
@@ -129,6 +129,15 @@ public interface ZtsClient extends AutoCloseable {
     AthenzAccessToken getAccessToken(List<AthenzRole> athenzRole);
 
     /**
+     * Fetch an access token for the target roles
+     *
+     * @param athenzRole List of athenz roles to get access token for
+     * @param proxyPrincipals List of principals to allow proxying token
+     * @return An Athenz access token
+     */
+    AthenzAccessToken getAccessToken(List<AthenzRole> athenzRole, List<AthenzIdentity> proxyPrincipals);
+
+    /**
      * Fetch role certificate for the target domain and role
      *
      * @param role Target role
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
index 623a8c856bc..a3916cbe080 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
@@ -20,6 +20,7 @@ import com.yahoo.security.SslContextBuilder;
 import com.yahoo.security.X509CertificateWithKey;
 import com.yahoo.vespa.athenz.api.AthenzAccessToken;
 import com.yahoo.vespa.athenz.api.AthenzDomain;
+import com.yahoo.vespa.athenz.api.AthenzIdentity;
 import com.yahoo.vespa.athenz.api.AthenzRole;
 import com.yahoo.vespa.athenz.api.AthenzService;
 import com.yahoo.vespa.athenz.api.ZToken;
@@ -28,6 +29,7 @@ import com.yahoo.vespa.athenz.client.zts.ZtsClient;
 import com.yahoo.vespa.athenz.identity.ServiceIdentityProvider;
 import com.yahoo.vespa.athenz.identityprovider.api.VespaUniqueInstanceId;
 import com.yahoo.vespa.athenz.tls.AthenzX509CertificateUtils;
+import com.yahoo.vespa.athenz.utils.AthenzIdentities;
 import com.yahoo.vespa.athenz.utils.SiaUtils;
 
 import javax.net.ssl.SSLContext;
@@ -43,6 +45,7 @@ import java.time.Instant;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
@@ -89,8 +92,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     private final Map<AthenzRole, MutableX509KeyManager> roleKeyManagerCache;
     private final LoadingCache<AthenzRole, ZToken> roleSpecificRoleTokenCache;
     private final LoadingCache<AthenzDomain, ZToken> domainSpecificRoleTokenCache;
-    private final LoadingCache<AthenzDomain, AthenzAccessToken> domainSpecificAccessTokenCache;
-    private final LoadingCache<List<AthenzRole>, AthenzAccessToken> roleSpecificAccessTokenCache;
+    private final LoadingCache<AccessTokenCacheKey, AthenzAccessToken> accessTokenCache;
     private final CsrGenerator csrGenerator;
 
     @Inject
@@ -115,8 +117,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         this.roleKeyManagerCache = new HashMap<>();
         this.roleSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
         this.domainSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
-        this.domainSpecificAccessTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createAccessToken);
-        this.roleSpecificAccessTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createAccessToken);
+        this.accessTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createAccessToken);
         this.csrGenerator = new CsrGenerator(config.athenzDnsSuffix(), config.configserverIdentityName());
         this.autoReloadingX509KeyManager = autoReloadingX509KeyManager;
         this.identitySslContext = createIdentitySslContext(autoReloadingX509KeyManager, trustStore);
@@ -221,7 +222,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     @Override
     public String getAccessToken(String domain) {
         try {
-            return domainSpecificAccessTokenCache.get(new AthenzDomain(domain)).value();
+            return accessTokenCache.get(AccessTokenCacheKey.from(domain, List.of(), List.of())).value();
         } catch (Exception e) {
             throw new AthenzIdentityProviderException("Could not retrieve access token: " + e.getMessage(), e);
         }
@@ -230,10 +231,16 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     @Override
     public String getAccessToken(String domain, List<String> roles) {
         try {
-            List<AthenzRole> roleList = roles.stream()
-                    .map(roleName -> new AthenzRole(domain, roleName))
-                    .toList();
-            return roleSpecificAccessTokenCache.get(roleList).value();
+            return accessTokenCache.get(AccessTokenCacheKey.from(domain, roles, List.of())).value();
+        } catch (Exception e) {
+            throw new AthenzIdentityProviderException("Could not retrieve access token: " + e.getMessage(), e);
+        }
+    }
+
+    @Override
+    public String getAccessToken(String domain, List<String> roles, List<String> proxyPrincipals) {
+        try {
+            return accessTokenCache.get(AccessTokenCacheKey.from(domain, roles, proxyPrincipals)).value();
         } catch (Exception e) {
             throw new AthenzIdentityProviderException("Could not retrieve access token: " + e.getMessage(), e);
         }
@@ -305,15 +312,15 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         }
     }
 
-    private AthenzAccessToken createAccessToken(AthenzDomain domain) {
+    private AthenzAccessToken createAccessToken(AccessTokenCacheKey cacheKey) {
+        List<AthenzRole> roles = Optional.ofNullable(cacheKey.roles()).orElse(List.of());
+        List<AthenzIdentity> proxyPrincipals = Optional.ofNullable(cacheKey.proxyPrincipals()).orElse(List.of());
         try (ZtsClient client = createZtsClient()) {
-            return client.getAccessToken(domain);
-        }
-    }
-
-    private AthenzAccessToken createAccessToken(List<AthenzRole> roles) {
-        try (ZtsClient client = createZtsClient()) {
-            return client.getAccessToken(roles);
+            if (roles.isEmpty()) {
+                return client.getAccessToken(cacheKey.domain(), proxyPrincipals);
+            } else {
+                return client.getAccessToken(roles, proxyPrincipals);
+            }
         }
     }
 
@@ -349,5 +356,15 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
             log.log(Level.WARNING, "Failed to update metrics: " + t.getMessage(), t);
         }
     }
+    private record AccessTokenCacheKey(AthenzDomain domain, List<AthenzRole> roles, List<AthenzIdentity> proxyPrincipals) {
+        static AccessTokenCacheKey from(String domain, List<String> roles, List<String> proxyPrincipals) {
+            List<AthenzRole> roleList = Optional.ofNullable(roles).orElse(List.of()).stream()
+                    .map(roleName -> new AthenzRole(domain, roleName))
+                    .toList();
+            List<AthenzIdentity> proxyPrincipalList = Optional.ofNullable(proxyPrincipals).orElse(List.of()).stream()
+                    .map(AthenzIdentities::from)
+                    .toList();
+            return new AccessTokenCacheKey(new AthenzDomain(domain), roleList, proxyPrincipalList);
+        }
+    }
 }
-
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/LegacyAthenzIdentityProviderImpl.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/LegacyAthenzIdentityProviderImpl.java
index 4058af31dde..7c375d4c0ff 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/LegacyAthenzIdentityProviderImpl.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/LegacyAthenzIdentityProviderImpl.java
@@ -261,6 +261,11 @@ public final class LegacyAthenzIdentityProviderImpl extends AbstractComponent im
     }
 
     @Override
+    public String getAccessToken(String domain, List<String> roles, List<String> proxyPrincipal) {
+        throw new UnsupportedOperationException("Not implemented in legacy client");
+    }
+
+    @Override
     public PrivateKey getPrivateKey() {
         return credentials.getKeyPair().getPrivate();
     }