vite 源码 - 创建 ws 服务

wss 配置

在 vite 中，对于 ws 服务相关配置指的是 server.hmr 。这个字段接收两个类型的参数：

• 布尔值（当值为 false 时则禁用热更新）
• 对象

其中对象支持以下几种属性：

export interface HmrOptions {
  protocol?: string // 协议 （`ws` 或 `wss`）
  host?: string // 主机
  port?: number // 端口
  clientPort?: number // 客户端端口
  path?: string // 路径
  timeout?: number // 超时时间
  overlay?: boolean // 是否在浏览器上显示 vite 错误覆盖层
  server?: HttpServer // node 服务
}

创建 ws 服务

vite 中调用 createWebSocketServer 函数来创建 ws 服务。首先学习这个函数的定义。

function createWebSocketServer(
  server: HttpServer | null,
  config: ResolvedConfig,
  httpsOptions?: HttpsServerOptions,
): WebSocketServer

根据这个函数的定义，可以知道这个函数接收三个参数，分别是：

• server: 之前创建的 node 服务
• config: 处理好的配置对象
• httpsOptions: https 配置对象即相关文件

返回一个 WebSocketServer。

了解了这个函数定义之后，再结合之前 vite 的配置，可以推测出需要根据不同的配置返回不同的 WebSocketServer。

首先需要处理的是用户禁用 hmr 即配置为 false 的情况。

按理来说，对于这种情况，可以直接返回一个 null，但是这里并没有返回一个 null 而是返回了一个对象，这个对象中存在 WebSocketServer 的方法和属性，其中的方法值为 noop ，用来模拟 WebSocketServer。这样做的目的有几个：

• 避免判空，在其他地方需要调用 ws 服务时，就不用写 ifelse 判断。
• 符合类型安全

// 返回的对象
{
      [isWebSocketServer]: true,
      get clients() {
        return new Set<WebSocketClient>()
      },
      async close() {},
      on: noop as any as WebSocketServer['on'],
      off: noop as any as WebSocketServer['off'],
      setInvokeHandler: noop,
      handleInvoke: async () => ({
        error: {
          name: 'TransportError',
          message: 'handleInvoke not implemented',
          stack: new Error().stack,
        },
      }),
      listen: noop,
      send: noop,
}

再处理了禁用 hmr 的情况后，需要真正的创建 ws 服务了。

这里也分为两种情况。

在分析这两种情况之前，需要了解 HTTP 和 WebSocket 协议。

核心关系：WebSocket 通过 HTTP 协议"升级"而来

• WebSocket 连接的建立必须经过一次 HTTP 请求 ，称为 "WebSocket 握手（Handshake）" 。
• 客户端发送一个特殊的 HTTP GET 请求，包含特定的头部；
• 服务器如果支持 WebSocket，就返回 101 Switching Protocols 响应；
• 之后，TCP 连接从 HTTP 协议"升级"为 WebSocket 协议，双方开始全双工通信。

所以，要完成 hmr ，需要一个 HTTP 服务和 WebSocket 服务。

这两个服务在代码中对应两个变量：

• wsServer: 底层 HTTP 服务器 ，用于接收upgrade请求
• wss: WebSocket 连接管理器，处理具体 WebSocket 逻辑

wsServer 的来源就是 hmr 配置中的 server 字段，如果传入一个 server，vite 会复用这个 server。

存在 wsServer

直接使用该 wsServer ，监听 upgrade 事件。

let hmrBase = config.base
const hmrPath = hmr ? hmr.path : undefined
if (hmrPath) {
  hmrBase = path.posix.join(hmrBase, hmrPath)
}
hmrServerWsListener = (req, socket, head) => {
  // 安全校验
  const protocol = req.headers['sec-websocket-protocol']!
  const parsedUrl = new URL(`http://example.com${req.url!}`)
  if (
    [HMR_HEADER, 'vite-ping'].includes(protocol) &&
    parsedUrl.pathname === hmrBase
  ) {
    handleUpgrade(req, socket as Socket, head, protocol === 'vite-ping')
  }
}
wsServer.on('upgrade', hmrServerWsListener)

不存在 wsServer

由之前的内容，可以推断出这里的处理逻辑。

1. 创建 wsServer
2. 使用 wsServer 监听 upgrade 事件
3. 当触发 upgrade 事件， 执行 handleUpgrade 函数

