Handle WebSocket messages when they come in the same frame as the Upgrade response, close AsyncHttpClient#1095

Author: slandelle

client/src/main/java/org/asynchttpclient/netty/channel/ChannelManager.java

@@ -471,9 +471,9 @@ public Bootstrap getBootstrap(Uri uri, ProxyServer proxy) {
 
     public void upgradePipelineForWebSockets(ChannelPipeline pipeline) {
         pipeline.addAfter(HTTP_CLIENT_CODEC, WS_ENCODER_HANDLER, new WebSocket08FrameEncoder(true));
-        pipeline.remove(HTTP_CLIENT_CODEC);
         pipeline.addBefore(AHC_WS_HANDLER, WS_DECODER_HANDLER, new WebSocket08FrameDecoder(false, false, config.getWebSocketMaxFrameSize()));
         pipeline.addAfter(WS_DECODER_HANDLER, WS_FRAME_AGGREGATOR, new WebSocketFrameAggregator(config.getWebSocketMaxBufferSize()));
+        pipeline.remove(HTTP_CLIENT_CODEC);
     }
 
     public final Callback newDrainCallback(final NettyResponseFuture<?> future, final Channel channel, final boolean keepAlive, final Object partitionKey) {

client/src/main/java/org/asynchttpclient/netty/handler/WebSocketHandler.java

@@ -54,18 +54,6 @@ public WebSocketHandler(AsyncHttpClientConfig config,//
         super(config, channelManager, requestSender);
     }
 
-    // We don't need to synchronize as replacing the "ws-decoder" will
-    // process using the same thread.
-    private void invokeOnSucces(Channel channel, WebSocketUpgradeHandler h) {
-        if (!h.touchSuccess()) {
-            try {
-                h.onSuccess(new NettyWebSocket(channel, config));
-            } catch (Exception ex) {
-                logger.warn("onSuccess unexpected exception", ex);
-            }
-        }
-    }
-
     private class UpgradeCallback extends Callback {
 
         private final Channel channel;
@@ -84,6 +72,18 @@ public UpgradeCallback(NettyResponseFuture<?> future, Channel channel, HttpRespo
             this.responseHeaders = responseHeaders;
         }
 
+        // We don't need to synchronize as replacing the "ws-decoder" will
+        // process using the same thread.
+        private void invokeOnSucces(Channel channel, WebSocketUpgradeHandler h) {
+            if (!h.touchSuccess()) {
+                try {
+                    h.onSuccess(new NettyWebSocket(channel, config));
+                } catch (Exception ex) {
+                    logger.warn("onSuccess unexpected exception", ex);
+                }
+            }
+        }
+
         @Override
         public void call() throws Exception {
 
@@ -116,14 +116,16 @@ public void call() throws Exception {
                 requestSender.abort(channel, future, new IOException(String.format("Invalid challenge. Actual: %s. Expected: %s", accept, key)));
             }
 
+            // set back the future so the protocol gets notified of frames
+            // removing the HttpClientCodec from the pipeline might trigger a read with a WebSocket message
+            // if it comes in the same frame as the HTTP Upgrade response
+            Channels.setAttribute(channel, future);
+
             channelManager.upgradePipelineForWebSockets(channel.pipeline());
 
             invokeOnSucces(channel, handler);
             future.done();
-            // set back the future so the protocol gets notified of frames
-            Channels.setAttribute(channel, future);
         }
-
     }
 
     @Override
@@ -144,43 +146,61 @@ public void handleRead(Channel channel, NettyResponseFuture<?> future, Object e)
                 Channels.setAttribute(channel, new UpgradeCallback(future, channel, response, handler, status, responseHeaders));
             }
 
-
         } else if (e instanceof WebSocketFrame) {
-
             final WebSocketFrame frame = (WebSocketFrame) e;
             WebSocketUpgradeHandler handler = (WebSocketUpgradeHandler) future.getAsyncHandler();
             NettyWebSocket webSocket = (NettyWebSocket) handler.onCompleted();
 
             if (webSocket != null) {
-                if (frame instanceof CloseWebSocketFrame) {
-                    Channels.setDiscard(channel);
-                    CloseWebSocketFrame closeFrame = (CloseWebSocketFrame) frame;
-                    webSocket.onClose(closeFrame.statusCode(), closeFrame.reasonText());
-                } else {
-                    ByteBuf buf = frame.content();
-                    if (buf != null && buf.readableBytes() > 0) {
-                        HttpResponseBodyPart part = config.getResponseBodyPartFactory().newResponseBodyPart(buf, frame.isFinalFragment());
-                        handler.onBodyPartReceived(part);
-
-                        if (frame instanceof BinaryWebSocketFrame) {
-                            webSocket.onBinaryFragment(part);
-                        } else if (frame instanceof TextWebSocketFrame) {
-                            webSocket.onTextFragment(part);
-                        } else if (frame instanceof PingWebSocketFrame) {
-                            webSocket.onPing(part);
-                        } else if (frame instanceof PongWebSocketFrame) {
-                            webSocket.onPong(part);
-                        }
-                    }
-                }
+                handleFrame(channel, frame, handler, webSocket);
             } else {
-                logger.debug("UpgradeHandler returned a null NettyWebSocket");
+                logger.debug("Frame received but WebSocket is not available yet, buffering frame");
+                frame.retain();
+                Runnable bufferedFrame = new Runnable() {
+                    public void run() {
+                        try {
+                            // WebSocket is now not null
+                            NettyWebSocket webSocket = (NettyWebSocket) handler.onCompleted();
+                            handleFrame(channel, frame, handler, webSocket);
+                        } catch (Exception e) {
+                            logger.debug("Failure while handling buffered frame", e);
+                            handler.onFailure(e);
+                        } finally {
+                            frame.release();
+                        }
+                    };
+                };
+                handler.bufferFrame(bufferedFrame);
             }
         } else {
             logger.error("Invalid message {}", e);
         }
     }
 
