fix(rsc): client element as bound arg encryption (#905)

Author: hi-ogawa

packages/rsc/examples/basic/src/routes/action-bind/server.tsx

@@ -44,13 +44,18 @@ export function TestServerActionBindSimple() {
 let testServerActionBindClientState = 0;
 
 export function TestServerActionBindClient() {
+  // client element as server action bound argument
+  const client = <ActionBindClient />;
+
+  const action = async () => {
+    "use server";
+    return client;
+  };
+
   return (
     <TestServerActionBindClientForm
       key={testServerActionBindClientState}
-      action={async () => {
-        "use server";
-        return <ActionBindClient />;
-      }}
+      action={action}
     />
   );
 }

packages/rsc/src/core/rsc.ts

@@ -1,12 +1,16 @@
 import { memoize, tinyassert } from "@hiogawa/utils";
-import type { BundlerConfig, ImportManifestEntry } from "../types";
+import type { BundlerConfig, ImportManifestEntry, ModuleMap } from "../types";
 import {
+  SERVER_DECODE_CLIENT_PREFIX,
   SERVER_REFERENCE_PREFIX,
   createReferenceCacheTag,
   removeReferenceCacheTag,
   setInternalRequire,
 } from "./shared";
 
+// @ts-ignore
+import * as ReactServer from "react-server-dom-webpack/server.edge";
+
 let init = false;
 let requireModule!: (id: string) => unknown;
 
@@ -23,6 +27,24 @@ export function setRequireModule(options: {
   // need memoize to return stable promise from __webpack_require__
   (globalThis as any).__vite_rsc_server_require__ = memoize(requireModule);
 
+  (globalThis as any).__vite_rsc_server_decode_client__ = memoize(
+    async (raw: string) => {
+      // restore client reference on server for decoding.
+      // learned from https://github.com/lazarv/react-server/blob/79e7acebc6f4a8c930ad8422e2a4a9fdacfcce9b/packages/react-server/server/module-loader.mjs#L19
+      const { id, name } = JSON.parse(raw);
+      const reference = ReactServer.registerClientReference(
+        () => {
+          throw new Error(
+            `Unexpectedly client reference export '${name}' is called on server`,
+          );
+        },
+        removeReferenceCacheTag(id),
+        name,
+      );
+      return { [name]: reference };
+    },
+  );
+
   setInternalRequire();
 }
 
@@ -54,6 +76,29 @@ export function createServerManifest(): BundlerConfig {
   );
 }
 
+export function createServerDecodeClientManifest(): ModuleMap {
+  return new Proxy(
+    {},
+    {
+      get(_target, id: string) {
+        return new Proxy(
+          {},
+          {
+            get(_target, name: string) {
+              return {
+                id: SERVER_DECODE_CLIENT_PREFIX + JSON.stringify({ id, name }),
+                name,
+                chunks: [],
+                async: true,
+              };
+            },
+          },
+        );
+      },
+    },
+  );
+}
+
 export function createClientManifest(): BundlerConfig {
   const cacheTag = import.meta.env.DEV ? createReferenceCacheTag() : "";
 

packages/rsc/src/core/shared.ts

@@ -1,6 +1,8 @@
 // use special prefix to switch client/server reference loading inside __webpack_require__
 export const SERVER_REFERENCE_PREFIX = "$$server:";
 
+export const SERVER_DECODE_CLIENT_PREFIX = "$$decode-client:";
+
 // cache bust memoized require promise during dev
 export function createReferenceCacheTag(): string {
   const cache = Math.random().toString(36).slice(2);
@@ -14,6 +16,10 @@ export function removeReferenceCacheTag(id: string): string {
 export function setInternalRequire(): void {
   // branch client and server require to support the case when ssr and rsc share the same global
   (globalThis as any).__webpack_require__ = (id: string) => {
+    if (id.startsWith(SERVER_DECODE_CLIENT_PREFIX)) {
+      id = id.slice(SERVER_DECODE_CLIENT_PREFIX.length);
+      return (globalThis as any).__vite_rsc_server_decode_client__(id);
+    }
     if (id.startsWith(SERVER_REFERENCE_PREFIX)) {
       id = id.slice(SERVER_REFERENCE_PREFIX.length);
       return (globalThis as any).__vite_rsc_server_require__(id);

packages/rsc/src/react/rsc.ts

@@ -3,7 +3,11 @@ import type { ReactFormState } from "react-dom/client";
 import * as ReactClient from "react-server-dom-webpack/client.edge";
 // @ts-ignore
 import * as ReactServer from "react-server-dom-webpack/server.edge";
-import { createClientManifest, createServerManifest } from "../core/rsc";
+import {
+  createClientManifest,
+  createServerDecodeClientManifest,
+  createServerManifest,
+} from "../core/rsc";
 
 export { loadServerAction, setRequireModule } from "../core/rsc";
 
@@ -27,6 +31,7 @@ export function createFromReadableStream<T>(
       // https://github.com/facebook/react/pull/31300
       // https://github.com/vercel/next.js/pull/71527
       serverModuleMap: createServerManifest(),
+      moduleMap: createServerDecodeClientManifest(),
     },
     ...options,
   });