Skip to content

Commit 223b164

Browse files
james-elicxjames-elicx
committed
fix(rsc): validate custom server directives
1 parent 142fb07 commit 223b164

4 files changed

Lines changed: 203 additions & 38 deletions

File tree

‎packages/plugin-rsc/src/plugins/server-function-directives.test.ts‎

Lines changed: 111 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import type { Rollup } from 'vite'
22
import { describe, expect, it, vi } from 'vitest'
33
import {
4-
SERVER_FUNCTION_DIRECTIVE_MARKER,
54
vitePluginServerFunctionDirectives,
65
type ServerFunctionDirective,
76
} from './server-function-directives'
@@ -75,6 +74,8 @@ function cacheDirective(
7574
directive: /^use cache(?:: .+)?$/,
7675
test: (code) => code.includes('use cache'),
7776
rejectNonAsyncFunction: true,
77+
clientError: ({ id, environment }) =>
78+
`inline use cache is not allowed in ${environment}: ${id}`,
7879
wrap: ({ value, directiveMatch, location }) =>
7980
`cache(${value}, ${JSON.stringify(directiveMatch[0])}, ${JSON.stringify(location)})`,
8081
...overrides,
@@ -90,10 +91,19 @@ export async function getData() {
9091
return 1;
9192
}
9293
`)
93-
expect(result?.code).toContain(SERVER_FUNCTION_DIRECTIVE_MARKER)
94-
expect(result?.code).toContain('cache($$hoist_')
95-
expect(result?.code).toContain('$$ReactServer.registerServerReference')
96-
expect(result?.code).toContain('/rsc-runtime.js')
94+
expect(result?.code).toMatchInlineSnapshot(`
95+
"/* __vite_rsc_server_function_directives__ */
96+
import * as $$ReactServer from "/rsc-runtime.js";
97+
98+
export const getData = /* #__PURE__ */ $$ReactServer.registerServerReference(cache($$hoist_e9c2205b6101_0_getData, "use cache", "inline"), "53eb073e2100", "$$hoist_e9c2205b6101_0_getData");
99+
100+
;export async function $$hoist_e9c2205b6101_0_getData() {
101+
"use cache";
102+
return 1;
103+
};
104+
/* #__PURE__ */ Object.defineProperty($$hoist_e9c2205b6101_0_getData, "name", { value: "getData" });
105+
"
106+
`)
97107
expect(
98108
manager.serverReferenceMetaMap['/src/example.ts']?.exportNames,
99109
).toEqual([expect.stringMatching(/^\$\$hoist_/)])
@@ -110,9 +120,22 @@ export async function outer(value) {
110120
};
111121
}
112122
`)
113-
expect(result?.code).toContain('encryptActionBoundArgs([value])')
114-
expect(result?.code).toContain('decryptActionBoundArgs($$encoded)')
115-
expect(result?.code).toContain('/encryption-runtime.js')
123+
expect(result?.code).toMatchInlineSnapshot(`
124+
"/* __vite_rsc_server_function_directives__ */
125+
import * as $$ReactServer from "/rsc-runtime.js";
126+
import * as __vite_rsc_encryption_runtime from "/encryption-runtime.js";
127+
128+
export async function outer(value) {
129+
return /* #__PURE__ */ $$ReactServer.registerServerReference((($$wrapped) => async ($$encoded, ...$$args) => $$wrapped(...await __vite_rsc_encryption_runtime.decryptActionBoundArgs($$encoded), ...$$args))(cache($$hoist_ab3ae7af371a_0_cached)), "53eb073e2100", "$$hoist_ab3ae7af371a_0_cached").bind(null, __vite_rsc_encryption_runtime.encryptActionBoundArgs([value]));
130+
}
131+
132+
;export async function $$hoist_ab3ae7af371a_0_cached(value) {
133+
"use cache";
134+
return value;
135+
};
136+
/* #__PURE__ */ Object.defineProperty($$hoist_ab3ae7af371a_0_cached, "name", { value: "cached" });
137+
"
138+
`)
116139
expect(wrap).toHaveBeenCalledWith(
117140
expect.objectContaining({ location: 'inline', hasBoundArgs: true }),
118141
)
@@ -131,8 +154,18 @@ export async function getData() { return 1 }
131154
export const metadata = { title: "test" };
132155
`)
133156
expect(expandExportAll).toHaveBeenCalledOnce()
134-
expect(result?.code).toContain('cache(getData, "use cache", "module")')
135-
expect(result?.code).not.toContain('cache(metadata')
157+
expect(result?.code).toMatchInlineSnapshot(`
158+
"/* __vite_rsc_server_function_directives__ */
159+
160+
161+
/* "use cache" */;
162+
async function getData() { return 1 }
163+
let metadata = { title: "test" };
164+
getData = /* #__PURE__ */ $$ReactServer.registerServerReference(cache(getData, "use cache", "module"), "53eb073e2100", "getData");
165+
export { getData };
166+
export { metadata };
167+
"
168+
`)
136169
expect(
137170
manager.serverReferenceMetaMap['/src/example.ts']?.exportNames,
138171
).toEqual(['getData'])
@@ -141,36 +174,63 @@ export const metadata = { title: "test" };
141174
)
142175
})
143176

144-
it.each([
145-
['client', '/browser-runtime.js'],
146-
['ssr', '/ssr-runtime.js'],
147-
] as const)('creates module proxies in %s', async (environment, runtime) => {
177+
it('creates module proxies in client', async () => {
178+
const { run } = createHarness([cacheDirective()])
179+
const result = await run(
180+
`"use cache"; export async function getData() { return 1 }`,
181+
'client',
182+
)
183+
expect(result?.code).toMatchInlineSnapshot(`
184+
"import * as $$ReactClient from "/browser-runtime.js";
185+
export const getData = /* #__PURE__ */ $$ReactClient.createServerReference("53eb073e2100#getData",$$ReactClient.callServer,undefined,undefined,"getData");
186+
"
187+
`)
188+
})
189+
190+
it('creates module proxies in SSR', async () => {
148191
const { run } = createHarness([cacheDirective()])
149192
const result = await run(
150193
`"use cache"; export async function getData() { return 1 }`,
151-
environment,
194+
'ssr',
152195
)
153-
expect(result?.code).toContain(runtime)
154-
expect(result?.code).toContain('$$ReactClient.createServerReference')
155-
expect(result?.code).toContain('#getData')
196+
expect(result?.code).toMatchInlineSnapshot(`
197+
"import * as $$ReactClient from "/ssr-runtime.js";
198+
export const getData = /* #__PURE__ */ $$ReactClient.createServerReference("53eb073e2100#getData",$$ReactClient.callServer,undefined,undefined,"getData");
199+
"
200+
`)
156201
})
157202

158-
it('uses clientError for non-RSC module boundaries', async () => {
203+
it('uses clientError for non-RSC inline directives', async () => {
159204
const { run } = createHarness([
160205
cacheDirective({
161206
clientError: ({ id, environment }) => `${environment}:${id}`,
162207
}),
163208
])
164209
await expect(
165-
run(`"use cache"; export async function getData() {}`, 'client'),
210+
run(`export async function getData() { "use cache" }`, 'client'),
166211
).rejects.toThrow('client:/src/example.ts')
167212
})
168213

169-
it('leaves non-server inline directives untouched', async () => {
214+
it.each(['client', 'ssr'] as const)(
215+
'rejects inline directives in %s when clientError is configured',
216+
async (environment) => {
217+
const { run } = createHarness([cacheDirective()])
218+
const code = `export async function getData() { "use cache" }`
219+
await expect(run(code, environment)).rejects.toThrow(
220+
`inline use cache is not allowed in ${environment}: /src/example.ts`,
221+
)
222+
},
223+
)
224+
225+
it('leaves non-server inline directives untouched without clientError', async () => {
170226
const { run } = createHarness([cacheDirective()])
227+
const { run: runWithoutError } = createHarness([
228+
cacheDirective({ clientError: undefined }),
229+
])
171230
const code = `export async function getData() { "use cache" }`
172-
await expect(run(code, 'client')).resolves.toBeUndefined()
173-
await expect(run(code, 'ssr')).resolves.toBeUndefined()
231+
await expect(run(code, 'client')).rejects.toThrow()
232+
await expect(runWithoutError(code, 'client')).resolves.toBeUndefined()
233+
await expect(runWithoutError(code, 'ssr')).resolves.toBeUndefined()
174234
})
175235

176236
it('wraps inline directives inside use-server modules without owning metadata', async () => {
@@ -187,8 +247,20 @@ export async function action() {
187247
return cached();
188248
}
189249
`)
190-
expect(result?.code).toContain('cache($$hoist_')
191-
expect(result?.code).not.toContain('$$ReactServer.registerServerReference')
250+
expect(result?.code).toMatchInlineSnapshot(`
251+
"/* __vite_rsc_server_function_directives__ */
252+
253+
254+
"use server";
255+
export async function action() {
256+
const cached = /* #__PURE__ */ cache($$hoist_bf311121ee97_0_cached, "use cache", "inline");
257+
return cached();
258+
}
259+
260+
;export async function $$hoist_bf311121ee97_0_cached() { "use cache"; return 1 };
261+
/* #__PURE__ */ Object.defineProperty($$hoist_bf311121ee97_0_cached, "name", { value: "cached" });
262+
"
263+
`)
192264
expect(manager.serverReferenceMetaMap['/src/example.ts']).toEqual({
193265
importId: '/src/example.ts',
194266
referenceKey: 'existing',
@@ -226,6 +298,20 @@ export async function action() {
226298
).rejects.toThrow('non async function')
227299
})
228300

