Don't upgrade pipeline for WebSocket prior to installing WebSocket on Handler, close AsyncHttpClient#1364

Motivation:

Installing the WebSocket handler can trigger a read, so if there's a
pending frame that was send along the HTTP upgrade response, we could
crash because the handler is not properly configured yet.

This is a regression caused by AsyncHttpClient#1348.

Modifications:

Buffer frames until WebSocket is fully open.

Result:

No more crash when frames are sent along the HTTP upgrade response.

Author: slandelle

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

@@ -75,12 +75,13 @@ private void upgrade(Channel channel, NettyResponseFuture<?> future, WebSocketUp
         // if it comes in the same frame as the HTTP Upgrade response
         Channels.setAttribute(channel, future);
 
+        handler.setWebSocket(new NettyWebSocket(channel, responseHeaders.getHeaders()));
         channelManager.upgradePipelineForWebSockets(channel.pipeline());
 
         // We don't need to synchronize as replacing the "ws-decoder" will
         // process using the same thread.
         try {
-            handler.openWebSocket(new NettyWebSocket(channel, responseHeaders.getHeaders()));
+            handler.onOpen();
         } catch (Exception ex) {
             logger.warn("onSuccess unexpected exception", ex);
         }
@@ -105,7 +106,7 @@ public void handleRead(Channel channel, NettyResponseFuture<?> future, Object e)
                 logger.debug("\n\nRequest {}\n\nResponse {}\n", httpRequest, response);
             }
 
-            WebSocketUpgradeHandler handler = WebSocketUpgradeHandler.class.cast(future.getAsyncHandler());
+            WebSocketUpgradeHandler handler = (WebSocketUpgradeHandler) future.getAsyncHandler();
             HttpResponseStatus status = new NettyResponseStatus(future.getUri(), response, channel);
             HttpResponseHeaders responseHeaders = new HttpResponseHeaders(response.headers());
 
@@ -123,15 +124,23 @@ public void handleRead(Channel channel, NettyResponseFuture<?> future, Object e)
             final WebSocketFrame frame = (WebSocketFrame) e;
             WebSocketUpgradeHandler handler = (WebSocketUpgradeHandler) future.getAsyncHandler();
             NettyWebSocket webSocket = (NettyWebSocket) handler.onCompleted();
-            handleFrame(channel, frame, handler, webSocket);
+            // retain because we might buffer the frame
+            frame.retain();
+            if (handler.isOpen()) {
+                handleFrame(channel, frame, handler, webSocket);
+            } else {
+                // WebSocket hasn't been open yet, but upgrading the pipeline triggered a read and a frame was sent along the HTTP upgrade response
+                // as we want to keep sequential order (but can't notify user of open before upgrading so he doesn't to try send immediately), we have to buffer
+                handler.bufferFrame(() -> handleFrame(channel, frame, handler, webSocket));
+            }
 
         } else if (!(e instanceof LastHttpContent)) {
             // ignore, end of handshake response
             logger.error("Invalid message {}", e);
         }
     }
 
-    private void handleFrame(Channel channel, WebSocketFrame frame, WebSocketUpgradeHandler handler, NettyWebSocket webSocket) throws Exception {
+    private void handleFrame(Channel channel, WebSocketFrame frame, WebSocketUpgradeHandler handler, NettyWebSocket webSocket) {
         if (frame instanceof TextWebSocketFrame) {
             webSocket.onTextFrame((TextWebSocketFrame) frame);
 
@@ -150,6 +159,8 @@ private void handleFrame(Channel channel, WebSocketFrame frame, WebSocketUpgrade
         } else if (frame instanceof PongWebSocketFrame) {
             webSocket.onPong((PongWebSocketFrame) frame);
         }
+        // release because we had to retain in case the frame had to be buffered
+        frame.release();
     }
 
     @Override

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

@@ -29,7 +29,9 @@ public class WebSocketUpgradeHandler implements AsyncHandler<WebSocket> {
     private static final int SWITCHING_PROTOCOLS = io.netty.handler.codec.http.HttpResponseStatus.SWITCHING_PROTOCOLS.code();
 
     private WebSocket webSocket;
+    private boolean open;
     private final List<WebSocketListener> listeners;
+    private List<Runnable> bufferedFrames;
 
     public WebSocketUpgradeHandler(List<WebSocketListener> listeners) {
         this.listeners = listeners;
@@ -65,12 +67,33 @@ public final void onThrowable(Throwable t) {
         }
     }
 
-    public final void openWebSocket(WebSocket webSocket) {
+    public final void setWebSocket(WebSocket webSocket) {
         this.webSocket = webSocket;
+    }
+    
+    public final void onOpen() {
+        open = true;
         for (WebSocketListener listener : listeners) {
             webSocket.addWebSocketListener(listener);
             listener.onOpen(webSocket);
         }
+        if (bufferedFrames != null) {
+            for (Runnable bufferedFrame : bufferedFrames) {
+                bufferedFrame.run();
+            }
+            bufferedFrames = null;
+        }
+    }
+
+    public final boolean isOpen() {
+        return open;
+    }
+
+    public final void bufferFrame(Runnable bufferedFrame) {
+        if (bufferedFrames == null) {
+            bufferedFrames = new ArrayList<>(1);
+        }
+        bufferedFrames.add(bufferedFrame);
     }
 
     /**