ZOOKEEPER-4453: NettyServerCnxnFactory: allow to configure the early TLS connection drop feature

- add new flag netty.server.earlyDropSecureConnectionHandshakes to turn on/off ZOOKEEPER-3682
- disable ZOOKEEPER-3682 by default
- add docs
- add tests for this patch and for ZOOKEEPER-3682

see https://issues.apache.org/jira/browse/ZOOKEEPER-4453 for more context

Author: Enrico Olivelli <eolivelli@apache.org>

Reviewers: Mate Szalay-Beko <symat@apache.org>

Closes #1800 from eolivelli/fix/ZOOKEEPER-4453-b

4 files changed, 83 insertions, 35 deletions

diff --git a/zookeeper-docs/src/main/resources/markdown/zookeeperAdmin.md b/zookeeper-docs/src/main/resources/markdown/zookeeperAdmin.md
index eaefdda82..ce5dcaa01 100644
--- a/zookeeper-docs/src/main/resources/markdown/zookeeperAdmin.md
+++ b/zookeeper-docs/src/main/resources/markdown/zookeeperAdmin.md
@@ -1161,6 +1161,15 @@ property, when available, is noted below.
     effect due to TLS handshake timeout when there are too many in-flight TLS
     handshakes. Set it to something like 250 is good enough to avoid herd effect.
 
