ZOOKEEPER-3160: Custom User SSLContext

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

@@ -28,6 +28,7 @@
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
 import java.security.KeyStore;
 
 import org.slf4j.Logger;
@@ -85,60 +86,74 @@ public static SSLContext createSSLContext() throws SSLContextException {
     }
 
     public static SSLContext createSSLContext(ZKConfig config) throws SSLContextException {
-        KeyManager[] keyManagers = null;
-        TrustManager[] trustManagers = null;
+        if (config.getProperty(ZKConfig.SSL_CLIENT_CONTEXT) != null) {
+            LOG.debug("Loading SSLContext from property '" + ZKConfig.SSL_CLIENT_CONTEXT + "'");
+            String sslClientContextClass = config.getProperty(ZKConfig.SSL_CLIENT_CONTEXT);
+            try {
+                Class<?> sslContextClass = Class.forName(sslClientContextClass);
+                ZKClientSSLContext sslContext = (ZKClientSSLContext) sslContextClass.getConstructor().newInstance();
+                return sslContext.getSSLContext();
+            } catch (ClassNotFoundException | ClassCastException | NoSuchMethodException | InvocationTargetException |
+                    InstantiationException | IllegalAccessException e) {
+                throw new SSLContextException("Could not retrieve the SSLContext from source '" + sslClientContextClass +
+                        "' provided in the property '" + ZKConfig.SSL_CLIENT_CONTEXT + "'", e);
+            }
+        } else {
+            KeyManager[] keyManagers = null;
+            TrustManager[] trustManagers = null;
 
-        String keyStoreLocationProp = config.getProperty(ZKConfig.SSL_KEYSTORE_LOCATION);
-        String keyStorePasswordProp = config.getProperty(ZKConfig.SSL_KEYSTORE_PASSWD);
+            String keyStoreLocationProp = config.getProperty(ZKConfig.SSL_KEYSTORE_LOCATION);
+            String keyStorePasswordProp = config.getProperty(ZKConfig.SSL_KEYSTORE_PASSWD);
 
-        // There are legal states in some use cases for null KeyManager or TrustManager.
-        // But if a user wanna specify one, location and password are required.
+            // There are legal states in some use cases for null KeyManager or TrustManager.
+            // But if a user wanna specify one, location and password are required.
 
-        if (keyStoreLocationProp == null && keyStorePasswordProp == null) {
-            LOG.warn("keystore not specified for client connection");
-        } else {
-            if (keyStoreLocationProp == null) {
-                throw new SSLContextException("keystore location not specified for client connection");
-            }
-            if (keyStorePasswordProp == null) {
-                throw new SSLContextException("keystore password not specified for client connection");
-            }
-            try {
-                keyManagers = new KeyManager[]{
-                        createKeyManager(keyStoreLocationProp, keyStorePasswordProp)};
-            } catch (KeyManagerException e) {
-                throw new SSLContextException("Failed to create KeyManager", e);
+            if (keyStoreLocationProp == null && keyStorePasswordProp == null) {
+                LOG.warn("keystore not specified for client connection");
+            } else {
+                if (keyStoreLocationProp == null) {
+                    throw new SSLContextException("keystore location not specified for client connection");
+                }
+                if (keyStorePasswordProp == null) {
+                    throw new SSLContextException("keystore password not specified for client connection");
+                }
+                try {
+                    keyManagers = new KeyManager[]{
+                            createKeyManager(keyStoreLocationProp, keyStorePasswordProp)};
+                } catch (KeyManagerException e) {
+                    throw new SSLContextException("Failed to create KeyManager", e);
+                }
             }
-        }
 
-        String trustStoreLocationProp = config.getProperty(ZKConfig.SSL_TRUSTSTORE_LOCATION);
-        String trustStorePasswordProp = config.getProperty(ZKConfig.SSL_TRUSTSTORE_PASSWD);
+            String trustStoreLocationProp = config.getProperty(ZKConfig.SSL_TRUSTSTORE_LOCATION);
+            String trustStorePasswordProp = config.getProperty(ZKConfig.SSL_TRUSTSTORE_PASSWD);
 
-        if (trustStoreLocationProp == null && trustStorePasswordProp == null) {
-            LOG.warn("Truststore not specified for client connection");
-        } else {
-            if (trustStoreLocationProp == null) {
-                throw new SSLContextException("Truststore location not specified for client connection");
-            }
-            if (trustStorePasswordProp == null) {
-                throw new SSLContextException("Truststore password not specified for client connection");
+            if (trustStoreLocationProp == null && trustStorePasswordProp == null) {
+                LOG.warn("Truststore not specified for client connection");
+            } else {
+                if (trustStoreLocationProp == null) {
+                    throw new SSLContextException("Truststore location not specified for client connection");
+                }
+                if (trustStorePasswordProp == null) {
+                    throw new SSLContextException("Truststore password not specified for client connection");
+                }
+                try {
+                    trustManagers = new TrustManager[]{
+                            createTrustManager(trustStoreLocationProp, trustStorePasswordProp)};
+                } catch (TrustManagerException e) {
+                    throw new SSLContextException("Failed to create TrustManager", e);
+                }
             }
+
+            SSLContext sslContext = null;
             try {
-                trustManagers = new TrustManager[]{
-                        createTrustManager(trustStoreLocationProp, trustStorePasswordProp)};
-            } catch (TrustManagerException e) {
-                throw new SSLContextException("Failed to create TrustManager", e);
+                sslContext = SSLContext.getInstance("TLSv1");
+                sslContext.init(keyManagers, trustManagers, null);
+            } catch (Exception e) {
+                throw new SSLContextException(e);
             }
+            return sslContext;
         }
-
-        SSLContext sslContext = null;
-        try {
-            sslContext = SSLContext.getInstance("TLSv1");
-            sslContext.init(keyManagers, trustManagers, null);
-        } catch (Exception e) {
-            throw new SSLContextException(e);
-        }
-        return sslContext;
     }
 
     public static X509KeyManager createKeyManager(String keyStoreLocation, String keyStorePassword)

src/java/main/org/apache/zookeeper/common/ZKClientSSLContext.java

@@ -0,0 +1,18 @@
+package org.apache.zookeeper.common;
+
+import javax.net.ssl.SSLContext;
+
+/**
+ * An interface for providing a custom {@link SSLContext} object to {@link X509Util} using {@link ZKConfig#SSL_CLIENT_CONTEXT}
+ */
+public interface ZKClientSSLContext {
+
+    /**
+     * Returns an {@link SSLContext} for use within the {@link X509Util}
+     *
+     * @return {@link SSLContext} for use within {@link X509Util}
+     * @throws X509Exception.SSLContextException if {@link SSLContext} cannot be created
+     */
+    SSLContext getSSLContext() throws X509Exception.SSLContextException;
+
+}

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

@@ -50,6 +50,7 @@ public class ZKConfig {
     public static final String SSL_TRUSTSTORE_LOCATION = X509Util.SSL_TRUSTSTORE_LOCATION;
     @SuppressWarnings("deprecation")
     public static final String SSL_TRUSTSTORE_PASSWD = X509Util.SSL_TRUSTSTORE_PASSWD;
+    public static final String SSL_CLIENT_CONTEXT = "zookeeper.ssl.client.context";
     @SuppressWarnings("deprecation")
     public static final String SSL_AUTHPROVIDER = X509Util.SSL_AUTHPROVIDER;
     public static final String JUTE_MAXBUFFER = "jute.maxbuffer";

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

@@ -0,0 +1,50 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.zookeeper.common;
+
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import org.junit.Test;
+
+import javax.net.ssl.SSLContext;
+
+public class X509UtilTest {
+
+    @Test
+    public void testCreateSSLContext_invalidCustomSSLContextClass() {
+        ZKConfig zkConfig = new ZKConfig();
+        zkConfig.setProperty(ZKConfig.SSL_CLIENT_CONTEXT, String.class.getCanonicalName());
+        try {
+            X509Util.createSSLContext(zkConfig);
+            fail("SSLContextException expected.");
+        } catch (X509Exception.SSLContextException e) {
+            assertTrue(e.getMessage().contains(ZKConfig.SSL_CLIENT_CONTEXT));
+        }
+    }
+
+    @Test
+    public void testCreateSSLContext_validCustomSSLContextClass() throws X509Exception.SSLContextException {
+        ZKConfig zkConfig = new ZKConfig();
+        zkConfig.setProperty(ZKConfig.SSL_CLIENT_CONTEXT, ZKTestClientSSLContext.class.getCanonicalName());
+        final SSLContext sslContext = X509Util.createSSLContext(zkConfig);
+        assertNull(sslContext);
+    }
+
+}

src/java/test/org/apache/zookeeper/common/ZKTestClientSSLContext.java

@@ -0,0 +1,12 @@
+package org.apache.zookeeper.common;
+
+import javax.net.ssl.SSLContext;
+
+public class ZKTestClientSSLContext implements ZKClientSSLContext {
+
+    @Override
+    public SSLContext getSSLContext() {
+        return null;
+    }
+
+}