prevent cookies from being set in the wrong phase

packages/next/src/server/app-render/action-handler.ts

@@ -43,6 +43,7 @@ import { fromNodeOutgoingHttpHeaders } from '../web/utils'
 import { selectWorkerForForwarding } from './action-utils'
 import { isNodeNextRequest, isWebNextRequest } from '../base-http/helpers'
 import { RedirectStatusCode } from '../../client/components/redirect-status-code'
+import { synchronizeMutableCookies } from '../async-storage/request-store'
 
 function formDataFromSearchQueryString(query: string) {
   const searchParams = new URLSearchParams(query)
@@ -459,6 +460,15 @@ export async function handleAction({
     )
   }
 
+  const finalizeAndGenerateFlight: GenerateFlight = (...args) => {
+    // When we switch to the render phase, cookies() will return
+    // `workUnitStore.cookies` instead of `workUnitStore.userspaceMutableCookies`.
+    // We want the render to see any cookie writes that we performed during the action,
+    // so we need to update the immutable cookies to reflect the changes.
+    synchronizeMutableCookies(requestStore)
+    return generateFlight(...args)
+  }
+
   requestStore.phase = 'action'
 
   // When running actions the default is no-store, you can still `cache: 'force-cache'`
@@ -548,7 +558,7 @@ export async function handleAction({
 
         return {
           type: 'done',
-          result: await generateFlight(req, ctx, {
+          result: await finalizeAndGenerateFlight(req, ctx, {
             actionResult: promise,
             // if the page was not revalidated, we can skip the rendering the flight tree
             skipFlight: !workStore.pathWasRevalidated,
@@ -841,7 +851,7 @@ export async function handleAction({
           requestStore,
         })
 
-        actionResult = await generateFlight(req, ctx, {
+        actionResult = await finalizeAndGenerateFlight(req, ctx, {
           actionResult: Promise.resolve(returnVal),
           // if the page was not revalidated, or if the action was forwarded from another worker, we can skip the rendering the flight tree
           skipFlight: !workStore.pathWasRevalidated || actionWasForwarded,
@@ -916,7 +926,7 @@ export async function handleAction({
         }
         return {
           type: 'done',
-          result: await generateFlight(req, ctx, {
+          result: await finalizeAndGenerateFlight(req, ctx, {
             skipFlight: false,
             actionResult: promise,
           }),

packages/next/src/server/app-render/work-unit-async-storage.external.ts

@@ -38,8 +38,11 @@ export type RequestStore = {
   }
 
   readonly headers: ReadonlyHeaders
-  readonly cookies: ReadonlyRequestCookies
+  // This is mutable because we need to reassign it when transitioning from the action phase to the render phase.
+  // The cookie object itself is deliberately read only and thus can't be updated.
+  cookies: ReadonlyRequestCookies
   readonly mutableCookies: ResponseCookies
+  readonly userspaceMutableCookies: ResponseCookies
   readonly draftMode: DraftModeProvider
   readonly isHmrRefresh?: boolean
   readonly serverComponentsHmrCache?: ServerComponentsHmrCache

packages/next/src/server/async-storage/request-store.ts

@@ -13,6 +13,8 @@ import {
 import {
   MutableRequestCookiesAdapter,
   RequestCookiesAdapter,
+  responseCookiesToRequestCookies,
+  wrapWithMutableAccessCheck,
   type ReadonlyRequestCookies,
 } from '../web/spec-extension/adapters/request-cookies'
 import { ResponseCookies, RequestCookies } from '../web/spec-extension/cookies'
@@ -165,6 +167,7 @@ function createRequestStoreImpl(
     headers?: ReadonlyHeaders
     cookies?: ReadonlyRequestCookies
     mutableCookies?: ResponseCookies
+    userspaceMutableCookies?: ResponseCookies
     draftMode?: DraftModeProvider
   } = {}
 
@@ -202,6 +205,9 @@ function createRequestStoreImpl(
 
       return cache.cookies
     },
+    set cookies(value: ReadonlyRequestCookies) {
+      cache.cookies = value
+    },
     get mutableCookies() {
       if (!cache.mutableCookies) {
         const mutableCookies = getMutableCookies(
@@ -215,6 +221,15 @@ function createRequestStoreImpl(
       }
       return cache.mutableCookies
     },
+    get userspaceMutableCookies() {
+      if (!cache.userspaceMutableCookies) {
+        const userspaceMutableCookies = wrapWithMutableAccessCheck(
+          this.mutableCookies
+        )
+        cache.userspaceMutableCookies = userspaceMutableCookies
+      }
+      return cache.userspaceMutableCookies
+    },
     get draftMode() {
       if (!cache.draftMode) {
         cache.draftMode = new DraftModeProvider(
@@ -234,3 +249,10 @@ function createRequestStoreImpl(
       (globalThis as any).__serverComponentsHmrCache,
   }
 }
+
+export function synchronizeMutableCookies(store: RequestStore) {
+  // TODO: does this need to update headers as well?
+  store.cookies = RequestCookiesAdapter.seal(
+    responseCookiesToRequestCookies(store.mutableCookies)
+  )
+}

packages/next/src/server/request/cookies.ts

@@ -1,6 +1,7 @@
 import {
   type ReadonlyRequestCookies,
   type ResponseCookies,
+  areCookiesMutableInCurrentPhase,
   RequestCookiesAdapter,
 } from '../web/spec-extension/adapters/request-cookies'
 import { RequestCookies } from '../web/spec-extension/cookies'
@@ -16,7 +17,6 @@ import {
   trackDynamicDataInDynamicRender,
 } from '../app-render/dynamic-rendering'
 import { getExpectedRequestStore } from '../app-render/work-unit-async-storage.external'
-import { actionAsyncStorage } from '../app-render/action-async-storage.external'
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
 import { makeResolvedReactPromise } from './utils'
 import { makeHangingPromise } from '../dynamic-rendering-utils'
@@ -113,22 +113,16 @@ export function cookies(): Promise<ReadonlyRequestCookies> {
   }
 
   // cookies is being called in a dynamic context
-  const actionStore = actionAsyncStorage.getStore()
 
   const requestStore = getExpectedRequestStore(callingExpression)
 
   let underlyingCookies: ReadonlyRequestCookies
 
-  // The current implementation of cookies will return Response cookies
-  // for a server action during the render phase of a server action.
-  // This is not correct b/c the type of cookies during render is ReadOnlyRequestCookies
-  // where as the type of cookies during action is ResponseCookies
-  // This was found because RequestCookies is iterable and ResponseCookies is not
-  if (actionStore?.isAction || actionStore?.isAppRoute) {
+  if (areCookiesMutableInCurrentPhase(requestStore)) {
     // We can't conditionally return different types here based on the context.
     // To avoid confusion, we always return the readonly type here.
     underlyingCookies =
-      requestStore.mutableCookies as unknown as ReadonlyRequestCookies
+      requestStore.userspaceMutableCookies as unknown as ReadonlyRequestCookies
   } else {
     underlyingCookies = requestStore.cookies
   }

packages/next/src/server/web/spec-extension/adapters/request-cookies.test.ts

@@ -1,7 +1,13 @@
-import { RequestCookies } from '../cookies'
+import {
+  workUnitAsyncStorage,
+  type RequestStore,
+} from '../../../app-render/work-unit-async-storage.external'
+import { RequestCookies, ResponseCookies } from '../cookies'
 import {
   ReadonlyRequestCookiesError,
   RequestCookiesAdapter,
+  MutableRequestCookiesAdapter,
+  wrapWithMutableAccessCheck,
 } from './request-cookies'
 
 describe('RequestCookiesAdapter', () => {
@@ -52,3 +58,97 @@ describe('RequestCookiesAdapter', () => {
     expect(sealed.get('bar')).toEqual(undefined)
   })
 })
+
+describe('MutableRequestCookiesAdapter', () => {
+  it('supports chained set calls and preserves wrapping', () => {
+    const headers = new Headers({})
+    const underlyingCookies = new RequestCookies(headers)
+    const onUpdateCookies = jest.fn<void, [string[]]>()
+
+    const wrappedCookies = MutableRequestCookiesAdapter.wrap(
+      underlyingCookies,
+      onUpdateCookies
+    )
+
+    const returned = wrappedCookies.set('foo', '1').set('bar', '2')
+
+    expect(returned).toBe(wrappedCookies)
+    expect(onUpdateCookies).toHaveBeenCalledWith([
+      expect.stringContaining('foo=1'),
+    ])
+    expect(onUpdateCookies).toHaveBeenCalledWith([
+      expect.stringContaining('foo=1'),
+      expect.stringContaining('bar=2'),
+    ])
+  })
+
+  it('supports chained delete calls and preserves wrapping', () => {
+    const headers = new Headers({})
+    const underlyingCookies = new RequestCookies(headers)
+    underlyingCookies.set('foo', '1').set('bar', '2')
+
+    const onUpdateCookies = jest.fn<void, [string[]]>()
+    const wrappedCookies = MutableRequestCookiesAdapter.wrap(
+      underlyingCookies,
+      onUpdateCookies
+    )
+
+    const returned = wrappedCookies.delete('foo').delete('bar')
+
+    expect(returned).toBe(wrappedCookies)
+    expect(onUpdateCookies).toHaveBeenCalledWith([
+      expect.stringContaining('foo=;'),
+    ])
+    expect(onUpdateCookies).toHaveBeenCalledWith([
+      expect.stringContaining('foo=;'),
+      expect.stringContaining('bar=;'),
+    ])
+  })
+})
+
+describe('wrapWithMutableAccessCheck', () => {
+  const createMockRequestStore = (phase: RequestStore['phase']) =>
+    ({ type: 'request', phase }) as RequestStore
+
+  it('prevents setting cookies in the render phase', () => {
+    const requestStore = createMockRequestStore('action')
+    workUnitAsyncStorage.run(requestStore, () => {
+      const headers = new Headers({})
+      const underlyingCookies = new ResponseCookies(headers)
+      const wrappedCookies = wrapWithMutableAccessCheck(underlyingCookies)
+
+      // simulate changing phases
+      requestStore.phase = 'render'
+
+      const EXPECTED_ERROR =
+        /Cookies can only be modified in a Server Action or Route Handler\./
+
+      expect(() => {
+        wrappedCookies.set('foo', '1')
+      }).toThrow(EXPECTED_ERROR)
+
+      expect(wrappedCookies.get('foo')).toBe(undefined)
+    })
+  })
+
+  it('prevents deleting cookies in the render phase', () => {
+    const requestStore = createMockRequestStore('action')
+    workUnitAsyncStorage.run(requestStore, () => {
+      const headers = new Headers({})
+      const underlyingCookies = new ResponseCookies(headers)
+      const wrappedCookies = wrapWithMutableAccessCheck(underlyingCookies)
+      wrappedCookies.set('foo', '1')
+
+      // simulate changing phases
+      requestStore.phase = 'render'
+
+      const EXPECTED_ERROR =
+        /Cookies can only be modified in a Server Action or Route Handler\./
+
+      expect(() => {
+        wrappedCookies.delete('foo')
+      }).toThrow(EXPECTED_ERROR)
+      expect(wrappedCookies.get('foo')?.value).toEqual('1')
+    })
+  })
+})

packages/next/src/server/web/spec-extension/adapters/request-cookies.ts

@@ -1,8 +1,12 @@
-import type { RequestCookies } from '../cookies'
+import { RequestCookies } from '../cookies'
 
 import { ResponseCookies } from '../cookies'
 import { ReflectAdapter } from './reflect'
 import { workAsyncStorage } from '../../../app-render/work-async-storage.external'
+import {
+  getExpectedRequestStore,
+  type RequestStore,
+} from '../../../app-render/work-unit-async-storage.external'
 
 /**
  * @internal
@@ -63,6 +67,10 @@ export function getModifiedCookieValues(
   return modified
 }
 
+type SetCookieArgs =
+  | [key: string, value: string, cookie?: Partial<ResponseCookie>]
+  | [options: ResponseCookie]
+
 export function appendMutableCookies(
   headers: Headers,
   mutableCookies: ResponseCookies
@@ -128,7 +136,7 @@ export class MutableRequestCookiesAdapter {
       }
     }
 
-    return new Proxy(responseCookies, {
+    const wrappedCookies = new Proxy(responseCookies, {
       get(target, prop, receiver) {
         switch (prop) {
           // A special symbol to get the modified cookie values
@@ -144,29 +152,86 @@ export class MutableRequestCookiesAdapter {
               )
               try {
                 target.delete(...args)
+                return wrappedCookies
               } finally {
                 updateResponseCookies()
               }
             }
           case 'set':
-            return function (
-              ...args:
-                | [key: string, value: string, cookie?: Partial<ResponseCookie>]
-                | [options: ResponseCookie]
-            ) {
+            return function (...args: SetCookieArgs) {
               modifiedCookies.add(
                 typeof args[0] === 'string' ? args[0] : args[0].name
               )
               try {
-                return target.set(...args)
+                target.set(...args)
+                return wrappedCookies
               } finally {
                 updateResponseCookies()
               }
             }
+
           default:
             return ReflectAdapter.get(target, prop, receiver)
         }
       },
     })
+
+    return wrappedCookies
+  }
+}
+
+export function wrapWithMutableAccessCheck(
+  responseCookies: ResponseCookies
+): ResponseCookies {
+  const wrappedCookies = new Proxy(responseCookies, {
+    get(target, prop, receiver) {
+      switch (prop) {
+        case 'delete':
+          return function (...args: [string] | [ResponseCookie]) {
+            ensureCookiesAreStillMutable('cookies().delete')
+            target.delete(...args)
+            return wrappedCookies
+          }
+        case 'set':
+          return function (...args: SetCookieArgs) {
+            ensureCookiesAreStillMutable('cookies().set')
+            target.set(...args)
+            return wrappedCookies
+          }
+
+        default:
+          return ReflectAdapter.get(target, prop, receiver)
+      }
+    },
+  })
+  return wrappedCookies
+}
+
+export function areCookiesMutableInCurrentPhase(requestStore: RequestStore) {
+  return requestStore.phase === 'action'
+}
+
+/** Ensure that cookies() starts throwing on mutation
+ * if we changed phases and can no longer mutate.
+ *
+ * This can happen when going:
+ *   'render' -> 'after'
+ *   'action' -> 'render'
+ * */
+function ensureCookiesAreStillMutable(callingExpression: string) {
+  const requestStore = getExpectedRequestStore(callingExpression)
+  if (!areCookiesMutableInCurrentPhase(requestStore)) {
+    // TODO: maybe we can give a more precise error message based on callingExpression?
+    throw new ReadonlyRequestCookiesError()
+  }
+}
+
+export function responseCookiesToRequestCookies(
+  responseCookies: ResponseCookies
+): RequestCookies {
+  const requestCookies = new RequestCookies(new Headers())
+  for (const cookie of responseCookies.getAll()) {
+    requestCookies.set(cookie)
   }
+  return requestCookies
 }