301+
it.each(['this', 'super', 'arguments'] as const)(
302+
'rejects %s inside inline directive functions',
303+
async (expression) => {
304+
const { run } = createHarness([cacheDirective()])
305+
const code =
306+
expression === 'super'
307+
? `class Base { value() {} } class Test extends Base { static async value() { "use cache"; return super.value() } }`
308+
: `export async function getData() { "use cache"; return ${expression} }`
309+
await expect(run(code)).rejects.toThrow(
310+
`"use cache" functions cannot use ${JSON.stringify(expression)}`,
311+
)
312+
},
313+
)
314+
229315
it('respects source and id prefilters and clears stale metadata', async () => {
230316
const test = vi.fn(() => false)
231317
const filter = vi.fn(() => false)

‎packages/plugin-rsc/src/plugins/server-function-directives.ts‎

Lines changed: 54 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { exactRegex } from '@rolldown/pluginutils'
22
import type { Literal, Program } from 'estree'
3+
import { walk } from 'estree-walker'
34
import type { Plugin, ResolvedConfig, Rollup, ViteDevServer } from 'vite'
45
import { parseAstAsync } from 'vite'
56
import {
@@ -56,7 +57,7 @@ export type ServerFunctionDirective = {
5657
id: string
5758
meta: Parameters<TransformWrapExportFilter>[1]
5859
}) => boolean
59-
/** Creates the error shown when a module boundary is imported outside RSC. */
60+
/** Creates the error shown for inline directives outside RSC. */
6061
clientError?: (context: { id: string; environment: string }) => string
6162
}
6263

@@ -118,6 +119,39 @@ function findModuleDirective(
118119
}
119120
}
120121