// **HTTP 服务器的请求回调**，用于处理**所有非 WebSocket 的普通 HTTP 请求**。
// 总是返回 426 升级协议
const route = ((_, res) => {
  const statusCode = 426
  const body = STATUS_CODES[statusCode]
  if (!body)
    throw new Error(`No body text found for the ${statusCode} status code`)

  res.writeHead(statusCode, {
    'Content-Length': body.length,
    'Content-Type': 'text/plain',
  })
  res.end(body)
}) as Parameters<typeof createHttpServer>[1]
if (httpsOptions) {
  wsHttpServer = createHttpsServer(httpsOptions, route)
} else {
  wsHttpServer = createHttpServer(route)
}
wsHttpServer.on('upgrade', (req, socket, head) => {
  const protocol = req.headers['sec-websocket-protocol']!
  // 防止服务未就绪时客户端连接
  if (protocol === 'vite-ping' && server && !server.listening) {
    req.destroy()
    return
  }
  handleUpgrade(req, socket as Socket, head, protocol === 'vite-ping')
})

handleUpgrade

这里的逻辑很简单，主要就是建立 websocket 连接。

const handleUpgrade = (
    req: IncomingMessage,
    socket: Duplex,
    head: Buffer,
    isPing: boolean,
) => {
    wss.handleUpgrade(req, socket as Socket, head, (ws) => {
      if (isPing) {
        ws.close(/* Normal Closure */ 1000)
        return
      }
      wss.emit('connection', ws, req)
    })
}

在 handleUpgrade 函数中，手动触发了 connection 事件，那么就需要实现 connection 事件的回调函数。

wss.on('connection', (socket) => {
    socket.on('message', (raw) => {
      if (!customListeners.size) return
      let parsed: any
      try {
        parsed = JSON.parse(String(raw))
      } catch {}
      if (!parsed || parsed.type !== 'custom' || !parsed.event) return
      const listeners = customListeners.get(parsed.event)
      if (!listeners?.size) return
      const client = getSocketClient(socket)
      listeners.forEach((listener) =>
        listener(parsed.data, client, parsed.invoke),
      )
    })
    socket.on('error', (err) => {
      // 错误处理，日志打印
    })
    socket.send(JSON.stringify({ type: 'connected' }))
    // bufferedError 处理
  })

从代码中可以看出，如果忽略对参数的校验和处理，核心逻辑其实是根据 parsed.event 从 customListeners 中取得所有的监听函数，并全部执行。那么该如何添加监听函数呢？

在最后的部分，vite 对 ws 的功能进行的封装和增强。通过一个对象放回，通过调用这些函数，可以添加监听函数。

// normalizeHotChannel 对**底层 HMR（热模块替换）通信通道（`HotChannel`）进行标准化封装**
const normalizedHotChannel = normalizeHotChannel(
    {
      send(payload) {
        if (payload.type === 'error' && !wss.clients.size) {
          bufferedError = payload
          return
        }

        const stringified = JSON.stringify(payload)
        wss.clients.forEach((client) => {
          // readyState 1 means the connection is open
          if (client.readyState === 1) {
            client.send(stringified)
          }
        })
      },
      on(event: string, fn: any) {
        if (!customListeners.has(event)) {
          customListeners.set(event, new Set())
        }
        customListeners.get(event)!.add(fn)
      },
      off(event: string, fn: any) {
        customListeners.get(event)?.delete(fn)
      },
      listen() {
        wsHttpServer?.listen(port, host)
      },
      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<void>((resolve, reject) => {
          wss.clients.forEach((client) => {
            client.terminate()
          })
          wss.close((err) => {
            if (err) {
              reject(err)
            } else {
              if (wsHttpServer) {
                wsHttpServer.close((err) => {
                  if (err) {
                    reject(err)
                  } else {
                    resolve()
                  }
                })
              } else {
                resolve()
              }
            }
          })
        })
      },
    },
    config.server.hmr !== false,
    false,
  )
  return {
    ...normalizedHotChannel,

    on: ((event: string, fn: any) => {
      if (wsServerEvents.includes(event)) {
        wss.on(event, fn)
        return
      }
      normalizedHotChannel.on(event, fn)
    }) as WebSocketServer['on'],
    off: ((event: string, fn: any) => {
      if (wsServerEvents.includes(event)) {
        wss.off(event, fn)
        return
      }
      normalizedHotChannel.off(event, fn)
    }) as WebSocketServer['off'],
    async close() {
      await normalizedHotChannel.close()
    },

    [isWebSocketServer]: true,
    get clients() {
      return new Set(Array.from(wss.clients).map(getSocketClient))
    },
  }