feat: allow custom websocket server (#2338)

Author: egoist

packages/vite/src/node/server/hmr.ts

@@ -9,6 +9,7 @@ import { CLIENT_DIR } from '../constants'
 import { RollupError } from 'rollup'
 import { prepareError } from './middlewares/error'
 import match from 'minimatch'
+import { Server } from 'http'
 
 export const debugHmr = createDebugger('vite:hmr')
 
@@ -21,6 +22,7 @@ export interface HmrOptions {
   path?: string
   timeout?: number
   overlay?: boolean
+  server?: Server
 }
 
 export interface HmrContext {

packages/vite/src/node/server/ws.ts

@@ -17,9 +17,12 @@ export function createWebSocketServer(
 ): WebSocketServer {
   let wss: WebSocket.Server
 
-  if (server) {
+  const hmr = typeof config.server.hmr === 'object' && config.server.hmr
+  const wsServer = (hmr && hmr.server) || server
+
+  if (wsServer) {
     wss = new WebSocket.Server({ noServer: true })
-    server.on('upgrade', (req, socket, head) => {
+    wsServer.on('upgrade', (req, socket, head) => {
       if (req.headers['sec-websocket-protocol'] === HMR_HEADER) {
         wss.handleUpgrade(req, socket, head, (ws) => {
           wss.emit('connection', ws, req)
@@ -29,9 +32,7 @@ export function createWebSocketServer(
   } else {
     // vite dev server in middleware mode
     wss = new WebSocket.Server({
-      port:
-        (typeof config.server.hmr === 'object' && config.server.hmr.port) ||
-        24678
+      port: (hmr && hmr.port) || 24678
     })
   }