ZOOKEEPER-3160: Custom User SSLContext

This is a master branch version of: apache#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 <[email protected]>
Author: Alex Rankin <[email protected]>

Reviewers: [email protected]

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

a20c62f [Alex Rankin] Merge branch 'master' into ZOOKEEPER-3160
5a9b8fc [Alex Rankin] Merge pull request #7 from apache/master
3c3dfdd [Alex Rankin] Re-ordering imports.
69e0b6c [Alex Rankin] Updating custom SSLContext supplier with review comments
874529b [Alex Rankin] Using supplier interface instead of custom interface, and renaming property
ec27260 [Alex Rankin] Merge branch 'master' into ZOOKEEPER-3160
75a010e [Alex Rankin] Merge pull request #6 from apache/master
838f61c [Alex Rankin] Merge branch 'master' into ZOOKEEPER-3160
f85d7e5 [Alex Rankin] Merge pull request #5 from apache/master
31d8dd5 [Alex Rankin] Extracting SSLContext creation from config to new method.
400839a [Alex Rankin] Adding ability to specify custom SSLContext for client
7ae7485 [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.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 static ClientAuth fromPropertyValue(String prop) {
     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 String getSslTruststoreTypeProperty() {
         return sslTruststoreTypeProperty;
     }
 
+    public String getSslContextSupplierClassProperty() {
+        return sslContextSupplierClassProperty;
+    }
+
     public String getSslHostnameVerificationEnabledProperty() {
         return sslHostnameVerificationEnabledProperty;
     }
@@ -282,7 +285,28 @@ public int getSslHandshakeTimeoutMillis() {
         }
     }
 
+    @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 @@ private void putSSLProperties(X509Util x509Util) {
                 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.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.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 void testGetSslHandshakeDetectionTimeoutMillisProperty() {
         }
     }
 
+    @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 @@ private void setCustomCipherSuites() {
         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);
+            }
+        }
+
+    }
+
 }