ZOOKEEPER-4849: Option to Provide Custom X509 Implementation of QuorumAuthServer and QuorumAuthLearner

zookeeper-server/src/main/java/org/apache/zookeeper/server/quorum/QuorumPeer.java

@@ -28,6 +28,7 @@
 import java.io.StringReader;
 import java.io.StringWriter;
 import java.io.Writer;
+import java.lang.reflect.InvocationTargetException;
 import java.net.DatagramPacket;
 import java.net.DatagramSocket;
 import java.net.InetAddress;
@@ -693,6 +694,13 @@ void setId(long id) {
     }
 
     private boolean sslQuorum;
+
+    /** Class name to instantiate for SSL AuthServerProvider */
+    private String sslAuthServerProvider;
+
+    /** Class name to instantiate for SSL AuthLearnerProvider */
+    private String sslAuthLearnerProvider;
+
     private boolean shouldUsePortUnification;
 
     public boolean isSslQuorum() {
@@ -1176,12 +1184,96 @@ public void initialize() throws SaslException {
             }
             authServer = new SaslQuorumAuthServer(isQuorumServerSaslAuthRequired(), quorumServerLoginContext, authzHosts);
             authLearner = new SaslQuorumAuthLearner(isQuorumLearnerSaslAuthRequired(), quorumServicePrincipal, quorumLearnerLoginContext);
+        } else if (isSslQuorum()) {
+            try {
+                authServer = getSslQuorumAuthServer();
+                authLearner = getSslQuorumAuthLearner();
+            }  catch (Exception e) {
+                LOG.error(e.getMessage(), e);
+                throw new SaslException(e.getMessage());
+            }
         } else {
             authServer = new NullQuorumAuthServer();
             authLearner = new NullQuorumAuthLearner();
         }
     }
 
+    /**
+     * Instantiate the configured SSL QuorumAuthServer implementation.
+     *
+     * <p>Checks for a JVM system property named by {@value QuorumAuth#QUORUM_SSL_AUTHPROVIDER}
+     * first; if unset, falls back to the value read from config.
+     * If still undefined or empty, returns {@link NullQuorumAuthServer}.</p>
+     *
+     * @return a {@link QuorumAuthServer} instance
+     * @throws SaslException if the class named cannot be loaded, instantiated,
+     *         or doesn’t implement {@link QuorumAuthServer}
+     */
+    private QuorumAuthServer getSslQuorumAuthServer() throws SaslException {
+        String providerClass = System.getProperty(
+                "zookeeper." + QuorumAuth.QUORUM_SSL_AUTHPROVIDER,
+                sslAuthServerProvider
+        );
+
+        if (providerClass == null || providerClass.isEmpty()) {
+            LOG.debug("{} not defined; using NullQuorumAuthServer", QuorumAuth.QUORUM_SSL_AUTHPROVIDER);
+            return new NullQuorumAuthServer();
+        }
+
+        try {
+            Class<?> cls = Class.forName(providerClass);
+            return (QuorumAuthServer) cls.getDeclaredConstructor().newInstance();
+
+        } catch (ClassNotFoundException e) {
+            throw new SaslException("SSL auth server provider class not found: " + providerClass, e);
+        } catch (NoSuchMethodException
+                 | InstantiationException
+                 | IllegalAccessException
+                 | InvocationTargetException e) {
+            throw new SaslException("Failed to instantiate SSL auth server provider: " + providerClass, e);
+        } catch (ClassCastException e) {
+            throw new SaslException("Configured class does not implement QuorumAuthServer: " + providerClass, e);
+        }
+    }
+
+    /**
+     * Instantiate the configured SSL QuorumAuthLearner implementation.
+     *
+     * <p>Checks for a JVM system property named by {@value QuorumAuth#QUORUM_SSL_LEARNER_AUTHPROVIDER}
+     * first; if unset, falls back to the value read from zoo.cfg.
+     * If still undefined or empty, returns {@link NullQuorumAuthLearner}.</p>
+     *
+     * @return a {@link QuorumAuthLearner} instance
+     * @throws SaslException if the class named cannot be loaded, instantiated,
+     *         or doesn’t implement {@link QuorumAuthLearner}
+     */
+    private QuorumAuthLearner getSslQuorumAuthLearner() throws SaslException {
+        String providerClass = System.getProperty(
+                "zookeeper." + QuorumAuth.QUORUM_SSL_LEARNER_AUTHPROVIDER,
+                sslAuthLearnerProvider
+        );
+
+        if (providerClass == null || providerClass.isEmpty()) {
+            LOG.debug("{} not defined; using NullQuorumAuthLearner", QuorumAuth.QUORUM_SSL_LEARNER_AUTHPROVIDER);
+            return new NullQuorumAuthLearner();
+        }
+
+        try {
+            Class<?> cls = Class.forName(providerClass);
+            return (QuorumAuthLearner) cls.getDeclaredConstructor().newInstance();
+
+        } catch (ClassNotFoundException e) {
+            throw new SaslException("SSL auth learner provider class not found: " + providerClass, e);
+        } catch (NoSuchMethodException
+                 | InstantiationException
+                 | IllegalAccessException
+                 | InvocationTargetException e) {
+            throw new SaslException("Failed to instantiate SSL auth learner provider: " + providerClass, e);
+        } catch (ClassCastException e) {
+            throw new SaslException("Configured class does not implement QuorumAuthLearner: " + providerClass, e);
+        }
+    }
+
     QuorumStats quorumStats() {
         return quorumStats;
     }
