Add initial implementation for retrieving role token/cert

3 files changed, 101 insertions, 1 deletions

diff --git a/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/AthenzIdentityProvider.java b/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/AthenzIdentityProvider.java
index bb6a033a4ab..9762d69ec31 100644
--- a/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/AthenzIdentityProvider.java
+++ b/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/AthenzIdentityProvider.java
@@ -10,4 +10,7 @@ public interface AthenzIdentityProvider {
     String domain();
     String service();
     SSLContext getIdentitySslContext();
+    SSLContext getRoleSslContext(String domain, String role);
+    String getRoleToken(String domain);
+    String getRoleToken(String domain, String 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 3773af6b808..eb0ae89fdcf 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
@@ -8,12 +8,17 @@ import com.yahoo.container.jdisc.athenz.AthenzIdentityProvider;
 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.AthenzService;
 import com.yahoo.vespa.athenz.identity.ServiceIdentityProvider;
+import com.yahoo.vespa.athenz.tls.KeyStoreType;
+import com.yahoo.vespa.athenz.tls.SslContextBuilder;
 import com.yahoo.vespa.defaults.Defaults;
 
 import javax.net.ssl.SSLContext;
 import java.io.File;
+import java.security.PrivateKey;
+import java.security.cert.X509Certificate;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
@@ -38,12 +43,14 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     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;
     private final Clock clock;
     private final AthenzService identity;
 
+    // TODO IdentityConfig should contain ZTS uri and dns suffix
     @Inject
     public AthenzIdentityProviderImpl(IdentityConfig config, Metric metric) {
         this(config,
@@ -51,7 +58,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
              new AthenzCredentialsService(config,
                                           new IdentityDocumentService(config.loadBalancerAddress()),
                                           new ZtsClient(),
-                                          new File(Defaults.getDefaults().underVespaHome("share/ssl/certs/yahoo_certificate_bundle.jks"))),
+                                          getDefaultTrustStoreLocation()),
              new ScheduledThreadPoolExecutor(1),
              Clock.systemUTC());
     }
@@ -101,6 +108,45 @@ 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.
+        PrivateKey privateKey = credentials.getKeyPair().getPrivate();
+        X509Certificate roleCertificate = ztsClient.getRoleCertificate(
+                new AthenzDomain(domain),
+                role,
+                credentials.getIdentityDocument().dnsSuffix,
+                credentials.getIdentityDocument().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),
+                        credentials.getIdentityDocument().ztsEndpoint,
+                        credentials.getIdentitySslContext())
+                .getRawToken();
+    }
+
+    @Override
+    public String getRoleToken(String domain, String role) {
+        return ztsClient
+                .getRoleToken(
+                        new AthenzDomain(domain),
+                        role,
+                        credentials.getIdentityDocument().ztsEndpoint,
+                        credentials.getIdentitySslContext())
+                .getRawToken();
+    }
+
+    @Override
     public void deconstruct() {
         try {
             scheduler.shutdownNow();
@@ -110,6 +156,10 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         }
     }
 
+    private static File getDefaultTrustStoreLocation() {
+        return new File(Defaults.getDefaults().underVespaHome("share/ssl/certs/yahoo_certificate_bundle.jks"));
+    }
+
     private boolean isExpired(AthenzCredentials credentials) {
         return clock.instant().isAfter(getExpirationTime(credentials));
     }
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 a4fde8270c6..c995bfba791 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
@@ -3,6 +3,13 @@ package com.yahoo.vespa.athenz.identityprovider.client;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
+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.AthenzService;
+import com.yahoo.vespa.athenz.api.ZToken;
+import com.yahoo.vespa.athenz.tls.X509CertificateUtils;
 import org.apache.http.client.HttpRequestRetryHandler;
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpUriRequest;
@@ -19,6 +26,9 @@ import javax.net.ssl.SSLContext;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.net.URI;
+import java.security.PrivateKey;
+import java.security.cert.X509Certificate;
+import java.time.Duration;
 
 /**
  * @author mortent
@@ -71,6 +81,43 @@ 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(AthenzDomain roleDomain,
+                                       String roleName,
+                                       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(), roleDomain.getName(), roleName, privateKey, dnsSuffix, (int) Duration.ofHours(1).getSeconds());
+        RoleToken pemCert = ztsClient.postRoleCertificateRequest(roleDomain.getName(), roleName, rcr);
+        return X509CertificateUtils.fromPem(pemCert.token);
+    }
+
     private InstanceIdentity getInstanceIdentity(CloseableHttpClient client, HttpUriRequest postRequest)
             throws IOException {
         try (CloseableHttpResponse response = client.execute(postRequest)) {