fix(cookies): prevent mutation outside action phase

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

@@ -1,14 +1,10 @@
 import { RequestCookies } from '../cookies'
-
 import { ResponseCookies } from '../cookies'
 import { ReflectAdapter } from './reflect'
 import { workAsyncStorage } from '../../../app-render/work-async-storage.external'
 import type { RequestStore } from '../../../app-render/work-unit-async-storage.external'
 import { ActionDidRevalidateStaticAndDynamic } from '../../../../shared/lib/action-revalidation-kind'
 
-/**
- * @internal
- */
 export class ReadonlyRequestCookiesError extends Error {
   constructor() {
     super(
@@ -21,12 +17,8 @@ export class ReadonlyRequestCookiesError extends Error {
   }
 }
 
-// We use this to type some APIs but we don't construct instances directly
 export type { ResponseCookies }
 
-// The `cookies()` API is a mix of request and response cookies. For `.get()` methods,
-// we want to return the request cookie if it exists. For mutative methods like `.set()`,
-// we want to return the response cookie.
 export type ReadonlyRequestCookies = Omit<
   RequestCookies,
   'set' | 'clear' | 'delete'
@@ -69,34 +61,6 @@ type SetCookieArgs =
   | [key: string, value: string, cookie?: Partial<ResponseCookie>]
   | [options: ResponseCookie]
 
-export function appendMutableCookies(
-  headers: Headers,
-  mutableCookies: ResponseCookies
-): boolean {
-  const modifiedCookieValues = getModifiedCookieValues(mutableCookies)
-  if (modifiedCookieValues.length === 0) {
-    return false
-  }
-
-  // Return a new response that extends the response with
-  // the modified cookies as fallbacks. `res` cookies
-  // will still take precedence.
-  const resCookies = new ResponseCookies(headers)
-  const returnedCookies = resCookies.getAll()
-
-  // Set the modified cookies as fallbacks.
-  for (const cookie of modifiedCookieValues) {
-    resCookies.set(cookie)
-  }
-
-  // Set the original cookies as the final values.
-  for (const cookie of returnedCookies) {
-    resCookies.set(cookie)
-  }
-
-  return true
-}
-
 type ResponseCookie = NonNullable<
   ReturnType<InstanceType<typeof ResponseCookies>['get']>
 >
@@ -107,23 +71,26 @@ export class MutableRequestCookiesAdapter {
     onUpdateCookies?: (cookies: string[]) => void
   ): ResponseCookies {
     const responseCookies = new ResponseCookies(new Headers())
+
     for (const cookie of cookies.getAll()) {
       responseCookies.set(cookie)
     }
 
     let modifiedValues: ResponseCookie[] = []
     const modifiedCookies = new Set<string>()
+
     const updateResponseCookies = () => {
-      // TODO-APP: change method of getting workStore
       const workStore = workAsyncStorage.getStore()
       if (workStore) {
         workStore.pathWasRevalidated = ActionDidRevalidateStaticAndDynamic
       }
 
       const allCookies = responseCookies.getAll()
       modifiedValues = allCookies.filter((c) => modifiedCookies.has(c.name))
+
       if (onUpdateCookies) {
         const serializedCookies: string[] = []
+
         for (const cookie of modifiedValues) {
           const tempCookies = new ResponseCookies(new Headers())
           tempCookies.set(cookie)
@@ -137,29 +104,45 @@ export class MutableRequestCookiesAdapter {
     const wrappedCookies = new Proxy(responseCookies, {
       get(target, prop, receiver) {
         switch (prop) {
-          // A special symbol to get the modified cookie values
           case SYMBOL_MODIFY_COOKIE_VALUES:
             return modifiedValues
 
-          // TODO: Throw error if trying to set a cookie after the response
-          // headers have been set.
           case 'delete':
             return function (...args: [string] | [ResponseCookie]) {
+              const requestStore = workAsyncStorage.getStore() as
+                | RequestStore
+                | undefined
+
+              if (requestStore) {
+                ensureCookiesAreStillMutable(requestStore, 'cookies().delete')
+              }
+
               modifiedCookies.add(
                 typeof args[0] === 'string' ? args[0] : args[0].name
               )
+
               try {
                 target.delete(...args)
                 return wrappedCookies
               } finally {
                 updateResponseCookies()
               }
             }
+
           case 'set':
             return function (...args: SetCookieArgs) {
+              const requestStore = workAsyncStorage.getStore() as
+                | RequestStore
+                | undefined
+
+              if (requestStore) {
+                ensureCookiesAreStillMutable(requestStore, 'cookies().set')
+              }
+
               modifiedCookies.add(
                 typeof args[0] === 'string' ? args[0] : args[0].name
               )
+
               try {
                 target.set(...args)
                 return wrappedCookies
@@ -190,6 +173,7 @@ export function createCookiesWithMutableAccessCheck(
             target.delete(...args)
             return wrappedCookies
           }
+
         case 'set':
           return function (...args: SetCookieArgs) {
             ensureCookiesAreStillMutable(requestStore, 'cookies().set')
@@ -202,26 +186,19 @@ export function createCookiesWithMutableAccessCheck(
       }
     },
   })
+
   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(
   requestStore: RequestStore,
   _callingExpression: string
 ) {
   if (!areCookiesMutableInCurrentPhase(requestStore)) {
-    // TODO: maybe we can give a more precise error message based on callingExpression?
     throw new ReadonlyRequestCookiesError()
   }
 }
@@ -230,8 +207,34 @@ export function responseCookiesToRequestCookies(
   responseCookies: ResponseCookies
 ): RequestCookies {
   const requestCookies = new RequestCookies(new Headers())
+
   for (const cookie of responseCookies.getAll()) {
     requestCookies.set(cookie)
   }
+
   return requestCookies
 }
+
+export function appendMutableCookies(
+  headers: Headers,
+  mutableCookies: ResponseCookies
+): boolean {
+  const modifiedCookieValues = getModifiedCookieValues(mutableCookies)
+
+  if (modifiedCookieValues.length === 0) {
+    return false
+  }
+
+  const resCookies = new ResponseCookies(headers)
+  const returnedCookies = resCookies.getAll()
+
+  for (const cookie of modifiedCookieValues) {
+    resCookies.set(cookie)
+  }
+
+  for (const cookie of returnedCookies) {
+    resCookies.set(cookie)
+  }
+
+  return true
+}