Skip to content

Commit 509a020

Browse files
antfuantfu
committed
Merge branch 'main' into feat/auth
2 parents ebedf38 + 50c65f5 commit 509a020

File tree

6 files changed

+113
-32
lines changed

6 files changed

+113
-32
lines changed

‎packages/core/src/node/host-functions.ts‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
1-
import type { DevToolsNodeContext, DevToolsRpcClientFunctions, DevToolsRpcServerFunctions, RpcFunctionsHost as RpcFunctionsHostType } from '@vitejs/devtools-kit'
1+
import type { DevToolsNodeContext, DevToolsNodeRpcSession, DevToolsRpcClientFunctions, DevToolsRpcServerFunctions, RpcFunctionsHost as RpcFunctionsHostType } from '@vitejs/devtools-kit'
22
import type { BirpcGroup } from 'birpc'
3+
import type { AsyncLocalStorage } from 'node:async_hooks'
34
import { RpcFunctionsCollectorBase } from 'birpc-x'
45

56
export class RpcFunctionsHost extends RpcFunctionsCollectorBase<DevToolsRpcServerFunctions, DevToolsNodeContext> implements RpcFunctionsHostType {
67
/**
78
* @internal
89
*/
9-
rpcGroup: BirpcGroup<DevToolsRpcClientFunctions, DevToolsRpcServerFunctions, false> = undefined!
10+
_rpcGroup: BirpcGroup<DevToolsRpcClientFunctions, DevToolsRpcServerFunctions, false> = undefined!
11+
_asyncStorage: AsyncLocalStorage<DevToolsNodeRpcSession> = undefined!
1012

1113
constructor(context: DevToolsNodeContext) {
1214
super(context)
@@ -19,9 +21,15 @@ export class RpcFunctionsHost extends RpcFunctionsCollectorBase<DevToolsRpcServe
1921
name: T,
2022
...args: Args
2123
): Promise<(Awaited<ReturnType<DevToolsRpcClientFunctions[T]>> | undefined)[]> {
22-
if (!this.rpcGroup)
23-
throw new Error('RpcFunctionsHost.rpcGroup is not set, it likely to be an internal bug of Vite DevTools')
24+
if (!this._rpcGroup)
25+
throw new Error('RpcFunctionsHost] RpcGroup is not set, it likely to be an internal bug of Vite DevTools')
2426
// @ts-expect-error - BirpcGroup.broadcast.$callOptional is not typed correctly
25-
return this.rpcGroup.broadcast.$callOptional<T>(name, ...args)
27+
return this._rpcGroup.broadcast.$callOptional<T>(name, ...args)
28+
}
29+
30+
getCurrentRpcSession(): DevToolsNodeRpcSession | undefined {
31+
if (!this._asyncStorage)
32+
throw new Error('RpcFunctionsHost] AsyncLocalStorage is not set, it likely to be an internal bug of Vite DevTools')
33+
return this._asyncStorage.getStore()
2634
}
2735
}

‎packages/core/src/node/rpc/internal/docks-list.ts‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,14 @@ export const docksList = defineRpcFunction({
1818
]
1919

2020
return {
21-
handler: (): DevToolsDockEntry[] => [
22-
...Array.from(context.docks.values()),
23-
...builtinDocksEntries,
24-
],
21+
handler(): DevToolsDockEntry[] {
22+
// const session = context.rpc.getCurrentRpcSession()
23+
// console.log('session', session?.meta.id)
24+
return [
25+
...Array.from(context.docks.values()),
26+
...builtinDocksEntries,
27+
]
28+
},
2529
}
2630
},
2731
})

