From 4c27f1c722323c2c24026cbbfe49e910272c3a96 Mon Sep 17 00:00:00 2001
From: David Berlin <47757213+davbree@users.noreply.github.com>
Date: Wed, 29 May 2024 14:52:33 +0300
Subject: [PATCH] fix: websocket redirects (#6672)

* fix: fix websocket redirects
---
 src/utils/proxy.ts | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/utils/proxy.ts b/src/utils/proxy.ts
index 4e14b239064..85b0476c228 100644
--- a/src/utils/proxy.ts
+++ b/src/utils/proxy.ts
@@ -661,7 +661,7 @@ const initializeProxy = async function ({
       return proxy.web(req, res, options)
     },
     // @ts-expect-error TS(7006) FIXME: Parameter 'req' implicitly has an 'any' type.
-    ws: (req, socket, head) => proxy.ws(req, socket, head),
+    ws: (req, socket, head, options) => proxy.ws(req, socket, head, options),
   }
 
   return handlers
@@ -876,8 +876,15 @@ export const startProxy = async function ({
   const primaryServer = settings.https
     ? https.createServer({ cert: settings.https.cert, key: settings.https.key }, onRequestWithOptions)
     : http.createServer(onRequestWithOptions)
-  const onUpgrade = function onUpgrade(req: http.IncomingMessage, socket: Duplex, head: Buffer) {
-    proxy.ws(req, socket, head)
+  const onUpgrade = async function onUpgrade(req: http.IncomingMessage, socket: Duplex, head: Buffer) {
+    const match = await rewriter(req)
+    if (match && !match.force404 && isExternal(match)) {
+      const reqUrl = reqToURL(req, req.url)
+      const dest = new URL(match.to, `${reqUrl.protocol}//${reqUrl.host}`)
+      const destURL = stripOrigin(dest)
+      return proxy.ws(req, socket, head, { target: dest.origin, changeOrigin: true, pathRewrite: () => destURL })
+    }
+    return proxy.ws(req, socket, head, {})
   }
 
   primaryServer.on('upgrade', onUpgrade)