fix: setting cookies in draftmode

Author: lubieowoce

packages/next/src/server/async-storage/draft-mode-provider.ts

@@ -1,5 +1,8 @@
 import type { IncomingMessage } from 'http'
-import type { ReadonlyRequestCookies } from '../web/spec-extension/adapters/request-cookies'
+import {
+  setMutableCookieUnchecked,
+  type ReadonlyRequestCookies,
+} from '../web/spec-extension/adapters/request-cookies'
 import type { ResponseCookies } from '../web/spec-extension/cookies'
 import type { BaseNextRequest } from '../base-http'
 import type { NextRequest } from '../web/spec-extension/request'
@@ -58,7 +61,7 @@ export class DraftModeProvider {
       )
     }
 
-    this._mutableCookies.set({
+    setMutableCookieUnchecked(this._mutableCookies, {
       name: COOKIE_NAME_PRERENDER_BYPASS,
       value: this._previewModeId,
       httpOnly: true,
@@ -72,7 +75,7 @@ export class DraftModeProvider {
     // To delete a cookie, set `expires` to a date in the past:
     // https://tools.ietf.org/html/rfc6265#section-4.1.1
     // `Max-Age: 0` is not valid, thus ignored, and the cookie is persisted.
-    this._mutableCookies.set({
+    setMutableCookieUnchecked(this._mutableCookies, {
       name: COOKIE_NAME_PRERENDER_BYPASS,
       value: '',
       httpOnly: true,

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

@@ -19,6 +19,7 @@ import {
   MutableRequestCookiesAdapter,
   RequestCookiesAdapter,
   responseCookiesToRequestCookies,
+  setMutableCookieUnchecked,
   type ReadonlyRequestCookies,
 } from '../web/spec-extension/adapters/request-cookies'
 import { ResponseCookies, RequestCookies } from '../web/spec-extension/cookies'
@@ -99,7 +100,11 @@ function mergeMiddlewareCookies(
 
     // Transfer cookies from ResponseCookies to RequestCookies
     for (const cookie of responseCookies.getAll()) {
-      existingCookies.set(cookie)
+      if (existingCookies instanceof ResponseCookies) {
+        setMutableCookieUnchecked(existingCookies, cookie)
+      } else {
+        existingCookies.set(cookie)
+      }
     }
   }
 }

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

@@ -53,6 +53,7 @@ export class RequestCookiesAdapter {
 }
 
 const SYMBOL_MODIFY_COOKIE_VALUES = Symbol.for('next.mutated.cookies')
+const SYMBOL_COOKIE_SET_UNCHECKED = Symbol.for('next.cookies.setUnchecked')
 
 export function getModifiedCookieValues(
   cookies: ResponseCookies
@@ -67,6 +68,25 @@ export function getModifiedCookieValues(
   return modified
 }
 
+type SetCookieArgs =
+  | [key: string, value: string, cookie?: Partial<ResponseCookie>]
+  | [options: ResponseCookie]
+
+/** Allows setting mutable cookies during render.  */
+export function setMutableCookieUnchecked(
+  mutableCookies: ResponseCookies,
+  ...args: SetCookieArgs
+) {
+  const setUnchecked: ResponseCookies['set'] =
+    // @ts-expect-error
+    mutableCookies[SYMBOL_COOKIE_SET_UNCHECKED]
+  if (setUnchecked) {
+    return setUnchecked(...args)
+  } else {
+    return mutableCookies.set(...args)
+  }
+}
+
 export function appendMutableCookies(
   headers: Headers,
   mutableCookies: ResponseCookies
@@ -132,12 +152,27 @@ export class MutableRequestCookiesAdapter {
       }
     }
 
+    const setUnchecked = (target: ResponseCookies, ...args: SetCookieArgs) => {
+      modifiedCookies.add(typeof args[0] === 'string' ? args[0] : args[0].name)
+      try {
+        return target.set(...args)
+      } finally {
+        updateResponseCookies()
+      }
+    }
+
     return 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
+          // A special symbol to get an unchecked version of .set()
+          // that allows setting cookies during render
+          case SYMBOL_COOKIE_SET_UNCHECKED:
+            return function (...args: SetCookieArgs) {
+              setUnchecked(target, ...args)
+            }
 
           // TODO: Throw error if trying to set a cookie after the response
           // headers have been set.
@@ -154,21 +189,11 @@ export class MutableRequestCookiesAdapter {
               }
             }
           case 'set':
-            return function (
-              ...args:
-                | [key: string, value: string, cookie?: Partial<ResponseCookie>]
-                | [options: ResponseCookie]
-            ) {
+            return function (...args: SetCookieArgs) {
               ensureCookiesAreStillMutable('cookies().set')
-              modifiedCookies.add(
-                typeof args[0] === 'string' ? args[0] : args[0].name
-              )
-              try {
-                return target.set(...args)
-              } finally {
-                updateResponseCookies()
-              }
+              setUnchecked(target, ...args)
             }
+
           default:
             return ReflectAdapter.get(target, prop, receiver)
         }