feat(core): do ws auth eagerly

Author: antfu

packages/core/src/node/rpc/anonymous/auth.ts

@@ -1,4 +1,3 @@
-import process from 'node:process'
 import * as p from '@clack/prompts'
 import { defineRpcFunction } from '@vitejs/devtools-kit'
 import c from 'ansis'
@@ -20,19 +19,13 @@ export const anonymousAuth = defineRpcFunction({
   setup: (context) => {
     const internal = getInternalContext(context)
     const storage = internal.storage.auth
-    const isClientAuthDisabled = context.mode === 'build' || context.viteConfig.devtools?.clientAuth === false || process.env.VITE_DEVTOOLS_DISABLE_CLIENT_AUTH === 'true'
-
-    if (isClientAuthDisabled) {
-      console.warn('[Vite DevTools] Client authentication is disabled. Any browser can connect to the devtools and access to your server and filesystem.')
-    }
-
     return {
       handler: async (query: DevToolsAuthInput): Promise<DevToolsAuthReturn> => {
         const session = context.rpc.getCurrentRpcSession()
         if (!session)
           throw new Error('Failed to retrieve the current RPC session')
 
-        if (isClientAuthDisabled || storage.get().trusted[query.authId]) {
+        if (session.meta.isTrusted || storage.get().trusted[query.authId]) {
           session.meta.clientAuthId = query.authId
           session.meta.isTrusted = true
           return {

packages/core/src/node/ws.ts

@@ -3,7 +3,6 @@ import type { ConnectionMeta, DevToolsNodeContext, DevToolsNodeRpcSession, DevTo
 import type { WebSocket } from 'ws'
 import type { RpcFunctionsHost } from './host-functions'
 import { AsyncLocalStorage } from 'node:async_hooks'
-import process from 'node:process'
 import { createRpcServer } from '@vitejs/devtools-rpc'
 import { createWsRpcPreset } from '@vitejs/devtools-rpc/presets/ws/server'
 import c from 'ansis'
@@ -30,6 +29,7 @@ export async function createWsServer(options: CreateWsServerOptions) {
   const preset = createWsRpcPreset({
     port,
     host,
+    context: options.context,
     onConnected: (ws, meta) => {
       wsClients.add(ws)
       console.log(c.green`${MARK_CHECK} Websocket client [${meta.id}] connected`)
@@ -42,8 +42,6 @@ export async function createWsServer(options: CreateWsServerOptions) {
 
   const asyncStorage = new AsyncLocalStorage<DevToolsNodeRpcSession>()
 
-  const isClientAuthDisabled = options.context.mode === 'build' || options.context.viteConfig.devtools?.clientAuth === false || process.env.VITE_DEVTOOLS_DISABLE_CLIENT_AUTH === 'true'
-
   const rpcGroup = createRpcServer<DevToolsRpcClientFunctions, DevToolsRpcServerFunctions>(
     rpcHost.functions,
     {
@@ -62,7 +60,7 @@ export async function createWsServer(options: CreateWsServerOptions) {
           const rpc = this
 
           // Block unauthorized access to non-anonymous methods
-          if (!name.startsWith(ANONYMOUS_SCOPE) && !rpc.$meta.isTrusted && !isClientAuthDisabled) {
+          if (!name.startsWith(ANONYMOUS_SCOPE) && !rpc.$meta.isTrusted) {
             return () => {
               throw new Error(`Unauthorized access to method ${JSON.stringify(name)} from client [${rpc.$meta.id}]`)
             }

packages/kit/src/client/rpc.ts

@@ -172,6 +172,7 @@ export async function getDevToolsRpcClient(
     {
       preset: createWsRpcPreset({
         url,
+        authId,
         ...options.wsOptions,
       }),
       rpcOptions,

packages/rpc/src/presets/ws/client.ts

@@ -8,12 +8,17 @@ export interface WebSocketRpcClientOptions {
   onConnected?: (e: Event) => void
   onError?: (e: Error) => void
   onDisconnected?: (e: CloseEvent) => void
+  authId?: string
 }
 
 function NOOP() {}
 
 export const createWsRpcPreset: RpcClientPreset<(options: WebSocketRpcClientOptions) => ChannelOptions> = defineRpcClientPreset((options: WebSocketRpcClientOptions) => {
-  const ws = new WebSocket(options.url)
+  let url = options.url
+  if (options.authId) {
+    url = `${url}?vite_devtools_auth_id=${encodeURIComponent(options.authId)}`
+  }
+  const ws = new WebSocket(url)
   const {
     onConnected = NOOP,
     onError = NOOP,

packages/rpc/src/presets/ws/server.ts

@@ -1,14 +1,17 @@
-import type { DevToolsNodeRpcSessionMeta } from '@vitejs/devtools-kit'
+import type { DevToolsNodeContext, DevToolsNodeRpcSessionMeta } from '@vitejs/devtools-kit'
 import type { BirpcGroup, BirpcOptions, ChannelOptions } from 'birpc'
 import type { WebSocket } from 'ws'
 import type { RpcServerPreset } from '..'
+import process from 'node:process'
 import { parse, stringify } from 'structured-clone-es'
 import { WebSocketServer } from 'ws'
 import { defineRpcServerPreset } from '..'
+import { getInternalContext } from '../../../../core/src/node/context-internal'
 
 export interface WebSocketRpcServerOptions {
   port: number
   host?: string
+  context: DevToolsNodeContext
   onConnected?: (ws: WebSocket, meta: DevToolsNodeRpcSessionMeta) => void
   onDisconnected?: (ws: WebSocket, meta: DevToolsNodeRpcSessionMeta) => void
 }
@@ -32,27 +35,48 @@ export const createWsRpcPreset: RpcServerPreset<
     host = 'localhost',
     onConnected = NOOP,
     onDisconnected = NOOP,
+    context,
   } = options
 
+  const isClientAuthDisabled = context.mode === 'build' || context.viteConfig.devtools?.clientAuth === false || process.env.VITE_DEVTOOLS_DISABLE_CLIENT_AUTH === 'true'
+  if (isClientAuthDisabled) {
+    console.warn('[Vite DevTools] Client authentication is disabled. Any browser can connect to the devtools and access to your server and filesystem.')
+  }
+
+  const internal = getInternalContext(context)
+
   const wss = new WebSocketServer({
     port,
     host,
   })
 
   return <ClientFunctions extends object, ServerFunctions extends object>(
-    rpc: BirpcGroup<ClientFunctions, ServerFunctions, false>,
+    rpcGroup: BirpcGroup<ClientFunctions, ServerFunctions, false>,
     options?: Pick<BirpcOptions<ClientFunctions, ServerFunctions, false>, 'serialize' | 'deserialize'>,
   ) => {
     const {
       serialize = stringify,
       deserialize = parse,
     } = options ?? {}
 
-    wss.on('connection', (ws) => {
+    wss.on('connection', (ws, req) => {
+      const url = new URL(req.url ?? '', 'http://localhost')
+      const authId = url.searchParams.get('vite_devtools_auth_id') ?? undefined
+      let isTrusted = false
+      if (isClientAuthDisabled) {
+        isTrusted = true
+      }
+      else if (authId && internal.storage.auth.get().trusted[authId]) {
+        isTrusted = true
+      }
+
       const meta: DevToolsNodeRpcSessionMeta = {
         id: id++,
         ws,
+        isTrusted,
+        clientAuthId: authId,
       }
+
       const channel: ChannelOptions = {
         post: (data) => {
           ws.send(data)
@@ -67,12 +91,14 @@ export const createWsRpcPreset: RpcServerPreset<
         meta,
       }
 
-      rpc.updateChannels((channels) => {
+      rpcGroup.updateChannels((channels) => {
         channels.push(channel)
       })
 
+      // const rpc = rpcGroup.clients.find(client => client.$meta.id === meta.id)
+
       ws.on('close', () => {
-        rpc.updateChannels((channels) => {
+        rpcGroup.updateChannels((channels) => {
           const index = channels.indexOf(channel)
           if (index >= 0)
             channels.splice(index, 1)