ensure cookies set in middleware can be read in a server action (#67924)

If middleware targets a server action handler and sets or updates a
cookie, the newly updated cookie would not be reflected in the
`cookies()` response of the action handler

In #65008 we fixed a bug where cookies set in middleware were not
reflected in the `cookies()` call in a server component from the same
request. We did this by introducing a `x-middleware-set-cookie` header,
that signaled to downstream handlers that middleware had run on the
request & set a cookie. However this handling was only applied to the
sealed/read-only cookies. Cookies accessed from a server action use
`mutableCookies`, since those aren't frozen as a server action is
allowed to modify cookies.

This pulls the cookie merge handling into a function and applies the
merge to `mutableCookies`.

Fixes #67814
Closes NDX-95

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

@@ -44,6 +44,35 @@ export type RequestContext = {
   renderOpts?: RenderOpts
 }
 
+/**
+ * If middleware set cookies in this request (indicated by `x-middleware-set-cookie`),
+ * then merge those into the existing cookie object, so that when `cookies()` is accessed
+ * it's able to read the newly set cookies.
+ */
+function mergeMiddlewareCookies(
+  req: RequestContext['req'],
+  existingCookies: RequestCookies | ResponseCookies
+) {
+  if (
+    'x-middleware-set-cookie' in req.headers &&
+    typeof req.headers['x-middleware-set-cookie'] === 'string'
+  ) {
+    const setCookieValue = req.headers['x-middleware-set-cookie']
+    const responseHeaders = new Headers()
+
+    for (const cookie of splitCookiesString(setCookieValue)) {
+      responseHeaders.append('set-cookie', cookie)
+    }
+
+    const responseCookies = new ResponseCookies(responseHeaders)
+
+    // Transfer cookies from ResponseCookies to RequestCookies
+    for (const cookie of responseCookies.getAll()) {
+      existingCookies.set(cookie.name, cookie.value ?? '')
+    }
+  }
+}
+
 export const RequestAsyncStorageWrapper: AsyncStorageWrapper<
   RequestStore,
   RequestContext
@@ -100,24 +129,7 @@ export const RequestAsyncStorageWrapper: AsyncStorageWrapper<
             HeadersAdapter.from(req.headers)
           )
 
-          if (
-            'x-middleware-set-cookie' in req.headers &&
-            typeof req.headers['x-middleware-set-cookie'] === 'string'
-          ) {
-            const setCookieValue = req.headers['x-middleware-set-cookie']
-            const responseHeaders = new Headers()
-
-            for (const cookie of splitCookiesString(setCookieValue)) {
-              responseHeaders.append('set-cookie', cookie)
-            }
-
-            const responseCookies = new ResponseCookies(responseHeaders)
-
-            // Transfer cookies from ResponseCookies to RequestCookies
-            for (const cookie of responseCookies.getAll()) {
-              requestCookies.set(cookie.name, cookie.value ?? '')
-            }
-          }
+          mergeMiddlewareCookies(req, requestCookies)
 
           // Seal the cookies object that'll freeze out any methods that could
           // mutate the underlying data.
@@ -128,11 +140,15 @@ export const RequestAsyncStorageWrapper: AsyncStorageWrapper<
       },
       get mutableCookies() {
         if (!cache.mutableCookies) {
-          cache.mutableCookies = getMutableCookies(
+          const mutableCookies = getMutableCookies(
             req.headers,
             renderOpts?.onUpdateCookies ||
               (res ? defaultOnUpdateCookies : undefined)
           )
+
+          mergeMiddlewareCookies(req, mutableCookies)
+
+          cache.mutableCookies = mutableCookies
         }
         return cache.mutableCookies
       },

test/e2e/app-dir/app-middleware/app-middleware.test.ts

@@ -171,6 +171,31 @@ createNextDescribe(
         )
       })
     })
+
+    it('should be possible to read cookies that are set during the middleware handling of a server action', async () => {
+      const browser = await next.browser('/rsc-cookies')
+      const initialRandom1 = await browser.elementById('rsc-cookie-1').text()
+      const initialRandom2 = await browser.elementById('rsc-cookie-2').text()
+      const totalCookies = await browser.elementById('total-cookies').text()
+
+      // cookies were set in middleware, assert they are present and match the Math.random() pattern
+      expect(initialRandom1).toMatch(/Cookie 1: \d+\.\d+/)
+      expect(initialRandom2).toMatch(/Cookie 2: \d+\.\d+/)
+      expect(totalCookies).toBe('Total Cookie Length: 2')
+
+      expect(await browser.eval('document.cookie')).toBeTruthy()
+
+      await browser.deleteCookies()
+
+      // assert that document.cookie is empty
+      expect(await browser.eval('document.cookie')).toBeFalsy()
+
+      await browser.elementById('submit-server-action').click()
+
+      await retry(() => {
+        expect(next.cliOutput).toMatch(/\[Cookie From Action\]: \d+\.\d+/)
+      })
+    })
   }
 )
 

test/e2e/app-dir/app-middleware/app/rsc-cookies/page.js

@@ -11,6 +11,20 @@ export default function Page() {
       <p id="rsc-cookie-2">Cookie 2: {rscCookie2}</p>
       <p id="total-cookies">Total Cookie Length: {cookies().size}</p>
       <Link href="/rsc-cookies-delete">To Delete Cookies Route</Link>
+
+      <form
+        action={async () => {
+          'use server'
+          console.log(
+            '[Cookie From Action]:',
+            cookies().get('rsc-cookie-value-1')?.value
+          )
+        }}
+      >
+        <button type="submit" id="submit-server-action">
+          server action
+        </button>
+      </form>
     </div>
   )
 }