Merge pull request #5296 from vespa-engine/bjorncs/keystore-builder

Add fluent api for building KeyStore

12 files changed, 342 insertions, 133 deletions

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 31e1a8519f4..e4e964c7088 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
@@ -1,4 +1,4 @@
-// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright 2018 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;
 
 import com.google.inject.Inject;
@@ -9,25 +9,19 @@ import com.yahoo.container.jdisc.athenz.AthenzIdentityProvider;
 import com.yahoo.jdisc.http.ssl.SslKeyStoreConfigurator;
 import com.yahoo.jdisc.http.ssl.SslKeyStoreContext;
 import com.yahoo.log.LogLevel;
+import com.yahoo.vespa.athenz.tls.KeyStoreBuilder;
+import com.yahoo.vespa.athenz.tls.KeyStoreType;
 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.io.BufferedInputStream;
-import java.io.BufferedOutputStream;
 import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileOutputStream;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.security.GeneralSecurityException;
 import java.security.KeyStore;
 import java.security.KeyStoreException;
 import java.security.PrivateKey;
-import java.security.cert.Certificate;
 import java.security.cert.X509Certificate;
 import java.time.Duration;
 import java.time.Instant;
@@ -37,6 +31,7 @@ import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
 
+import static com.yahoo.vespa.athenz.tls.KeyStoreUtils.writeKeyStoreToFile;
 import static com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.Utils.getZoneConfig;
 
 /**
@@ -87,15 +82,14 @@ public class AthenzSslKeyStoreConfigurator extends AbstractComponent implements
     private static Optional<KeyStore> tryReadKeystoreFile(File certificateFile, Duration updatePeriod) {
         try {
             if (!certificateFile.exists()) return Optional.empty();
-            KeyStore keyStore = KeyStore.getInstance("JKS");
-            try (InputStream in = new BufferedInputStream(new FileInputStream(certificateFile))) {
-                keyStore.load(in, new char[0]);
-            }
+            KeyStore keyStore = KeyStoreBuilder.withType(KeyStoreType.JKS)
+                    .fromFile(certificateFile)
+                    .build();
             Instant minimumExpiration = Instant.now().plus(updatePeriod).plus(EXPIRATION_MARGIN);
             boolean isExpired = getCertificateExpiry(keyStore).isBefore(minimumExpiration);
             if (isExpired) return Optional.empty();
             return Optional.of(keyStore);
-        } catch (IOException | GeneralSecurityException e) {
+        } catch (GeneralSecurityException e) {
             log.log(LogLevel.ERROR, "Failed to read keystore from disk: " + e.getMessage(), e);
             return Optional.empty();
         }
@@ -139,28 +133,23 @@ public class AthenzSslKeyStoreConfigurator extends AbstractComponent implements
                                                 AthenzCertificateClient certificateClient,
                                                 AthenzProviderServiceConfig.Zones zoneConfig,
                                                 Path keystoreCachePath) {
-        try {
-            PrivateKey privateKey = keyProvider.getPrivateKey(zoneConfig.secretVersion());
-            X509Certificate certificate = certificateClient.updateCertificate(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));
-
-            KeyStore keyStore = KeyStore.getInstance("JKS");
-            keyStore.load(null);
-            keyStore.setKeyEntry(
-                    CERTIFICATE_ALIAS, privateKey, CERTIFICATE_PASSWORD.toCharArray(), new Certificate[]{certificate});
-            tryWriteKeystore(keyStore, keystoreCachePath);
-            return keyStore;
-        } catch (IOException | GeneralSecurityException e) {
-            throw new RuntimeException(e);
-        }
+        PrivateKey privateKey = keyProvider.getPrivateKey(zoneConfig.secretVersion());
+        X509Certificate certificate = certificateClient.updateCertificate(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));
+
+        KeyStore keyStore = KeyStoreBuilder.withType(KeyStoreType.JKS)
+                .withKeyEntry(CERTIFICATE_ALIAS, privateKey, CERTIFICATE_PASSWORD.toCharArray(), certificate)
+                .build();
+        tryWriteKeystore(keyStore, keystoreCachePath);
+        return keyStore;
     }
 
     private static void tryWriteKeystore(KeyStore keyStore, Path keystoreCachePath) {
-        try (OutputStream out = new BufferedOutputStream(new FileOutputStream(keystoreCachePath.toFile()))) {
-            keyStore.store(out, new char[0]);
-        } catch (IOException | GeneralSecurityException e) {
+        try  {
+            writeKeyStoreToFile(keyStore, keystoreCachePath.toFile());
+        } catch (Exception e) {
             log.log(LogLevel.ERROR, "Failed to write keystore to disk: " + e.getMessage(), e);
         }
     }
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzSslTrustStoreConfigurator.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzSslTrustStoreConfigurator.java
index 7e24109a197..376dd2ed4ac 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzSslTrustStoreConfigurator.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzSslTrustStoreConfigurator.java
@@ -1,4 +1,4 @@
-// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright 2018 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;
 
 import com.google.inject.Inject;
@@ -6,6 +6,8 @@ import com.yahoo.cloud.config.ConfigserverConfig;
 import com.yahoo.jdisc.http.ssl.SslTrustStoreConfigurator;
 import com.yahoo.jdisc.http.ssl.SslTrustStoreContext;
 import com.yahoo.log.LogLevel;
+import com.yahoo.vespa.athenz.tls.KeyStoreBuilder;
+import com.yahoo.vespa.athenz.tls.KeyStoreType;
 import com.yahoo.vespa.hosted.athenz.instanceproviderservice.config.AthenzProviderServiceConfig;
 import org.bouncycastle.asn1.x500.X500Name;
 import org.bouncycastle.asn1.x509.BasicConstraints;
@@ -20,7 +22,7 @@ import org.bouncycastle.operator.ContentSigner;
 import org.bouncycastle.operator.OperatorCreationException;
 import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
 
-import java.io.FileInputStream;
+import java.io.File;
 import java.io.IOException;
 import java.math.BigInteger;
 import java.security.KeyPair;
@@ -70,12 +72,10 @@ public class AthenzSslTrustStoreConfigurator implements SslTrustStoreConfigurato
             KeyPair keyPair = getKeyPair(keyProvider, configserverConfig, athenzProviderServiceConfig);
             X509Certificate selfSignedCertificate = createSelfSignedCertificate(keyPair, configserverConfig);
             log.log(LogLevel.FINE, "Generated self-signed certificate: " + selfSignedCertificate);
-            KeyStore trustStore = KeyStore.getInstance("JKS");
-            try (FileInputStream in = new FileInputStream(athenzProviderServiceConfig.athenzCaTrustStore())) {
-                trustStore.load(in, "changeit".toCharArray());
-            }
-            trustStore.setCertificateEntry(CERTIFICATE_ALIAS, selfSignedCertificate);
-            return trustStore;
+            return KeyStoreBuilder.withType(KeyStoreType.JKS)
+                    .fromFile(new File(athenzProviderServiceConfig.athenzCaTrustStore()), "changeit".toCharArray())
+                    .withCertificateEntry(CERTIFICATE_ALIAS, selfSignedCertificate)
+                    .build();
         } catch (Exception e) {
             throw new RuntimeException(e);
         }
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/AthenzSslContextProviderImpl.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/AthenzSslContextProviderImpl.java
index 261c2c0f2ad..ca939a9f862 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/AthenzSslContextProviderImpl.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/AthenzSslContextProviderImpl.java
@@ -13,7 +13,7 @@ import java.time.Duration;
 import java.time.Instant;
 import java.util.concurrent.atomic.AtomicReference;
 
-import static com.yahoo.vespa.athenz.tls.AthenzSslContextBuilder.KeyStoreType.JKS;
+import static com.yahoo.vespa.athenz.tls.KeyStoreType.JKS;
 
 /**
  * @author bjorncs
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/SslConfigServerApiImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/SslConfigServerApiImpl.java
index 7abe9bce718..04b222875c3 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/SslConfigServerApiImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/SslConfigServerApiImpl.java
@@ -3,7 +3,7 @@ package com.yahoo.vespa.hosted.node.admin.configserver;
 
 import com.yahoo.vespa.athenz.tls.AthenzIdentityVerifier;
 import com.yahoo.vespa.athenz.tls.AthenzSslContextBuilder;
-import com.yahoo.vespa.athenz.tls.AthenzSslContextBuilder.KeyStoreType;
+import com.yahoo.vespa.athenz.tls.KeyStoreType;
 import com.yahoo.vespa.hosted.node.admin.component.Environment;
 import com.yahoo.vespa.hosted.node.admin.configserver.certificate.ConfigServerKeyStoreRefresher;
 import com.yahoo.vespa.hosted.node.admin.util.KeyStoreOptions;
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImpl.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImpl.java
index 2bfcaae79e6..bf7a9eb1c31 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImpl.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImpl.java
@@ -17,14 +17,12 @@ import java.io.File;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Logger;
 
-import static com.yahoo.vespa.athenz.tls.AthenzSslContextBuilder.KeyStoreType.JKS;
+import static com.yahoo.vespa.athenz.tls.KeyStoreType.JKS;
 
 /**
  * @author mortent
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/AthenzSslContextBuilder.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/AthenzSslContextBuilder.java
index fdf58f9e64b..57fc7d1b581 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/AthenzSslContextBuilder.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/AthenzSslContextBuilder.java
@@ -2,7 +2,6 @@
 package com.yahoo.vespa.athenz.tls;
 
 import com.yahoo.vespa.athenz.api.AthenzIdentityCertificate;
-import org.bouncycastle.jce.provider.BouncyCastleProvider;
 
 import javax.net.ssl.KeyManager;
 import javax.net.ssl.KeyManagerFactory;
@@ -10,14 +9,11 @@ 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.io.UncheckedIOException;
 import java.security.GeneralSecurityException;
 import java.security.KeyStore;
-import java.security.KeyStoreException;
 import java.security.PrivateKey;
-import java.security.cert.Certificate;
 import java.security.cert.X509Certificate;
 
 /**
@@ -25,22 +21,6 @@ import java.security.cert.X509Certificate;
  */
 public class AthenzSslContextBuilder {
 
-    public enum KeyStoreType {
-        JKS {
-            KeyStore createKeystore() throws KeyStoreException {
-                return KeyStore.getInstance("JKS");
-            }
-        },
-        PKCS12 {
-            private final BouncyCastleProvider bouncyCastleProvider = new BouncyCastleProvider();
-
-            KeyStore createKeystore() throws KeyStoreException {
-                return KeyStore.getInstance("PKCS12", bouncyCastleProvider);
-            }
-        };
-        abstract KeyStore createKeystore() throws GeneralSecurityException;
-    }
-
     private KeyStoreSupplier trustStoreSupplier;
     private KeyStoreSupplier keyStoreSupplier;
     private char[] keyStorePassword;
@@ -48,7 +28,7 @@ public class AthenzSslContextBuilder {
     public AthenzSslContextBuilder() {}
 
     public AthenzSslContextBuilder withTrustStore(File file, KeyStoreType trustStoreType) {
-        this.trustStoreSupplier = () -> loadKeyStoreFromFile(file, null, trustStoreType);
+        this.trustStoreSupplier = () -> KeyStoreBuilder.withType(trustStoreType).fromFile(file).build();
         return this;
     }
 
@@ -63,7 +43,7 @@ public class AthenzSslContextBuilder {
 
     public AthenzSslContextBuilder withKeyStore(PrivateKey privateKey, X509Certificate certificate) {
         char[] pwd = new char[0];
-        this.keyStoreSupplier = () -> createJksKeyStore(privateKey, certificate, pwd);
+        this.keyStoreSupplier = () -> KeyStoreBuilder.withType(KeyStoreType.JKS).withKeyEntry("athenz", privateKey, certificate).build();
         this.keyStorePassword = pwd;
         return this;
     }
@@ -75,7 +55,7 @@ public class AthenzSslContextBuilder {
     }
 
     public AthenzSslContextBuilder withKeyStore(File file, char[] password, KeyStoreType keyStoreType) {
-        this.keyStoreSupplier = () -> loadKeyStoreFromFile(file, password, keyStoreType);
+        this.keyStoreSupplier = () -> KeyStoreBuilder.withType(keyStoreType).fromFile(file, password).build();
         this.keyStorePassword = password;
         return this;
     }
@@ -112,23 +92,6 @@ public class AthenzSslContextBuilder {
         return keyManagerFactory.getKeyManagers();
     }
 
-    private static KeyStore loadKeyStoreFromFile(File file, char[] password, KeyStoreType keyStoreType)
-            throws IOException, GeneralSecurityException{
-        KeyStore keyStore = keyStoreType.createKeystore();
-        try (FileInputStream in = new FileInputStream(file)) {
-            keyStore.load(in, password);
-        }
-        return keyStore;
-    }
-
-    private KeyStore createJksKeyStore(PrivateKey privateKey, X509Certificate certificate, char[] password)
-            throws GeneralSecurityException, IOException{
-        KeyStore keyStore = KeyStoreType.JKS.createKeystore();
-        keyStore.load(null);
-        keyStore.setKeyEntry("athenz-identity", privateKey, password, new Certificate[]{certificate});
-        return keyStore;
-    }
-
     private interface KeyStoreSupplier {
         KeyStore get() throws IOException, GeneralSecurityException;
     }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/KeyStoreBuilder.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/KeyStoreBuilder.java
new file mode 100644
index 00000000000..a9279f45129
--- /dev/null
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/KeyStoreBuilder.java
@@ -0,0 +1,121 @@
+// 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.tls;
+
+import java.io.BufferedInputStream;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.UncheckedIOException;
+import java.security.GeneralSecurityException;
+import java.security.KeyStore;
+import java.security.PrivateKey;
+import java.security.cert.Certificate;
+import java.security.cert.X509Certificate;
+import java.util.ArrayList;
+import java.util.List;
+
+import static java.util.Collections.singletonList;
+
+/**
+ * @author bjorncs
+ */
+public class KeyStoreBuilder {
+
+    private final List<KeyEntry> keyEntries = new ArrayList<>();
+    private final List<CertificateEntry> certificateEntries = new ArrayList<>();
+
+    private final KeyStoreType keyStoreType;
+    private File inputFile;
+    private char[] inputFilePassword;
+
+    private KeyStoreBuilder(KeyStoreType keyStoreType) {
+        this.keyStoreType = keyStoreType;
+    }
+
+    public static KeyStoreBuilder withType(KeyStoreType type) {
+        return new KeyStoreBuilder(type);
+    }
+
+    public KeyStoreBuilder fromFile(File file, char[] password) {
+        this.inputFile = file;
+        this.inputFilePassword = password;
+        return this;
+    }
+
+    public KeyStoreBuilder fromFile(File file) {
+        return fromFile(file, null);
+    }
+
+    public KeyStoreBuilder withKeyEntry(String alias, PrivateKey privateKey, char[] password, List<X509Certificate> certificateChain) {
+        keyEntries.add(new KeyEntry(alias, privateKey, certificateChain, password));
+        return this;
+    }
+
+    public KeyStoreBuilder withKeyEntry(String alias, PrivateKey privateKey, char[] password, X509Certificate certificate) {
+        return withKeyEntry(alias, privateKey, password, singletonList(certificate));
+    }
+
+    public KeyStoreBuilder withKeyEntry(String alias, PrivateKey privateKey, X509Certificate certificate) {
+        return withKeyEntry(alias, privateKey, null, certificate);
+    }
+
+    public KeyStoreBuilder withKeyEntry(String alias, PrivateKey privateKey, List<X509Certificate> certificateChain) {
+        return withKeyEntry(alias, privateKey, null, certificateChain);
+    }
+
+    public KeyStoreBuilder withCertificateEntry(String alias, X509Certificate certificate) {
+        certificateEntries.add(new CertificateEntry(alias, certificate));
+        return this;
+    }
+
+    public KeyStore build() {
+        try {
+            KeyStore keystore = this.keyStoreType.createKeystore();
+            if (this.inputFile != null) {
+                try (InputStream in = new BufferedInputStream(new FileInputStream(this.inputFile))) {
+                    keystore.load(in, this.inputFilePassword);
+                }
+            } else {
+                keystore.load(null);
+            }
+            for (KeyEntry entry : keyEntries) {
+                char[] password = entry.password != null ? entry.password : new char[0];
+                Certificate[] certificateChain = entry.certificateChain.toArray(new Certificate[entry.certificateChain.size()]);
+                keystore.setKeyEntry(entry.alias, entry.privateKey, password, certificateChain);
+            }
+            for (CertificateEntry entry : certificateEntries) {
+                keystore.setCertificateEntry(entry.alias, entry.certificate);
+            }
+            return keystore;
+        } catch (GeneralSecurityException e) {
+            throw new RuntimeException(e);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    private static class KeyEntry {
+        final String alias;
+        final PrivateKey privateKey;
+        final List<X509Certificate> certificateChain;
+        final char[] password;
+
+        KeyEntry(String alias, PrivateKey privateKey, List<X509Certificate> certificateChain, char[] password) {
+            this.alias = alias;
+            this.privateKey = privateKey;
+            this.certificateChain = certificateChain;
+            this.password = password;
+        }
+    }
+
+    private static class CertificateEntry {
+        final String alias;
+        final X509Certificate certificate;
+
+        CertificateEntry(String alias, X509Certificate certificate) {
+            this.alias = alias;
+            this.certificate = certificate;
+        }
+    }
+}
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/KeyStoreType.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/KeyStoreType.java
new file mode 100644
index 00000000000..a5c35549540
--- /dev/null
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/KeyStoreType.java
@@ -0,0 +1,27 @@
+// 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.tls;
+
+import org.bouncycastle.jce.provider.BouncyCastleProvider;
+
+import java.security.GeneralSecurityException;
+import java.security.KeyStore;
+import java.security.KeyStoreException;
+
+/**
+ * @author bjorncs
+ */
+public enum KeyStoreType {
+    JKS {
+        KeyStore createKeystore() throws KeyStoreException {
+            return KeyStore.getInstance("JKS");
+        }
+    },
+    PKCS12 {
+        private final BouncyCastleProvider bouncyCastleProvider = new BouncyCastleProvider();
+
+        KeyStore createKeystore() throws KeyStoreException {
+            return KeyStore.getInstance("PKCS12", bouncyCastleProvider);
+        }
+    };
+    abstract KeyStore createKeystore() throws GeneralSecurityException;
+}
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/KeyStoreUtils.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/KeyStoreUtils.java
new file mode 100644
index 00000000000..12aaa40cce4
--- /dev/null
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/tls/KeyStoreUtils.java
@@ -0,0 +1,34 @@
+// 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.tls;
+
+import java.io.BufferedOutputStream;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.io.UncheckedIOException;
+import java.security.GeneralSecurityException;
+import java.security.KeyStore;
+
+/**
+ * @author bjorncs
+ */
+public class KeyStoreUtils {
+    private KeyStoreUtils() {}
+
+    public static void writeKeyStoreToFile(KeyStore keyStore, File file, char[] password) {
+        try (OutputStream out = new BufferedOutputStream(new FileOutputStream(file))) {
+            keyStore.store(out, password);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        } catch (GeneralSecurityException e) {
+            throw new RuntimeException(e);
+        }
+
+    }
+
+    public static void writeKeyStoreToFile(KeyStore keyStore, File file) {
+        writeKeyStoreToFile(keyStore, file, new char[0]);
+    }
+
+}
diff --git a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/tls/AthenzSslContextBuilderTest.java b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/tls/AthenzSslContextBuilderTest.java
index 38f3b12b16b..20ac8791863 100644
--- a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/tls/AthenzSslContextBuilderTest.java
+++ b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/tls/AthenzSslContextBuilderTest.java
@@ -1,27 +1,19 @@
 // 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.tls;
 
-import com.yahoo.athenz.auth.util.Crypto;
-import com.yahoo.vespa.athenz.tls.AthenzSslContextBuilder.KeyStoreType;
-import org.bouncycastle.asn1.x500.X500Name;
-import org.bouncycastle.operator.OperatorCreationException;
-import org.bouncycastle.pkcs.PKCS10CertificationRequest;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
 
-import java.io.BufferedOutputStream;
 import java.io.File;
-import java.io.FileOutputStream;
-import java.io.IOException;
-import java.io.OutputStream;
 import java.security.KeyPair;
-import java.security.KeyPairGenerator;
-import java.security.KeyStore;
-import java.security.NoSuchAlgorithmException;
-import java.security.cert.Certificate;
 import java.security.cert.X509Certificate;
 
+import static com.yahoo.vespa.athenz.tls.TestUtils.createCertificate;
+import static com.yahoo.vespa.athenz.tls.TestUtils.createKeyPair;
+import static com.yahoo.vespa.athenz.tls.TestUtils.createKeystore;
+import static com.yahoo.vespa.athenz.tls.TestUtils.createKeystoreFile;
+
 /**
  * @author bjorncs
  */
@@ -35,22 +27,22 @@ public class AthenzSslContextBuilderTest {
     @Test
     public void can_build_sslcontext_with_truststore_only() throws Exception {
         new AthenzSslContextBuilder()
-                .withTrustStore(createKeystore("JKS"))
+                .withTrustStore(createKeystore(KeyStoreType.JKS, PASSWORD))
                 .build();
     }
 
     @Test
     public void can_build_sslcontext_with_keystore_only() throws Exception {
         new AthenzSslContextBuilder()
-                .withKeyStore(createKeystore("JKS"), PASSWORD)
+                .withKeyStore(createKeystore(KeyStoreType.JKS, PASSWORD), PASSWORD)
                 .build();
     }
 
     @Test
     public void can_build_sslcontext_with_truststore_and_keystore() throws Exception {
         new AthenzSslContextBuilder()
-                .withKeyStore(createKeystore("JKS"), PASSWORD)
-                .withTrustStore(createKeystore("JKS"))
+                .withKeyStore(createKeystore(KeyStoreType.JKS, PASSWORD), PASSWORD)
+                .withTrustStore(createKeystore(KeyStoreType.JKS, PASSWORD))
                 .build();
     }
 
@@ -65,11 +57,9 @@ public class AthenzSslContextBuilderTest {
 
     @Test
     public void can_build_sslcontext_with_jks_keystore_from_file() throws Exception {
-        KeyStore keystore = createKeystore("JKS");
         File keystoreFile = tempDirectory.newFile();
-        try (OutputStream out = new BufferedOutputStream(new FileOutputStream(keystoreFile))) {
-            keystore.store(out, PASSWORD);
-        }
+        createKeystoreFile(keystoreFile, KeyStoreType.JKS, PASSWORD);
+
         new AthenzSslContextBuilder()
                 .withKeyStore(keystoreFile, PASSWORD, KeyStoreType.JKS)
                 .build();
@@ -77,36 +67,12 @@ public class AthenzSslContextBuilderTest {
 
     @Test
     public void can_build_sslcontext_with_pcks12_keystore_from_file() throws Exception {
-        KeyStore keystore = createKeystore("PKCS12");
         File keystoreFile = tempDirectory.newFile();
-        try (OutputStream out = new BufferedOutputStream(new FileOutputStream(keystoreFile))) {
-            keystore.store(out, PASSWORD);
-        }
+        createKeystoreFile(keystoreFile, KeyStoreType.PKCS12, PASSWORD);
+
         new AthenzSslContextBuilder()
                 .withKeyStore(keystoreFile, PASSWORD, KeyStoreType.PKCS12)
                 .build();
     }
 
-    private static KeyStore createKeystore(String type) throws Exception {
-        KeyPair keyPair = createKeyPair();
-        KeyStore keystore = KeyStore.getInstance(type);
-        keystore.load(null);
-        keystore.setKeyEntry("entry-name", keyPair.getPrivate(), PASSWORD, new Certificate[]{createCertificate(keyPair)});
-        return keystore;
-    }
-
-    private static X509Certificate createCertificate(KeyPair keyPair) throws
-            OperatorCreationException, IOException {
-        String x500Principal = "CN=mysubject";
-        PKCS10CertificationRequest csr =
-                Crypto.getPKCS10CertRequest(
-                        Crypto.generateX509CSR(keyPair.getPrivate(), x500Principal, null));
-        return Crypto.generateX509Certificate(csr, keyPair.getPrivate(), new X500Name(x500Principal), 3600, false);
-    }
-
-    private static KeyPair createKeyPair() throws NoSuchAlgorithmException {
-        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
-        keyGen.initialize(512);
-        return keyGen.genKeyPair();
-    }
 }
diff --git a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/tls/KeyStoreBuilderTest.java b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/tls/KeyStoreBuilderTest.java
new file mode 100644
index 00000000000..1b6fa8bcbf1
--- /dev/null
+++ b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/tls/KeyStoreBuilderTest.java
@@ -0,0 +1,54 @@
+package com.yahoo.vespa.athenz.tls;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.io.File;
+import java.security.KeyPair;
+import java.security.cert.X509Certificate;
+
+import static com.yahoo.vespa.athenz.tls.TestUtils.createCertificate;
+import static com.yahoo.vespa.athenz.tls.TestUtils.createKeyPair;
+import static com.yahoo.vespa.athenz.tls.TestUtils.createKeystoreFile;
+
+/**
+ * @author bjorncs
+ */
+public class KeyStoreBuilderTest {
+
+    private static final char[] PASSWORD = new char[0];
+
+    @Rule
+    public TemporaryFolder tempDirectory = new TemporaryFolder();
+
+    @Test
+    public void can_create_jks_keystore_from_privatekey_and_certificate() throws Exception {
+        KeyPair keyPair = createKeyPair();
+        X509Certificate certificate = createCertificate(keyPair);
+        KeyStoreBuilder.withType(KeyStoreType.JKS)
+                .withKeyEntry("key", keyPair.getPrivate(), certificate)
+                .build();
+    }
+
+    @Test
+    public void can_build_jks_keystore_from_file() throws Exception {
+        File keystoreFile = tempDirectory.newFile();
+        createKeystoreFile(keystoreFile, KeyStoreType.JKS, PASSWORD);
+
+        KeyStoreBuilder.withType(KeyStoreType.JKS)
+                .fromFile(keystoreFile, PASSWORD)
+                .build();
+    }
+
+    @Test
+    public void can_build_pcks12_keystore_from_file() throws Exception {
+        File keystoreFile = tempDirectory.newFile();
+        createKeystoreFile(keystoreFile, KeyStoreType.PKCS12, PASSWORD);
+
+        KeyStoreBuilder.withType(KeyStoreType.PKCS12)
+                .fromFile(keystoreFile, PASSWORD)
+                .build();
+    }
+
+}
\ No newline at end of file
diff --git a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/tls/TestUtils.java b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/tls/TestUtils.java
new file mode 100644
index 00000000000..54601c04514
--- /dev/null
+++ b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/tls/TestUtils.java
@@ -0,0 +1,57 @@
+// 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.tls;
+
+import com.yahoo.athenz.auth.util.Crypto;
+import org.bouncycastle.asn1.x500.X500Name;
+import org.bouncycastle.operator.OperatorCreationException;
+import org.bouncycastle.pkcs.PKCS10CertificationRequest;
+
+import java.io.BufferedOutputStream;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.security.GeneralSecurityException;
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.KeyStore;
+import java.security.NoSuchAlgorithmException;
+import java.security.cert.Certificate;
+import java.security.cert.X509Certificate;
+
+import static com.yahoo.vespa.athenz.tls.KeyStoreUtils.writeKeyStoreToFile;
+
+/**
+ * @author bjorncs
+ */
+class TestUtils {
+
+    static KeyStore createKeystore(KeyStoreType type, char[] password)
+            throws GeneralSecurityException, IOException, OperatorCreationException {
+        KeyPair keyPair = createKeyPair();
+        KeyStore keystore = type.createKeystore();
+        keystore.load(null);
+        keystore.setKeyEntry("entry-name", keyPair.getPrivate(), password, new Certificate[]{createCertificate(keyPair)});
+        return keystore;
+    }
+
+    static X509Certificate createCertificate(KeyPair keyPair)
+            throws OperatorCreationException, IOException {
+        String x500Principal = "CN=mysubject";
+        PKCS10CertificationRequest csr =
+                Crypto.getPKCS10CertRequest(
+                        Crypto.generateX509CSR(keyPair.getPrivate(), x500Principal, null));
+        return Crypto.generateX509Certificate(csr, keyPair.getPrivate(), new X500Name(x500Principal), 3600, false);
+    }
+
+    static KeyPair createKeyPair() throws NoSuchAlgorithmException {
+        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
+        keyGen.initialize(4096);
+        return keyGen.genKeyPair();
+    }
+
+    static void createKeystoreFile(File file, KeyStoreType type, char[] password)
+            throws IOException, GeneralSecurityException, OperatorCreationException {
+        writeKeyStoreToFile(createKeystore(type, password), file, password);
+    }
+}