+    private void handleFrame(Channel channel, WebSocketFrame frame, WebSocketUpgradeHandler handler, NettyWebSocket webSocket) throws Exception {
+        if (frame instanceof CloseWebSocketFrame) {
+            Channels.setDiscard(channel);
+            CloseWebSocketFrame closeFrame = (CloseWebSocketFrame) frame;
+            webSocket.onClose(closeFrame.statusCode(), closeFrame.reasonText());
+        } else {
+            ByteBuf buf = frame.content();
+            if (buf != null && buf.readableBytes() > 0) {
+                HttpResponseBodyPart part = config.getResponseBodyPartFactory().newResponseBodyPart(buf, frame.isFinalFragment());
+                handler.onBodyPartReceived(part);
+
+                if (frame instanceof BinaryWebSocketFrame) {
+                    webSocket.onBinaryFragment(part);
+                } else if (frame instanceof TextWebSocketFrame) {
+                    webSocket.onTextFragment(part);
+                } else if (frame instanceof PingWebSocketFrame) {
+                    webSocket.onPing(part);
+                } else if (frame instanceof PongWebSocketFrame) {
+                    webSocket.onPong(part);
+                }
+            }
+        }
+    }
+
     @Override
     public void handleException(NettyResponseFuture<?> future, Throwable e) {
         logger.warn("onError {}", e);

client/src/main/java/org/asynchttpclient/ws/WebSocketUpgradeHandler.java

@@ -12,7 +12,7 @@
  */
 package org.asynchttpclient.ws;
 
-import static org.asynchttpclient.util.Assertions.*;
+import static org.asynchttpclient.util.MiscUtils.isNonEmpty;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -28,16 +28,26 @@
  */
 public class WebSocketUpgradeHandler implements UpgradeHandler<WebSocket>, AsyncHandler<WebSocket> {
 
+    private static final int SWITCHING_PROTOCOLS = io.netty.handler.codec.http.HttpResponseStatus.SWITCHING_PROTOCOLS.code();
+
     private WebSocket webSocket;
     private final List<WebSocketListener> listeners;
     private final AtomicBoolean ok = new AtomicBoolean(false);
     private boolean onSuccessCalled;
     private int status;
+    private List<Runnable> bufferedFrames;
 
     public WebSocketUpgradeHandler(List<WebSocketListener> listeners) {
         this.listeners = listeners;
     }
 
+    public void bufferFrame(Runnable bufferedFrame) {
+        if (bufferedFrames == null) {
+            bufferedFrames = new ArrayList<>();
+        }
+        bufferedFrames.add(bufferedFrame);
+    }
+
     /**
      * {@inheritDoc}
      */
@@ -66,11 +76,7 @@ public final State onBodyPartReceived(HttpResponseBodyPart bodyPart) throws Exce
     @Override
     public final State onStatusReceived(HttpResponseStatus responseStatus) throws Exception {
         status = responseStatus.getStatusCode();
-        if (responseStatus.getStatusCode() == 101) {
-            return State.UPGRADE;
-        } else {
-            return State.ABORT;
-        }
+        return status == SWITCHING_PROTOCOLS ? State.UPGRADE : State.ABORT;
     }
 
     /**
@@ -87,15 +93,15 @@ public final State onHeadersReceived(HttpResponseHeaders headers) throws Excepti
     @Override
     public final WebSocket onCompleted() throws Exception {
 
-        if (status != 101) {
+        if (status != SWITCHING_PROTOCOLS) {
             IllegalStateException e = new IllegalStateException("Invalid Status Code " + status);
             for (WebSocketListener listener : listeners) {
                 listener.onError(e);
             }
             throw e;
         }
 
-        return assertNotNull(webSocket, "webSocket");
+        return webSocket;
     }
 
     /**
@@ -108,6 +114,12 @@ public final void onSuccess(WebSocket webSocket) {
             webSocket.addWebSocketListener(listener);
             listener.onOpen(webSocket);
         }
+        if (isNonEmpty(bufferedFrames)) {
+            for (Runnable bufferedFrame : bufferedFrames) {
+                bufferedFrame.run();
+            }
+            bufferedFrames = null;
+        }
         ok.set(true);
     }