+* *netty.server.earlyDropSecureConnectionHandshakes*
+  (Java system property: **zookeeper.netty.server.earlyDropSecureConnectionHandshakes**)
+    If the ZooKeeper server is not fully started, drop TCP connections before performing the TLS handshake.
+    This is useful in order to prevent flooding the server with many concurrent TLS handshakes after a restart.
+    Please note that if you enable this flag the server won't answer to 'ruok' commands if it is not fully started.
+
+    The behaviour of dropping the connection has been introduced in ZooKeeper 3.7 and it was not possible to disable it.
+    Since 3.7.1 and 3.8.0 this feature is disabled by default.
+
 * *throttledOpWaitTime*
     (Java system property: **zookeeper.throttled_op_wait_time**)
     The time in the RequestThrottler queue longer than which a request will be marked as throttled.
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/server/NettyServerCnxn.java b/zookeeper-server/src/main/java/org/apache/zookeeper/server/NettyServerCnxn.java
index 269fcd942..8937039bc 100644
--- a/zookeeper-server/src/main/java/org/apache/zookeeper/server/NettyServerCnxn.java
+++ b/zookeeper-server/src/main/java/org/apache/zookeeper/server/NettyServerCnxn.java
@@ -522,7 +522,10 @@ public class NettyServerCnxn extends ServerCnxn {
                         }
                         ZooKeeperServer zks = this.zkServer;
                         if (zks == null || !zks.isRunning()) {
-                            throw new IOException("ZK down");
+                            LOG.info("Closing connection to {} because the server is not ready",
+                                    getRemoteSocketAddress());
+                            close(DisconnectReason.IO_EXCEPTION);
+                            return;
                         }
                         // checkRequestSize will throw IOException if request is rejected
                         zks.checkRequestSizeWhenReceivingMessage(len);
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/server/NettyServerCnxnFactory.java b/zookeeper-server/src/main/java/org/apache/zookeeper/server/NettyServerCnxnFactory.java
index 6b2713269..040650e4a 100644
--- a/zookeeper-server/src/main/java/org/apache/zookeeper/server/NettyServerCnxnFactory.java
+++ b/zookeeper-server/src/main/java/org/apache/zookeeper/server/NettyServerCnxnFactory.java
@@ -85,6 +85,7 @@ public class NettyServerCnxnFactory extends ServerCnxnFactory {
      * Allow client-server sockets to accept both SSL and plaintext connections
      */
     public static final String PORT_UNIFICATION_KEY = "zookeeper.client.portUnification";
+    public static final String EARLY_DROP_SECURE_CONNECTION_HANDSHAKES = "zookeeper.netty.server.earlyDropSecureConnectionHandshakes";
     private final boolean shouldUsePortUnification;
 
     /**
@@ -227,11 +228,15 @@ public class NettyServerCnxnFactory extends ServerCnxnFactory {
 
             // Check the zkServer assigned to the cnxn is still running,
             // close it before starting the heavy TLS handshake
-            if (!cnxn.isZKServerRunning()) {
-                LOG.warn("Zookeeper server is not running, close the connection before starting the TLS handshake");
-                ServerMetrics.getMetrics().CNXN_CLOSED_WITHOUT_ZK_SERVER_RUNNING.add(1);
-                channel.close();
-                return;
+            if (secure && !cnxn.isZKServerRunning()) {
+                boolean earlyDropSecureConnectionHandshakes = Boolean.getBoolean(EARLY_DROP_SECURE_CONNECTION_HANDSHAKES);
+                if (earlyDropSecureConnectionHandshakes) {
+                    LOG.info("Zookeeper server is not running, close the connection to {} before starting the TLS handshake",
+                            cnxn.getChannel().remoteAddress());
+                    ServerMetrics.getMetrics().CNXN_CLOSED_WITHOUT_ZK_SERVER_RUNNING.add(1);
+                    channel.close();
+                    return;
+                }
             }
 
             if (handshakeThrottlingEnabled) {
@@ -510,6 +515,7 @@ public class NettyServerCnxnFactory extends ServerCnxnFactory {
         x509Util = new ClientX509Util();
 
         boolean usePortUnification = Boolean.getBoolean(PORT_UNIFICATION_KEY);
+
         LOG.info("{}={}", PORT_UNIFICATION_KEY, usePortUnification);
         if (usePortUnification) {
             try {
diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/server/NettyServerCnxnTest.java b/zookeeper-server/src/test/java/org/apache/zookeeper/server/NettyServerCnxnTest.java
index a02d5d755..e27361c37 100644
--- a/zookeeper-server/src/test/java/org/apache/zookeeper/server/NettyServerCnxnTest.java
+++ b/zookeeper-server/src/test/java/org/apache/zookeeper/server/NettyServerCnxnTest.java
@@ -29,6 +29,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.when;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelFuture;
@@ -180,42 +181,71 @@ public class NettyServerCnxnTest extends ClientBase {
         when(zks.isRunning()).thenReturn(true);
         ServerStats.Provider providerMock = mock(ServerStats.Provider.class);
         when(zks.serverStats()).thenReturn(new ServerStats(providerMock));
-        testNonMTLSRemoteConn(zks);
+        testNonMTLSRemoteConn(zks, false, false);
     }
 
     @Test
     public void testNonMTLSRemoteConnZookKeeperServerNotReady() throws Exception {
-        testNonMTLSRemoteConn(null);
+        testNonMTLSRemoteConn(null, false, false);
+    }
+
+    @Test
+    public void testNonMTLSRemoteConnZookKeeperServerNotReadyEarlyDropEnabled() throws Exception {
+        testNonMTLSRemoteConn(null, false, true);
+    }
+
+    @Test
+    public void testMTLSRemoteConnZookKeeperServerNotReadyEarlyDropEnabled() throws Exception {
+        testNonMTLSRemoteConn(null, true, true);
+    }
+
+    @Test
+    public void testMTLSRemoteConnZookKeeperServerNotReadyEarlyDropDisabled() throws Exception {
+        testNonMTLSRemoteConn(null, true, true);
     }
 
     @SuppressWarnings("unchecked")
-    private void testNonMTLSRemoteConn(ZooKeeperServer zks) throws Exception {
-        Channel channel = mock(Channel.class);
-        ChannelId id = mock(ChannelId.class);
-        ChannelFuture success = mock(ChannelFuture.class);
-        ChannelHandlerContext context = mock(ChannelHandlerContext.class);
-        ChannelPipeline channelPipeline = mock(ChannelPipeline.class);
-
-        when(context.channel()).thenReturn(channel);
-        when(channel.pipeline()).thenReturn(channelPipeline);
-        when(success.channel()).thenReturn(channel);
-        when(channel.closeFuture()).thenReturn(success);
-
-        InetSocketAddress address = new InetSocketAddress(0);
-        when(channel.remoteAddress()).thenReturn(address);
-        when(channel.id()).thenReturn(id);
-        NettyServerCnxnFactory factory = new NettyServerCnxnFactory();
-        factory.setZooKeeperServer(zks);
-        Attribute atr = mock(Attribute.class);
-        Mockito.doReturn(atr).when(channel).attr(
-                Mockito.any()
-        );
-        doNothing().when(atr).set(Mockito.any());
-        factory.channelHandler.channelActive(context);
-
-        if (zks != null) {
-            assertEquals(0, zks.serverStats().getNonMTLSLocalConnCount());
-            assertEquals(1, zks.serverStats().getNonMTLSRemoteConnCount());
+    private void testNonMTLSRemoteConn(ZooKeeperServer zks, boolean secure, boolean earlyDrop) throws Exception {
+        try {
+            System.setProperty(NettyServerCnxnFactory.EARLY_DROP_SECURE_CONNECTION_HANDSHAKES, earlyDrop + "");
+
+            Channel channel = mock(Channel.class);
+            ChannelId id = mock(ChannelId.class);
+            ChannelFuture success = mock(ChannelFuture.class);
+            ChannelHandlerContext context = mock(ChannelHandlerContext.class);
+            ChannelPipeline channelPipeline = mock(ChannelPipeline.class);
+
+            when(context.channel()).thenReturn(channel);
+            when(channel.pipeline()).thenReturn(channelPipeline);
+            when(success.channel()).thenReturn(channel);
+            when(channel.closeFuture()).thenReturn(success);
+
+            InetSocketAddress address = new InetSocketAddress(0);
+            when(channel.remoteAddress()).thenReturn(address);
+            when(channel.id()).thenReturn(id);
+            NettyServerCnxnFactory factory = new NettyServerCnxnFactory();
+            factory.setSecure(secure);
+            factory.setZooKeeperServer(zks);
+            Attribute atr = mock(Attribute.class);
+            Mockito.doReturn(atr).when(channel).attr(
+                    Mockito.any()
+            );
+            doNothing().when(atr).set(Mockito.any());
+            factory.channelHandler.channelActive(context);
+
+            if (zks != null)  {
+                assertEquals(0, zks.serverStats().getNonMTLSLocalConnCount());
+                assertEquals(1, zks.serverStats().getNonMTLSRemoteConnCount());
+            } else {
+                if (earlyDrop && secure) {
+                    // the channel must have been forcibly closed
+                    Mockito.verify(channel, times(1)).close();
+                } else {
+                    Mockito.verify(channel, times(0)).close();
+                }
+            }
+        } finally {
+            System.clearProperty(NettyServerCnxnFactory.EARLY_DROP_SECURE_CONNECTION_HANDSHAKES);
         }
     }