From 959a49707d8e6915817045f932d9f458308d8a0c Mon Sep 17 00:00:00 2001
From: Hendrik Liebau <mail@hendrik-liebau.de>
Date: Wed, 23 Oct 2024 20:15:09 +0200
Subject: [PATCH] Fix race condition when setting client reference manifests

Similar to how #71669 fixed a race condition when setting the manifests
singleton regarding the server action manifests, this PR fixes the same
race condition for the client reference manifests.
---
 .../src/build/templates/edge-app-route.ts     |   1 +
 .../next/src/build/templates/edge-ssr-app.ts  |   1 +
 .../webpack/plugins/flight-manifest-plugin.ts |  19 ++--
 .../next/src/server/app-render/app-render.tsx |   1 +
 .../src/server/app-render/encryption-utils.ts | 107 +++++++++++++++---
 .../next/src/server/app-render/encryption.ts  |  17 ++-
 packages/next/src/server/load-components.ts   |   1 +
 .../src/server/use-cache/use-cache-wrapper.ts |  14 +--
 test/e2e/app-dir/actions/app-action.test.ts   |  10 ++
 .../app-dir/actions/app/encryption/client.js  |   5 +
 .../app-dir/actions/app/encryption/form.js    |  14 +++
 .../app-dir/actions/app/encryption/page.js    |  24 ++--
 .../app/with-server-action/layout.tsx         |  14 ---
 13 files changed, 162 insertions(+), 66 deletions(-)
 create mode 100644 test/e2e/app-dir/actions/app/encryption/client.js
 create mode 100644 test/e2e/app-dir/actions/app/encryption/form.js
 delete mode 100644 test/e2e/app-dir/use-cache/app/with-server-action/layout.tsx

