From 441642e848b7f9c8435a76b5b8b33988210fe9c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=BF=A0=20/=20green?= 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, + socket: Duplex, + head: Buffer, + ) => void const customListeners = new Map>>() const clientsMap = new WeakMap() 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()