From 441642e848b7f9c8435a76b5b8b33988210fe9c6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=BF=A0=20/=20green?= <green@sapphi.red>
Date: Thu, 17 Aug 2023 03:04:06 +0900
Subject: [PATCH] fix: ws never connects after restarting server if
 server.hmr.server is set (#14127)

---
 packages/vite/src/node/server/ws.ts | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/packages/vite/src/node/server/ws.ts b/packages/vite/src/node/server/ws.ts
index 4a298cac1eebd2..7df59370f9e1d8 100644
--- a/packages/vite/src/node/server/ws.ts
+++ b/packages/vite/src/node/server/ws.ts
@@ -1,9 +1,10 @@
 import path from 'node:path'
-import type { Server } from 'node:http'
+import type { IncomingMessage, Server } from 'node:http'
 import { STATUS_CODES, createServer as createHttpServer } from 'node:http'
 import type { ServerOptions as HttpsServerOptions } from 'node:https'
 import { createServer as createHttpsServer } from 'node:https'
 import type { Socket } from 'node:net'
+import type { Duplex } from 'node:stream'
 import colors from 'picocolors'
 import type { WebSocket as WebSocketRaw } from 'ws'
 import { WebSocketServer as WebSocketServerRaw_ } from 'ws'
@@ -104,6 +105,11 @@ export function createWebSocketServer(
   // TODO: the main server port may not have been chosen yet as it may use the next available
   const portsAreCompatible = !hmrPort || hmrPort === config.server.port
   const wsServer = hmrServer || (portsAreCompatible && server)
+  let hmrServerWsListener: (
+    req: InstanceType<typeof IncomingMessage>,
+    socket: Duplex,
+    head: Buffer,
+  ) => void
   const customListeners = new Map<string, Set<WebSocketCustomListener<any>>>()
   const clientsMap = new WeakMap<WebSocketRaw, WebSocketClient>()
   const port = hmrPort || 24678
@@ -116,7 +122,7 @@ export function createWebSocketServer(
       hmrBase = path.posix.join(hmrBase, hmrPath)
     }
     wss = new WebSocketServerRaw({ noServer: true })
-    wsServer.on('upgrade', (req, socket, head) => {
+    hmrServerWsListener = (req, socket, head) => {
       if (
         req.headers['sec-websocket-protocol'] === HMR_HEADER &&
         req.url === hmrBase
@@ -125,7 +131,8 @@ export function createWebSocketServer(
           wss.emit('connection', ws, req)
         })
       }
-    })
+    }
+    wsServer.on('upgrade', hmrServerWsListener)
   } else {
     // http server request handler keeps the same with
     // https://github.com/websockets/ws/blob/45e17acea791d865df6b255a55182e9c42e5877a/lib/websocket-server.js#L88-L96
@@ -273,6 +280,11 @@ export function createWebSocketServer(
     },
 
     close() {
+      // should remove listener if hmr.server is set
+      // otherwise the old listener swallows all WebSocket connections
+      if (hmrServerWsListener && wsServer) {
+        wsServer.off('upgrade', hmrServerWsListener)
+      }
       return new Promise((resolve, reject) => {
         wss.clients.forEach((client) => {
           client.terminate()