Ensure upgrade request has request meta (#51590)

x-ref: [slack thread](https://vercel.slack.com/archives/C03S8ED1DKM/p1686906407341689)
vercel · Jun 21, 2023 · 9d043f6 · 9d043f6
1 parent 4a50743
commit 9d043f6
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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(?:\\/)?$'),