ZOOKEEPER-3239: Adding EnsembleAuthProvider to verify the ensemble name

Author: Jie Huang <jiehuang@fb.com>

Reviewers: fangmin@apache.org, andor@apache.org

Closes #768 from jhuan31/ZOOKEEPER-3239 and squashes the following commits:

7e17e1b8a [Jie Huang] add doc for EnsembleAuthProvider
878bfb543 [Jie Huang] ZOOKEEPER-3239: Adding EnsembleAuthProvider to verify the ensemble name

zookeeper-docs/src/main/resources/markdown/zookeeperAdmin.md

@@ -1044,6 +1044,15 @@ encryption/authentication/authorization performed by the service.
     Then set this property **zookeeper.ssl.authProvider=[scheme]** and that provider
     will be used for secure authentication.
 
+* *zookeeper.ensembleAuthName* :
+    (Java system property only: **zookeeper.ensembleAuthName**)
+    **New in 3.6.0:**
+    Specify a list of comma-separated valid names/aliases of an ensemble. A client
+    can provide the ensemble name it intends to connect as the credential for scheme "ensemble". The EnsembleAuthenticationProvider will check the credential against
+    the list of names/aliases of the ensemble that receives the connection request.
+    If the credential is not in the list, the connection request will be refused.
+    This prevents a client accidentally connecting to a wrong ensemble.
+
 <a name="Experimental+Options%2FFeatures"></a>
 
 #### Experimental Options/Features

zookeeper-server/src/main/java/org/apache/zookeeper/server/DumbWatcher.java

@@ -54,7 +54,7 @@ public class DumbWatcher extends ServerCnxn {
     int getSessionTimeout() { return 0; }
 
     @Override
-    void close() { }
+    public void close() { }
 
     @Override
     public void sendResponse(ReplyHeader h, Record r, String tag, String cacheKey, Stat stat) throws IOException { }

zookeeper-server/src/main/java/org/apache/zookeeper/server/ServerCnxn.java

@@ -103,7 +103,7 @@ public abstract class ServerCnxn implements Stats, Watcher {
         }
     }
 
-    abstract void close();
+    public abstract void close();
 
     public abstract void sendResponse(ReplyHeader h, Record r,
             String tag, String cacheKey, Stat stat) throws IOException;

zookeeper-server/src/main/java/org/apache/zookeeper/server/ServerMetrics.java