‎packages/core/src/node/ws.ts‎

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
/* eslint-disable no-console */
2-
import type { ConnectionMeta, DevToolsNodeContext, DevToolsRpcClientFunctions, DevToolsRpcServerFunctions } from '@vitejs/devtools-kit'
2+
import type { ConnectionMeta, DevToolsNodeContext, DevToolsNodeRpcSession, DevToolsRpcClientFunctions, DevToolsRpcServerFunctions } from '@vitejs/devtools-kit'
33
import type { WebSocket } from 'ws'
44
import type { RpcFunctionsHost } from './host-functions'
5+
import { AsyncLocalStorage } from 'node:async_hooks'
56
import { createRpcServer } from '@vitejs/devtools-rpc'
67
import { createWsRpcPreset } from '@vitejs/devtools-rpc/presets/ws/server'
78
import c from 'ansis'
@@ -25,16 +26,18 @@ export async function createWsServer(options: CreateWsServerOptions) {
2526

2627
const preset = createWsRpcPreset({
2728
port: port!,
28-
onConnected: (ws) => {
29+
onConnected: (ws, meta) => {
2930
wsClients.add(ws)
30-
console.log(c.green`${MARK_CHECK} Websocket client connected`)
31+
console.log(c.green`${MARK_CHECK} Websocket client [${meta.id}] connected`)
3132
},
32-
onDisconnected: (ws) => {
33+
onDisconnected: (ws, meta) => {
3334
wsClients.delete(ws)
34-
console.log(c.red`${MARK_CHECK} Websocket client disconnected`)
35+
console.log(c.red`${MARK_CHECK} Websocket client [${meta.id}] disconnected`)
3536
},
3637
})
3738

39+
const asyncStorage = new AsyncLocalStorage<DevToolsNodeRpcSession>()
40+
3841
const rpcGroup = createRpcServer<DevToolsRpcClientFunctions, DevToolsRpcServerFunctions>(
3942
rpcHost.functions,
4043
{
@@ -49,21 +52,35 @@ export async function createWsServer(options: CreateWsServerOptions) {
4952
console.error(error)
5053
},
5154
resolver(name, fn) {
52-
if (name.startsWith(ANONYMOUS_SCOPE))
53-
return fn
54-
55-
if (!this.$meta.isTrusted) {
55+
// Block unauthorized access to non-anonymous methods
56+
if (!name.startsWith(ANONYMOUS_SCOPE) && !this.$meta.isTrusted) {
5657
return () => {
5758
throw new Error(`Unauthorized access to method ${JSON.stringify(name)}, please trust this client first`)
5859
}
5960
}
60-
return fn
61+
62+
// If the function is not found, return undefined
63+
if (!fn)
64+
return undefined
65+
66+
// Register AsyncContext for the current RPC call
67+
// eslint-disable-next-line ts/no-this-alias
68+
const rpc = this
69+
return async function (this: any, ...args) {
70+
return await asyncStorage.run({
71+
rpc,
72+
meta: rpc.$meta,
73+
}, async () => {
74+
return (await fn).apply(this, args)
75+
})
76+
}
6177
},
6278
},
6379
},
6480
)
6581

66-
rpcHost.rpcGroup = rpcGroup
82+
rpcHost._rpcGroup = rpcGroup
83+
rpcHost._asyncStorage = asyncStorage
6784

6885
const getConnectionMeta = async (): Promise<ConnectionMeta> => {
6986
return {

‎packages/kit/src/types/rpc.ts‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,37 @@
1+
import type { BirpcReturn } from 'birpc'
12
import type { RpcFunctionsCollectorBase } from 'birpc-x'
3+
import type { WebSocket } from 'ws'
24
import type { DevToolsRpcClientFunctions, DevToolsRpcServerFunctions } from './rpc-augments'
35
import type { DevToolsNodeContext } from './vite-plugin'
46

7+
export interface DevToolsNodeRpcSessionMeta {
8+
id: number
9+
ws?: WebSocket
10+
clientAuthId?: string
11+
isTrusted?: boolean
12+
}
13+
14+
export interface DevToolsNodeRpcSession {
15+
meta: DevToolsNodeRpcSessionMeta
16+
rpc: BirpcReturn<DevToolsRpcClientFunctions, DevToolsRpcServerFunctions, false>
17+
}
18+
519
export type RpcFunctionsHost = RpcFunctionsCollectorBase<DevToolsRpcServerFunctions, DevToolsNodeContext> & {
20+
/**
21+
* Broadcast a message to all connected clients
22+
*/
623
boardcast: <
724
T extends keyof DevToolsRpcClientFunctions,
825
Args extends Parameters<DevToolsRpcClientFunctions[T]>,
926
>(
1027
name: T,
1128
...args: Args
1229
) => Promise<(Awaited<ReturnType<DevToolsRpcClientFunctions[T]>> | undefined)[]>
30+
31+
/**
32+
* Get the current RPC client
33+
*
34+
* Available in RPC functions to get the current RPC client
35+
*/
36+
getCurrentRpcSession: () => DevToolsNodeRpcSession | undefined
1337
}

‎packages/kit/src/types/vite-plugin.ts‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,43 @@ export interface DevToolsPluginOptions {
1919
}
2020

2121
export interface DevToolsNodeContext {
22+
/**
23+
* Current working directory of Vite DevTools
24+
*/
2225
readonly cwd: string
26+
/**
27+
* Current mode of Vite DevTools
28+
* - 'dev' - when Vite DevTools is running in dev mode
29+
* - 'build' - when Vite DevTools is running in build mode (no server)
30+
*/
2331
readonly mode: 'dev' | 'build'
32+
/**
33+
* Resolved Vite configuration
34+
*/
2435
readonly viteConfig: ResolvedConfig
36+
/**
37+
* Vite dev server instance (only available in dev mode)
38+
*/
2539
readonly viteServer?: ViteDevServer
40+
/**
41+
* RPC functions host, for registering server-side RPC functions and calling client-side RPC functions
42+
*/
2643
rpc: RpcFunctionsHost
44+
/**
45+
* Docks host, for registering dock entries
46+
*/
2747
docks: DevToolsDockHost
48+
/**
49+
* Views host, for registering static views
50+
*/
2851
views: DevToolsViewHost
52+
/**
53+
* Utils for the node context
54+
*/
2955
utils: DevToolsNodeUtils
56+
/**
57+
* Terminals host, for registering terminal sessions and streaming terminal output
58+
*/
3059
terminals: DevToolsTerminalHost
3160
}
3261

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

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { DevToolsNodeRpcSessionMeta } from '@vitejs/devtools-kit'
12
import type { BirpcGroup, BirpcOptions, ChannelOptions } from 'birpc'
23
import type { WebSocket } from 'ws'
34
import type { RpcServerPreset } from '..'
@@ -7,14 +8,11 @@ import { defineRpcServerPreset } from '..'
78

89
export interface WebSocketRpcServerOptions {
910
port: number
10-
onConnected?: (ws: WebSocket) => void
11-
onDisconnected?: (ws: WebSocket) => void
11+
onConnected?: (ws: WebSocket, meta: DevToolsNodeRpcSessionMeta) => void
12+
onDisconnected?: (ws: WebSocket, meta: DevToolsNodeRpcSessionMeta) => void
1213
}
1314

14-
export interface WebSocketRpcServerChannelMeta {
15-
uid: string
16-
isTrusted?: boolean
17-
}
15+
let id = 0
1816

1917
function NOOP() {}
2018

@@ -48,6 +46,10 @@ export const createWsRpcPreset: RpcServerPreset<
4846
} = options ?? {}
4947

5048
wss.on('connection', (ws) => {
49+
const meta: DevToolsNodeRpcSessionMeta = {
50+
id: id++,
51+
ws,
52+
}
5153
const channel: ChannelOptions = {
5254
post: (data) => {
5355
ws.send(data)
@@ -59,10 +61,7 @@ export const createWsRpcPreset: RpcServerPreset<
5961
},
6062
serialize,
6163
deserialize,
62-
meta: <WebSocketRpcServerChannelMeta>{
63-
uid: crypto.randomUUID(),
64-
isTrusted: false,
65-
},
64+
meta,
6665
}
6766

6867
rpc.updateChannels((channels) => {
@@ -75,9 +74,9 @@ export const createWsRpcPreset: RpcServerPreset<
7574
if (index >= 0)
7675
channels.splice(index, 1)
7776
})
78-
onDisconnected(ws)
77+
onDisconnected(ws, meta)
7978
})
80-
onConnected(ws)
79+
onConnected(ws, meta)
8180
})
8281
}
8382
})

0 commit comments

Comments
 (0)