ZOOKEEPER-4947: Stop the send loop after SASL authentication fails

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

@@ -1217,6 +1217,7 @@ public void run() {
                                 eventThread.queueEvent(new WatchedEvent(Watcher.Event.EventType.None, authState, null));
                                 if (state == States.AUTH_FAILED) {
                                     eventThread.queueEventOfDeath();
+                                    break;
                                 }
                             }
                         }

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

@@ -1319,7 +1319,7 @@ public synchronized void register(Watcher watcher) {
      * @throws InterruptedException
      */
     public synchronized void close() throws InterruptedException {
-        if (!cnxn.getState().isAlive()) {
+        if (cnxn.getState() == States.CLOSED) {
             LOG.debug("Close called on already closed client");
             return;
         }

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

@@ -85,6 +85,11 @@ public void run() {
         }
     }
 
+    @Override
+    public boolean testableWaitForShutdown(int wait) throws InterruptedException {
+        return super.testableWaitForShutdown(wait);
+    }
+
     public SocketAddress testableLocalSocketAddress() {
         return super.testableLocalSocketAddress();
     }

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

@@ -18,12 +18,18 @@
 
 package org.apache.zookeeper.test;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import org.apache.zookeeper.CreateMode;
+import org.apache.zookeeper.KeeperException;
+import org.apache.zookeeper.TestableZooKeeper;
 import org.apache.zookeeper.WatchedEvent;
 import org.apache.zookeeper.Watcher.Event.KeeperState;
 import org.apache.zookeeper.ZooDefs.Ids;
@@ -77,7 +83,7 @@ public synchronized void process(WatchedEvent event) {
 
     @Test
     public void testAuthFail() {
-        try (ZooKeeper zk = createClient()) {
+        try (ZooKeeper zk = new ZooKeeper(hostPort, CONNECTION_TIMEOUT, new CountdownWatcher())) {
             zk.create("/path1", null, Ids.CREATOR_ALL_ACL, CreateMode.PERSISTENT);
             fail("Should have gotten exception.");
         } catch (Exception e) {
@@ -88,9 +94,18 @@ public void testAuthFail() {
 
     @Test
     public void testBadSaslAuthNotifiesWatch() throws Exception {
-        try (ZooKeeper ignored = createClient(new MyWatcher(), hostPort)) {
+        try (TestableZooKeeper zk = new TestableZooKeeper(hostPort, CONNECTION_TIMEOUT, new MyWatcher())) {
             // wait for authFailed event from client's EventThread.
-            authFailed.await();
+            assertTrue(authFailed.await(CONNECTION_TIMEOUT, TimeUnit.MILLISECONDS));
+            boolean threadsStopped = zk.testableWaitForShutdown(1000);
+            LOG.info("SASL failure shutdown without close: threadsStopped={}, state={}",
+                     threadsStopped, zk.getState());
+            assertTrue(threadsStopped, "Client threads should stop after SASL authentication fails");
+            assertEquals(ZooKeeper.States.AUTH_FAILED, zk.getState());
+            assertThrows(KeeperException.AuthFailedException.class, () -> zk.exists("/", false));
+            zk.close();
+            assertEquals(ZooKeeper.States.CLOSED, zk.getState());
+            assertTrue(zk.close(CONNECTION_TIMEOUT));
         }
     }
 

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

@@ -19,10 +19,15 @@
 package org.apache.zookeeper.test;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import javax.security.auth.login.Configuration;
 import org.apache.zookeeper.CreateMode;
 import org.apache.zookeeper.KeeperException;
+import org.apache.zookeeper.Watcher.Event.KeeperState;
 import org.apache.zookeeper.ZooDefs.Ids;
 import org.apache.zookeeper.ZooKeeper;
 import org.junit.jupiter.api.AfterAll;
@@ -55,12 +60,7 @@ public void testClientOpWithInvalidSASLUserAuthAfterSuccessLogin() throws Except
         }
 
         resetJaasConfiguration("jaas.conf", "super_wrong", "test");
-        try  (ZooKeeper wrongUserZk = createClient()) {
-            wrongUserZk.create("/bar", null, Ids.CREATOR_ALL_ACL, CreateMode.PERSISTENT);
-            fail("Client with wrong SASL config should not pass SASL authentication.");
-        } catch (KeeperException e) {
-            assertEquals(KeeperException.Code.AUTHFAILED, e.code());
-        }
+        assertClientAuthFailed();
     }
 
     @Test
@@ -73,11 +73,21 @@ public void testClientOpWithInvalidSASLPasswordAuthAfterSuccessLogin() throws Ex
         }
 
         resetJaasConfiguration("jaas.conf", "super", "test_wrongong");
-        try (ZooKeeper wrongPasswordZk = createClient()) {
-            wrongPasswordZk.create("/bar", null, Ids.CREATOR_ALL_ACL, CreateMode.PERSISTENT);
-            fail("Client with wrong SASL config should not pass SASL authentication.");
-        } catch (KeeperException e) {
-            assertEquals(KeeperException.Code.AUTHFAILED, e.code());
+        assertClientAuthFailed();
+    }
+
+    private void assertClientAuthFailed() throws Exception {
+        CountDownLatch authFailed = new CountDownLatch(1);
+        // A rejected connection may disappear before createClient's JMX check.
+        try (ZooKeeper zk = new ZooKeeper(hostPort, CONNECTION_TIMEOUT, event -> {
+            if (event.getState() == KeeperState.AuthFailed) {
+                authFailed.countDown();
+            }
+        })) {
+            assertTrue(authFailed.await(CONNECTION_TIMEOUT, TimeUnit.MILLISECONDS));
+            assertEquals(ZooKeeper.States.AUTH_FAILED, zk.getState());
+            assertThrows(KeeperException.AuthFailedException.class,
+                         () -> zk.create("/bar", null, Ids.CREATOR_ALL_ACL, CreateMode.PERSISTENT));
         }
     }