1 files changed, 31 insertions, 35 deletions

diff --git a/security-utils/src/main/java/com/yahoo/security/SslContextBuilder.java b/security-utils/src/main/java/com/yahoo/security/SslContextBuilder.java
index 09a5a87138f..1ef4df9c7bc 100644
--- a/security-utils/src/main/java/com/yahoo/security/SslContextBuilder.java
+++ b/security-utils/src/main/java/com/yahoo/security/SslContextBuilder.java
@@ -1,11 +1,14 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.security;
 
+import com.yahoo.security.tls.KeyManagerUtils;
+import com.yahoo.security.tls.TrustManagerUtils;
+
 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 javax.net.ssl.X509ExtendedKeyManager;
+import javax.net.ssl.X509ExtendedTrustManager;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.nio.file.Files;
@@ -19,14 +22,17 @@ import java.util.List;
 import static java.util.Collections.singletonList;
 
 /**
+ * A builder for {@link SSLContext}.
+ *
  * @author bjorncs
  */
 public class SslContextBuilder {
 
-    private KeyStoreSupplier trustStoreSupplier;
-    private KeyStoreSupplier keyStoreSupplier;
+    private KeyStoreSupplier trustStoreSupplier = () -> null;
+    private KeyStoreSupplier keyStoreSupplier = () -> null;
     private char[] keyStorePassword;
-    private TrustManagersFactory trustManagersFactory = SslContextBuilder::createDefaultTrustManagers;
+    private TrustManagerFactory trustManagerFactory = TrustManagerUtils::createDefaultX509TrustManager;
+    private KeyManagerFactory keyManagerFactory = KeyManagerUtils::createDefaultX509KeyManager;
 
     public SslContextBuilder() {}
 
@@ -94,18 +100,21 @@ public class SslContextBuilder {
         return this;
     }
 
-    public SslContextBuilder withTrustManagerFactory(TrustManagersFactory trustManagersFactory) {
-        this.trustManagersFactory = trustManagersFactory;
+    public SslContextBuilder withTrustManagerFactory(TrustManagerFactory trustManagersFactory) {
+        this.trustManagerFactory = trustManagersFactory;
+        return this;
+    }
+
+    public SslContextBuilder withKeyManagerFactory(KeyManagerFactory keyManagerFactory) {
+        this.keyManagerFactory = keyManagerFactory;
         return this;
     }
 
     public SSLContext build() {
         try {
             SSLContext sslContext = SSLContext.getInstance("TLSv1.2");
-            TrustManager[] trustManagers =
-                    trustStoreSupplier != null ? createTrustManagers(trustManagersFactory, trustStoreSupplier) : null;
-            KeyManager[] keyManagers =
-                    keyStoreSupplier != null ? createKeyManagers(keyStoreSupplier, keyStorePassword) : null;
+            TrustManager[] trustManagers = new TrustManager[] { trustManagerFactory.createTrustManager(trustStoreSupplier.get()) };
+            KeyManager[] keyManagers = new KeyManager[] { keyManagerFactory.createKeyManager(keyStoreSupplier.get(), keyStorePassword) };
             sslContext.init(keyManagers, trustManagers, null);
             return sslContext;
         } catch (GeneralSecurityException e) {
@@ -115,27 +124,6 @@ public class SslContextBuilder {
         }
     }
 
-    private static TrustManager[] createTrustManagers(TrustManagersFactory trustManagersFactory, KeyStoreSupplier trustStoreSupplier)
-            throws GeneralSecurityException, IOException {
-        KeyStore truststore = trustStoreSupplier.get();
-        return trustManagersFactory.createTrustManagers(truststore);
-    }
-
-    private static TrustManager[] createDefaultTrustManagers(KeyStore truststore) throws GeneralSecurityException {
-        TrustManagerFactory trustManagerFactory =
-                TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
-        trustManagerFactory.init(truststore);
-        return trustManagerFactory.getTrustManagers();
-    }
-
-    private static KeyManager[] createKeyManagers(KeyStoreSupplier keyStoreSupplier, char[] password)
-            throws GeneralSecurityException, IOException {
-        KeyManagerFactory keyManagerFactory =
-                KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
-        keyManagerFactory.init(keyStoreSupplier.get(), password);
-        return keyManagerFactory.getKeyManagers();
-    }
-
     private static KeyStore createTrustStore(List<X509Certificate> caCertificates) {
         KeyStoreBuilder trustStoreBuilder = KeyStoreBuilder.withType(KeyStoreType.JKS);
         for (int i = 0; i < caCertificates.size(); i++) {
@@ -149,11 +137,19 @@ public class SslContextBuilder {
     }
 
     /**
-     * A factory interface that is similar to {@link TrustManagerFactory}, but is an interface instead of a class.
+     * A factory interface for creating {@link X509ExtendedTrustManager}.
+     */
+    @FunctionalInterface
+    public interface TrustManagerFactory {
+        X509ExtendedTrustManager createTrustManager(KeyStore truststore) throws GeneralSecurityException;
+    }
+
+    /**
+     * A factory interface for creating {@link X509ExtendedKeyManager}.
      */
     @FunctionalInterface
-    public interface TrustManagersFactory {
-        TrustManager[] createTrustManagers(KeyStore truststore) throws GeneralSecurityException;
+    public interface KeyManagerFactory {
+        X509ExtendedKeyManager createKeyManager(KeyStore truststore, char[] password) throws GeneralSecurityException;
     }
 
 }