Fix Websocket race condition while protocol switching

src/main/java/org/littleshoot/proxy/impl/ClientToProxyConnection.java

@@ -4,6 +4,7 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.Channel;
+import io.netty.channel.ChannelFuture;
 import io.netty.channel.ChannelPipeline;
 import io.netty.handler.codec.haproxy.HAProxyMessage;
 import io.netty.handler.codec.haproxy.HAProxyMessageDecoder;
@@ -302,11 +303,6 @@ private ConnectionState doReadHTTPInitial(HttpRequest httpRequest) {
                     "Not reusing existing ProxyToServerConnection because request is a CONNECT for: {}",
                     serverHostAndPort);
             newConnectionRequired = true;
-        } else if (ProxyUtils.isSwitchingToWebSocketProtocol(httpRequest)) {
-            LOG.debug(
-                    "Not reusing existing ProxyToServerConnection because request is an upgrade to websocket for: {}",
-                    serverHostAndPort);
-            newConnectionRequired = true;
         } else if (currentServerConnection == null) {
             LOG.debug("Didn't find existing ProxyToServerConnection for: {}",
                     serverHostAndPort);
@@ -455,7 +451,7 @@ void respond(ProxyToServerConnection serverConnection, HttpFilters filters,
             return;
         }
 
-        boolean isSwitchingToWebSocketProtocol = false;
+        final boolean isSwitchingToWebSocketProtocol;
         if (httpObject instanceof HttpResponse) {
             HttpResponse httpResponse = (HttpResponse) httpObject;
 
@@ -484,25 +480,31 @@ void respond(ProxyToServerConnection serverConnection, HttpFilters filters,
 
             fixHttpVersionHeaderIfNecessary(httpResponse);
             modifyResponseHeadersToReflectProxying(httpResponse);
+        } else {
+            isSwitchingToWebSocketProtocol = false;
         }
 
-        httpObject = filters.proxyToClientResponse(httpObject);
-        if (httpObject == null) {
+        final HttpObject filteredhttpObject = filters.proxyToClientResponse(httpObject);
+        if (filteredhttpObject == null) {
             forceDisconnect(serverConnection);
             return;
         }
 
-        write(httpObject);
-
-        if (ProxyUtils.isLastChunk(httpObject)) {
-            writeEmptyBuffer();
-        }
-        else if (isSwitchingToWebSocketProtocol) {
-            switchToWebSocketProtocol(serverConnection);
+        if (isSwitchingToWebSocketProtocol) {
+            serverConnection.switchToWebSocketProtocol();
         }
+        write(filteredhttpObject).addListener(l -> {
 
-        closeConnectionsAfterWriteIfNecessary(serverConnection,
-                currentHttpRequest, currentHttpResponse, httpObject);
+	        if (ProxyUtils.isLastChunk(filteredhttpObject)) {
+	            writeEmptyBuffer();
+	        }
+	        else if (isSwitchingToWebSocketProtocol) {
+	            switchToWebSocketProtocol(serverConnection);
+	        }
+
+	        closeConnectionsAfterWriteIfNecessary(serverConnection,
+	                currentHttpRequest, currentHttpResponse, filteredhttpObject);
+        });
     }
 
     private void resetCurrentRequest() {
@@ -520,7 +522,6 @@ private void switchToWebSocketProtocol(final ProxyToServerConnection serverConne
                     new ProxyConnectionPipeHandler(serverConnection));
         }
         orderedHandlersToRemove.forEach(this::removeHandlerIfPresent);
-        serverConnection.switchToWebSocketProtocol();
     }
 
     /* *************************************************************************
@@ -1388,8 +1389,8 @@ private String identifyHostAndPort(HttpRequest httpRequest) {
      * and using the empty buffer's future instead to handle any operations we
      * need to when responses are fully written back to clients.
      */
-    private void writeEmptyBuffer() {
-        write(Unpooled.EMPTY_BUFFER);
+    private ChannelFuture writeEmptyBuffer() {
+        return write(Unpooled.EMPTY_BUFFER);
     }
 
     public boolean isMitming() {

src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java

@@ -71,7 +71,7 @@ abstract class ProxyConnection<I extends HttpObject> extends
     protected volatile Channel channel;
 
     private volatile ConnectionState currentState;
-    private volatile boolean tunneling = false;
+    protected volatile boolean tunneling = false;
     protected volatile long lastReadTime = 0;
 
     /**
@@ -211,23 +211,23 @@ private void readHTTP(HttpObject httpObject) {
      * This method is called by users of the ProxyConnection to send stuff out
      * over the socket.
      */
-    void write(Object msg) {
+    ChannelFuture write(Object msg) {
         if (msg instanceof ReferenceCounted) {
             LOG.debug("Retaining reference counted message");
             ((ReferenceCounted) msg).retain();
         }
 
-        doWrite(msg);
+        return doWrite(msg);
     }
 
-    void doWrite(Object msg) {
+    ChannelFuture doWrite(Object msg) {
         LOG.debug("Writing: {}", msg);
 
         try {
             if (msg instanceof HttpObject) {
-                writeHttp((HttpObject) msg);
+                return writeHttp((HttpObject) msg);
             } else {
-                writeRaw((ByteBuf) msg);
+                return writeRaw((ByteBuf) msg);
             }
         } finally {
             LOG.debug("Wrote: {}", msg);
@@ -237,25 +237,29 @@ void doWrite(Object msg) {
     /**
      * Writes HttpObjects to the connection asynchronously.
      */
-    protected void writeHttp(HttpObject httpObject) {
+    protected ChannelFuture writeHttp(HttpObject httpObject) {
         if (ProxyUtils.isLastChunk(httpObject)) {
             channel.write(httpObject);
             LOG.debug("Writing an empty buffer to signal the end of our chunked transfer");
-            writeToChannel(Unpooled.EMPTY_BUFFER);
+            return writeToChannel(Unpooled.EMPTY_BUFFER);
         } else {
-            writeToChannel(httpObject);
+            return writeToChannel(httpObject);
         }
     }
 
     /**
      * Writes raw buffers to the connection.
      */
-    protected void writeRaw(ByteBuf buf) {
-        writeToChannel(buf);
+    protected ChannelFuture writeRaw(ByteBuf buf) {
+        return writeToChannel(buf);
     }
 
     protected ChannelFuture writeToChannel(final Object msg) {
-        return channel.writeAndFlush(msg);
+        return channel.writeAndFlush(msg).addListener(l-> {
+            if (!l.isSuccess()) {
+                LOG.debug("writeToChannel failed sending message {}", msg, l.cause());
+            }
+        });
     }
 
     /* *************************************************************************

src/main/java/org/littleshoot/proxy/impl/ProxyToServerConnection.java

@@ -351,7 +351,7 @@ void write(Object msg, HttpFilters filters) {
     }
 
     @Override
-    void write(Object msg) {
+    ChannelFuture write(Object msg) {
         LOG.debug("Requested write of {}", msg);
 
         if (msg instanceof ReferenceCounted) {
@@ -362,6 +362,7 @@ void write(Object msg) {
         if (is(DISCONNECTED) && msg instanceof HttpRequest) {
             LOG.debug("Currently disconnected, connect and then write the message");
             connectAndWrite((HttpRequest) msg);
+            return this.clientConnection.channel.newSucceededFuture();
         } else {
             if (isConnecting()) {
                 synchronized (connectLock) {
@@ -381,24 +382,24 @@ void write(Object msg) {
             // already disconnected
             if (isConnecting() || getCurrentState().isDisconnectingOrDisconnected()) {
                 LOG.debug("Connection failed or timed out while waiting to write message to server. Message will be discarded: {}", msg);
-                return;
+                return channel.newFailedFuture(new Exception("Connection failed or timed out while waiting to write message to server. Message will be discarded."));
             }
 
             LOG.debug("Using existing connection to: {}", remoteAddress);
-            doWrite(msg);
+            return doWrite(msg);
         }
     }
 
     @Override
-    protected void writeHttp(HttpObject httpObject) {
+    protected ChannelFuture writeHttp(HttpObject httpObject) {
         if (chainedProxy != null) {
             chainedProxy.filterRequest(httpObject);
         }
         if (httpObject instanceof HttpRequest) {
             // Remember that we issued this HttpRequest for later
             currentHttpRequest = (HttpRequest) httpObject;
         }
-        super.writeHttp(httpObject);
+        return super.writeHttp(httpObject);
     }
 
     /* *************************************************************************
@@ -1217,6 +1218,7 @@ void switchToWebSocketProtocol() {
                     new ProxyConnectionPipeHandler(clientConnection));
         }
         orderedHandlersToRemove.forEach(this::removeHandlerIfPresent);
+        tunneling = true;
     }
 
     /* *************************************************************************