From 839aab855844e23f0f0d394155baa2d0d63507ff Mon Sep 17 00:00:00 2001
From: Norman Maurer <norman_maurer@apple.com>
Date: Sat, 17 Aug 2019 10:00:01 +0200
Subject: [PATCH] Ensure we replace WebSocketServerProtocolHandshakeHandler
 before doing the handshake (#9472)

Motivation:

We need to ensure we replace WebSocketServerProtocolHandshakeHandler before doing the actual handshake as the handshake itself may complete directly and so forward pending bytes through the pipeline.

Modifications:

Replace the handler before doing the actual handshake.

Result:

Fixes https://github.com/netty/netty/issues/9471.
---
 ...bSocketServerProtocolHandshakeHandler.java | 12 +++++++---
 .../WebSocketServerProtocolHandlerTest.java   | 22 +++++++++++++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/codec-http/src/main/java/io/netty/handler/codec/http/websocketx/WebSocketServerProtocolHandshakeHandler.java b/codec-http/src/main/java/io/netty/handler/codec/http/websocketx/WebSocketServerProtocolHandshakeHandler.java
index 2362e60d1b..3c866eec78 100644
--- a/codec-http/src/main/java/io/netty/handler/codec/http/websocketx/WebSocketServerProtocolHandshakeHandler.java
+++ b/codec-http/src/main/java/io/netty/handler/codec/http/websocketx/WebSocketServerProtocolHandshakeHandler.java
@@ -88,6 +88,15 @@ class WebSocketServerProtocolHandshakeHandler extends ChannelInboundHandlerAdapt
             if (handshaker == null) {
                 WebSocketServerHandshakerFactory.sendUnsupportedVersionResponse(ctx.channel());
             } else {
+                // Ensure we set the handshaker and replace this handler before we
+                // trigger the actual handshake. Otherwise we may receive websocket bytes in this handler
+                // before we had a chance to replace it.
+                //
+                // See https://github.com/netty/netty/issues/9471.
+                WebSocketServerProtocolHandler.setHandshaker(ctx.channel(), handshaker);
+                ctx.pipeline().replace(this, "WS403Responder",
+                        WebSocketServerProtocolHandler.forbiddenHttpRequestResponder());
+
                 final ChannelFuture handshakeFuture = handshaker.handshake(ctx.channel(), req);
                 handshakeFuture.addListener(new ChannelFutureListener() {
                     @Override
@@ -107,9 +116,6 @@ class WebSocketServerProtocolHandshakeHandler extends ChannelInboundHandlerAdapt
                     }
                 });
                 applyHandshakeTimeout();
-                WebSocketServerProtocolHandler.setHandshaker(ctx.channel(), handshaker);
-                ctx.pipeline().replace(this, "WS403Responder",
-                        WebSocketServerProtocolHandler.forbiddenHttpRequestResponder());
             }
         } finally {
             req.release();
diff --git a/codec-http/src/test/java/io/netty/handler/codec/http/websocketx/WebSocketServerProtocolHandlerTest.java b/codec-http/src/test/java/io/netty/handler/codec/http/websocketx/WebSocketServerProtocolHandlerTest.java
index f6b748fba5..02c7d4046d 100644
--- a/codec-http/src/test/java/io/netty/handler/codec/http/websocketx/WebSocketServerProtocolHandlerTest.java
+++ b/codec-http/src/test/java/io/netty/handler/codec/http/websocketx/WebSocketServerProtocolHandlerTest.java
@@ -63,6 +63,28 @@ public class WebSocketServerProtocolHandlerTest {
         assertFalse(ch.finish());
     }
 
+    @Test
+    public void testWebSocketServerProtocolHandshakeHandlerReplacedBeforeHandshake() throws Exception {
+        EmbeddedChannel ch = createChannel(new MockOutboundHandler());
+        ChannelHandlerContext handshakerCtx = ch.pipeline().context(WebSocketServerProtocolHandshakeHandler.class);
+        ch.pipeline().addLast(new ChannelInboundHandlerAdapter() {
+            @Override
+            public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
+                if (evt instanceof WebSocketServerProtocolHandler.HandshakeComplete) {
+                    // We should have removed the handler already.
+                    assertNull(ctx.pipeline().context(WebSocketServerProtocolHandshakeHandler.class));
+                }
+            }
+        });
+        writeUpgradeRequest(ch);
+
+        FullHttpResponse response = responses.remove();
+        assertEquals(SWITCHING_PROTOCOLS, response.status());
+        response.release();
+        assertNotNull(WebSocketServerProtocolHandler.getHandshaker(handshakerCtx.channel()));
+        assertFalse(ch.finish());
+    }
+
     @Test
     public void testSubsequentHttpRequestsAfterUpgradeShouldReturn403() throws Exception {
         EmbeddedChannel ch = createChannel();