Add fluent api for building KeyStore

7 files changed, 277 insertions, 90 deletions

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/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..2a2b74ab949
--- /dev/null
+++ b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/tls/TestUtils.java
@@ -0,0 +1,58 @@
+// 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;
+
+/**
+ * @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 {
+        try (OutputStream out = new BufferedOutputStream(new FileOutputStream(file))) {
+            KeyStore keystore = createKeystore(type, password);
+            keystore.store(out, password);
+        }
+    }
+}