Improve SSL initializer tests

Got rid of the ugly use of locks and consolidate synchronization between I/O thread and test thread using count down latches. I believe this makes the code much cleaner and easy to reason about.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=208739770

javatests/google/registry/proxy/handler/SslClientInitializerTest.java

@@ -28,8 +28,6 @@ import google.registry.proxy.Protocol;
 import google.registry.proxy.Protocol.BackendProtocol;
 import google.registry.proxy.handler.SslInitializerTestUtils.DumpHandler;
 import google.registry.proxy.handler.SslInitializerTestUtils.EchoHandler;
-import io.netty.buffer.ByteBuf;
-import io.netty.buffer.Unpooled;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelInitializer;
@@ -38,6 +36,7 @@ import io.netty.channel.EventLoopGroup;
 import io.netty.channel.embedded.EmbeddedChannel;
 import io.netty.channel.local.LocalAddress;
 import io.netty.channel.local.LocalChannel;
+import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.handler.ssl.OpenSsl;
 import io.netty.handler.ssl.SniHandler;
 import io.netty.handler.ssl.SslContext;
@@ -50,9 +49,8 @@ import java.security.KeyPair;
 import java.security.PrivateKey;
 import java.security.cert.CertificateException;
 import java.security.cert.X509Certificate;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 import javax.net.ssl.SSLException;
