Cache role ssl token

5 files changed, 86 insertions, 20 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
new file mode 100644
index 00000000000..b011cc8fbcc
--- /dev/null
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/AthenzRole.java
@@ -0,0 +1,47 @@
+// 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.api;
+
+import java.util.Objects;
+
+/**
+ * @author tokle
+ */
+public class AthenzRole {
+    private final AthenzDomain domain;
+    private final String roleName;
+
+    public AthenzRole(AthenzDomain domain, String roleName) {
+        this.domain = domain;
+        this.roleName = roleName;
+    }
+
+    public AthenzDomain domain() {
+        return domain;
+    }
+
+    public String roleName() {
+        return roleName;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        AthenzRole that = (AthenzRole) o;
+        return Objects.equals(domain, that.domain) &&
+               Objects.equals(roleName, that.roleName);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(domain, roleName);
+    }
+
+    @Override
+    public String toString() {
+        return "AthenzRole{" +
+               "domain=" + domain +
+               ", roleName='" + roleName + '\'' +
+               '}';
+    }
+}
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 3f1531ad991..9d1b16da42a 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
@@ -6,6 +6,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
 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.NToken;
 import com.yahoo.vespa.athenz.api.ZToken;
@@ -147,18 +148,17 @@ public class DefaultZtsClient implements ZtsClient {
     }
 
     @Override
