feat: allow custom websocket server (#2338)

vitejs · Mar 15, 2021 · 9243cc9 · 9243cc9
1 parent ee7eb82
commit 9243cc9
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/packages/vite/src/node/server/hmr.ts b/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 {

diff --git a/packages/vite/src/node/server/ws.ts b/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
  })
  }