ZOOKEEPER-3160: Custom User SSLContext

This is a master branch version of: https://github.com/apache/zookeeper/pull/654

The previous PR was for branch 3.5, and couldn't be merged as that branch is closed for new features.

The Zookeeper libraries currently allow you to set up your SSL Context via system properties such as "zookeeper.ssl.keyStore.location" in the X509Util. This covers most simple use cases, where users have software keystores on their harddrive.

There are, however, a few additional scenarios that this doesn't cover. Two possible ones would be:

1. The user has a hardware keystore, loaded in using PKCS11 or something similar.
2. The user has no access to the software keystore, but can retrieve an already-constructed SSLContext from their container.

For this, I would propose that the X509Util be extended to allow a user to set a property "zookeeper.ssl.client.context" to provide a class which supplies a custom SSL context. This gives a lot more flexibility to the ZK client, and allows the user to construct the SSLContext in whatever way they please (which also future proofs the implementation somewhat).

I added a few simple tests to this class around setting the SSLContext, and setting an invalid one. I'm not testing the actual functionality of the SSLContext, etc.

Author: Alex Rankin <davelister@gmail.com>
Author: Alex Rankin <alex.rankin@mastercard.com>

Reviewers: andor@apache.org

Closes #728 from arankin-irl/ZOOKEEPER-3160 and squashes the following commits:

a20c62feb [Alex Rankin] Merge branch 'master' into ZOOKEEPER-3160
5a9b8fcb4 [Alex Rankin] Merge pull request #7 from apache/master
3c3dfdd70 [Alex Rankin] Re-ordering imports.
69e0b6c93 [Alex Rankin] Updating custom SSLContext supplier with review comments
874529ba0 [Alex Rankin] Using supplier interface instead of custom interface, and renaming property
ec272607d [Alex Rankin] Merge branch 'master' into ZOOKEEPER-3160
75a010e34 [Alex Rankin] Merge pull request #6 from apache/master
838f61c1a [Alex Rankin] Merge branch 'master' into ZOOKEEPER-3160
f85d7e5cb [Alex Rankin] Merge pull request #5 from apache/master
31d8dd5a4 [Alex Rankin] Extracting SSLContext creation from config to new method.
400839a60 [Alex Rankin] Adding ability to specify custom SSLContext for client
7ae74851b [Alex Rankin] Merge pull request #4 from apache/master

zookeeper-server/src/main/java/org/apache/zookeeper/common/X509Util.java

@@ -17,10 +17,9 @@
  */
 package org.apache.zookeeper.common;
 
-
-import java.io.ByteArrayInputStream;
 import java.io.Closeable;
 import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
 import java.net.Socket;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -33,15 +32,14 @@ import java.security.NoSuchAlgorithmException;
 import java.security.Security;
 import java.security.cert.PKIXBuilderParameters;
 import java.security.cert.X509CertSelector;
-import java.util.Arrays;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
 import javax.net.ssl.CertPathTrustManagerParameters;
 import javax.net.ssl.KeyManager;
 import javax.net.ssl.KeyManagerFactory;
 import javax.net.ssl.SSLContext;
-import javax.net.ssl.SSLParameters;
 import javax.net.ssl.SSLServerSocket;
 import javax.net.ssl.SSLSocket;
 import javax.net.ssl.TrustManager;
@@ -137,6 +135,7 @@ public abstract class X509Util implements Closeable, AutoCloseable {
     private String sslTruststoreLocationProperty = getConfigPrefix() + "trustStore.location";
     private String sslTruststorePasswdProperty = getConfigPrefix() + "trustStore.password";
     private String sslTruststoreTypeProperty = getConfigPrefix() + "trustStore.type";
+    private String sslContextSupplierClassProperty = getConfigPrefix() + "context.supplier.class";
     private String sslHostnameVerificationEnabledProperty = getConfigPrefix() + "hostnameVerification";
     private String sslCrlEnabledProperty = getConfigPrefix() + "crl";
     private String sslOcspEnabledProperty = getConfigPrefix() + "ocsp";
@@ -202,6 +201,10 @@ public abstract class X509Util implements Closeable, AutoCloseable {
         return sslTruststoreTypeProperty;
     }
 
+    public String getSslContextSupplierClassProperty() {
+        return sslContextSupplierClassProperty;
+    }
+
     public String getSslHostnameVerificationEnabledProperty() {
         return sslHostnameVerificationEnabledProperty;
     }
@@ -282,7 +285,28 @@ public abstract class X509Util implements Closeable, AutoCloseable {
         }
     }
 