-    public X509Certificate getRoleCertificate(AthenzDomain domain,
-                                              String roleName,
+    public X509Certificate getRoleCertificate(AthenzRole role,
                                               Duration expiry,
                                               KeyPair keyPair,
                                               String cloud) {
-        X500Principal principal = new X500Principal(String.format("cn=%s:role.%s", domain.getName(), roleName));
+        X500Principal principal = new X500Principal(String.format("cn=%s:role.%s", role.domain().getName(), role.roleName()));
         Pkcs10Csr csr = Pkcs10CsrBuilder.fromKeypair(principal, keyPair, SHA256_WITH_RSA)
                 .addSubjectAlternativeName(DNS_NAME, String.format("%s.%s.%s", identity.getName(), identity.getDomainName().replace('.', '-'), cloud))
                 .addSubjectAlternativeName(RFC822_NAME, String.format("%s.%s@%s", identity.getDomainName(), identity.getName(), cloud))
                 .build();
         RoleCertificateRequestEntity requestEntity = new RoleCertificateRequestEntity(csr, expiry);
-        URI uri = ztsUrl.resolve(String.format("domain/%s/role/%s/token", domain.getName(), roleName));
+        URI uri = ztsUrl.resolve(String.format("domain/%s/role/%s/token", role.domain().getName(), role.roleName()));
         HttpUriRequest request = RequestBuilder.post(uri)
                 .setEntity(toJsonStringEntity(requestEntity))
                 .build();
@@ -171,11 +171,10 @@ public class DefaultZtsClient implements ZtsClient {
     }
 
     @Override
-    public X509Certificate getRoleCertificate(AthenzDomain domain,
-                                              String roleName,
+    public X509Certificate getRoleCertificate(AthenzRole role,
                                               KeyPair keyPair,
                                               String cloud) {
-        return getRoleCertificate(domain, roleName, null, keyPair, cloud);
+        return getRoleCertificate(role, null, keyPair, cloud);
     }
 
     private static InstanceIdentity getInstanceIdentity(HttpResponse response) throws IOException {
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 9f3a3b33f25..270954c73b2 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
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.athenz.client.zts;
 
 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.tls.Pkcs10Csr;
@@ -61,15 +62,13 @@ public interface ZtsClient extends AutoCloseable {
     /**
      * Fetch role certificate for the target domain and role
      *
-     * @param domain Target domain
-     * @param roleName Target role
+     * @param role Target role
      * @param expiry Certificate expiry
      * @param keyPair Key pair which will be used to generate CSR (certificate signing request)
      * @param cloud The cloud suffix used in DNS SAN entries
      * @return A role certificate
      */
-    X509Certificate getRoleCertificate(AthenzDomain domain,
-                                       String roleName,
+    X509Certificate getRoleCertificate(AthenzRole role,
                                        Duration expiry,
                                        KeyPair keyPair,
                                        String cloud);
@@ -77,14 +76,12 @@ public interface ZtsClient extends AutoCloseable {
     /**
      * Fetch role certificate for the target domain and role
      *
-     * @param domain Target domain
-     * @param roleName Target role
+     * @param role Target role
      * @param keyPair Key pair which will be used to generate CSR (certificate signing request)
      * @param cloud The cloud suffix used in DNS SAN entries
      * @return A role certificate
      */
-    X509Certificate getRoleCertificate(AthenzDomain domain,
-                                       String roleName,
+    X509Certificate getRoleCertificate(AthenzRole role,
                                        KeyPair keyPair,
                                        String cloud);
 
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 db949929115..813941ac9b2 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
@@ -1,6 +1,9 @@
 // Copyright 2019 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.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
 import com.google.inject.Inject;
 import com.yahoo.component.AbstractComponent;
 import com.yahoo.container.core.identity.IdentityConfig;
@@ -9,6 +12,7 @@ import com.yahoo.container.jdisc.athenz.AthenzIdentityProviderException;
 import com.yahoo.jdisc.Metric;
 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.identity.ServiceIdentityProvider;
 import com.yahoo.vespa.athenz.identity.ServiceIdentityProviderListenerHelper;
@@ -52,6 +56,9 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     private final AthenzService identity;
     private final ServiceIdentityProviderListenerHelper listenerHelper;
 
+    private final LoadingCache<AthenzRole, SSLContext> roleSslContextCache;
+    private final static Duration roleSslContextExpiry = Duration.ofHours(24);
+
     // TODO IdentityConfig should contain ZTS uri and dns suffix
     @Inject
     public AthenzIdentityProviderImpl(IdentityConfig config, Metric metric) {
@@ -78,6 +85,15 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         this.identity = new AthenzService(config.domain(), config.service());
         this.listenerHelper = new ServiceIdentityProviderListenerHelper(this.identity);
         registerInstance();
+        roleSslContextCache = CacheBuilder.newBuilder()
+                .refreshAfterWrite(roleSslContextExpiry.dividedBy(2).toMinutes(), TimeUnit.MINUTES)
+                .expireAfterWrite(roleSslContextExpiry.toMinutes(), TimeUnit.MINUTES)
+                .build(new CacheLoader<AthenzRole, SSLContext>() {
+                    @Override
+                    public SSLContext load(AthenzRole key) throws Exception {
+                        return createRoleSslContext(key);
+                    }
+                });
     }
 
     private void registerInstance() {
@@ -123,9 +139,16 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     @Override
     public SSLContext getRoleSslContext(String domain, String role) {
         // This ssl context should ideally be cached as it is quite expensive to create.
+        try {
+            return roleSslContextCache.get(new AthenzRole(new AthenzDomain(domain), role));
+        } catch (Exception e) {
+            throw new AthenzIdentityProviderException("Could not retrieve role certificate.", e);
+        }
+    }
+
+    private SSLContext createRoleSslContext(AthenzRole role) {
         PrivateKey privateKey = credentials.getKeyPair().getPrivate();
         X509Certificate roleCertificate = ztsClient.getRoleCertificate(
-                new AthenzDomain(domain),
                 role,
                 credentials.getIdentityDocument().dnsSuffix(),
                 credentials.getIdentityDocument().ztsEndpoint(),
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
index 47539d5d60a..afdccac62cf 100644
--- 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
@@ -7,6 +7,7 @@ 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.tls.X509CertificateUtils;
@@ -102,8 +103,7 @@ class ZtsClient {
                         .getRoleToken(domain.getName(), roleName).getToken());
     }
 
-    X509Certificate getRoleCertificate(AthenzDomain roleDomain,
-                                       String roleName,
+    X509Certificate getRoleCertificate(AthenzRole role,
                                        String dnsSuffix,
                                        URI ztsEndpoint,
                                        AthenzService identity,
@@ -113,8 +113,8 @@ class ZtsClient {
         URI correctedZtsEndpoint = ztsEndpoint.resolve("/zts/v1");
         ZTSClient ztsClient = new ZTSClient(correctedZtsEndpoint.toString(), sslContext);
         RoleCertificateRequest rcr = ZTSClient.generateRoleCertificateRequest(
-                identity.getDomain().getName(), identity.getName(), roleDomain.getName(), roleName, privateKey, dnsSuffix, (int) Duration.ofHours(1).getSeconds());
-        RoleToken pemCert = ztsClient.postRoleCertificateRequest(roleDomain.getName(), roleName, rcr);
+                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);
     }