diff --git a/packages/next/src/build/templates/edge-app-route.ts b/packages/next/src/build/templates/edge-app-route.ts
index b3514afe44c78..3401fdf7348fe 100644
--- a/packages/next/src/build/templates/edge-app-route.ts
+++ b/packages/next/src/build/templates/edge-app-route.ts
@@ -17,6 +17,7 @@ const rscServerManifest = maybeJSONParse(self.__RSC_SERVER_MANIFEST)
 
 if (rscManifest && rscServerManifest) {
   setReferenceManifestsSingleton({
+    page: 'VAR_PAGE',
     clientReferenceManifest: rscManifest,
     serverActionsManifest: rscServerManifest,
     serverModuleMap: createServerModuleMap({
diff --git a/packages/next/src/build/templates/edge-ssr-app.ts b/packages/next/src/build/templates/edge-ssr-app.ts
index 4eb8a04de9a45..c30af2f3cbec7 100644
--- a/packages/next/src/build/templates/edge-ssr-app.ts
+++ b/packages/next/src/build/templates/edge-ssr-app.ts
@@ -56,6 +56,7 @@ const interceptionRouteRewrites =
 
 if (rscManifest && rscServerManifest) {
   setReferenceManifestsSingleton({
+    page: 'VAR_PAGE',
     clientReferenceManifest: rscManifest,
     serverActionsManifest: rscServerManifest,
     serverModuleMap: createServerModuleMap({
diff --git a/packages/next/src/build/webpack/plugins/flight-manifest-plugin.ts b/packages/next/src/build/webpack/plugins/flight-manifest-plugin.ts
index 594bef23030a7..14bc8e920b52e 100644
--- a/packages/next/src/build/webpack/plugins/flight-manifest-plugin.ts
+++ b/packages/next/src/build/webpack/plugins/flight-manifest-plugin.ts
@@ -72,12 +72,21 @@ export interface ManifestNode {
   }
 }
 
-export type ClientReferenceManifest = {
+export interface ClientReferenceManifestForRsc {
+  clientModules: ManifestNode
+  rscModuleMapping: {
+    [moduleId: string]: ManifestNode
+  }
+  edgeRscModuleMapping: {
+    [moduleId: string]: ManifestNode
+  }
+}
+
+export interface ClientReferenceManifest extends ClientReferenceManifestForRsc {
   readonly moduleLoading: {
     prefix: string
     crossOrigin: string | null
   }
-  clientModules: ManifestNode
   ssrModuleMapping: {
     [moduleId: string]: ManifestNode
   }
@@ -90,12 +99,6 @@ export type ClientReferenceManifest = {
   entryJSFiles?: {
     [entry: string]: string[]
   }
-  rscModuleMapping: {
-    [moduleId: string]: ManifestNode
-  }
-  edgeRscModuleMapping: {
-    [moduleId: string]: ManifestNode
-  }
 }
 
 function getAppPathRequiredChunks(
diff --git a/packages/next/src/server/app-render/app-render.tsx b/packages/next/src/server/app-render/app-render.tsx
index 766c4d487da81..1dd03b858cfbf 100644
--- a/packages/next/src/server/app-render/app-render.tsx
+++ b/packages/next/src/server/app-render/app-render.tsx
@@ -1060,6 +1060,7 @@ async function renderToHTMLOrFlightImpl(
   const serverModuleMap = createServerModuleMap({ serverActionsManifest })
 
   setReferenceManifestsSingleton({
+    page: workStore.page,
     clientReferenceManifest,
     serverActionsManifest,
     serverModuleMap,
diff --git a/packages/next/src/server/app-render/encryption-utils.ts b/packages/next/src/server/app-render/encryption-utils.ts
index 7a59dbf0ace63..287138245afaa 100644
--- a/packages/next/src/server/app-render/encryption-utils.ts
+++ b/packages/next/src/server/app-render/encryption-utils.ts
@@ -1,6 +1,11 @@
 import type { ActionManifest } from '../../build/webpack/plugins/flight-client-entry-plugin'
-import type { ClientReferenceManifest } from '../../build/webpack/plugins/flight-manifest-plugin'
+import type {
+  ClientReferenceManifest,
+  ClientReferenceManifestForRsc,
+} from '../../build/webpack/plugins/flight-manifest-plugin'
 import type { DeepReadonly } from '../../shared/lib/deep-readonly'
+import { InvariantError } from '../../shared/lib/invariant-error'
+import { workAsyncStorage } from './work-async-storage.external'
 
 let __next_loaded_action_key: CryptoKey
 
@@ -64,10 +69,12 @@ const SERVER_ACTION_MANIFESTS_SINGLETON = Symbol.for(
 )
 
 export function setReferenceManifestsSingleton({
+  page,
   clientReferenceManifest,
   serverActionsManifest,
   serverModuleMap,
 }: {
+  page: string
   clientReferenceManifest: DeepReadonly<ClientReferenceManifest>
   serverActionsManifest: DeepReadonly<ActionManifest>
   serverModuleMap: {
@@ -78,9 +85,19 @@ export function setReferenceManifestsSingleton({
     }
   }
 }) {
-  // @ts-ignore
+  // @ts-expect-error
+  const clientReferenceManifestsPerPage = globalThis[
+    SERVER_ACTION_MANIFESTS_SINGLETON
+  ]?.clientReferenceManifestsPerPage as
+    | undefined
+    | DeepReadonly<Record<string, ClientReferenceManifest>>
+
+  // @ts-expect-error
   globalThis[SERVER_ACTION_MANIFESTS_SINGLETON] = {
-    clientReferenceManifest,
+    clientReferenceManifestsPerPage: {
+      ...clientReferenceManifestsPerPage,
+      [normalizePage(page)]: clientReferenceManifest,
+    },
     serverActionsManifest,
     serverModuleMap,
   }
@@ -100,29 +117,49 @@ export function getServerModuleMap() {
   }
 
   if (!serverActionsManifestSingleton) {
-    throw new Error(
-      'Missing manifest for Server Actions. This is a bug in Next.js'
-    )
+    throw new InvariantError('Missing manifest for Server Actions.')
   }
 
   return serverActionsManifestSingleton.serverModuleMap
 }
 
-export function getClientReferenceManifestSingleton() {
+export function getClientReferenceManifestForRsc(): DeepReadonly<ClientReferenceManifestForRsc> {
   const serverActionsManifestSingleton = (globalThis as any)[
     SERVER_ACTION_MANIFESTS_SINGLETON
   ] as {
-    clientReferenceManifest: DeepReadonly<ClientReferenceManifest>
-    serverActionsManifest: DeepReadonly<ActionManifest>
+    clientReferenceManifestsPerPage: DeepReadonly<
+      Record<string, ClientReferenceManifest>
+    >
   }
 
   if (!serverActionsManifestSingleton) {
-    throw new Error(
-      'Missing manifest for Server Actions. This is a bug in Next.js'
-    )
+    throw new InvariantError('Missing manifest for Server Actions.')
+  }
+
+  const { clientReferenceManifestsPerPage } = serverActionsManifestSingleton
+  const workStore = workAsyncStorage.getStore()
+
+  if (!workStore) {
+    // If there's no work store defined, we can assume that a client reference
+    // manifest is needed during module evaluation, e.g. to create a server
+    // action using a higher-order function. This might also use client
+    // components which need to be serialized by Flight, and therefore client
+    // references need to be resolvable. To make this work, we're returning a
+    // merged manifest across all pages. This is fine as long as the module IDs
+    // are not page specific, which they are not for Webpack. TODO: Fix this in
+    // Turbopack.
+    return mergeClientReferenceManifests(clientReferenceManifestsPerPage)
   }
 
-  return serverActionsManifestSingleton.clientReferenceManifest
+  const page = normalizePage(workStore.page)
+
+  const clientReferenceManifest = clientReferenceManifestsPerPage[page]
+
+  if (!clientReferenceManifest) {
+    throw new InvariantError(`Missing Client Reference Manifest for ${page}.`)
+  }
+
+  return clientReferenceManifest
 }
 
 export async function getActionEncryptionKey() {
@@ -133,14 +170,11 @@ export async function getActionEncryptionKey() {
   const serverActionsManifestSingleton = (globalThis as any)[
     SERVER_ACTION_MANIFESTS_SINGLETON
   ] as {
-    clientReferenceManifest: DeepReadonly<ClientReferenceManifest>
     serverActionsManifest: DeepReadonly<ActionManifest>
   }
 
   if (!serverActionsManifestSingleton) {
-    throw new Error(
-      'Missing manifest for Server Actions. This is a bug in Next.js'
-    )
+    throw new InvariantError('Missing manifest for Server Actions.')
   }
 
   const rawKey =
@@ -148,7 +182,7 @@ export async function getActionEncryptionKey() {
     serverActionsManifestSingleton.serverActionsManifest.encryptionKey
 
   if (rawKey === undefined) {
-    throw new Error('Missing encryption key for Server Actions')
+    throw new InvariantError('Missing encryption key for Server Actions')
   }
 
   __next_loaded_action_key = await crypto.subtle.importKey(
@@ -161,3 +195,40 @@ export async function getActionEncryptionKey() {
 
   return __next_loaded_action_key
 }
+
+function normalizePage(page: string): string {
+  return page.replace(/\/(page|route)$/, '')
+}
+
+function mergeClientReferenceManifests(
+  clientReferenceManifestsPerPage: DeepReadonly<
+    Record<string, ClientReferenceManifest>
+  >
+): ClientReferenceManifestForRsc {
+  const clientReferenceManifests = Object.values(
+    clientReferenceManifestsPerPage as Record<string, ClientReferenceManifest>
+  )
+
+  const mergedClientReferenceManifest: ClientReferenceManifestForRsc = {
+    clientModules: {},
+    edgeRscModuleMapping: {},
+    rscModuleMapping: {},
+  }
+
+  for (const clientReferenceManifest of clientReferenceManifests) {
+    mergedClientReferenceManifest.clientModules = {
+      ...mergedClientReferenceManifest.clientModules,
+      ...clientReferenceManifest.clientModules,
+    }
+    mergedClientReferenceManifest.edgeRscModuleMapping = {
+      ...mergedClientReferenceManifest.edgeRscModuleMapping,
+      ...clientReferenceManifest.edgeRscModuleMapping,
+    }
+    mergedClientReferenceManifest.rscModuleMapping = {
+      ...mergedClientReferenceManifest.rscModuleMapping,
+      ...clientReferenceManifest.rscModuleMapping,
+    }
+  }
+
+  return mergedClientReferenceManifest
+}
diff --git a/packages/next/src/server/app-render/encryption.ts b/packages/next/src/server/app-render/encryption.ts
index f5ee59db8c569..7649ef2dfe206 100644
--- a/packages/next/src/server/app-render/encryption.ts
+++ b/packages/next/src/server/app-render/encryption.ts
@@ -12,7 +12,7 @@ import {
   decrypt,
   encrypt,
   getActionEncryptionKey,
-  getClientReferenceManifestSingleton,
+  getClientReferenceManifestForRsc,
   getServerModuleMap,
   stringToUint8Array,
 } from './encryption-utils'
@@ -70,11 +70,11 @@ async function encodeActionBoundArg(actionId: string, arg: string) {
 
 // Encrypts the action's bound args into a string.
 export async function encryptActionBoundArgs(actionId: string, args: any[]) {
-  const clientReferenceManifestSingleton = getClientReferenceManifestSingleton()
+  const { clientModules } = getClientReferenceManifestForRsc()
 
   // Using Flight to serialize the args into a string.
   const serialized = await streamToString(
-    renderToReadableStream(args, clientReferenceManifestSingleton.clientModules)
+    renderToReadableStream(args, clientModules)
   )
 
   // Encrypt the serialized string with the action id as the salt.
@@ -90,16 +90,17 @@ export async function decryptActionBoundArgs(
   actionId: string,
   encrypted: Promise<string>
 ) {
-  const clientReferenceManifestSingleton = getClientReferenceManifestSingleton()
+  const { edgeRscModuleMapping, rscModuleMapping } =
+    getClientReferenceManifestForRsc()
 
   // Decrypt the serialized string with the action id as the salt.
-  const decryped = await decodeActionBoundArg(actionId, await encrypted)
+  const decrypted = await decodeActionBoundArg(actionId, await encrypted)
 
   // Using Flight to deserialize the args from the string.
   const deserialized = await createFromReadableStream(
     new ReadableStream({
       start(controller) {
-        controller.enqueue(textEncoder.encode(decryped))
+        controller.enqueue(textEncoder.encode(decrypted))
         controller.close()
       },
     }),
@@ -109,9 +110,7 @@ export async function decryptActionBoundArgs(
         // to be added to the current execution. Instead, we'll wait for any ClientReference
         // to be emitted which themselves will handle the preloading.
         moduleLoading: null,
-        moduleMap: isEdgeRuntime
-          ? clientReferenceManifestSingleton.edgeRscModuleMapping
-          : clientReferenceManifestSingleton.rscModuleMapping,
+        moduleMap: isEdgeRuntime ? edgeRscModuleMapping : rscModuleMapping,
         serverModuleMap: getServerModuleMap(),
       },
     }
diff --git a/packages/next/src/server/load-components.ts b/packages/next/src/server/load-components.ts
index 5c826205a16eb..907d68f3b85b5 100644
--- a/packages/next/src/server/load-components.ts
+++ b/packages/next/src/server/load-components.ts
@@ -177,6 +177,7 @@ async function loadComponentsImpl<N = any>({
   // to them at the top level of the page module.
   if (serverActionsManifest && clientReferenceManifest) {
     setReferenceManifestsSingleton({
+      page,
       clientReferenceManifest,
       serverActionsManifest,
       serverModuleMap: createServerModuleMap({
diff --git a/packages/next/src/server/use-cache/use-cache-wrapper.ts b/packages/next/src/server/use-cache/use-cache-wrapper.ts
index 9597aaf1bc0b4..b681d8aa4588a 100644
--- a/packages/next/src/server/use-cache/use-cache-wrapper.ts
+++ b/packages/next/src/server/use-cache/use-cache-wrapper.ts
@@ -25,10 +25,10 @@ import { makeHangingPromise } from '../dynamic-rendering-utils'
 
 import { cacheScopeAsyncLocalStorage } from '../async-storage/cache-scope.external'
 
-import type { ClientReferenceManifest } from '../../build/webpack/plugins/flight-manifest-plugin'
+import type { ClientReferenceManifestForRsc } from '../../build/webpack/plugins/flight-manifest-plugin'
 
 import {
-  getClientReferenceManifestSingleton,
+  getClientReferenceManifestForRsc,
   getServerModuleMap,
 } from '../app-render/encryption-utils'
 import type { CacheScopeStore } from '../async-storage/cache-scope.external'
@@ -66,7 +66,7 @@ function generateCacheEntry(
   workStore: WorkStore,
   outerWorkUnitStore: WorkUnitStore | undefined,
   cacheScope: undefined | CacheScopeStore,
-  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
   fn: any
 ): Promise<[ReadableStream, Promise<CacheEntry>]> {
@@ -90,7 +90,7 @@ function generateCacheEntryWithRestoredWorkStore(
   workStore: WorkStore,
   outerWorkUnitStore: WorkUnitStore | undefined,
   cacheScope: undefined | CacheScopeStore,
-  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
   fn: any
 ) {
@@ -128,7 +128,7 @@ function generateCacheEntryWithRestoredWorkStore(
 function generateCacheEntryWithCacheContext(
   workStore: WorkStore,
   outerWorkUnitStore: WorkUnitStore | undefined,
-  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
   fn: any
 ) {
@@ -298,7 +298,7 @@ async function generateCacheEntryImpl(
   workStore: WorkStore,
   outerWorkUnitStore: WorkUnitStore | undefined,
   innerCacheStore: UseCacheStore,
-  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
   fn: any
 ): Promise<[ReadableStream, Promise<CacheEntry>]> {
@@ -455,7 +455,7 @@ export function cache(kind: string, id: string, fn: any) {
 
       // Get the clientReferenceManifest while we're still in the outer Context.
       // In case getClientReferenceManifestSingleton is implemented using AsyncLocalStorage.
-      const clientReferenceManifest = getClientReferenceManifestSingleton()
+      const clientReferenceManifest = getClientReferenceManifestForRsc()
 
       // Because the Action ID is not yet unique per implementation of that Action we can't
       // safely reuse the results across builds yet. In the meantime we add the buildId to the
diff --git a/test/e2e/app-dir/actions/app-action.test.ts b/test/e2e/app-dir/actions/app-action.test.ts
index 21968dcc4b9aa..f37c9db5f9d87 100644
--- a/test/e2e/app-dir/actions/app-action.test.ts
+++ b/test/e2e/app-dir/actions/app-action.test.ts
@@ -1587,6 +1587,16 @@ describe('app-dir action handling', () => {
       expect(html).not.toContain('qwerty123')
       expect(html).not.toContain('some-module-level-encryption-value')
     })
+
+    it('should be able to resolve other server actions and client components', async () => {
+      const browser = await next.browser('/encryption')
+      expect(await browser.elementByCss('p').text()).toBe('initial')
+      await browser.elementByCss('button').click()
+
+      await retry(async () => {
+        expect(await browser.elementByCss('p').text()).toBe('hello from client')
+      })
+    })
   })
 
   describe('redirects', () => {
diff --git a/test/e2e/app-dir/actions/app/encryption/client.js b/test/e2e/app-dir/actions/app/encryption/client.js
new file mode 100644
index 0000000000000..0ebd41fc1b60f
--- /dev/null
+++ b/test/e2e/app-dir/actions/app/encryption/client.js
@@ -0,0 +1,5 @@
+'use client'
+
+export function Client() {
+  return 'hello from client'
+}
diff --git a/test/e2e/app-dir/actions/app/encryption/form.js b/test/e2e/app-dir/actions/app/encryption/form.js
new file mode 100644
index 0000000000000..bad0045173c5b
--- /dev/null
+++ b/test/e2e/app-dir/actions/app/encryption/form.js
@@ -0,0 +1,14 @@
+'use client'
+
+import { useActionState } from 'react'
+
+export default function Form({ action }) {
+  const [result, formAction] = useActionState(action, 'initial')
+
+  return (
+    <form action={formAction}>
+      <button>Submit</button>
+      <p>{result}</p>
+    </form>
+  )
+}
diff --git a/test/e2e/app-dir/actions/app/encryption/page.js b/test/e2e/app-dir/actions/app/encryption/page.js
index 2fc280d49a43e..2b640845d5cb4 100644
--- a/test/e2e/app-dir/actions/app/encryption/page.js
+++ b/test/e2e/app-dir/actions/app/encryption/page.js
@@ -1,36 +1,40 @@
+import { Client } from './client'
+import Form from './form'
+
 async function otherAction() {
   'use server'
   return 'hi'
 }
+
 // Test top-level encryption (happens during the module load phase)
-function wrapAction(value, action) {
+function wrapAction(value, action, element) {
   return async function () {
     'use server'
     const v = await action()
     if (v === 'hi') {
       console.log(value)
     }
+    return element
   }
 }
 
-const action = wrapAction('some-module-level-encryption-value', otherAction)
+const action = wrapAction(
+  'some-module-level-encryption-value',
+  otherAction,
+  <Client />
+)
 
 // Test runtime encryption (happens during the rendering phase)
 export default function Page() {
   const secret = 'my password is qwerty123'
 
   return (
-    <form
+    <Form
       action={async () => {
         'use server'
         console.log(secret)
-        await action()
-        return 'success'
+        return action()
       }}
-    >
-      <button type="submit" id="submit">
-        Submit
-      </button>
-    </form>
+    />
   )
 }
diff --git a/test/e2e/app-dir/use-cache/app/with-server-action/layout.tsx b/test/e2e/app-dir/use-cache/app/with-server-action/layout.tsx
deleted file mode 100644
index 7ca1569e06a67..0000000000000
--- a/test/e2e/app-dir/use-cache/app/with-server-action/layout.tsx
+++ /dev/null
@@ -1,14 +0,0 @@
-import { connection } from 'next/server'
-
-export default async function DynamicLayout({
-  children,
-}: {
-  children: React.ReactNode
-}) {
-  // TODO: This is a workaround for Turbopack. Figure out why this fails during
-  // prerendering with:
-  // TypeError: Cannot read properties of undefined (reading 'Form')
-  await connection()
-
-  return children
-}