+    @SuppressWarnings("unchecked")
     public SSLContextAndOptions createSSLContextAndOptions(ZKConfig config) throws SSLContextException {
+        final String supplierContextClassName = config.getProperty(sslContextSupplierClassProperty);
+        if (supplierContextClassName != null) {
+            if (LOG.isDebugEnabled()) {
+                LOG.debug("Loading SSLContext supplier from property '{}'", sslContextSupplierClassProperty);
+            }
+            try {
+                Class<?> sslContextClass = Class.forName(supplierContextClassName);
+                Supplier<SSLContext> sslContextSupplier = (Supplier<SSLContext>) sslContextClass.getConstructor().newInstance();
+                return new SSLContextAndOptions(this, config, sslContextSupplier.get());
+            } catch (ClassNotFoundException | ClassCastException | NoSuchMethodException | InvocationTargetException |
+                    InstantiationException | IllegalAccessException e) {
+                throw new SSLContextException("Could not retrieve the SSLContext from supplier source '" + supplierContextClassName +
+                        "' provided in the property '" + sslContextSupplierClassProperty + "'", e);
+            }
+        } else {
+            return createSSLContextAndOptionsFromConfig(config);
+        }
+    }
+
+    public SSLContextAndOptions createSSLContextAndOptionsFromConfig(ZKConfig config) throws SSLContextException {
         KeyManager[] keyManagers = null;
         TrustManager[] trustManagers = null;
 

zookeeper-server/src/main/java/org/apache/zookeeper/common/ZKConfig.java

@@ -133,6 +133,8 @@ public class ZKConfig {
                 System.getProperty(x509Util.getSslTruststorePasswdProperty()));
         properties.put(x509Util.getSslTruststoreTypeProperty(),
                 System.getProperty(x509Util.getSslTruststoreTypeProperty()));
+        properties.put(x509Util.getSslContextSupplierClassProperty(),
+                System.getProperty(x509Util.getSslContextSupplierClassProperty()));
         properties.put(x509Util.getSslHostnameVerificationEnabledProperty(),
                 System.getProperty(x509Util.getSslHostnameVerificationEnabledProperty()));
         properties.put(x509Util.getSslCrlEnabledProperty(),

zookeeper-server/src/test/java/org/apache/zookeeper/common/X509UtilTest.java

@@ -22,6 +22,7 @@ import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.ServerSocket;
 import java.net.Socket;
+import java.security.NoSuchAlgorithmException;
 import java.security.Security;
 import java.util.Collection;
 import java.util.concurrent.Callable;
@@ -30,6 +31,7 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Supplier;
 
 import javax.net.ssl.HandshakeCompletedEvent;
 import javax.net.ssl.HandshakeCompletedListener;
@@ -403,6 +405,23 @@ public class X509UtilTest extends BaseX509ParameterizedTestCase {
         }
     }
 
+    @Test(expected = X509Exception.SSLContextException.class)
+    public void testCreateSSLContext_invalidCustomSSLContextClass() throws Exception {
+        ZKConfig zkConfig = new ZKConfig();
+        ClientX509Util clientX509Util = new ClientX509Util();
+        zkConfig.setProperty(clientX509Util.getSslContextSupplierClassProperty(), String.class.getCanonicalName());
+        clientX509Util.createSSLContext(zkConfig);
+    }
+
+    @Test
+    public void testCreateSSLContext_validCustomSSLContextClass() throws Exception {
+        ZKConfig zkConfig = new ZKConfig();
+        ClientX509Util clientX509Util = new ClientX509Util();
+        zkConfig.setProperty(clientX509Util.getSslContextSupplierClassProperty(), SslContextSupplier.class.getName());
+        final SSLContext sslContext = clientX509Util.createSSLContext(zkConfig);
+        Assert.assertEquals(SSLContext.getDefault(), sslContext);
+    }
+
     private static void forceClose(Socket s) {
         if (s == null || s.isClosed()) {
             return;
@@ -528,4 +547,18 @@ public class X509UtilTest extends BaseX509ParameterizedTestCase {
         x509Util.close(); // remember to close old instance before replacing it
         x509Util = new ClientX509Util();
     }
+
+    public static class SslContextSupplier implements Supplier<SSLContext> {
+
+        @Override
+        public SSLContext get() {
+            try {
+                return SSLContext.getDefault();
+            } catch (NoSuchAlgorithmException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+    }
+
 }