Replace use of AthenzCertificateClient with DefaultZtsClient

3 files changed, 28 insertions, 83 deletions

diff --git a/athenz-identity-provider-service/pom.xml b/athenz-identity-provider-service/pom.xml
index 982cb89f2bf..16e18e99e1b 100644
--- a/athenz-identity-provider-service/pom.xml
+++ b/athenz-identity-provider-service/pom.xml
@@ -63,38 +63,6 @@
             <groupId>org.apache.httpcomponents</groupId>
             <artifactId>httpcore</artifactId>
         </dependency>
-        <dependency>
-            <groupId>com.yahoo.athenz</groupId>
-            <artifactId>athenz-zts-java-client</artifactId>
-            <scope>compile</scope>
-            <exclusions>
-                <!--Exclude all bundles provided by JDisc -->
-                <exclusion>
-                    <groupId>org.bouncycastle</groupId>
-                    <artifactId>bcpkix-jdk15on</artifactId>
-                </exclusion>
-                <exclusion>
-                    <groupId>org.bouncycastle</groupId>
-                    <artifactId>bcprov-jdk15on</artifactId>
-                </exclusion>
-                <exclusion>
-                    <groupId>org.slf4j</groupId>
-                    <artifactId>slf4j-api</artifactId>
-                </exclusion>
-                <exclusion>
-                    <groupId>com.fasterxml.jackson.core</groupId>
-                    <artifactId>jackson-core</artifactId>
-                </exclusion>
-                <exclusion>
-                    <groupId>com.fasterxml.jackson.core</groupId>
-                    <artifactId>jackson-databind</artifactId>
-                </exclusion>
-                <exclusion>
-                    <groupId>com.fasterxml.jackson.core</groupId>
-                    <artifactId>jackson-annotations</artifactId>
-                </exclusion>
-            </exclusions>
-        </dependency>
 
         <!-- TEST -->
         <dependency>
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzSslKeyStoreConfigurator.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzSslKeyStoreConfigurator.java
index 2e0cf04ae4f..4a060c7284b 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzSslKeyStoreConfigurator.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzSslKeyStoreConfigurator.java
@@ -9,22 +9,30 @@ import com.yahoo.jdisc.http.ssl.SslKeyStoreConfigurator;
 import com.yahoo.jdisc.http.ssl.SslKeyStoreContext;
 import com.yahoo.log.LogLevel;
 import com.yahoo.vespa.athenz.api.AthenzService;
+import com.yahoo.vespa.athenz.client.zts.DefaultZtsClient;
+import com.yahoo.vespa.athenz.client.zts.Identity;
+import com.yahoo.vespa.athenz.client.zts.ZtsClient;
 import com.yahoo.vespa.athenz.identity.ServiceIdentityProvider;
 import com.yahoo.vespa.athenz.tls.KeyStoreBuilder;
 import com.yahoo.vespa.athenz.tls.KeyStoreType;
+import com.yahoo.vespa.athenz.tls.KeyUtils;
 import com.yahoo.vespa.athenz.utils.SiaUtils;
 import com.yahoo.vespa.defaults.Defaults;
 import com.yahoo.vespa.hosted.athenz.instanceproviderservice.config.AthenzProviderServiceConfig;
-import com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.AthenzCertificateClient;
 
+import java.net.URI;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.security.KeyPair;
 import java.security.KeyStore;
 import java.security.KeyStoreException;
 import java.security.PrivateKey;
+import java.security.PublicKey;
 import java.security.cert.X509Certificate;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.Executors;
@@ -48,7 +56,7 @@ public class AthenzSslKeyStoreConfigurator extends AbstractComponent implements
     private static final Path VESPA_SIA_DIRECTORY = Paths.get(Defaults.getDefaults().underVespaHome("var/vespa/sia"));
 
     private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
-    private final AthenzCertificateClient certificateClient;
+    private final ZtsClient ztsClient;
     private final KeyProvider keyProvider;
     private final AthenzProviderServiceConfig.Zones zoneConfig;
     private final Duration updatePeriod;
