Add withKeyManagerFactory() to specify custom key manager

- Introduce an interface for key manager factory.
- Change SslContextBuilder to call trust/key manager factory even when no truststore/keystore has been specified.
- Change trust manager factory to be specific for x509.
- Use TrustManagerUtils/KeyManagerUtil to construct default managers.

3 files changed, 37 insertions, 59 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;
     }
 
 }
diff --git a/security-utils/src/main/java/com/yahoo/security/tls/authz/PeerAuthorizerTrustManager.java b/security-utils/src/main/java/com/yahoo/security/tls/authz/PeerAuthorizerTrustManager.java
index 80acc940a99..eee2e502183 100644
--- a/security-utils/src/main/java/com/yahoo/security/tls/authz/PeerAuthorizerTrustManager.java
+++ b/security-utils/src/main/java/com/yahoo/security/tls/authz/PeerAuthorizerTrustManager.java
@@ -3,14 +3,12 @@ package com.yahoo.security.tls.authz;
 
 import com.yahoo.security.X509CertificateUtils;
 import com.yahoo.security.tls.AuthorizationMode;
+import com.yahoo.security.tls.TrustManagerUtils;
 import com.yahoo.security.tls.policy.AuthorizedPeers;
 
 import javax.net.ssl.SSLEngine;
-import javax.net.ssl.TrustManager;
-import javax.net.ssl.TrustManagerFactory;
 import javax.net.ssl.X509ExtendedTrustManager;
 import java.net.Socket;
-import java.security.GeneralSecurityException;
 import java.security.KeyStore;
 import java.security.cert.CertificateException;
 import java.security.cert.X509Certificate;
@@ -39,22 +37,8 @@ public class PeerAuthorizerTrustManager extends X509ExtendedTrustManager {
         this.defaultTrustManager = defaultTrustManager;
     }
 
-    public static TrustManager[] wrapTrustManagersFromKeystore(AuthorizedPeers authorizedPeers, AuthorizationMode mode, KeyStore keystore) throws GeneralSecurityException {
-        TrustManagerFactory factory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
-        factory.init(keystore);
-        return wrapTrustManagers(authorizedPeers, mode, factory.getTrustManagers());
-    }
-
-    public static TrustManager[] wrapTrustManagers(AuthorizedPeers authorizedPeers, AuthorizationMode mode, TrustManager[] managers) {
-        TrustManager[] wrappedManagers = new TrustManager[managers.length];
-        for (int i = 0; i < managers.length; i++) {
-            if (managers[i] instanceof X509ExtendedTrustManager) {
-                wrappedManagers[i] = new PeerAuthorizerTrustManager(authorizedPeers, mode, (X509ExtendedTrustManager) managers[i]);
-            } else {
-                wrappedManagers[i] = managers[i];
-            }
-        }
-        return wrappedManagers;
+    public PeerAuthorizerTrustManager(AuthorizedPeers authorizedPeers, AuthorizationMode mode, KeyStore truststore) {
+        this(authorizedPeers, mode, TrustManagerUtils.createDefaultX509TrustManager(truststore));
     }
 
     @Override
diff --git a/security-utils/src/main/java/com/yahoo/security/tls/authz/PeerAuthorizerTrustManagersFactory.java b/security-utils/src/main/java/com/yahoo/security/tls/authz/PeerAuthorizerTrustManagersFactory.java
index c0a3b4e41a5..6ec8450c035 100644
--- a/security-utils/src/main/java/com/yahoo/security/tls/authz/PeerAuthorizerTrustManagersFactory.java
+++ b/security-utils/src/main/java/com/yahoo/security/tls/authz/PeerAuthorizerTrustManagersFactory.java
@@ -5,14 +5,12 @@ import com.yahoo.security.SslContextBuilder;
 import com.yahoo.security.tls.AuthorizationMode;
 import com.yahoo.security.tls.policy.AuthorizedPeers;
 
-import javax.net.ssl.TrustManager;
-import java.security.GeneralSecurityException;
 import java.security.KeyStore;
 
 /**
  * @author bjorncs
  */
-public class PeerAuthorizerTrustManagersFactory implements SslContextBuilder.TrustManagersFactory {
+public class PeerAuthorizerTrustManagersFactory implements SslContextBuilder.TrustManagerFactory {
     private final AuthorizedPeers authorizedPeers;
     private AuthorizationMode mode;
 
@@ -22,7 +20,7 @@ public class PeerAuthorizerTrustManagersFactory implements SslContextBuilder.Tru
     }
 
     @Override
-    public TrustManager[] createTrustManagers(KeyStore truststore) throws GeneralSecurityException {
-        return PeerAuthorizerTrustManager.wrapTrustManagersFromKeystore(authorizedPeers, mode, truststore);
+    public PeerAuthorizerTrustManager createTrustManager(KeyStore truststore) {
+        return new PeerAuthorizerTrustManager(authorizedPeers, mode, truststore);
     }
 }