s2a: Add S2AStub cleanup handler.

netty/src/main/java/io/grpc/netty/InternalProtocolNegotiators.java

@@ -24,6 +24,7 @@
 import io.netty.channel.ChannelHandler;
 import io.netty.handler.ssl.SslContext;
 import io.netty.util.AsciiString;
+import java.util.Optional;
 import java.util.concurrent.Executor;
 
 /**
@@ -40,9 +41,10 @@
    * @param executorPool a dedicated {@link Executor} pool for time-consuming TLS tasks
    */
   public static InternalProtocolNegotiator.ProtocolNegotiator tls(SslContext sslContext,
-          ObjectPool<? extends Executor> executorPool) {
+          ObjectPool<? extends Executor> executorPool,
+          Optional<Runnable> handshakeCompleteRunnable) {
     final io.grpc.netty.ProtocolNegotiator negotiator = ProtocolNegotiators.tls(sslContext,
-        executorPool);
+        executorPool, handshakeCompleteRunnable);
     final class TlsNegotiator implements InternalProtocolNegotiator.ProtocolNegotiator {
 
       @Override
@@ -70,7 +72,7 @@
    * may happen immediately, even before the TLS Handshake is complete.
    */
   public static InternalProtocolNegotiator.ProtocolNegotiator tls(SslContext sslContext) {
-    return tls(sslContext, null);
+    return tls(sslContext, null, Optional.empty());
   }
 
   /**
@@ -167,7 +169,8 @@
   public static ChannelHandler clientTlsHandler(
       ChannelHandler next, SslContext sslContext, String authority,
       ChannelLogger negotiationLogger) {
-    return new ClientTlsHandler(next, sslContext, authority, null, negotiationLogger);
+    return new ClientTlsHandler(next, sslContext, authority, null, negotiationLogger,
+        Optional.empty());
   }
 
   public static class ProtocolNegotiationHandler

netty/src/main/java/io/grpc/netty/NettyChannelBuilder.java

@@ -63,6 +63,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -604,7 +605,7 @@ static ProtocolNegotiator createProtocolNegotiatorByType(
       case PLAINTEXT_UPGRADE:
         return ProtocolNegotiators.plaintextUpgrade();
       case TLS:
-        return ProtocolNegotiators.tls(sslContext, executorPool);
+        return ProtocolNegotiators.tls(sslContext, executorPool, Optional.empty());
       default:
         throw new IllegalArgumentException("Unsupported negotiationType: " + negotiationType);
     }

netty/src/main/java/io/grpc/netty/ProtocolNegotiators.java

@@ -72,6 +72,7 @@
 import java.nio.channels.ClosedChannelException;
 import java.util.Arrays;
 import java.util.EnumSet;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.logging.Level;
@@ -543,16 +544,18 @@
   static final class ClientTlsProtocolNegotiator implements ProtocolNegotiator {
 
     public ClientTlsProtocolNegotiator(SslContext sslContext,
-        ObjectPool<? extends Executor> executorPool) {
+        ObjectPool<? extends Executor> executorPool, Optional<Runnable> handshakeCompleteRunnable) {
       this.sslContext = checkNotNull(sslContext, "sslContext");
       this.executorPool = executorPool;
       if (this.executorPool != null) {
         this.executor = this.executorPool.getObject();
       }
+      this.handshakeCompleteRunnable = handshakeCompleteRunnable;
     }
 
     private final SslContext sslContext;
     private final ObjectPool<? extends Executor> executorPool;
+    private final Optional<Runnable> handshakeCompleteRunnable;
     private Executor executor;
 
     @Override
@@ -565,7 +568,7 @@
       ChannelHandler gnh = new GrpcNegotiationHandler(grpcHandler);
       ChannelLogger negotiationLogger = grpcHandler.getNegotiationLogger();
       ChannelHandler cth = new ClientTlsHandler(gnh, sslContext, grpcHandler.getAuthority(),
-          this.executor, negotiationLogger);
+          this.executor, negotiationLogger, handshakeCompleteRunnable);
       return new WaitUntilActiveHandler(cth, negotiationLogger);
     }
 
@@ -583,15 +586,18 @@
     private final String host;
     private final int port;
     private Executor executor;
+    private final Optional<Runnable> handshakeCompleteRunnable;
 
     ClientTlsHandler(ChannelHandler next, SslContext sslContext, String authority,
-        Executor executor, ChannelLogger negotiationLogger) {
+        Executor executor, ChannelLogger negotiationLogger,
+        Optional<Runnable> handshakeCompleteRunnable) {
       super(next, negotiationLogger);
       this.sslContext = checkNotNull(sslContext, "sslContext");
       HostPort hostPort = parseAuthority(authority);
       this.host = hostPort.host;
       this.port = hostPort.port;
       this.executor = executor;
+      this.handshakeCompleteRunnable = handshakeCompleteRunnable;
     }
 
     @Override
@@ -634,6 +640,9 @@
                 .withCause(t)
                 .asRuntimeException();
           }
+          if (handshakeCompleteRunnable.isPresent()) {
+            handshakeCompleteRunnable.get().run();
+          }
           ctx.fireExceptionCaught(t);
         }
       } else {
@@ -649,6 +658,9 @@
           .set(Grpc.TRANSPORT_ATTR_SSL_SESSION, session)
           .build();
       replaceProtocolNegotiationEvent(existingPne.withAttributes(attrs).withSecurity(security));
+      if (handshakeCompleteRunnable.isPresent()) {
+        handshakeCompleteRunnable.get().run();
+      }
       fireProtocolNegotiationEvent(ctx);
     }
   }
@@ -683,8 +695,8 @@
    * @param executorPool a dedicated {@link Executor} pool for time-consuming TLS tasks
    */
   public static ProtocolNegotiator tls(SslContext sslContext,
-      ObjectPool<? extends Executor> executorPool) {
-    return new ClientTlsProtocolNegotiator(sslContext, executorPool);
+      ObjectPool<? extends Executor> executorPool, Optional<Runnable> handshakeCompleteRunnable) {
+    return new ClientTlsProtocolNegotiator(sslContext, executorPool, handshakeCompleteRunnable);
   }
 
   /**
@@ -693,7 +705,7 @@
    * may happen immediately, even before the TLS Handshake is complete.
    */
   public static ProtocolNegotiator tls(SslContext sslContext) {
-    return tls(sslContext, null);
+    return tls(sslContext, null, Optional.empty());
   }
 
   public static ProtocolNegotiator.ClientFactory tlsClientFactory(SslContext sslContext) {

netty/src/test/java/io/grpc/netty/NettyClientTransportTest.java

@@ -105,6 +105,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -766,7 +767,8 @@ public void tlsNegotiationServerExecutorShouldSucceed() throws Exception {
         .trustManager(caCert)
         .keyManager(clientCert, clientKey)
         .build();
-    ProtocolNegotiator negotiator = ProtocolNegotiators.tls(clientContext, clientExecutorPool);
+    ProtocolNegotiator negotiator = ProtocolNegotiators.tls(clientContext, clientExecutorPool,
+        Optional.empty());
     // after starting the client, the Executor in the client pool should be used
     assertEquals(true, clientExecutorPool.isInUse());
     final NettyClientTransport transport = newTransport(negotiator);

netty/src/test/java/io/grpc/netty/ProtocolNegotiatorsTest.java

@@ -120,6 +120,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.Queue;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -876,7 +877,7 @@ public String applicationProtocol() {
     DefaultEventLoopGroup elg = new DefaultEventLoopGroup(1);
 
     ClientTlsHandler handler = new ClientTlsHandler(grpcHandler, sslContext,
-        "authority", elg, noopLogger);
+        "authority", elg, noopLogger, Optional.empty());
     pipeline.addLast(handler);
     pipeline.replace(SslHandler.class, null, goodSslHandler);
     pipeline.fireUserEventTriggered(ProtocolNegotiationEvent.DEFAULT);
@@ -914,7 +915,7 @@ public String applicationProtocol() {
         .applicationProtocolConfig(apn).build();
 
     ClientTlsHandler handler = new ClientTlsHandler(grpcHandler, sslContext,
-        "authority", elg, noopLogger);
+        "authority", elg, noopLogger, Optional.empty());
     pipeline.addLast(handler);
     pipeline.replace(SslHandler.class, null, goodSslHandler);
     pipeline.fireUserEventTriggered(ProtocolNegotiationEvent.DEFAULT);
@@ -938,7 +939,7 @@ public String applicationProtocol() {
     DefaultEventLoopGroup elg = new DefaultEventLoopGroup(1);
 
     ClientTlsHandler handler = new ClientTlsHandler(grpcHandler, sslContext,
-        "authority", elg, noopLogger);
+        "authority", elg, noopLogger, Optional.empty());
     pipeline.addLast(handler);
 
     final AtomicReference<Throwable> error = new AtomicReference<>();
@@ -966,7 +967,7 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
   @Test
   public void clientTlsHandler_closeDuringNegotiation() throws Exception {
     ClientTlsHandler handler = new ClientTlsHandler(grpcHandler, sslContext,
-        "authority", null, noopLogger);
+        "authority", null, noopLogger, Optional.empty());
     pipeline.addLast(new WriteBufferingAndExceptionHandler(handler));
     ChannelFuture pendingWrite = channel.writeAndFlush(NettyClientHandler.NOOP_MESSAGE);
 
@@ -1228,7 +1229,8 @@ public void clientTlsHandler_firesNegotiation() throws Exception {
       serverSslContext = GrpcSslContexts.forServer(server1Chain, server1Key).build();
     }
     FakeGrpcHttp2ConnectionHandler gh = FakeGrpcHttp2ConnectionHandler.newHandler();
-    ClientTlsProtocolNegotiator pn = new ClientTlsProtocolNegotiator(clientSslContext, null);
+    ClientTlsProtocolNegotiator pn = new ClientTlsProtocolNegotiator(clientSslContext,
+        null, Optional.empty());
     WriteBufferingAndExceptionHandler clientWbaeh =
         new WriteBufferingAndExceptionHandler(pn.newHandler(gh));
 

s2a/src/main/java/io/grpc/s2a/S2AChannelCredentials.java

@@ -31,6 +31,7 @@
 import io.grpc.s2a.internal.channel.S2AHandshakerServiceChannel;
 import io.grpc.s2a.internal.handshaker.S2AIdentity;
 import io.grpc.s2a.internal.handshaker.S2AProtocolNegotiatorFactory;
+import io.grpc.s2a.internal.handshaker.S2AStub;
 import javax.annotation.concurrent.NotThreadSafe;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -59,6 +60,7 @@ public static final class Builder {
     private final String s2aAddress;
     private final ChannelCredentials s2aChannelCredentials;
     private @Nullable S2AIdentity localIdentity = null;
+    private @Nullable S2AStub stub = null;
 
     Builder(String s2aAddress, ChannelCredentials s2aChannelCredentials) {
       this.s2aAddress = s2aAddress;
@@ -104,6 +106,16 @@ public Builder setLocalUid(String localUid) {
       return this;
     }
 
+    /**
+     * Sets the stub to use to communicate with S2A. This is only used for testing that the
+     * stream to S2A gets closed.
+     */
+    public Builder setStub(S2AStub stub) {
+      checkNotNull(stub);
+      this.stub = stub;
+      return this;
+    }
+
     public ChannelCredentials build() {
       return InternalNettyChannelCredentials.create(buildProtocolNegotiatorFactory());
     }
@@ -113,7 +125,7 @@ InternalProtocolNegotiator.ClientFactory buildProtocolNegotiatorFactory() {
           SharedResourcePool.forResource(
               S2AHandshakerServiceChannel.getChannelResource(s2aAddress, s2aChannelCredentials));
       checkNotNull(s2aChannelPool, "s2aChannelPool");
-      return S2AProtocolNegotiatorFactory.createClientFactory(localIdentity, s2aChannelPool);
+      return S2AProtocolNegotiatorFactory.createClientFactory(localIdentity, s2aChannelPool, stub);
     }
   }