ws.ts

/* eslint-disable no-console */
import type { ConnectionMeta, DevToolsNodeRpcSession, DevToolsRpcClientFunctions, DevToolsRpcServerFunctions, ViteDevToolsNodeContext } from '@vitejs/devtools-kit'
import type { RpcFunctionsHost } from 'devframe/node'
import type { WebSocket } from 'ws'
import { AsyncLocalStorage } from 'node:async_hooks'
import process from 'node:process'
import { getInternalContext } from 'devframe/node/internal'
import { createRpcServer } from 'devframe/rpc/server'
import { attachWsRpcTransport } from 'devframe/rpc/transports/ws-server'
import { colors as c } from 'devframe/utils/colors'
import { getPort } from 'get-port-please'
import { createDebug } from 'obug'
import { MARK_INFO } from './constants'
import { diagnostics } from './diagnostics'

const debugInvoked = createDebug('vite:devtools:rpc:invoked')

export interface CreateWsServerOptions {
  cwd: string
  websocket: {
    port?: number
    host: string
    https?: ViteDevToolsNodeContext['viteConfig']['server']['https'] | false
  }
  base?: string
  context: ViteDevToolsNodeContext
}

const ANONYMOUS_SCOPE = 'vite:anonymous:'

function buildWsUrl({ host, port, https }: { host: string, port: number, https: boolean }): string {
  // 0.0.0.0 / :: / unspecified bindings listen on all interfaces, but a hosted
  // page can't dial into an unspecified address — substitute localhost so the
  // descriptor we hand out is dialable.
  const reachableHost = host === '0.0.0.0' || host === '::' || host === '' ? 'localhost' : host
  const protocol = https ? 'wss' : 'ws'
  return `${protocol}://${reachableHost}:${port}`
}

export async function createWsServer(options: CreateWsServerOptions) {
  const rpcHost = options.context.rpc as unknown as RpcFunctionsHost
  const host = options.websocket.host
  const https = options.websocket.https === false ? undefined : (options.websocket.https ?? options.context.viteConfig.server.https)
  const port = options.websocket.port ?? await getPort({ port: 7812, host, random: true })!

  const wsClients = new Set<WebSocket>()

  const context = options.context
  const contextInternal = getInternalContext(context)

  const isClientAuthDisabled = context.mode === 'build' || context.viteConfig.devtools?.config?.clientAuth === false || process.env.VITE_DEVTOOLS_DISABLE_CLIENT_AUTH === 'true'
  if (isClientAuthDisabled) {
    diagnostics.DTK0008()
  }

  contextInternal.wsEndpoint = {
    url: buildWsUrl({ host, port, https: !!https }),
  }

  // `docks.register()` runs before the WS port is allocated, so any remote
  // docks registered during plugin setup were projected without a connection
  // descriptor. Now that the endpoint is resolved, re-emit their update events
  // so the shared-state consumers pick up the enriched URLs.
  for (const view of context.docks.views.values()) {
    if (view.type === 'iframe' && view.remote) {
      context.docks.events.emit('dock:entry:updated', view)
    }
  }

  const asyncStorage = new AsyncLocalStorage<DevToolsNodeRpcSession>()

  const rpcGroup = createRpcServer<DevToolsRpcClientFunctions, DevToolsRpcServerFunctions>(
    rpcHost.functions,
    {
      rpcOptions: {
        onFunctionError(error, name) {
          diagnostics.DTK0011({ name, cause: error })
        },
        onGeneralError(error) {
          diagnostics.DTK0012({ cause: error })
        },
        resolver(name, fn) {
          // eslint-disable-next-line ts/no-this-alias
          const rpc = this

          // Block unauthorized access to non-anonymous methods
          if (!name.startsWith(ANONYMOUS_SCOPE) && !rpc.$meta.isTrusted) {
            return () => {
              throw diagnostics.DTK0013({ name, clientId: rpc.$meta.id })
            }
          }

          // If the function is not found, return undefined
          if (!fn)
            return undefined

          // Register AsyncContext for the current RPC call
          return async function (this: any, ...args) {
            debugInvoked(`${JSON.stringify(name)} from #${rpc.$meta.id}`)
            return await asyncStorage.run({
              rpc,
              meta: rpc.$meta,
            }, async () => {
              return (await fn).apply(this, args)
            })
          }
        },
      },
    },
  )

  attachWsRpcTransport(rpcGroup, {
    port,
    host,
    https,
    definitions: rpcHost.definitions,
    onConnected: (ws, req, meta) => {
      const url = new URL(req.url ?? '', 'http://localhost')
      const authToken = url.searchParams.get('vite_devtools_auth_token') ?? undefined
      const requestOrigin = req.headers.origin
      if (isClientAuthDisabled) {
        meta.isTrusted = true
      }
      else if (authToken && contextInternal.isRemoteTokenTrusted(authToken, requestOrigin)) {
        meta.isTrusted = true
        meta.clientAuthToken = authToken
      }
      else if (authToken && contextInternal.storage.auth.value().trusted[authToken]) {
        meta.isTrusted = true
        meta.clientAuthToken = authToken
      }
      else if (authToken && (context.viteConfig.devtools?.config?.clientAuthTokens ?? []).includes(authToken)) {
        meta.isTrusted = true
        meta.clientAuthToken = authToken
      }

      wsClients.add(ws)
      const color = meta.isTrusted ? c.green : c.yellow
      console.log(color`${MARK_INFO} Websocket client connected. [${meta.id}] [${meta.clientAuthToken}] (${meta.isTrusted ? 'trusted' : 'untrusted'})`)
    },
    onDisconnected: (ws, meta) => {
      wsClients.delete(ws)
      rpcHost._emitSessionDisconnected(meta)
      console.log(c.red`${MARK_INFO} Websocket client disconnected. [${meta.id}]`)
    },
  })

  rpcHost._rpcGroup = rpcGroup
  rpcHost._asyncStorage = asyncStorage

  const getConnectionMeta = async (): Promise<ConnectionMeta> => {
    const jsonSerializableMethods: string[] = []
    for (const def of rpcHost.definitions.values()) {
      if (def.jsonSerializable === true)
        jsonSerializableMethods.push(def.name)
    }
    return {
      backend: 'websocket',
      websocket: port,
      jsonSerializableMethods,
    }
  }

  return {
    port,
    rpc: rpcGroup,
    rpcHost,
    getConnectionMeta,
  }
}