Function setupWebSocket

src/websocket.ts:122–207

(options: {
  server: ServerType
  fetchCallback: FetchCallback
  wss: WebSocketServerLike
})

}

export const setupWebSocket = (options: {
  server: ServerType
  fetchCallback: FetchCallback
  wss: WebSocketServerLike
}): void => {
  const { server, fetchCallback, wss } = options

  const waiterMap = new Map<
    IncomingMessage,
    { resolve: (ws: WebSocketLike) => void; connectionSymbol: symbol }
  >()

  wss.on('connection', (ws, request) => {
    const waiter = waiterMap.get(request)
    if (waiter) {
      waiter.resolve(ws)
      waiterMap.delete(request)
    }
  })

  const waitForWebSocket: WaitForWebSocket = (request, connectionSymbol) => {
    return new Promise<WebSocketLike>((resolve) => {
      waiterMap.set(request, { resolve, connectionSymbol })
    })
  }

  server.on('upgrade', async (request, socket: Duplex, head) => {
    if (request.headers.upgrade?.toLowerCase() !== 'websocket') {
      return
    }

    const env: UpgradeBindings = {
      incoming: request,
      outgoing: undefined,
      wss,
      [WAIT_FOR_WEBSOCKET_SYMBOL]: waitForWebSocket,
    }

    let status = 400
    let responseHeaders: Headers | undefined
    try {
      const response = (await fetchCallback(
        createUpgradeRequest(request),
        env as unknown as Parameters<FetchCallback>[1]
      )) as Response
      if (response instanceof Response) {
        status = response.status
        responseHeaders = response.headers
      }
    } catch {
      if (server.listenerCount('upgrade') === 1) {
        rejectUpgradeRequest(socket, 500)
      }
      return
    }

    const waiter = waiterMap.get(request)