From 9d043f6e359eb7f0bc94982c200aacc6e1fae3ba Mon Sep 17 00:00:00 2001
From: JJ Kasper <jj@jjsweb.site>
Date: Wed, 21 Jun 2023 14:52:12 -0400
Subject: [PATCH] Ensure upgrade request has request meta (#51590)

x-ref: [slack
thread](https://vercel.slack.com/archives/C03S8ED1DKM/p1686906407341689)
---
 packages/next/src/server/next-server.ts       | 26 +++++++++++++--
 test/integration/custom-routes/next.config.js |  4 +++
 .../custom-routes/test/index.test.js          | 32 +++++++++++++++++++
 3 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/packages/next/src/server/next-server.ts b/packages/next/src/server/next-server.ts
index 9a4d8d2d16646..d71d3c6bf30be 100644
--- a/packages/next/src/server/next-server.ts
+++ b/packages/next/src/server/next-server.ts
@@ -808,7 +808,14 @@ export default class NextNodeServer extends BaseServer {
   }
 
   protected async handleUpgrade(req: NodeNextRequest, socket: any, head: any) {
-    await this.router.execute(req, socket, nodeParseUrl(req.url, true), head)
+    try {
+      const parsedUrl = nodeParseUrl(req.url, true)
+      this.attachRequestMeta(req, parsedUrl, true)
+      await this.router.execute(req, socket, parsedUrl, head)
+    } catch (err) {
+      console.error(err)
+      socket.end('Internal Server Error')
+    }
   }
 
   protected async proxyRequest(
@@ -1486,6 +1493,15 @@ export default class NextNodeServer extends BaseServer {
               },
               getRequestMeta(req, '__NEXT_CLONABLE_BODY')?.cloneBodyStream()
             )
+
+            // if this is an upgrade request just pipe body back
+            if (!res.setHeader) {
+              invokeRes.pipe(res as any as ServerResponse)
+              return {
+                finished: true,
+              }
+            }
+
             const noFallback = invokeRes.headers['x-no-fallback']
 
             if (noFallback) {
@@ -2748,7 +2764,8 @@ export default class NextNodeServer extends BaseServer {
 
   protected attachRequestMeta(
     req: BaseNextRequest,
-    parsedUrl: NextUrlWithParsedQuery
+    parsedUrl: NextUrlWithParsedQuery,
+    isUpgradeReq?: boolean
   ) {
     const protocol = (
       (req as NodeNextRequest).originalRequest?.socket as TLSSocket
@@ -2767,7 +2784,10 @@ export default class NextNodeServer extends BaseServer {
     addRequestMeta(req, '__NEXT_INIT_URL', initUrl)
     addRequestMeta(req, '__NEXT_INIT_QUERY', { ...parsedUrl.query })
     addRequestMeta(req, '_protocol', protocol)
-    addRequestMeta(req, '__NEXT_CLONABLE_BODY', getCloneableBody(req.body))
+
+    if (!isUpgradeReq) {
+      addRequestMeta(req, '__NEXT_CLONABLE_BODY', getCloneableBody(req.body))
+    }
   }
 
   protected async runEdgeFunction(params: {
diff --git a/test/integration/custom-routes/next.config.js b/test/integration/custom-routes/next.config.js
index 58a057f9983be..c90d3b17e8164 100644
--- a/test/integration/custom-routes/next.config.js
+++ b/test/integration/custom-routes/next.config.js
@@ -19,6 +19,10 @@ module.exports = {
           destination:
             'http://localhost:__EXTERNAL_PORT__/_next/webpack-hmr?page=/about',
         },
+        {
+          source: '/websocket-to-page',
+          destination: '/hello',
+        },
         {
           source: '/to-nowhere',
           destination: 'http://localhost:12233',
diff --git a/test/integration/custom-routes/test/index.test.js b/test/integration/custom-routes/test/index.test.js
index e80e08ef55cdc..eda2ebfc65d66 100644
--- a/test/integration/custom-routes/test/index.test.js
+++ b/test/integration/custom-routes/test/index.test.js
@@ -103,6 +103,33 @@ const runTests = (isDev = false, isTurbo = false) => {
     expect([...externalServerHits]).toEqual(['/_next/webpack-hmr?page=/about'])
   })
 
+  it('should successfully rewrite a WebSocket request to a page', async () => {
+    // TODO: remove once test failure has been fixed
+    if (isTurbo) return
+
+    const messages = []
+    try {
+      const ws = await new Promise((resolve, reject) => {
+        let socket = new WebSocket(
+          `ws://localhost:${appPort}/websocket-to-page`
+        )
+        socket.on('message', (data) => {
+          messages.push(data.toString())
+        })
+        socket.on('open', () => resolve(socket))
+        socket.on('error', (err) => {
+          console.error(err)
+          socket.close()
+          reject()
+        })
+      })
+      ws.close()
+    } catch (err) {
+      messages.push(err)
+    }
+    expect(stderr).not.toContain('unhandledRejection')
+  })
+
   it('should not rewrite for _next/data route when a match is found', async () => {
     const initial = await fetchViaHTTP(appPort, '/overridden/first')
     expect(initial.status).toBe(200)
@@ -2210,6 +2237,11 @@ const runTests = (isDev = false, isTurbo = false) => {
               regex: normalizeRegEx('^\\/to-websocket(?:\\/)?$'),
               source: '/to-websocket',
             },
+            {
+              destination: '/hello',
+              regex: normalizeRegEx('^\\/websocket-to-page(?:\\/)?$'),
+              source: '/websocket-to-page',
+            },
             {
               destination: 'http://localhost:12233',
               regex: normalizeRegEx('^\\/to-nowhere(?:\\/)?$'),