+import org.junit.After;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
@@ -83,6 +81,7 @@ public class SslClientInitializerTest {
   @Parameter(0)
   public SslProvider sslProvider;
 
+  // We do our best effort to test all available SSL providers.
   @Parameters(name = "{0}")
   public static SslProvider[] data() {
     return OpenSsl.isAvailable()
@@ -90,8 +89,20 @@ public class SslClientInitializerTest {
         : new SslProvider[] {SslProvider.JDK};
   }
 
+  /** Saves the SNI hostname received by the server, if sent by the client. */
   private String sniHostReceived;
 
+  // All I/O operations are done inside the single thread within this event loop group, which is
+  // different from the main test thread. Therefore synchronizations are required to make sure that
+  // certain I/O activities are finished when assertions are performed.
+  private final EventLoopGroup eventLoopGroup = new NioEventLoopGroup(1);
+
+  // Handler attached to server's channel to record the request received.
+  private final EchoHandler echoHandler = new EchoHandler();
+
+  // Handler attached to client's channel to record the response received.
+  private final DumpHandler dumpHandler = new DumpHandler();
+
   /** Fake protocol saved in channel attribute. */
   private static final BackendProtocol PROTOCOL =
       Protocol.backendBuilder()
@@ -101,16 +112,17 @@ public class SslClientInitializerTest {
           .handlerProviders(ImmutableList.of())
           .build();
 
+  @After
+  public void shutDown() {
+    Future<?> unusedFuture = eventLoopGroup.shutdownGracefully();
+  }
+
   private ChannelInitializer<LocalChannel> getServerInitializer(
-      PrivateKey privateKey,
+      PrivateKey privateKey, X509Certificate certificate) throws Exception {
-      X509Certificate certificate,
-      Lock serverLock,
-      Exception serverException)
-      throws Exception {
     SslContext sslContext = SslContextBuilder.forServer(privateKey, certificate).build();
     return new ChannelInitializer<LocalChannel>() {
       @Override
-      protected void initChannel(LocalChannel ch) throws Exception {
+      protected void initChannel(LocalChannel ch) {
         ch.pipeline()
             .addLast(
                 new SniHandler(
@@ -118,21 +130,17 @@ public class SslClientInitializerTest {
                       sniHostReceived = hostname;
                       return sslContext;
                     }),
-                new EchoHandler(serverLock, serverException));
+                echoHandler);
       }
     };
   }
 
   private ChannelInitializer<LocalChannel> getClientInitializer(
-      SslClientInitializer<LocalChannel> sslClientInitializer,
+      SslClientInitializer<LocalChannel> sslClientInitializer) {
-      Lock clientLock,
-      ByteBuf buffer,
-      Exception clientException) {
     return new ChannelInitializer<LocalChannel>() {
       @Override
-      protected void initChannel(LocalChannel ch) throws Exception {
+      protected void initChannel(LocalChannel ch) {
-        ch.pipeline()
+        ch.pipeline().addLast(sslClientInitializer, dumpHandler);
-            .addLast(sslClientInitializer, new DumpHandler(clientLock, buffer, clientException));
       }
     };
   }
@@ -169,45 +177,27 @@ public class SslClientInitializerTest {
     SelfSignedCertificate ssc = new SelfSignedCertificate(SSL_HOST);
     LocalAddress localAddress =
         new LocalAddress("DEFAULT_TRUST_MANAGER_REJECT_SELF_SIGNED_CERT_" + sslProvider);
-    Lock clientLock = new ReentrantLock();
+    setUpServer(eventLoopGroup, getServerInitializer(ssc.key(), ssc.cert()), localAddress);
-    Lock serverLock = new ReentrantLock();
-    ByteBuf buffer = Unpooled.buffer();
-    Exception clientException = new Exception();
-    Exception serverException = new Exception();
-    EventLoopGroup eventLoopGroup =
-        setUpServer(
-            getServerInitializer(ssc.key(), ssc.cert(), serverLock, serverException), localAddress);
     SslClientInitializer<LocalChannel> sslClientInitializer =
         new SslClientInitializer<>(sslProvider);
     Channel channel =
         setUpClient(
-            eventLoopGroup,
+            eventLoopGroup, getClientInitializer(sslClientInitializer), localAddress, PROTOCOL);
-            getClientInitializer(sslClientInitializer, clientLock, buffer, clientException),
-            localAddress,
-            PROTOCOL);
     // Wait for handshake exception to throw.
-    clientLock.lock();
+    echoHandler.waitTillReady();
-    serverLock.lock();
+    dumpHandler.waitTillReady();
     // The connection is now terminated, both the client side and the server side should get
-    // exceptions (caught in the caughtException method in EchoHandler and DumpHandler,
+    // exceptions.
-    // respectively).
+    assertThat(Throwables.getRootCause(dumpHandler.getCause()))
-    assertThat(Throwables.getRootCause(clientException))
         .isInstanceOf(SunCertPathBuilderException.class);
-    assertThat(Throwables.getRootCause(serverException)).isInstanceOf(SSLException.class);
+    assertThat(Throwables.getRootCause(echoHandler.getCause())).isInstanceOf(SSLException.class);
     assertThat(channel.isActive()).isFalse();
-
-    Future<?> unusedFuture = eventLoopGroup.shutdownGracefully().syncUninterruptibly();
   }
 
   @Test
   public void testSuccess_customTrustManager_acceptCertSignedByTrustedCa() throws Exception {
     LocalAddress localAddress =
         new LocalAddress("CUSTOM_TRUST_MANAGER_ACCEPT_CERT_SIGNED_BY_TRUSTED_CA_" + sslProvider);
-    Lock clientLock = new ReentrantLock();
-    Lock serverLock = new ReentrantLock();
-    ByteBuf buffer = Unpooled.buffer();
-    Exception clientException = new Exception();
-    Exception serverException = new Exception();
 
     // Generate a new key pair.
     KeyPair keyPair = getKeyPair();
@@ -218,37 +208,25 @@ public class SslClientInitializerTest {
 
     // Set up the server to use the signed cert and private key to perform handshake;
     PrivateKey privateKey = keyPair.getPrivate();
-    EventLoopGroup eventLoopGroup =
+    setUpServer(eventLoopGroup, getServerInitializer(privateKey, cert), localAddress);
-        setUpServer(
-            getServerInitializer(privateKey, cert, serverLock, serverException), localAddress);
 
     // Set up the client to trust the self signed cert used to sign the cert that server provides.
     SslClientInitializer<LocalChannel> sslClientInitializer =
         new SslClientInitializer<>(sslProvider, new X509Certificate[] {ssc.cert()});
     Channel channel =
         setUpClient(
-            eventLoopGroup,
+            eventLoopGroup, getClientInitializer(sslClientInitializer), localAddress, PROTOCOL);
-            getClientInitializer(sslClientInitializer, clientLock, buffer, clientException),
-            localAddress,
-            PROTOCOL);
 
-    verifySslChannel(channel, ImmutableList.of(cert), clientLock, serverLock, buffer);
+    verifySslChannel(channel, ImmutableList.of(cert), echoHandler, dumpHandler);
 
     // Verify that the SNI extension is sent during handshake.
     assertThat(sniHostReceived).isEqualTo(SSL_HOST);
-
-    Future<?> unusedFuture = eventLoopGroup.shutdownGracefully().syncUninterruptibly();
   }
 
   @Test
   public void testFailure_customTrustManager_wrongHostnameInCertificate() throws Exception {
     LocalAddress localAddress =
         new LocalAddress("CUSTOM_TRUST_MANAGER_WRONG_HOSTNAME_" + sslProvider);
-    Lock clientLock = new ReentrantLock();
-    Lock serverLock = new ReentrantLock();
-    ByteBuf buffer = Unpooled.buffer();
-    Exception clientException = new Exception();
-    Exception serverException = new Exception();
 
     // Generate a new key pair.
     KeyPair keyPair = getKeyPair();
@@ -259,31 +237,24 @@ public class SslClientInitializerTest {
 
     // Set up the server to use the signed cert and private key to perform handshake;
     PrivateKey privateKey = keyPair.getPrivate();
-    EventLoopGroup eventLoopGroup =
+    setUpServer(eventLoopGroup, getServerInitializer(privateKey, cert), localAddress);
-        setUpServer(
-            getServerInitializer(privateKey, cert, serverLock, serverException), localAddress);
 
     // Set up the client to trust the self signed cert used to sign the cert that server provides.
     SslClientInitializer<LocalChannel> sslClientInitializer =
         new SslClientInitializer<>(sslProvider, new X509Certificate[] {ssc.cert()});
     Channel channel =
         setUpClient(
-            eventLoopGroup,
+            eventLoopGroup, getClientInitializer(sslClientInitializer), localAddress, PROTOCOL);
-            getClientInitializer(sslClientInitializer, clientLock, buffer, clientException),
-            localAddress,
-            PROTOCOL);
 
-    serverLock.lock();
+    echoHandler.waitTillReady();
-    clientLock.lock();
+    dumpHandler.waitTillReady();
 
-    // When the client rejects the server cert due to wrong hostname, the client error is wrapped
+    // When the client rejects the server cert due to wrong hostname, both the client and server
-    // several layers in the exception. The server also throws an exception.
+    // should throw exceptions.
-    Throwable rootCause = Throwables.getRootCause(clientException);
+    Throwable rootCause = Throwables.getRootCause(dumpHandler.getCause());
     assertThat(rootCause).isInstanceOf(CertificateException.class);
     assertThat(rootCause).hasMessageThat().contains(SSL_HOST);
-    assertThat(Throwables.getRootCause(serverException)).isInstanceOf(SSLException.class);
+    assertThat(Throwables.getRootCause(echoHandler.getCause())).isInstanceOf(SSLException.class);
     assertThat(channel.isActive()).isFalse();
-
-    Future<?> unusedFuture = eventLoopGroup.shutdownGracefully().syncUninterruptibly();
   }
 }

javatests/google/registry/proxy/handler/SslInitializerTestUtils.java

@@ -16,6 +16,7 @@ package google.registry.proxy.handler;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.proxy.Protocol.PROTOCOL_KEY;
+import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.collect.ImmutableList;
 import google.registry.proxy.Protocol.BackendProtocol;
@@ -32,9 +33,9 @@ import io.netty.channel.EventLoopGroup;
 import io.netty.channel.local.LocalAddress;
 import io.netty.channel.local.LocalChannel;
 import io.netty.channel.local.LocalServerChannel;
-import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.handler.ssl.SslHandler;
 import io.netty.handler.ssl.util.SelfSignedCertificate;
+import io.netty.util.ReferenceCountUtil;
 import java.math.BigInteger;
 import java.nio.charset.StandardCharsets;
 import java.security.KeyPair;
@@ -45,7 +46,7 @@ import java.security.cert.X509Certificate;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Date;
-import java.util.concurrent.locks.Lock;
+import java.util.concurrent.CountDownLatch;
 import javax.net.ssl.SSLSession;
 import javax.security.auth.x500.X500Principal;
 import org.bouncycastle.jce.provider.BouncyCastleProvider;
@@ -61,43 +62,25 @@ public class SslInitializerTestUtils {
     Security.addProvider(new BouncyCastleProvider());
   }
 
-  /**
+  /** Sets up a server channel bound to the given local address. */
-   * Sets up a server channel bound to the given local address.
+  static void setUpServer(
-   *
+      EventLoopGroup eventLoopGroup,
-   * @return the event loop group used to process incoming connections.
+      ChannelInitializer<LocalChannel> serverInitializer,
-   */
+      LocalAddress localAddress) {
-  static EventLoopGroup setUpServer(
-      ChannelInitializer<LocalChannel> serverInitializer, LocalAddress localAddress)
-      throws Exception {
-    // Only use one thread in the event loop group. The same event loop group will be used to
-    // register client channels during setUpClient as well. This ensures that all I/O activities
-    // in both channels happen in the same thread, making debugging easier (i. e. no need to jump
-    // between threads when debugging, everything happens synchronously within the only I/O thread
-    // effectively). Note that the main thread is still separate from the I/O thread and
-    // synchronization (using the lock field) is still needed when the main thread needs to verify
-    // properties calculated by the I/O thread.
-    EventLoopGroup eventLoopGroup = new NioEventLoopGroup(1);
     ServerBootstrap sb =
         new ServerBootstrap()
             .group(eventLoopGroup)
             .channel(LocalServerChannel.class)
             .childHandler(serverInitializer);
     ChannelFuture unusedFuture = sb.bind(localAddress).syncUninterruptibly();
-    return eventLoopGroup;
   }
 
-  /**
+  /** Sets up a client channel connecting to the give local address. */
-   * Sets up a client channel connecting to the give local address.
-   *
-   * @param eventLoopGroup the same {@link EventLoopGroup} that is used to bootstrap server.
-   * @return the connected client channel.
-   */
   static Channel setUpClient(
       EventLoopGroup eventLoopGroup,
       ChannelInitializer<LocalChannel> clientInitializer,
       LocalAddress localAddress,
-      BackendProtocol protocol)
+      BackendProtocol protocol) {
-      throws Exception {
     Bootstrap b =
         new Bootstrap()
             .group(eventLoopGroup)
@@ -107,98 +90,102 @@ public class SslInitializerTestUtils {
     return b.connect(localAddress).syncUninterruptibly().channel();
   }
 
-  /** A handler that echoes back its inbound message. Used in test server. */
+  /**
+   * A handler that echoes back its inbound message. The message is also saved in a promise for
+   * inspection later.
+   */
   static class EchoHandler extends ChannelInboundHandlerAdapter {
 
-    /**
+    private final CountDownLatch latch = new CountDownLatch(1);
-     * A lock that synchronizes server I/O activity with the main thread. Acquired by the server I/O
+    private String request;
-     * thread when the handler is constructed, released when the server echoes back, or when an
+    private Throwable cause;
-     * exception is caught (during SSH handshake for example).
-     */
-    private final Lock lock;
 
-    /**
+    void waitTillReady() throws InterruptedException {
-     * Exception that would be initialized with the exception caught during SSL handshake. This
+      latch.await();
-     * field is constructed in the main thread and passed in the constructor. After a failure the
+    }
-     * main thread can inspect this object to assert the cause of the failure.
-     */
-    private final Exception serverException;
 
-    EchoHandler(Lock lock, Exception serverException) {
+    String getRequest() {
-      // This handler is constructed within getClientInitializer, which is called in the I/O thread.
+      return request;
-      // The server lock is therefore locked by the I/O thread.
+    }
-      lock.lock();
+
-      this.lock = lock;
+    Throwable getCause() {
-      this.serverException = serverException;
+      return cause;
     }
 
     @Override
     public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
-      // Always unlock regardless of whether the write is successful.
+      // In the test we only send messages of type ByteBuf.
-      ctx.writeAndFlush(msg).addListener(future -> lock.unlock());
+      assertThat(msg).isInstanceOf(ByteBuf.class);
+      request = ((ByteBuf) msg).toString(UTF_8);
+      // After the message is written back to the client, fulfill the promise.
+      ChannelFuture unusedFuture = ctx.writeAndFlush(msg).addListener(f -> latch.countDown());
     }
 
-    /** Saves any inbound error into the server exception field. */
+    /** Saves any inbound error as the cause of the promise failure. */
     @Override
     public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
-      serverException.initCause(cause);
+      this.cause = cause;
-      // If an exception is caught, we should also release the lock after the channel is closed
+      ChannelFuture unusedFuture =
-      // so that the main thread knows there is an exception to inspect now.
+          ctx.channel()
-      ctx.channel().closeFuture().addListener(f -> lock.unlock());
+              .closeFuture()
+              .addListener(
+                  // Apparently the JDK SSL provider will call #exceptionCaught twice with the same
+                  // exception when the handshake fails. In this case the second listener should not
+                  // set the promise again.
+                  f -> {
+                    if (latch.getCount() == 1) {
+                      latch.countDown();
+                    }
+                  });
     }
   }
 
-  /** A handler that dumps its inbound message in to {@link ByteBuf}. */
+  /** A handler that dumps its inbound message to a promise that can be inspected later. */
   static class DumpHandler extends ChannelInboundHandlerAdapter {
 
-    /**
+    private final CountDownLatch latch = new CountDownLatch(1);
-     * A lock that synchronizes client I/O activity with the main thread. Acquired by the client I/O
+    private String response;
-     * thread when the handler is constructed, released when the client receives an response, or
+    private Throwable cause;
-     * when an exception is caught (during SSH handshake for example).
-     */
-    private final Lock lock;
 
-    /**
+    void waitTillReady() throws InterruptedException {
-     * A Buffer that is used to store incoming message. Constructed in the main thread and passed in
+      latch.await();
-     * the constructor. The main thread can inspect this object to assert that the incoming message
+    }
-     * is as expected.
-     */
-    private final ByteBuf buffer;
 
-    /**
+    String getResponse() {
-     * Exception that would be initialized with the exception caught during SSL handshake. This
+      return response;
-     * field is constructed in the main thread and passed in the constructor. After a failure the
+    }
-     * main thread can inspect this object to assert the cause of the failure.
-     */
-    private final Exception clientException;
 
-    DumpHandler(Lock lock, ByteBuf buffer, Exception clientException) {
+    Throwable getCause() {
-      super();
+      return cause;
-      // This handler is constructed within getClientInitializer, which is called in the I/O thread.
-      // The client lock is therefore locked by the I/O thread.
-      lock.lock();
-      this.lock = lock;
-      this.buffer = buffer;
-      this.clientException = clientException;
     }
 
     @Override
     public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
-      buffer.writeBytes((ByteBuf) msg);
+      // In the test we only send messages of type ByteBuf.
-      // If a message is received here, the main thread must be waiting to acquire the lock from
+      assertThat(msg).isInstanceOf(ByteBuf.class);
-      // the I/O thread in order to verify it. Releasing the lock to notify the main thread it can
+      response = ((ByteBuf) msg).toString(UTF_8);
-      // continue now that the message has been written.
+      // There is no more use of this message, we should release its reference count so that it
-      lock.unlock();
+      // can be more effectively garbage collected by Netty.
+      ReferenceCountUtil.release(msg);
+      // Save the string in the promise and make it as complete.
+      latch.countDown();
     }
 
-    /** Saves any inbound error into clientException. */
+    /** Saves any inbound error into the failure cause of the promise. */
     @Override
     public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
-      clientException.initCause(cause);
+      this.cause = cause;
-      // If an exception is caught here, the main thread must be waiting to acquire the lock from
+      ctx.channel()
-      // the I/O thread in order to verify it. Releasing the lock after the channel is closed to
+          .closeFuture()
-      // notify the main thread it can continue now that the exception has been written.
+          .addListener(
-      ctx.channel().closeFuture().addListener(f -> lock.unlock());
+              f -> {
+                // Apparently the JDK SSL provider will call #exceptionCaught twice with the same
+                // exception when the handshake fails. In this case the second listener should not
+                // set the promise again.
+                if (latch.getCount() == 1) {
+                  latch.countDown();
+                }
+              });
     }
   }
 
@@ -237,9 +224,8 @@ public class SslInitializerTestUtils {
   static SSLSession verifySslChannel(
       Channel channel,
       ImmutableList<X509Certificate> certs,
-      Lock clientLock,
+      EchoHandler echoHandler,
-      Lock serverLock,
+      DumpHandler dumpHandler)
-      ByteBuf buffer)
       throws Exception {
     SslHandler sslHandler = channel.pipeline().get(SslHandler.class);
     // Wait till the handshake is complete.
@@ -254,18 +240,24 @@ public class SslInitializerTestUtils {
 
     // Test that message can go through, bound inbound and outbound.
     String inputString = "Hello, world!";
-    // The client writes the message to the server, which echos it back. The client receives the
+    // The client writes the message to the server, which echos it back and saves the string in its
-    // echo and writes to BUFFER. All these activities happens in the I/O thread, and this call
+    // promise. The client receives the echo and saves it in its promise. All these activities
-    // returns immediately.
+    // happens in the I/O thread, and this call itself returns immediately.
     ChannelFuture unusedFuture =
         channel.writeAndFlush(
             Unpooled.wrappedBuffer(inputString.getBytes(StandardCharsets.US_ASCII)));
-    // The lock is acquired by the I/O thread when the client's DumpHandler is constructed.
+
-    // Attempting to acquire it here blocks the main thread, until the I/O thread releases the lock
+    // Wait for both the server and the client to finish processing.
-    // after the DumpHandler writes the echo back to the buffer.
+    echoHandler.waitTillReady();
-    clientLock.lock();
+    dumpHandler.waitTillReady();
-    serverLock.lock();
+
-    assertThat(buffer.toString(StandardCharsets.US_ASCII)).isEqualTo(inputString);
+    // Checks that the message is transmitted faithfully.
+    String requestReceived = echoHandler.getRequest();
+    String responseReceived = dumpHandler.getResponse();
+    assertThat(inputString).isEqualTo(requestReceived);
+    assertThat(inputString).isEqualTo(responseReceived);
+
+    // Returns the SSL session for further assertion.
     return sslHandler.engine().getSession();
   }
 }

javatests/google/registry/proxy/handler/SslServerInitializerTest.java

@@ -28,8 +28,6 @@ import google.registry.proxy.Protocol;
 import google.registry.proxy.Protocol.BackendProtocol;
 import google.registry.proxy.handler.SslInitializerTestUtils.DumpHandler;
 import google.registry.proxy.handler.SslInitializerTestUtils.EchoHandler;
-import io.netty.buffer.ByteBuf;
-import io.netty.buffer.Unpooled;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelInitializer;
@@ -38,23 +36,22 @@ import io.netty.channel.EventLoopGroup;
 import io.netty.channel.embedded.EmbeddedChannel;
 import io.netty.channel.local.LocalAddress;
 import io.netty.channel.local.LocalChannel;
+import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.handler.ssl.OpenSsl;
 import io.netty.handler.ssl.SslContextBuilder;
 import io.netty.handler.ssl.SslHandler;
 import io.netty.handler.ssl.SslProvider;
 import io.netty.handler.ssl.util.SelfSignedCertificate;
-import io.netty.util.concurrent.Future;
 import java.security.KeyPair;
 import java.security.PrivateKey;
 import java.security.cert.CertificateException;
 import java.security.cert.X509Certificate;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 import javax.net.ssl.SSLEngine;
 import javax.net.ssl.SSLException;
 import javax.net.ssl.SSLHandshakeException;
 import javax.net.ssl.SSLParameters;
 import javax.net.ssl.SSLSession;
+import org.junit.After;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
@@ -93,6 +90,7 @@ public class SslServerInitializerTest {
   @Parameter(0)
   public SslProvider sslProvider;
 
+  // We do our best effort to test all available SSL providers.
   @Parameters(name = "{0}")
   public static SslProvider[] data() {
     return OpenSsl.isAvailable()
@@ -100,16 +98,27 @@ public class SslServerInitializerTest {
         : new SslProvider[] {SslProvider.JDK};
   }
 
+  // All I/O operations are done inside the single thread within this event loop group, which is
+  // different from the main test thread. Therefore synchronizations are required to make sure that
+  // certain I/O activities are finished when assertions are performed.
+  private final EventLoopGroup eventLoopGroup = new NioEventLoopGroup(1);
+
+  // Handler attached to server's channel to record the request received.
+  private final EchoHandler echoHandler = new EchoHandler();
+
+  // Handler attached to client's channel to record the response received.
+  private final DumpHandler dumpHandler = new DumpHandler();
+
+  @After
+  public void shutDown() {
+    eventLoopGroup.shutdownGracefully().getNow();
+  }
+
   private ChannelInitializer<LocalChannel> getServerInitializer(
-      boolean requireClientCert,
+      boolean requireClientCert, PrivateKey privateKey, X509Certificate... certificates) {
-      Lock serverLock,
-      Exception serverException,
-      PrivateKey privateKey,
-      X509Certificate... certificates)
-      throws Exception {
     return new ChannelInitializer<LocalChannel>() {
       @Override
-      protected void initChannel(LocalChannel ch) throws Exception {
+      protected void initChannel(LocalChannel ch) {
         ch.pipeline()
             .addLast(
                 new SslServerInitializer<LocalChannel>(
@@ -117,27 +126,18 @@ public class SslServerInitializerTest {
                     sslProvider,
                     Suppliers.ofInstance(privateKey),
                     Suppliers.ofInstance(certificates)),
-                new EchoHandler(serverLock, serverException));
+                echoHandler);
       }
     };
   }
 
   private ChannelInitializer<LocalChannel> getServerInitializer(
-      Lock serverLock,
+      PrivateKey privateKey, X509Certificate... certificates) {
-      Exception serverException,
+    return getServerInitializer(true, privateKey, certificates);
-      PrivateKey privateKey,
-      X509Certificate... certificates)
-      throws Exception {
-    return getServerInitializer(true, serverLock, serverException, privateKey, certificates);
   }
 
   private ChannelInitializer<LocalChannel> getClientInitializer(
-      X509Certificate trustedCertificate,
+      X509Certificate trustedCertificate, PrivateKey privateKey, X509Certificate certificate) {
-      PrivateKey privateKey,
-      X509Certificate certificate,
-      Lock clientLock,
-      ByteBuf buffer,
-      Exception clientException) {
     return new ChannelInitializer<LocalChannel>() {
       @Override
       protected void initChannel(LocalChannel ch) throws Exception {
@@ -155,8 +155,7 @@ public class SslServerInitializerTest {
         sslParameters.setEndpointIdentificationAlgorithm("HTTPS");
         sslEngine.setSSLParameters(sslParameters);
 
-        ch.pipeline().addLast("Client SSL Handler", sslHandler);
+        ch.pipeline().addLast(sslHandler, dumpHandler);
-        ch.pipeline().addLast(new DumpHandler(clientLock, buffer, clientException));
       }
     };
   }
@@ -184,72 +183,49 @@ public class SslServerInitializerTest {
   public void testSuccess_trustAnyClientCert() throws Exception {
     SelfSignedCertificate serverSsc = new SelfSignedCertificate(SSL_HOST);
     LocalAddress localAddress = new LocalAddress("TRUST_ANY_CLIENT_CERT_" + sslProvider);
-    Lock clientLock = new ReentrantLock();
+
-    Lock serverLock = new ReentrantLock();
-    ByteBuf buffer = Unpooled.buffer();
-    Exception clientException = new Exception();
-    Exception serverException = new Exception();
-    EventLoopGroup eventLoopGroup =
     setUpServer(
-            getServerInitializer(serverLock, serverException, serverSsc.key(), serverSsc.cert()),
+        eventLoopGroup, getServerInitializer(serverSsc.key(), serverSsc.cert()), localAddress);
-            localAddress);
     SelfSignedCertificate clientSsc = new SelfSignedCertificate();
     Channel channel =
         setUpClient(
             eventLoopGroup,
-            getClientInitializer(
+            getClientInitializer(serverSsc.cert(), clientSsc.key(), clientSsc.cert()),
-                serverSsc.cert(),
-                clientSsc.key(),
-                clientSsc.cert(),
-                clientLock,
-                buffer,
-                clientException),
             localAddress,
             PROTOCOL);
 
     SSLSession sslSession =
-        verifySslChannel(
+        verifySslChannel(channel, ImmutableList.of(serverSsc.cert()), echoHandler, dumpHandler);
-            channel, ImmutableList.of(serverSsc.cert()), clientLock, serverLock, buffer);
     // Verify that the SSL session gets the client cert. Note that this SslSession is for the client
     // channel, therefore its local certificates are the remote certificates of the SslSession for
     // the server channel, and vice versa.
     assertThat(sslSession.getLocalCertificates()).asList().containsExactly(clientSsc.cert());
     assertThat(sslSession.getPeerCertificates()).asList().containsExactly(serverSsc.cert());
-
-    Future<?> unusedFuture = eventLoopGroup.shutdownGracefully().syncUninterruptibly();
   }
 
   @Test
   public void testSuccess_doesNotRequireClientCert() throws Exception {
     SelfSignedCertificate serverSsc = new SelfSignedCertificate(SSL_HOST);
     LocalAddress localAddress = new LocalAddress("DOES_NOT_REQUIRE_CLIENT_CERT_" + sslProvider);
-    Lock clientLock = new ReentrantLock();
+
-    Lock serverLock = new ReentrantLock();
-    ByteBuf buffer = Unpooled.buffer();
-    Exception clientException = new Exception();
-    Exception serverException = new Exception();
-    EventLoopGroup eventLoopGroup =
     setUpServer(
-            getServerInitializer(
+        eventLoopGroup,
-                false, serverLock, serverException, serverSsc.key(), serverSsc.cert()),
+        getServerInitializer(false, serverSsc.key(), serverSsc.cert()),
         localAddress);
     Channel channel =
         setUpClient(
             eventLoopGroup,
-            getClientInitializer(serverSsc.cert(), null, null, clientLock, buffer, clientException),
+            getClientInitializer(serverSsc.cert(), null, null),
             localAddress,
             PROTOCOL);
 
     SSLSession sslSession =
-        verifySslChannel(
+        verifySslChannel(channel, ImmutableList.of(serverSsc.cert()), echoHandler, dumpHandler);
-            channel, ImmutableList.of(serverSsc.cert()), clientLock, serverLock, buffer);
     // Verify that the SSL session does not contain any client cert. Note that this SslSession is
     // for the client channel, therefore its local certificates are the remote certificates of the
     // SslSession for the server channel, and vice versa.
     assertThat(sslSession.getLocalCertificates()).isNull();
     assertThat(sslSession.getPeerCertificates()).asList().containsExactly(serverSsc.cert());
-
-    Future<?> unusedFuture = eventLoopGroup.shutdownGracefully().syncUninterruptibly();
   }
 
   @Test
@@ -259,16 +235,10 @@ public class SslServerInitializerTest {
     KeyPair keyPair = getKeyPair();
     X509Certificate serverCert = signKeyPair(caSsc, keyPair, SSL_HOST);
     LocalAddress localAddress = new LocalAddress("CERT_SIGNED_BY_OTHER_CA_" + sslProvider);
-    Lock clientLock = new ReentrantLock();
+
-    Lock serverLock = new ReentrantLock();
-    ByteBuf buffer = Unpooled.buffer();
-    Exception clientException = new Exception();
-    Exception serverException = new Exception();
-    EventLoopGroup eventLoopGroup =
     setUpServer(
+        eventLoopGroup,
         getServerInitializer(
-                serverLock,
-                serverException,
             keyPair.getPrivate(),
             // Serving both the server cert, and the CA cert
             serverCert,
@@ -280,41 +250,28 @@ public class SslServerInitializerTest {
             eventLoopGroup,
             getClientInitializer(
                 // Client trusts the CA cert
-                caSsc.cert(),
+                caSsc.cert(), clientSsc.key(), clientSsc.cert()),
-                clientSsc.key(),
-                clientSsc.cert(),
-                clientLock,
-                buffer,
-                clientException),
             localAddress,
             PROTOCOL);
 
     SSLSession sslSession =
         verifySslChannel(
-            channel, ImmutableList.of(serverCert, caSsc.cert()), clientLock, serverLock, buffer);
+            channel, ImmutableList.of(serverCert, caSsc.cert()), echoHandler, dumpHandler);
 
     assertThat(sslSession.getLocalCertificates()).asList().containsExactly(clientSsc.cert());
     assertThat(sslSession.getPeerCertificates())
         .asList()
         .containsExactly(serverCert, caSsc.cert())
         .inOrder();
-
-    Future<?> unusedFuture = eventLoopGroup.shutdownGracefully().syncUninterruptibly();
   }
 
   @Test
   public void testFailure_requireClientCertificate() throws Exception {
     SelfSignedCertificate serverSsc = new SelfSignedCertificate(SSL_HOST);
     LocalAddress localAddress = new LocalAddress("REQUIRE_CLIENT_CERT_" + sslProvider);
-    Lock clientLock = new ReentrantLock();
+
-    Lock serverLock = new ReentrantLock();
-    ByteBuf buffer = Unpooled.buffer();
-    Exception clientException = new Exception();
-    Exception serverException = new Exception();
-    EventLoopGroup eventLoopGroup =
     setUpServer(
-            getServerInitializer(serverLock, serverException, serverSsc.key(), serverSsc.cert()),
+        eventLoopGroup, getServerInitializer(serverSsc.key(), serverSsc.cert()), localAddress);
-            localAddress);
     Channel channel =
         setUpClient(
             eventLoopGroup,
@@ -322,62 +279,45 @@ public class SslServerInitializerTest {
                 serverSsc.cert(),
                 // No client cert/private key used.
                 null,
-                null,
+                null),
-                clientLock,
-                buffer,
-                clientException),
             localAddress,
             PROTOCOL);
 
-    serverLock.lock();
+    echoHandler.waitTillReady();
-    clientLock.lock();
+    dumpHandler.waitTillReady();
 
-    // When the server rejects the client during handshake due to lack of client certificate, only
+    // When the server rejects the client during handshake due to lack of client certificate, both
-    // the server throws an exception.
+    // should throw exceptions.
-    assertThat(Throwables.getRootCause(serverException)).isInstanceOf(SSLHandshakeException.class);
+    assertThat(Throwables.getRootCause(echoHandler.getCause()))
+        .isInstanceOf(SSLHandshakeException.class);
+    assertThat(Throwables.getRootCause(dumpHandler.getCause())).isInstanceOf(SSLException.class);
     assertThat(channel.isActive()).isFalse();
-
-    Future<?> unusedFuture = eventLoopGroup.shutdownGracefully().syncUninterruptibly();
   }
 
   @Test
   public void testFailure_wrongHostnameInCertificate() throws Exception {
     SelfSignedCertificate serverSsc = new SelfSignedCertificate("wrong.com");
     LocalAddress localAddress = new LocalAddress("WRONG_HOSTNAME_" + sslProvider);
-    Lock clientLock = new ReentrantLock();
+
-    Lock serverLock = new ReentrantLock();
-    ByteBuf buffer = Unpooled.buffer();
-    Exception clientException = new Exception();
-    Exception serverException = new Exception();
-    EventLoopGroup eventLoopGroup =
     setUpServer(
-            getServerInitializer(serverLock, serverException, serverSsc.key(), serverSsc.cert()),
+        eventLoopGroup, getServerInitializer(serverSsc.key(), serverSsc.cert()), localAddress);
-            localAddress);
     SelfSignedCertificate clientSsc = new SelfSignedCertificate();
     Channel channel =
         setUpClient(
             eventLoopGroup,
-            getClientInitializer(
+            getClientInitializer(serverSsc.cert(), clientSsc.key(), clientSsc.cert()),
-                serverSsc.cert(),
-                clientSsc.key(),
-                clientSsc.cert(),
-                clientLock,
-                buffer,
-                clientException),
             localAddress,
             PROTOCOL);
 
-    serverLock.lock();
+    echoHandler.waitTillReady();
-    clientLock.lock();
+    dumpHandler.waitTillReady();
 
-    // When the client rejects the server cert due to wrong hostname, the client error is wrapped
+    // When the client rejects the server cert due to wrong hostname, both the server and the client
-    // several layers in the exception. The server also throws an exception.
+    // throw exceptions.
-    Throwable rootCause = Throwables.getRootCause(clientException);
+    Throwable rootCause = Throwables.getRootCause(dumpHandler.getCause());
     assertThat(rootCause).isInstanceOf(CertificateException.class);
     assertThat(rootCause).hasMessageThat().contains(SSL_HOST);
-    assertThat(Throwables.getRootCause(serverException)).isInstanceOf(SSLException.class);
+    assertThat(Throwables.getRootCause(echoHandler.getCause())).isInstanceOf(SSLException.class);
     assertThat(channel.isActive()).isFalse();
-
-    Future<?> unusedFuture = eventLoopGroup.shutdownGracefully().syncUninterruptibly();
   }
 }