Use AthenzSslContextBuilder to build ssl context

2 files changed, 15 insertions, 42 deletions

diff --git a/container-disc/pom.xml b/container-disc/pom.xml
index 38b4bfc2ff5..f970de0b328 100644
--- a/container-disc/pom.xml
+++ b/container-disc/pom.xml
@@ -106,6 +106,12 @@
       <artifactId>vespalog</artifactId>
       <version>${project.version}</version>
     </dependency>
+    <dependency>
+      <groupId>com.yahoo.vespa</groupId>
+      <artifactId>vespa-athenz</artifactId>
+      <version>${project.version}</version>
+      <scope>provided</scope>
+    </dependency>
     <!-- WARNING: These are only here to make bundlification work -->
     <dependency>
       <groupId>com.yahoo.vespa</groupId>
diff --git a/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/impl/AthenzIdentityProviderImpl.java b/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/impl/AthenzIdentityProviderImpl.java
index b1cdbe50ae9..b6072b03e03 100644
--- a/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/impl/AthenzIdentityProviderImpl.java
+++ b/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/impl/AthenzIdentityProviderImpl.java
@@ -8,12 +8,15 @@ 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.AthenzIdentityCertificate;
+import com.yahoo.vespa.athenz.tls.AthenzSslContextBuilder;
 
 import javax.net.ssl.KeyManager;
 import javax.net.ssl.KeyManagerFactory;
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.TrustManager;
 import javax.net.ssl.TrustManagerFactory;
+import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.security.KeyManagementException;
@@ -114,48 +117,12 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     @Override
     public SSLContext getIdentitySslContext() {
-        try {
-            SSLContext sslContext = SSLContext.getInstance("TLSv1.2");
-            sslContext.init(createKeyManagersWithServiceCertificate(),
-                            createTrustManagersWithAthenzCa(),
-                            null);
-            return sslContext;
-        } catch (NoSuchAlgorithmException | KeyManagementException e) {
-            throw new RuntimeException(e);
-        }
-    }
-
-    private KeyManager[] createKeyManagersWithServiceCertificate() {
-        try {
-            credentialsRetrievedSignal.await();
-            KeyStore keyStore = KeyStore.getInstance("JKS");
-            keyStore.load(null);
-            keyStore.setKeyEntry("instance-key",
-                                 credentials.get().getKeyPair().getPrivate(),
-                                 new char[0],
-                                 new Certificate[]{credentials.get().getCertificate()});
-            KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
-            keyManagerFactory.init(keyStore, new char[0]);
-            return keyManagerFactory.getKeyManagers();
-        } catch (KeyStoreException | NoSuchAlgorithmException | UnrecoverableKeyException | CertificateException | IOException e) {
-            throw new RuntimeException(e);
-        } catch (InterruptedException e) {
-            throw new AthenzIdentityProviderException("Failed to register instance credentials", lastThrowable.get());
-        }
-    }
-
-    private static TrustManager[] createTrustManagersWithAthenzCa() {
-        try {
-            KeyStore trustStore = KeyStore.getInstance("JKS");
-            try (FileInputStream in = new FileInputStream("/opt/yahoo/share/ssl/certs/yahoo_certificate_bundle.jks")) {
-                trustStore.load(in, null);
-            }
-            TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
-            trustManagerFactory.init(trustStore);
-            return trustManagerFactory.getTrustManagers();
-        } catch (CertificateException | IOException | KeyStoreException | NoSuchAlgorithmException e) {
-            throw new RuntimeException(e);
-        }
+        return new AthenzSslContextBuilder()
+                .withIdentityCertificate(new AthenzIdentityCertificate(
+                        credentials.get().getCertificate(),
+                        credentials.get().getKeyPair().getPrivate()))
+                .withTrustStore(new File("/opt/yahoo/share/ssl/certs/yahoo_certificate_bundle.jks"), "JKS")
+                .build();
     }
 
     @Override