@@ -2190,6 +2282,14 @@ public void setSslQuorum(boolean sslQuorum) {
         this.sslQuorum = sslQuorum;
     }
 
+    public void setSslAuthServerProvider(String sslAuthServerProvider) {
+        this.sslAuthServerProvider = sslAuthServerProvider;
+    }
+
+    public void setSslAuthLearnerProvider(String sslAuthLearnerProvider) {
+        this.sslAuthLearnerProvider = sslAuthLearnerProvider;
+    }
+
     public void setUsePortUnification(boolean shouldUsePortUnification) {
         LOG.info("Port unification {}", shouldUsePortUnification ? "enabled" : "disabled");
         this.shouldUsePortUnification = shouldUsePortUnification;
@@ -2740,6 +2840,14 @@ boolean isQuorumSaslAuthEnabled() {
         return quorumSaslEnableAuth;
     }
 
+    public QuorumAuthServer getQuorumAuthServer() {
+        return authServer;
+    }
+
+    public QuorumAuthLearner getQuorumAuthLearner() {
+        return authLearner;
+    }
+
     private boolean isQuorumServerSaslAuthRequired() {
         return quorumServerSaslAuthRequired;
     }

zookeeper-server/src/main/java/org/apache/zookeeper/server/quorum/QuorumPeerConfig.java

@@ -77,6 +77,9 @@ public class QuorumPeerConfig {
     protected boolean shouldUsePortUnification = false;
     protected int observerMasterPort;
     protected boolean sslQuorumReloadCertFiles = false;
+    private String sslAuthServerProvider;
+    private String sslAuthLearnerProvider;
+
     protected File dataDir;
     protected File dataLogDir;
     protected String dynamicConfigFileStr = null;
@@ -390,6 +393,10 @@ public void parseProperties(Properties zkProp) throws IOException, ConfigExcepti
                 multiAddressReachabilityCheckEnabled = parseBoolean(key, value);
             } else if (key.equals("oraclePath")) {
                 oraclePath = value;
+            } else if (key.equals(QuorumAuth.QUORUM_SSL_AUTHPROVIDER)) {
+                sslAuthServerProvider = value;
+            } else if (key.equals(QuorumAuth.QUORUM_SSL_LEARNER_AUTHPROVIDER)) {
+                sslAuthLearnerProvider = value;
             } else {
                 System.setProperty("zookeeper." + key, value);
             }
@@ -875,7 +882,12 @@ public boolean isLocalSessionsUpgradingEnabled() {
     public boolean isSslQuorum() {
         return sslQuorum;
     }
-
+    public String getSslAuthServerProvider() {
+        return sslAuthServerProvider;
+    }
+    public String getSslAuthLearnerProvider() {
+        return sslAuthLearnerProvider;
+    }
     public boolean shouldUsePortUnification() {
         return shouldUsePortUnification;
     }

zookeeper-server/src/main/java/org/apache/zookeeper/server/quorum/QuorumPeerMain.java

@@ -201,6 +201,8 @@ public void runFromConfig(QuorumPeerConfig config) throws IOException, AdminServ
             quorumPeer.setSecureCnxnFactory(secureCnxnFactory);
             quorumPeer.setSslQuorum(config.isSslQuorum());
             quorumPeer.setUsePortUnification(config.shouldUsePortUnification());
+            quorumPeer.setSslAuthServerProvider(config.getSslAuthServerProvider());
+            quorumPeer.setSslAuthLearnerProvider(config.getSslAuthLearnerProvider());
             quorumPeer.setLearnerType(config.getPeerType());
             quorumPeer.setSyncEnabled(config.getSyncEnabled());
             quorumPeer.setQuorumListenOnAllIPs(config.getQuorumListenOnAllIPs());

zookeeper-server/src/main/java/org/apache/zookeeper/server/quorum/auth/QuorumAuth.java

@@ -42,6 +42,11 @@ public class QuorumAuth {
     public static final String QUORUM_SERVER_SASL_LOGIN_CONTEXT = "quorum.auth.server.saslLoginContext";
     public static final String QUORUM_SERVER_SASL_LOGIN_CONTEXT_DFAULT_VALUE = "QuorumServer";
 
+    /** Property key for custom SSL QuorumAuthServer provider (Class name) */
+    public static final String QUORUM_SSL_AUTHPROVIDER = "ssl.quorum.authProvider";
+    /** Property key for custom SSL QuorumAuthLearner provider (Class name) */
+    public static final String QUORUM_SSL_LEARNER_AUTHPROVIDER = "ssl.quorum.learner.authProvider";
+
     static final String QUORUM_SERVER_PROTOCOL_NAME = "zookeeper-quorum";
     static final String QUORUM_SERVER_SASL_DIGEST = "zk-quorum-sasl-md5";
     static final String QUORUM_AUTH_MESSAGE_TAG = "qpconnect";

zookeeper-server/src/test/java/org/apache/zookeeper/server/quorum/QuorumPeerAuthProviderTest.java

@@ -0,0 +1,171 @@
+/*
+ * 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.server.quorum;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import java.io.File;
+import java.lang.reflect.Method;
+import java.util.Properties;
+import org.apache.zookeeper.common.QuorumX509Util;
+import org.apache.zookeeper.server.quorum.auth.MockSSLQuorumAuthLearner;
+import org.apache.zookeeper.server.quorum.auth.MockSslQuorumAuthServer;
+import org.apache.zookeeper.server.quorum.auth.NullQuorumAuthLearner;
+import org.apache.zookeeper.server.quorum.auth.NullQuorumAuthServer;
+import org.apache.zookeeper.server.quorum.auth.QuorumAuth;
+import org.apache.zookeeper.server.quorum.auth.QuorumAuthLearner;
+import org.apache.zookeeper.server.quorum.auth.QuorumAuthServer;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Unit tests for pluggable SSL quorum auth providers in {@link QuorumPeer}.
+ */
+public class QuorumPeerAuthProviderTest {
+
+    private  QuorumX509Util quorumX509Util;
+
+    private static final String SSL_AUTH_PROVIDER_PROPERTY =
+            "zookeeper." + QuorumAuth.QUORUM_SSL_AUTHPROVIDER;
+    private static final String SSL_LEARNER_AUTH_PROVIDER_PROPERTY =
+            "zookeeper." + QuorumAuth.QUORUM_SSL_LEARNER_AUTHPROVIDER;
+
+    @BeforeEach
+    public void setup() {
+        quorumX509Util = new QuorumX509Util();
+    }
+    /**
+     * When sslAuthServerProvider is set to a custom provider, ensure it instantiates correctly.
+     */
+    @Test
+    public void testCustomSslAuthServerProvider() throws Exception {
+        // Prepare config with custom server auth provider
+        QuorumPeerConfig config = new QuorumPeerConfig();
+        Properties zkProp = getDefaultZKProperties();
+        zkProp.setProperty("sslQuorum", "true");
+        zkProp.setProperty(QuorumAuth.QUORUM_SSL_AUTHPROVIDER,
+                MockSslQuorumAuthServer.class.getName());
+        config.parseProperties(zkProp);
+
+        // Set on peer and invoke private method
+        QuorumPeer peer = new QuorumPeer();
+        peer.setSslAuthServerProvider(config.getSslAuthServerProvider());
+        QuorumAuthServer authServer = invokeGetSslQuorumAuthServer(peer);
+
+        assertTrue(authServer instanceof MockSslQuorumAuthServer,
+                "Expected MockSSLQuorumAuthServer when provider is configured");
+    }
+
+    /**
+     * When sslAuthServerProvider is set via JVM system property
+     */
+    @Test
+    public void testSslAuthServerProviderSystemProperty() throws Exception {
+        System.setProperty(SSL_AUTH_PROVIDER_PROPERTY,
+                MockSslQuorumAuthServer.class.getName());
+        try {
+            QuorumPeer peer = new QuorumPeer();
+            peer.setSslAuthServerProvider("some.invalid.Class");
+            QuorumAuthServer authServer = invokeGetSslQuorumAuthServer(peer);
+            assertTrue(authServer instanceof MockSslQuorumAuthServer,
+                    "Expected system property for server auth provider");
+        } finally {
+            System.clearProperty(SSL_AUTH_PROVIDER_PROPERTY);
+        }
+    }
+    /**
+     * Without any provider configured, default should be NullQuorumAuthServer.
+     */
+    @Test
+    public void testDefaultSslAuthServerProvider() throws Exception {
+        QuorumPeer peer = new QuorumPeer();
+        QuorumAuthServer authServer = invokeGetSslQuorumAuthServer(peer);
+        assertTrue(authServer instanceof NullQuorumAuthServer,
+                "Expected NullQuorumAuthServer when no provider is configured");
+    }
+
+    /**
+     * When sslAuthLearnerProvider is set to a custom provider, ensure it instantiates correctly.
+     */
+    @Test
+    public void testCustomSslAuthLearnerProvider() throws Exception {
+        QuorumPeerConfig config = new QuorumPeerConfig();
+        Properties zkProp = getDefaultZKProperties();
+        zkProp.setProperty("sslQuorum", "true");
+        zkProp.setProperty(QuorumAuth.QUORUM_SSL_LEARNER_AUTHPROVIDER,
+                MockSSLQuorumAuthLearner.class.getName());
+        config.parseProperties(zkProp);
+
+        QuorumPeer peer = new QuorumPeer();
+        peer.setSslAuthLearnerProvider(config.getSslAuthLearnerProvider());
+        QuorumAuthLearner authLearner = invokeGetSslQuorumAuthLearner(peer);
+
+        assertTrue(authLearner instanceof MockSSLQuorumAuthLearner,
+                "Expected MockSSLQuorumAuthLearner when learner provider is configured");
+    }
+
+    /**
+     * When sslAuthLearnerProvider is set via JVM system property
+     */
+    @Test
+    public void testSslAuthLearnerProviderSystemProperty() throws Exception {
+        System.setProperty(SSL_LEARNER_AUTH_PROVIDER_PROPERTY,
+                MockSSLQuorumAuthLearner.class.getName());
+        try {
+            QuorumPeer peer = new QuorumPeer();
+            QuorumAuthLearner authLearner = invokeGetSslQuorumAuthLearner(peer);
+            Properties props = System.getProperties();
+            props.forEach((key, value) -> System.out.println(key + " = " + value));
+
+            assertTrue(authLearner instanceof MockSSLQuorumAuthLearner,
+                    "Expected system property for learner auth provider");
+        } finally {
+            System.clearProperty(SSL_LEARNER_AUTH_PROVIDER_PROPERTY);
+        }
+    }
+    /**
+     * Without any learner provider configured, default should be NullQuorumAuthLearner.
+     */
+    @Test
+    public void testDefaultSslAuthLearnerProvider() throws Exception {
+        QuorumPeer peer = new QuorumPeer();
+        QuorumAuthLearner authLearner = invokeGetSslQuorumAuthLearner(peer);
+        assertTrue(authLearner instanceof NullQuorumAuthLearner,
+                "Expected NullQuorumAuthLearner when no learner provider is configured");
+    }
+
+    // Reflection helpers to access private methods
+
+    private QuorumAuthServer invokeGetSslQuorumAuthServer(QuorumPeer peer) throws Exception {
+        Method m = QuorumPeer.class.getDeclaredMethod("getSslQuorumAuthServer");
+        m.setAccessible(true);
+        return (QuorumAuthServer) m.invoke(peer);
+    }
+
+    private QuorumAuthLearner invokeGetSslQuorumAuthLearner(QuorumPeer peer) throws Exception {
+        Method m = QuorumPeer.class.getDeclaredMethod("getSslQuorumAuthLearner");
+        m.setAccessible(true);
+        return (QuorumAuthLearner) m.invoke(peer);
+    }
+    private Properties getDefaultZKProperties() {
+        Properties zkProp = new Properties();
+        zkProp.setProperty("dataDir", new File("myDataDir").getAbsolutePath());
+        zkProp.setProperty("oraclePath", new File("mastership").getAbsolutePath());
+        return zkProp;
+    }
+}