@@ -63,23 +71,23 @@ public class AthenzSslKeyStoreConfigurator extends AbstractComponent implements
                                          ConfigserverConfig configserverConfig) {
         AthenzProviderServiceConfig.Zones zoneConfig = getZoneConfig(config, zone);
         AthenzService configserverIdentity = new AthenzService(zoneConfig.domain(), zoneConfig.serviceName());
-        AthenzCertificateClient certificateClient = new AthenzCertificateClient(bootstrapIdentity, zoneConfig);
         Duration updatePeriod = Duration.ofDays(config.updatePeriodDays());
-        this.certificateClient = certificateClient;
+        DefaultZtsClient ztsClient = new DefaultZtsClient(URI.create(zoneConfig.ztsUrl()), bootstrapIdentity);
+        this.ztsClient = ztsClient;
         this.keyProvider = keyProvider;
         this.zoneConfig = zoneConfig;
-        this.currentKeyStore = initializeKeystore(configserverIdentity, keyProvider, certificateClient, zoneConfig, updatePeriod);
+        this.currentKeyStore = initializeKeystore(configserverIdentity, keyProvider, ztsClient, zoneConfig, updatePeriod);
         this.updatePeriod = updatePeriod;
         this.configserverIdentity = configserverIdentity;
     }
 
     private static KeyStoreAndPassword initializeKeystore(AthenzService configserverIdentity,
                                                           KeyProvider keyProvider,
-                                                          AthenzCertificateClient zoneConfig,
+                                                          ZtsClient ztsClient,
                                                           AthenzProviderServiceConfig.Zones keystoreCacheDirectory,
                                                           Duration updatePeriod) {
         return tryReadKeystoreFile(configserverIdentity, updatePeriod)
-                .orElseGet(() -> downloadCertificate(configserverIdentity, keyProvider, zoneConfig, keystoreCacheDirectory));
+                .orElseGet(() -> downloadCertificate(configserverIdentity, keyProvider, ztsClient, keystoreCacheDirectory));
     }
 
     private static Optional<KeyStoreAndPassword> tryReadKeystoreFile(AthenzService configserverIdentity,
@@ -113,6 +121,7 @@ public class AthenzSslKeyStoreConfigurator extends AbstractComponent implements
         try {
             scheduler.shutdownNow();
             scheduler.awaitTermination(30, TimeUnit.SECONDS);
+            ztsClient.close();
         } catch (InterruptedException e) {
             throw new RuntimeException("Failed to shutdown Athenz certificate updater on time", e);
         }
@@ -129,18 +138,26 @@ public class AthenzSslKeyStoreConfigurator extends AbstractComponent implements
 
     private static KeyStoreAndPassword downloadCertificate(AthenzService configserverIdentity,
                                                            KeyProvider keyProvider,
-                                                           AthenzCertificateClient certificateClient,
+                                                           ZtsClient ztsClient,
                                                            AthenzProviderServiceConfig.Zones zoneConfig) {
         PrivateKey privateKey = keyProvider.getPrivateKey(zoneConfig.secretVersion());
-        X509Certificate certificate = certificateClient.updateCertificate(privateKey);
+        PublicKey publicKey = KeyUtils.extractPublicKey(privateKey);
+        Identity serviceIdentity = ztsClient.getServiceIdentity(configserverIdentity,
+                                                                Integer.toString(zoneConfig.secretVersion()),
+                                                                new KeyPair(publicKey, privateKey),
+                                                                zoneConfig.certDnsSuffix());
+        X509Certificate certificate = serviceIdentity.certificate();
         writeCredentials(configserverIdentity, certificate, privateKey);
         Instant expirationTime = certificate.getNotAfter().toInstant();
         Duration expiry = Duration.between(certificate.getNotBefore().toInstant(), expirationTime);
         log.log(LogLevel.INFO, String.format("Got Athenz x509 certificate with expiry %s (expires %s)", expiry, expirationTime));
 
+        List<X509Certificate> chain = new ArrayList<>();
+        chain.add(certificate);
+        chain.addAll(serviceIdentity.caCertificates());
         char[] keystorePassword = generateKeystorePassword();
         KeyStore keyStore = KeyStoreBuilder.withType(KeyStoreType.JKS)
-                .withKeyEntry(CERTIFICATE_ALIAS, privateKey, keystorePassword, certificate)
+                .withKeyEntry(CERTIFICATE_ALIAS, privateKey, keystorePassword, chain)
                 .build();
         return new KeyStoreAndPassword(keyStore, keystorePassword);
     }
@@ -168,7 +185,7 @@ public class AthenzSslKeyStoreConfigurator extends AbstractComponent implements
         public void run() {
             try {
                 log.log(LogLevel.INFO, "Updating Athenz certificate from ZTS");
-                currentKeyStore = downloadCertificate(configserverIdentity, keyProvider, certificateClient, zoneConfig);
+                currentKeyStore = downloadCertificate(configserverIdentity, keyProvider, ztsClient, zoneConfig);
                 sslKeyStoreContext.updateKeyStore(currentKeyStore.keyStore, new String(currentKeyStore.password));
                 log.log(LogLevel.INFO, "Athenz certificate reload successfully completed");
             } catch (Throwable e) {
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/AthenzCertificateClient.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/AthenzCertificateClient.java
deleted file mode 100644
index 193a573c98d..00000000000
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/AthenzCertificateClient.java
+++ /dev/null
@@ -1,40 +0,0 @@
-// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl;
-
-import com.yahoo.athenz.zts.InstanceRefreshRequest;
-import com.yahoo.athenz.zts.ZTSClient;
-import com.yahoo.vespa.athenz.identity.ServiceIdentityProvider;
-import com.yahoo.vespa.athenz.tls.X509CertificateUtils;
-import com.yahoo.vespa.hosted.athenz.instanceproviderservice.config.AthenzProviderServiceConfig;
-
-import javax.net.ssl.SSLContext;
-import java.security.PrivateKey;
-import java.security.cert.X509Certificate;
-
-/**
- * @author bjorncs
- */
-public class AthenzCertificateClient {
-
-    private final AthenzProviderServiceConfig.Zones zoneConfig;
-    private final ServiceIdentityProvider bootstrapIdentity;
-
-    public AthenzCertificateClient(ServiceIdentityProvider bootstrapIdentity,
-                                   AthenzProviderServiceConfig.Zones zoneConfig) {
-        this.bootstrapIdentity = bootstrapIdentity;
-        this.zoneConfig = zoneConfig;
-    }
-
-    public X509Certificate updateCertificate(PrivateKey privateKey) {
-        SSLContext bootstrapSslContext = bootstrapIdentity.getIdentitySslContext();
-        ZTSClient ztsClient = new ZTSClient(zoneConfig.ztsUrl(), bootstrapSslContext);
-        InstanceRefreshRequest req =
-                ZTSClient.generateInstanceRefreshRequest(
-                        zoneConfig.domain(), zoneConfig.serviceName(), privateKey, zoneConfig.certDnsSuffix(), /*expiryTime*/0);
-        req.setKeyId(Integer.toString(zoneConfig.secretVersion()));
-        String pemEncoded = ztsClient.postInstanceRefreshRequest(zoneConfig.domain(), zoneConfig.serviceName(), req)
-                .getCertificate();
-        return X509CertificateUtils.fromPem(pemEncoded);
-    }
-
-}