feat(rsc): expose directive function validation

packages/plugin-rsc/examples/custom-server-function/custom-server-function-plugin.ts

@@ -4,6 +4,7 @@ import {
   transformDirectiveProxyExport,
   transformHoistInlineDirective,
   transformModuleExportEffect,
+  validateDirectiveFunction,
 } from '@vitejs/plugin-rsc/transforms'
 import { parseAstAsync, type Plugin } from 'vite'
 
@@ -39,12 +40,17 @@ export function customServerFunctionPlugin(): Plugin {
         const result = hasDirective(ast.body, directive)
           ? transformModuleExportEffect(code, ast, {
               rejectNonAsyncFunction: true,
-              generate: ({ binding, exportName }) =>
-                runtime(binding, exportName),
+              generate: ({ binding, exportName, meta }) => {
+                validateDirectiveFunction(meta.valueNode, directive)
+                return runtime(binding, exportName)
+              },
             })
           : transformHoistInlineDirective(code, ast, {
               directive,
-              runtime,
+              runtime: (value, name, meta) => {
+                validateDirectiveFunction(meta.valueNode, directive)
+                return runtime(value, name)
+              },
               rejectNonAsyncFunction: true,
             })
         if (!result.output.hasChanged()) {

packages/plugin-rsc/examples/use-cache-callable/callable-cache-plugin.ts

@@ -5,6 +5,7 @@ import {
   transformDirectiveProxyExport,
   transformHoistInlineDirective,
   transformWrapExport,
+  validateDirectiveFunction,
 } from '@vitejs/plugin-rsc/transforms'
 import { parseAstAsync, type Plugin } from 'vite'
 import type { CacheWrapperOptions } from './src/framework/use-cache-runtime'
@@ -40,10 +41,13 @@ export function callableCachePlugin(): Plugin {
           `$$cacheWrapper(${value}, ${JSON.stringify(options)}),` +
           `${JSON.stringify(reference.referenceKey)},` +
           `${JSON.stringify(name)})`
+        const wrap = (value: string, name: string, meta: ModuleExportMeta) => {
+          validateDirectiveFunction(meta.valueNode, directive)
+          return runtime(value, name, getCacheWrapperOptions(meta))
+        }
         const result = hasDirective(ast.body, directive)
           ? transformWrapExport(code, ast, {
-              runtime: (value, name, meta) =>
-                runtime(value, name, getCacheWrapperOptions(meta)),
+              runtime: wrap,
               // Next.js calls rejecting primitive literals while permitting
               // objects and arrays arbitrary, but keeps the latter for metadata
               // and viewport exports.
@@ -57,8 +61,7 @@ export function callableCachePlugin(): Plugin {
               directive,
               rejectNonAsyncFunction: true,
               hoistRuntime: true,
-              runtime: (value, name, meta) =>
-                runtime(value, name, getCacheWrapperOptions(meta)),
+              runtime: wrap,
               encode: (value) => `$$encryptCacheCaptures(${value})`,
               // The cache runtime replaces the envelope with decoded captures
               // before invoking this private implementation.

packages/plugin-rsc/examples/use-cache/vite.config.ts

@@ -1,6 +1,9 @@
 import react from '@vitejs/plugin-react'
 import rsc from '@vitejs/plugin-rsc'
-import { transformHoistInlineDirective } from '@vitejs/plugin-rsc/transforms'
+import {
+  transformHoistInlineDirective,
+  validateDirectiveFunction,
+} from '@vitejs/plugin-rsc/transforms'
 import { defineConfig, parseAstAsync, type Plugin } from 'vite'
 
 export default defineConfig({
@@ -25,7 +28,10 @@ function vitePluginUseCache(): Plugin[] {
         if (!code.includes('use cache')) return
         const ast = await parseAstAsync(code)
         const result = transformHoistInlineDirective(code, ast, {
-          runtime: (value) => `__vite_rsc_cache(${value})`,
+          runtime: (value, _name, meta) => {
+            validateDirectiveFunction(meta.valueNode, 'use cache')
+            return `__vite_rsc_cache(${value})`
+          },
           directive: 'use cache',
           rejectNonAsyncFunction: true,
           noExport: true,

packages/plugin-rsc/src/transforms/directive-function.test.ts

@@ -0,0 +1,111 @@
+import type { Node, Program } from 'estree'
+import { parseAstAsync } from 'vite'
+import { describe, expect, test } from 'vitest'
+import { validateDirectiveFunction } from './directive-function'
+
+async function getFunction(input: string): Promise<Node> {
+  const ast = (await parseAstAsync(input)) as unknown as Program
+  const statement = ast.body[0]!
+  if (statement.type === 'FunctionDeclaration') return statement
+  if (statement.type === 'ClassDeclaration') {
+    const member = statement.body.body[0]!
+    if (member.type === 'MethodDefinition') return member.value
+  }
+  throw new Error('unsupported test input')
+}
+
+describe(validateDirectiveFunction, () => {
+  test.each([
+    [`async function action() { this.value }`, 'this'],
+    [`async function action() { arguments }`, 'arguments'],
+    [`async function action(value = arguments) {}`, 'arguments'],
+    [`class A extends B { static async action() { super.action() } }`, 'super'],
+  ])('rejects %s', async (input, expression) => {
+    const node = await getFunction(input)
+    expect(() => validateDirectiveFunction(node, 'use server')).toThrow(
+      `"use server" functions cannot use "${expression}".`,
+    )
+  })
+
+  test('follows arrow lexical bindings', async () => {
+    const node = await getFunction(`
+async function action() {
+  const nested = () => arguments
+}
+`)
+    expect(() => validateDirectiveFunction(node, 'use server')).toThrow(
+      /arguments/,
+    )
+  })
+
+  test('stops at nested normal functions', async () => {
+    const node = await getFunction(`
+async function action() {
+  function nested() {
+    this.value
+    arguments
+    const arrow = () => this.value
+  }
+}
+`)
+    expect(() => validateDirectiveFunction(node, 'use server')).not.toThrow()
+  })
+
+  test('ignores syntax-only arguments identifiers and jsxDEV metadata', async () => {
+    const node = await getFunction(`
+async function action() {
+  const value = { arguments: 1 }
+  return _jsxDEV(Component, {}, undefined, false, undefined, this)
+}
+`)
+    expect(() => validateDirectiveFunction(node, 'use server')).not.toThrow()
+  })
+
+  test('ignores unknown and non-function values', async () => {
+    const ast = (await parseAstAsync(
+      'export default action',
+    )) as unknown as Program
+    const statement = ast.body[0]!
+    expect(statement.type).toBe('ExportDefaultDeclaration')
+    if (statement.type !== 'ExportDefaultDeclaration') return
+    expect(() =>
+      validateDirectiveFunction(statement.declaration, 'use server'),
+    ).not.toThrow()
+    expect(() =>
+      validateDirectiveFunction(undefined, 'use server'),
+    ).not.toThrow()
+  })
+
+  test('ignores arguments labels', async () => {
+    const node = await getFunction(`
+async function action() {
+  arguments: while (true) {
+    break arguments
+  }
+}
+`)
+    expect(() => validateDirectiveFunction(node, 'use server')).not.toThrow()
+  })
+
+  test.each([
+    `custom_jsxDEV(this)`,
+    `_jsxDEV(this)`,
+    `_jsxDEV(Component, {}, undefined, false, this, undefined)`,
+  ])('does not exempt user this expressions: %s', async (expression) => {
+    const node = await getFunction(`async function action() { ${expression} }`)
+    expect(() => validateDirectiveFunction(node, 'use server')).toThrow(/this/)
+  })
+
+  test('stops at nested classes', async () => {
+    const node = await getFunction(`
+async function action() {
+  class Nested extends Base {
+    value = this.value
+    method() { return super.method() }
+    static { this.initialize() }
+  }
+}
+`)
+    expect(() => validateDirectiveFunction(node, 'use server')).not.toThrow()
+  })
+})

packages/plugin-rsc/src/transforms/directive-function.ts

@@ -0,0 +1,110 @@
+import type {
+  ArrowFunctionExpression,
+  FunctionDeclaration,
+  FunctionExpression,
+  Identifier,
+  Node,
+} from 'estree'
+import { walk } from 'estree-walker'
+
+export type DirectiveFunction =
+  | ArrowFunctionExpression
+  | FunctionDeclaration
+  | FunctionExpression
+
+/**
+ * Applies the function restrictions used by Next.js Server Functions.
+ * Unknown and non-function values are ignored so transform callbacks can pass
+ * `meta.valueNode` directly.
+ */
+export function validateDirectiveFunction(
+  valueNode: { type: string } | null | undefined,
+  directive: string,
+): void {
+  if (!valueNode) return
+  const node = valueNode as Node
+  if (!isFunction(node)) return
+
+  walk(node, {
+    enter(child, parent) {
+      if (
+        child !== node &&
+        (child.type === 'FunctionDeclaration' ||
+          child.type === 'FunctionExpression' ||
+          child.type === 'ClassDeclaration' ||
+          child.type === 'ClassExpression')
+      ) {
+        this.skip()
+        return
+      }
+
+      const expression =
+        child.type === 'ThisExpression' && !isJsxDevSourceThis(child, parent)
+          ? 'this'
+          : child.type === 'Super'
+            ? 'super'
+            : child.type === 'Identifier' &&
+                child.name === 'arguments' &&
+                isReferenceIdentifier(child, parent)
+              ? 'arguments'
+              : undefined
+      if (expression) {
+        throw Object.assign(
+          new Error(
+            `${JSON.stringify(directive)} functions cannot use ${JSON.stringify(expression)}.`,
+          ),
+          { pos: child.start },
+        )
+      }
+    },
+  })
+}
+
+function isFunction(node: Node): node is DirectiveFunction {
+  return (
+    node.type === 'ArrowFunctionExpression' ||
+    node.type === 'FunctionDeclaration' ||
+    node.type === 'FunctionExpression'
+  )
+}
+
+function isJsxDevSourceThis(node: Node, parent: Node | null): boolean {
+  return (
+    node.type === 'ThisExpression' &&
+    parent?.type === 'CallExpression' &&
+    parent.arguments.length === 6 &&
+    parent.arguments[5] === node &&
+    parent.callee.type === 'Identifier' &&
+    (parent.callee.name === 'jsxDEV' || parent.callee.name === '_jsxDEV')
+  )
+}
+
+function isReferenceIdentifier(node: Identifier, parent: Node | null): boolean {
+  if (!parent) return true
+  if (
+    parent.type === 'MemberExpression' &&
+    parent.property === node &&
+    !parent.computed
+  ) {
+    return false
+  }
+  if (
+    (parent.type === 'Property' ||
+      parent.type === 'MethodDefinition' ||
+      parent.type === 'PropertyDefinition') &&
+    parent.key === node &&
+    !parent.computed &&
+    !(parent.type === 'Property' && parent.shorthand)
+  ) {
+    return false
+  }
+  if (
+    (parent.type === 'LabeledStatement' ||
+      parent.type === 'BreakStatement' ||
+      parent.type === 'ContinueStatement') &&
+    parent.label === node
+  ) {
+    return false
+  }
+  return true
+}

packages/plugin-rsc/src/transforms/index.ts

@@ -6,3 +6,4 @@ export * from './proxy-export'
 export * from './utils'
 export * from './server-action'
 export * from './expand-export-all'
+export * from './directive-function'

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

@@ -30,3 +30,30 @@ test('normalizes inline server reference names', async () => {
     referenceNames: ['$$hoist_0_anonymous_server_function'],
   })
 })
+
+test.each([
+  `'use server'; export async function action() { arguments }`,
+  `export function App() { return async function action() { 'use server'; this.value } }`,
+])('validates server functions: %s', async (input) => {
+  await expect(transform(input)).rejects.toThrow(/cannot use/)
+})
+
+test('does not apply a file directive to non-exported functions', async () => {
+  await expect(
+    transform(`
+'use server'
+async function helper() { arguments }
+export async function action() {}
+`),
+  ).resolves.toBeDefined()
+})
+
+test('leaves unresolved local exports for local binding resolution', async () => {
+  await expect(
+    transform(`
+'use server'
+async function action() { arguments }
+export { action }
+`),
+  ).resolves.toBeDefined()
+})

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

@@ -1,5 +1,6 @@
 import type MagicString from 'magic-string'
 import type { ESTree } from 'vite'
+import { validateDirectiveFunction } from './directive-function'
 import { transformHoistInlineDirective } from './hoist'
 import { transformModuleExportEffect } from './module-export-effect'
 import { hasDirective } from './utils'
@@ -32,14 +33,20 @@ export function transformServerActionServer(
   if (hasDirective(ast.body, 'use server')) {
     const result = transformModuleExportEffect(input, ast, {
       rejectNonAsyncFunction: options.rejectNonAsyncFunction,
-      generate: ({ binding, exportName }) =>
-        options.runtime(binding, exportName),
+      generate: ({ binding, exportName, meta }) => {
+        validateDirectiveFunction(meta.valueNode, 'use server')
+        return options.runtime(binding, exportName)
+      },
     })
     return { ...result, exportNames: result.referenceNames }
   }
   const result = transformHoistInlineDirective(input, ast, {
     ...options,
     directive: 'use server',
+    runtime: (value, name, meta) => {
+      validateDirectiveFunction(meta.valueNode, 'use server')
+      return options.runtime(value, name)
+    },
   })
   return { ...result, referenceNames: result.names }
 }