From b487f71821e948ecff1a7093d114bf0a3b5f6102 Mon Sep 17 00:00:00 2001
From: Trustin Lee <t@motd.kr>
Date: Tue, 18 May 2021 18:42:38 +0900
Subject: [PATCH] Provide a way to pass through a certain HTTP upgrade request
 (#11267)

Motivation:

A user might want to handle a certain HTTP upgrade request differently
than what `HttpServerUpgradeHandler` does by default. For example, a
user could let `HttpServerUpgradeHandler` handle HTTP/2 upgrades but
not WebSocket upgrades.

Modifications:

- Added `HttpServerUpgradeHandler.isUpgrade(HttpRequest)` so a user can
  tell `HttpServerUpgradeHandler` to pass the request as it is to the
  next handler.

Result:

- A user can handle a certain upgrade request specially.
---
 .../codec/http/HttpServerUpgradeHandler.java  | 39 ++++++++++++----
 .../http/HttpServerUpgradeHandlerTest.java    | 44 +++++++++++++++++++
 2 files changed, 74 insertions(+), 9 deletions(-)

diff --git a/codec-http/src/main/java/io/netty/handler/codec/http/HttpServerUpgradeHandler.java b/codec-http/src/main/java/io/netty/handler/codec/http/HttpServerUpgradeHandler.java
index 9266a8ef77..800f8f435e 100644
--- a/codec-http/src/main/java/io/netty/handler/codec/http/HttpServerUpgradeHandler.java
+++ b/codec-http/src/main/java/io/netty/handler/codec/http/HttpServerUpgradeHandler.java
@@ -224,13 +224,24 @@ public class HttpServerUpgradeHandler extends HttpObjectAggregator {
     @Override
     protected void decode(final ChannelHandlerContext ctx, HttpObject msg)
             throws Exception {
-        // Determine if we're already handling an upgrade request or just starting a new one.
-        handlingUpgrade |= isUpgradeRequest(msg);
+
         if (!handlingUpgrade) {
-            // Not handling an upgrade request, just pass it to the next handler.
-            ReferenceCountUtil.retain(msg);
-            ctx.fireChannelRead(msg);
-            return;
+            // Not handling an upgrade request yet. Check if we received a new upgrade request.
+            if (msg instanceof HttpRequest) {
+                HttpRequest req = (HttpRequest) msg;
+                if (req.headers().contains(HttpHeaderNames.UPGRADE) &&
+                    shouldHandleUpgradeRequest(req)) {
+                    handlingUpgrade = true;
+                } else {
+                    ReferenceCountUtil.retain(msg);
+                    ctx.fireChannelRead(msg);
+                    return;
+                }
+            } else {
+                ReferenceCountUtil.retain(msg);
+                ctx.fireChannelRead(msg);
+                return;
+            }
         }
 
         FullHttpRequest fullRequest;
@@ -261,10 +272,20 @@ public class HttpServerUpgradeHandler extends HttpObjectAggregator {
     }
 
     /**
-     * Determines whether or not the message is an HTTP upgrade request.
+     * Determines whether the specified upgrade {@link HttpRequest} should be handled by this handler or not.
+     * This method will be invoked only when the request contains an {@code Upgrade} header.
+     * It always returns {@code true} by default, which means any request with an {@code Upgrade} header
+     * will be handled. You can override this method to ignore certain {@code Upgrade} headers, for example:
+     * <pre>{@code
+     * @Override
+     * protected boolean isUpgradeRequest(HttpRequest req) {
+     *   // Do not handle WebSocket upgrades.
+     *   return !req.headers().contains(HttpHeaderNames.UPGRADE, "websocket", false);
+     * }
+     * }</pre>
      */
-    private static boolean isUpgradeRequest(HttpObject msg) {
-        return msg instanceof HttpRequest && ((HttpRequest) msg).headers().get(HttpHeaderNames.UPGRADE) != null;
+    protected boolean shouldHandleUpgradeRequest(HttpRequest req) {
+        return true;
     }
 
     /**
diff --git a/codec-http/src/test/java/io/netty/handler/codec/http/HttpServerUpgradeHandlerTest.java b/codec-http/src/test/java/io/netty/handler/codec/http/HttpServerUpgradeHandlerTest.java
index a52c8c663d..da773f6f42 100644
--- a/codec-http/src/test/java/io/netty/handler/codec/http/HttpServerUpgradeHandlerTest.java
+++ b/codec-http/src/test/java/io/netty/handler/codec/http/HttpServerUpgradeHandlerTest.java
@@ -123,4 +123,48 @@ public class HttpServerUpgradeHandlerTest {
         assertTrue(upgradeMessage.release());
         assertFalse(channel.finishAndReleaseAll());
     }
+
+    @Test
+    public void skippedUpgrade() {
+        final HttpServerCodec httpServerCodec = new HttpServerCodec();
+        final UpgradeCodecFactory factory = new UpgradeCodecFactory() {
+            @Override
+            public UpgradeCodec newUpgradeCodec(CharSequence protocol) {
+                fail("Should never be invoked");
+                return null;
+            }
+        };
+
+        HttpServerUpgradeHandler upgradeHandler = new HttpServerUpgradeHandler(httpServerCodec, factory) {
+            @Override
+            protected boolean shouldHandleUpgradeRequest(HttpRequest req) {
+                return !req.headers().contains(HttpHeaderNames.UPGRADE, "do-not-upgrade", false);
+            }
+        };
+
+        EmbeddedChannel channel = new EmbeddedChannel(httpServerCodec, upgradeHandler);
+
+        String upgradeString = "GET / HTTP/1.1\r\n" +
+                               "Host: example.com\r\n" +
+                               "Connection: Upgrade\r\n" +
+                               "Upgrade: do-not-upgrade\r\n\r\n";
+        ByteBuf upgrade = Unpooled.copiedBuffer(upgradeString, CharsetUtil.US_ASCII);
+
+        // The upgrade request should not be passed to the next handler without any processing.
+        assertTrue(channel.writeInbound(upgrade));
+        assertNotNull(channel.pipeline().get(HttpServerCodec.class));
+        assertNull(channel.pipeline().get("marker"));
+
+        HttpRequest req = channel.readInbound();
+        assertFalse(req instanceof FullHttpRequest); // Should not be aggregated.
+        assertTrue(req.headers().contains(HttpHeaderNames.CONNECTION, "Upgrade", false));
+        assertTrue(req.headers().contains(HttpHeaderNames.UPGRADE, "do-not-upgrade", false));
+        assertTrue(channel.readInbound() instanceof LastHttpContent);
+        assertNull(channel.readInbound());
+
+        // No response should be written because we're just passing through.
+        channel.flushOutbound();
+        assertNull(channel.readOutbound());
+        assertFalse(channel.finishAndReleaseAll());
+    }
 }