122+
function findInlineDirective(
123+
ast: Program,
124+
directive: string | RegExp,
125+
): StringLiteral | undefined {
126+
let result: StringLiteral | undefined
127+
walk(ast, {
128+
enter(node) {
129+
if (
130+
result ||
131+
(node.type !== 'FunctionDeclaration' &&
132+
node.type !== 'FunctionExpression' &&
133+
node.type !== 'ArrowFunctionExpression') ||
134+
node.body.type !== 'BlockStatement'
135+
) {
136+
return
137+
}
138+
for (const statement of node.body.body) {
139+
if (
140+
statement.type === 'ExpressionStatement' &&
141+
statement.expression.type === 'Literal' &&
142+
isStringLiteral(statement.expression) &&
143+
matchDirective(statement.expression.value, directive)
144+
) {
145+
result = statement.expression
146+
this.skip()
147+
return
148+
}
149+
}
150+
},
151+
})
152+
return result
153+
}
154+
121155
export function vitePluginServerFunctionDirectives(options: Options): Plugin {
122156
const { definitions, manager } = options
123157
return {
@@ -154,6 +188,23 @@ export function vitePluginServerFunctionDirectives(options: Options): Plugin {
154188
}
155189

156190
if (!isServer) {
191+
for (const definition of active) {
192+
const inlineDirective = findInlineDirective(
193+
ast,
194+
definition.directive,
195+
)
196+
if (inlineDirective && definition.clientError) {
197+
throw Object.assign(
198+
new Error(
199+
definition.clientError({
200+
id,
201+
environment: this.environment.name,
202+
}),
203+
),
204+
{ pos: inlineDirective.start },
205+
)
206+
}
207+
}
157208
const matches = active.flatMap((definition) => {
158209
const moduleDirective = findModuleDirective(
159210
ast,
@@ -175,18 +226,7 @@ export function vitePluginServerFunctionDirectives(options: Options): Plugin {
175226
}
176227
const match = matches[0]
177228
if (!match) return
178-
const [definition, moduleDirective] = match
179-
if (definition.clientError) {
180-
throw Object.assign(
181-
new Error(
182-
definition.clientError({
183-
id,
184-
environment: this.environment.name,
185-
}),
186-
),
187-
{ pos: moduleDirective.start },
188-
)
189-
}
229+
const [, moduleDirective] = match
190230

191231
const result = transformDirectiveProxyExport(ast, {
192232
code,
@@ -288,6 +328,7 @@ export function vitePluginServerFunctionDirectives(options: Options): Plugin {
288328
},
289329
stableName: true,
290330
detectUseServerModule: false,
331+
rejectForbiddenExpressions: true,
291332
})
292333
if (!result.output.hasChanged()) continue
293334

‎packages/plugin-rsc/src/transforms/hoist.ts‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { tinyassert } from '@hiogawa/utils'
22
import type {
33
Program,
4+
BlockStatement,
45
Literal,
56
Node,
67
MemberExpression,
@@ -30,6 +31,7 @@ export function transformHoistInlineDirective(
3031
decode?: (value: string) => string
3132
noExport?: boolean
3233
stableName?: boolean
34+
rejectForbiddenExpressions?: boolean
3335
},
3436
): {
3537
output: MagicString
@@ -67,6 +69,9 @@ export function transformHoistInlineDirective(
6769
},
6870
)
6971
}
72+
if (options.rejectForbiddenExpressions) {
73+
validateForbiddenExpressions(node.body, match[0])
74+
}
7075

7176
const isObjectMethod =
7277
node.type === 'FunctionExpression' &&
@@ -206,6 +211,37 @@ export function transformHoistInlineDirective(
206211
return { output, names }
207212
}
208213

214+
function validateForbiddenExpressions(body: BlockStatement, directive: string) {
215+
walk(body, {
216+
enter(node) {
217+
if (
218+
node !== body &&
219+
(node.type === 'FunctionDeclaration' ||
220+
node.type === 'FunctionExpression')
221+
) {
222+
this.skip()
223+
return
224+
}
225+
const expression =
226+
node.type === 'ThisExpression'
227+
? 'this'
228+
: node.type === 'Super'
229+
? 'super'
230+
: node.type === 'Identifier' && node.name === 'arguments'
231+
? 'arguments'
232+
: undefined
233+
if (expression) {
234+
throw Object.assign(
235+
new Error(
236+
`${JSON.stringify(directive)} functions cannot use ${JSON.stringify(expression)}.`,
237+
),
238+
{ pos: node.start },
239+
)
240+
}
241+
},
242+
})
243+
}
244+
209245
const exactRegex = (s: string): RegExp =>
210246
new RegExp('^' + s.replace(/[-[\]{}()*+?.,\\^$|#\s]/g, '\\$&') + '$')
211247

‎packages/plugin-rsc/src/transforms/server-action.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ export function transformServerActionServer(
2828
stableName?: boolean
2929
preserveModuleDirective?: boolean
3030
detectUseServerModule?: boolean
31+
rejectForbiddenExpressions?: boolean
3132
},
3233
):
3334
| {
@@ -78,5 +79,6 @@ export function transformServerActionServer(
7879
encode: options.encode,
7980
decode: options.decode,
8081
stableName: options.stableName,
82+
rejectForbiddenExpressions: options.rejectForbiddenExpressions,
8183
})
8284
}

0 commit comments

Comments
 (0)