@@ -77,7 +77,22 @@ public enum ServerMetrics {
     BYTES_RECEIVED_COUNT(new SimpleCounter("bytes_received_count")),
 
     RESPONSE_PACKET_CACHE_HITS(new SimpleCounter("response_packet_cache_hits")),
-    RESPONSE_PACKET_CACHE_MISSING(new SimpleCounter("response_packet_cache_misses"));
+    RESPONSE_PACKET_CACHE_MISSING(new SimpleCounter("response_packet_cache_misses")),
+    
+    /*
+     * Number of successful matches of expected ensemble name in EnsembleAuthenticationProvider.
+     */
+    ENSEMBLE_AUTH_SUCCESS(new SimpleCounter("ensemble_auth_success")),
+
+    /*
+     * Number of unsuccessful matches of expected ensemble name in EnsembleAuthenticationProvider.
+     */
+    ENSEMBLE_AUTH_FAIL(new SimpleCounter("ensemble_auth_fail")),
+
+    /*
+     * Number of client auth requests with no ensemble set in EnsembleAuthenticationProvider.
+     */
+    ENSEMBLE_AUTH_SKIP(new SimpleCounter("ensemble_auth_skip"));
 
     private final Metric metric;
 

zookeeper-server/src/main/java/org/apache/zookeeper/server/ZooKeeperServer.java

@@ -1203,6 +1203,8 @@ public class ZooKeeperServer implements SessionExpirer, ServerStats.Provider {
             Code authReturn = KeeperException.Code.AUTHFAILED;
             if(ap != null) {
                 try {
+                    // handleAuthentication may close the connection, to allow the client to choose
+                    // a different server to connect to.
                     authReturn = ap.handleAuthentication(new ServerAuthenticationProvider.ServerObjs(this, cnxn), authPacket.getAuth());
                 } catch(RuntimeException e) {
                     LOG.warn("Caught runtime exception from AuthenticationProvider: " + scheme + " due to " + e);

zookeeper-server/src/main/java/org/apache/zookeeper/server/auth/EnsembleAuthenticationProvider.java

@@ -0,0 +1,137 @@
+/**
+ * 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.auth;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+import org.apache.zookeeper.KeeperException;
+import org.apache.zookeeper.ZooKeeper;
+import org.apache.zookeeper.data.Id;
+import org.apache.zookeeper.jmx.MBeanRegistry;
+import org.apache.zookeeper.server.ServerCnxn;
+import org.apache.zookeeper.server.ServerMetrics;
+import org.apache.zookeeper.server.ZooKeeperServer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.zookeeper.jmx.ZKMBeanInfo;
+
+import javax.management.JMException;
+
+/**
+ * This is not a true AuthenticationProvider in the strict sense. it does
+ * handle add auth requests, but rather than authenticate the client, it checks
+ * to make sure that the ensemble name the client intends to connect to
+ * matches the name that the server thinks it belongs to. if the name does not match,
+ * this provider will close the connection.
+ */
+
+public class EnsembleAuthenticationProvider implements AuthenticationProvider {
+    private static final Logger LOG = LoggerFactory
+            .getLogger(EnsembleAuthenticationProvider.class);
+
+    public static final String ENSEMBLE_PROPERTY = "zookeeper.ensembleAuthName";
+    private static final int MIN_LOGGING_INTERVAL_MS = 1000;
+    private Set<String> ensembleNames;
+
+    public EnsembleAuthenticationProvider() {
+        String namesCSV = System.getProperty(ENSEMBLE_PROPERTY);
+        if (namesCSV != null) {
+            LOG.info("Set expected ensemble names to {}", namesCSV);
+            setEnsembleNames(namesCSV);
+        }
+    }
+
+    public void setEnsembleNames(String namesCSV) {
+        ensembleNames = new HashSet<String>();
+        for (String name: namesCSV.split(",")) {
+            ensembleNames.add(name.trim());    
+        }
+    }
+
+    /* provider methods */
+    @Override
+    public String getScheme() {
+        return "ensemble";
+    }
+
+    /**
+     * if things go bad, we don't want to freak out with the logging, so track
+     * the last time we logged something here.
+     */
+    private long lastFailureLogged;
+
+    @Override
+    public KeeperException.Code
+    handleAuthentication(ServerCnxn cnxn, byte[] authData)
+    {
+        if (authData == null || authData.length == 0) {
+            ServerMetrics.ENSEMBLE_AUTH_SKIP.add(1);
+            return KeeperException.Code.OK;
+        }
+
+        String receivedEnsembleName = new String(authData);
+
+        if (ensembleNames == null) {
+            ServerMetrics.ENSEMBLE_AUTH_SKIP.add(1);
+            return KeeperException.Code.OK;
+        }
+
+        if (ensembleNames.contains(receivedEnsembleName)) {
+            ServerMetrics.ENSEMBLE_AUTH_SUCCESS.add(1);
+            return KeeperException.Code.OK;
+        }
+
+        long currentTime = System.currentTimeMillis();
+        if (lastFailureLogged + MIN_LOGGING_INTERVAL_MS < currentTime) {
+            String id = cnxn.getRemoteSocketAddress().getAddress().getHostAddress();
+            LOG.warn("Unexpected ensemble name: ensemble name: {} client ip: {}", receivedEnsembleName, id);
+            lastFailureLogged = currentTime;
+        }
+        /*
+         * we are doing a close here rather than returning some other error
+         * since we want the client to choose another server to connect to. if
+         * we return an error, the client will get a fatal auth error and
+         * shutdown.
+         */
+        ServerMetrics.ENSEMBLE_AUTH_FAIL.add(1);
+        cnxn.close();
+        return KeeperException.Code.BADARGUMENTS;
+    }
+
+    /*
+     * since we aren't a true provider we return false for everything so that
+     * it isn't used in ACLs.
+     */
+    @Override
+    public boolean matches(String id, String aclExpr) {
+        return false;
+    }
+
+    @Override
+    public boolean isAuthenticated() {
+        return false;
+    }
+
+    @Override
+    public boolean isValid(String id) {
+        return false;
+    }
+}

zookeeper-server/src/main/java/org/apache/zookeeper/server/auth/ProviderRegistry.java

@@ -44,8 +44,6 @@ public class ProviderRegistry {
 
     public static void initialize() {
         synchronized (ProviderRegistry.class) {
-            if (initialized)
-                return;
             IPAuthenticationProvider ipp = new IPAuthenticationProvider();
             DigestAuthenticationProvider digp = new DigestAuthenticationProvider();
             authenticationProviders.put(ipp.getScheme(), ipp);
@@ -80,6 +78,10 @@ public class ProviderRegistry {
         return authenticationProviders.get(scheme);
     }
 
+    public static void removeProvider(String scheme) {
+        authenticationProviders.remove(scheme);
+    }
+
     public static String listProviders() {
         StringBuilder sb = new StringBuilder();
         for(String s: authenticationProviders.keySet()) {

zookeeper-server/src/test/java/org/apache/zookeeper/server/MockServerCnxn.java

@@ -40,7 +40,7 @@ public class MockServerCnxn extends ServerCnxn {
     }
 
     @Override
-    void close() {
+    public void close() {
     }
 
     @Override

zookeeper-server/src/test/java/org/apache/zookeeper/test/EnsembleAuthTest.java

@@ -0,0 +1,121 @@
+/**
+ * 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
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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.test;
+
+import org.apache.zookeeper.KeeperException;
+import org.apache.zookeeper.ZooKeeper;
+import org.apache.zookeeper.server.auth.EnsembleAuthenticationProvider;
+import org.apache.zookeeper.server.auth.ProviderRegistry;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.IOException;
+
+public class EnsembleAuthTest extends ClientBase {
+
+    @Before
+    public void setUp() throws Exception {
+        System.setProperty("zookeeper.authProvider.1", "org.apache.zookeeper.server.auth.EnsembleAuthenticationProvider");
+        super.setUp();
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        super.tearDown();
+        System.clearProperty("zookeeper.authProvider.1");
+        System.clearProperty(EnsembleAuthenticationProvider.ENSEMBLE_PROPERTY);
+        ProviderRegistry.removeProvider("ensemble");
+    }
+
+
+    @Test
+    public void noAuth() throws Exception {
+        resetEnsembleAuth(null, false);
+        connectToEnsemble(null);
+    }
+
+    @Test
+    public void emptyAuth() throws Exception {
+        resetEnsembleAuth(null, true);
+        connectToEnsemble("foo");
+    }
+
+    @Test
+    public void skipAuth() throws Exception {
+        resetEnsembleAuth("woo", true);
+        connectToEnsemble(null);
+    }
+
+    @Test
+    public void passAuth() throws Exception {
+        resetEnsembleAuth("woo", true);
+        connectToEnsemble("woo");
+    }
+
+    @Test
+    public void passAuthCSV() throws Exception {
+        resetEnsembleAuth(" foo,bar, baz ", true);
+
+        connectToEnsemble("foo");
+        connectToEnsemble("bar");
+        connectToEnsemble("baz");
+    }
+
+    @Test(expected = KeeperException.ConnectionLossException.class)
+    public void failAuth() throws Exception {
+        resetEnsembleAuth("woo", true);
+        connectToEnsemble("goo");
+    }
+
+    @Test(expected = KeeperException.AuthFailedException.class)
+    public void removeEnsembleAuthProvider() throws Exception {
+        resetEnsembleAuth(null, false);
+        connectToEnsemble("goo");
+    }
+
+
+    private void connectToEnsemble(final String auth) throws IOException, InterruptedException, KeeperException {
+        try (ZooKeeper zk = createClient()) {
+            // pass auth check
+            if (auth != null) {
+                zk.addAuthInfo("ensemble", auth.getBytes());
+            }
+            zk.getData("/", false, null);
+        }
+    }
+
+    private void resetEnsembleAuth(final String auth, final boolean useAuth) throws Exception {
+        stopServer();
+        if (auth == null) {
+            System.clearProperty(EnsembleAuthenticationProvider.ENSEMBLE_PROPERTY);
+        } else {
+            System.setProperty(EnsembleAuthenticationProvider.ENSEMBLE_PROPERTY, auth);
+        }
+        if (useAuth) {
+            System.setProperty("zookeeper.authProvider.1",
+                    "org.apache.zookeeper.server.auth.EnsembleAuthenticationProvider");
+        } else {
+            System.clearProperty("zookeeper.authProvider.1");
+        }
+        ProviderRegistry.removeProvider("ensemble");
+        ProviderRegistry.initialize();
+        startServer();
+    }
+}