Remove deprecated ZtsClient

- Replace use of old ZtsClient with DefaultZtsClient
- Add caching of role tokens
- Add constructor to AthenzRole taking only strings
- Change new ZtsClient interface to use AthenzRole for getRoleToken

5 files changed, 71 insertions, 106 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/AthenzRole.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/AthenzRole.java
index b011cc8fbcc..3a81e4a5e17 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/AthenzRole.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/AthenzRole.java
@@ -15,6 +15,11 @@ public class AthenzRole {
         this.roleName = roleName;
     }
 
+    public AthenzRole(String domain, String roleName) {
+        this.domain = new AthenzDomain(domain);
+        this.roleName = roleName;
+    }
+
     public AthenzDomain domain() {
         return domain;
     }
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 dfe49d1b407..dd97b20055d 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
@@ -132,7 +132,11 @@ public class DefaultZtsClient implements ZtsClient {
     }
 
     @Override
-    public ZToken getRoleToken(AthenzDomain domain, String roleName) {
+    public ZToken getRoleToken(AthenzRole athenzRole) {
+        return getRoleToken(athenzRole.domain(), athenzRole.roleName());
+    }
+
+    private ZToken getRoleToken(AthenzDomain domain, String roleName) {
         URI uri = ztsUrl.resolve(String.format("domain/%s/token", domain.getName()));
         RequestBuilder requestBuilder = RequestBuilder.get(uri)
                 .addHeader("Content-Type", "application/json");
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 43378b6507a..87a204095d9 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
@@ -51,13 +51,12 @@ public interface ZtsClient extends AutoCloseable {
     ZToken getRoleToken(AthenzDomain domain);
 
     /**
-     * Fetch a role token for the target domain and role
+     * Fetch a role token for the target role
      *
-     * @param domain Target domain
-     * @param roleName Target role
+     * @param athenzRole Target role
      * @return A role token
      */
-    ZToken getRoleToken(AthenzDomain domain, String roleName);
+    ZToken getRoleToken(AthenzRole athenzRole);
 
     /**
      * Fetch role certificate for the target domain and 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 f994316c22d..5f7f7d490fb 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
@@ -14,6 +14,9 @@ import com.yahoo.log.LogLevel;
 import com.yahoo.vespa.athenz.api.AthenzDomain;
 import com.yahoo.vespa.athenz.api.AthenzRole;
 import com.yahoo.vespa.athenz.api.AthenzService;
+import com.yahoo.vespa.athenz.api.ZToken;
+import com.yahoo.vespa.athenz.client.zts.DefaultZtsClient;
+import com.yahoo.vespa.athenz.client.zts.ZtsClient;
 import com.yahoo.vespa.athenz.identity.ServiceIdentityProvider;
 import com.yahoo.vespa.athenz.identity.ServiceIdentityProviderListenerHelper;
 import com.yahoo.vespa.athenz.identity.SiaIdentityProvider;
@@ -25,7 +28,6 @@ import com.yahoo.vespa.defaults.Defaults;
 import javax.net.ssl.SSLContext;
 import java.io.File;
 import java.net.URI;
-import java.security.PrivateKey;
 import java.security.cert.X509Certificate;
 import java.time.Clock;
 import java.time.Duration;
@@ -33,9 +35,12 @@ import java.time.Instant;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.logging.Logger;
 
 /**
+ * A {@link AthenzIdentityProvider} / {@link ServiceIdentityProvider} component that provides the tenant identity.
+ *
  * @author mortent
  * @author bjorncs
  */
@@ -48,11 +53,11 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     static final Duration UPDATE_PERIOD = Duration.ofDays(1);
     static final Duration AWAIT_TERMINTATION_TIMEOUT = Duration.ofSeconds(90);
     private final static Duration ROLE_SSL_CONTEXT_EXPIRY = Duration.ofHours(24);
+    private final static Duration ROLE_TOKEN_EXPIRY = Duration.ofMinutes(30);
 
     public static final String CERTIFICATE_EXPIRY_METRIC_NAME = "athenz-tenant-cert.expiry.seconds";
 
     private volatile AthenzCredentials credentials;
-    private final ZtsClient ztsClient = new ZtsClient();
     private final Metric metric;
     private final AthenzCredentialsService athenzCredentialsService;
     private final ScheduledExecutorService scheduler;
@@ -63,6 +68,8 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     private final URI ztsEndpoint;
 
     private final LoadingCache<AthenzRole, SSLContext> roleSslContextCache;
+    private final LoadingCache<AthenzRole, ZToken> roleSpecificRoleTokenCache;
+    private final LoadingCache<AthenzDomain, ZToken> domainSpecificRoleTokenCache;
 
     @Inject
     public AthenzIdentityProviderImpl(IdentityConfig config, Metric metric) {
@@ -91,14 +98,20 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         this.listenerHelper = new ServiceIdentityProviderListenerHelper(this.identity);
         this.dnsSuffix = config.athenzDnsSuffix();
         this.ztsEndpoint = URI.create(config.ztsUrl());
+        roleSslContextCache = createCache(ROLE_SSL_CONTEXT_EXPIRY, this::createRoleSslContext);
+        roleSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
+        domainSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
         registerInstance();
-        roleSslContextCache = CacheBuilder.newBuilder()
-                .refreshAfterWrite(ROLE_SSL_CONTEXT_EXPIRY.dividedBy(2).toMinutes(), TimeUnit.MINUTES)
-                .expireAfterWrite(ROLE_SSL_CONTEXT_EXPIRY.toMinutes(), TimeUnit.MINUTES)
-                .build(new CacheLoader<AthenzRole, SSLContext>() {
+    }
+
+    private static <KEY, VALUE> LoadingCache<KEY, VALUE> createCache(Duration expiry, Function<KEY, VALUE> cacheLoader) {
+        return CacheBuilder.newBuilder()
+                .refreshAfterWrite(expiry.dividedBy(2).toMinutes(), TimeUnit.MINUTES)
+                .expireAfterWrite(expiry.toMinutes(), TimeUnit.MINUTES)
+                .build(new CacheLoader<KEY, VALUE>() {
                     @Override
-                    public SSLContext load(AthenzRole key) throws Exception {
-                        return createRoleSslContext(key);
+                    public VALUE load(KEY key) {
+                        return cacheLoader.apply(key);
                     }
                 });
     }
@@ -149,44 +162,52 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         try {
             return roleSslContextCache.get(new AthenzRole(new AthenzDomain(domain), role));
         } catch (Exception e) {
-            throw new AthenzIdentityProviderException("Could not retrieve role certificate.", e);
+            throw new AthenzIdentityProviderException("Could not retrieve role certificate: " + e.getMessage(), e);
         }
     }
 
-    private SSLContext createRoleSslContext(AthenzRole role) {
-        PrivateKey privateKey = credentials.getKeyPair().getPrivate();
-        X509Certificate roleCertificate = ztsClient.getRoleCertificate(
-                role,
-                dnsSuffix,
-                ztsEndpoint,
-                identity,
-                privateKey,
-                credentials.getIdentitySslContext());
-        return new SslContextBuilder()
-                .withKeyStore(privateKey, roleCertificate)
-                .withTrustStore(getDefaultTrustStoreLocation(), KeyStoreType.JKS)
-                .build();
-    }
-
     @Override
     public String getRoleToken(String domain) {
-        return ztsClient
-                .getRoleToken(
-                        new AthenzDomain(domain),
-                        ztsEndpoint,
-                        credentials.getIdentitySslContext())
-                .getRawToken();
+        try {
+            return domainSpecificRoleTokenCache.get(new AthenzDomain(domain)).getRawToken();
+        } catch (Exception e) {
+            throw new AthenzIdentityProviderException("Could not retrieve role token: " + e.getMessage(), e);
+        }
     }
 
     @Override
     public String getRoleToken(String domain, String role) {
-        return ztsClient
-                .getRoleToken(
-                        new AthenzDomain(domain),
-                        role,
-                        ztsEndpoint,
-                        credentials.getIdentitySslContext())
-                .getRawToken();
+        try {
+            return roleSpecificRoleTokenCache.get(new AthenzRole(domain, role)).getRawToken();
+        } catch (Exception e) {
+            throw new AthenzIdentityProviderException("Could not retrieve role token: " + e.getMessage(), e);
+        }
+    }
+
+    private SSLContext createRoleSslContext(AthenzRole role) {
+        try (ZtsClient client = createZtsClient()) {
+            X509Certificate roleCertificate = client.getRoleCertificate(role, credentials.getKeyPair(), dnsSuffix);
+            return new SslContextBuilder()
+                    .withKeyStore(credentials.getKeyPair().getPrivate(), roleCertificate)
+                    .withTrustStore(getDefaultTrustStoreLocation(), KeyStoreType.JKS)
+                    .build();
+        }
+    }
+
+    private ZToken createRoleToken(AthenzRole athenzRole) {
+        try (ZtsClient client = createZtsClient()) {
+            return client.getRoleToken(athenzRole);
+        }
+    }
+
+    private ZToken createRoleToken(AthenzDomain domain) {
+        try (ZtsClient client = createZtsClient()) {
+            return client.getRoleToken(domain);
+        }
+    }
+
+    private DefaultZtsClient createZtsClient() {
+        return new DefaultZtsClient(ztsEndpoint, identity(), getIdentitySslContext());
     }
 
     @Override
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/ZtsClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/ZtsClient.java
deleted file mode 100644
index a3ec55eb815..00000000000
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/ZtsClient.java
+++ /dev/null
@@ -1,64 +0,0 @@
-// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.vespa.athenz.identityprovider.client;
-
-import com.yahoo.athenz.zts.RoleCertificateRequest;
-import com.yahoo.athenz.zts.RoleToken;
-import com.yahoo.athenz.zts.ZTSClient;
-import com.yahoo.vespa.athenz.api.AthenzDomain;
-import com.yahoo.vespa.athenz.api.AthenzRole;
-import com.yahoo.vespa.athenz.api.AthenzService;
-import com.yahoo.vespa.athenz.api.ZToken;
-import com.yahoo.vespa.athenz.client.zts.DefaultZtsClient;
-import com.yahoo.vespa.athenz.tls.X509CertificateUtils;
-
-import javax.net.ssl.SSLContext;
-import java.net.URI;
-import java.security.PrivateKey;
-import java.security.cert.X509Certificate;
-import java.time.Duration;
-
-/**
- * @author mortent
- * @author bjorncs
- * @deprecated Will be replaced by {@link DefaultZtsClient} once role token/certificate caching is ready.
- */
-@Deprecated
-class ZtsClient {
-
-    ZToken getRoleToken(AthenzDomain domain,
-                        URI ztsEndpoint,
-                        SSLContext sslContext) {
-         // TODO ztsEndpoint should contain '/zts/v1' as path
-        URI correctedZtsEndpoint = ztsEndpoint.resolve("/zts/v1");
-        return new ZToken(
-                new ZTSClient(correctedZtsEndpoint.toString(), sslContext)
-                        .getRoleToken(domain.getName()).getToken());
-    }
-
-    ZToken getRoleToken(AthenzDomain domain,
-                        String roleName,
-                        URI ztsEndpoint,
-                        SSLContext sslContext) {
-        // TODO ztsEndpoint should contain '/zts/v1' as path
-        URI correctedZtsEndpoint = ztsEndpoint.resolve("/zts/v1");
-        return new ZToken(
-                new ZTSClient(correctedZtsEndpoint.toString(), sslContext)
-                        .getRoleToken(domain.getName(), roleName).getToken());
-    }
-
-    X509Certificate getRoleCertificate(AthenzRole role,
-                                       String dnsSuffix,
-                                       URI ztsEndpoint,
-                                       AthenzService identity,
-                                       PrivateKey privateKey,
-                                       SSLContext sslContext) {
-        // TODO ztsEndpoint should contain '/zts/v1' as path
-        URI correctedZtsEndpoint = ztsEndpoint.resolve("/zts/v1");
-        ZTSClient ztsClient = new ZTSClient(correctedZtsEndpoint.toString(), sslContext);
-        RoleCertificateRequest rcr = ZTSClient.generateRoleCertificateRequest(
-                identity.getDomain().getName(), identity.getName(), role.domain().getName(), role.roleName(), privateKey, dnsSuffix, (int) Duration.ofHours(1).getSeconds());
-        RoleToken pemCert = ztsClient.postRoleCertificateRequest(role.domain().getName(), role.roleName(), rcr);
-        return X509CertificateUtils.fromPem(pemCert.token);
-    }
-
-}