initialize ALS with cookies in middleware (#65008)

### What Cookies set/updated/removed in middleware won't be accessible during the render in which they were set ### Why Middleware will properly set a `set-cookie` header to inform the client of the cookie change, but this means the `AsyncLocalStorage` context containing the cookies value wouldn't be updated until the next time the request headers were parsed. In other words, on the first request the cookie would be sent but wouldn't be available in the `cookies()` context. And then the following request would properly have the cookie values. ### How This uses a proxy on the `ResponseCookies` used in middleware to add a middleware override header with the cookie value. When we instantiate the cached cookies, we merge in whatever headers would have been set by middleware, so that they're available in the same render that invoked middleware. ### Test Plan This changeset adds a test to confirm cookies set/deleted in middleware are available in a single pass. Verified with a deployment [here](https://vtest314-e2e-tests-ldx7olfl1-ztanner.vercel.app/rsc-cookies). Fixes #49442 Closes NEXT-1126
vercel · Aug 28, 2024 · 4727137 · 4727137
1 parent a1c3a03
commit 4727137
Show file tree

Hide file tree

Showing 7 changed files with 130 additions and 4 deletions.
diff --git a/packages/next/src/build/index.ts b/packages/next/src/build/index.ts
@@ -2813,6 +2813,10 @@ export default async function build(
                   // normalize header values as initialHeaders
                   // must be Record<string, string>
                   for (const key of headerKeys) {
+                    // set-cookie is already handled - the middleware cookie setting case
+                    // isn't needed for the prerender manifest since it can't read cookies
+                    if (key === 'x-middleware-set-cookie') continue
+
                     let value = exportHeaders[key]
 
                     if (Array.isArray(value)) {

diff --git a/packages/next/src/server/async-storage/request-async-storage-wrapper.ts b/packages/next/src/server/async-storage/request-async-storage-wrapper.ts
@@ -101,9 +101,26 @@ export const RequestAsyncStorageWrapper: AsyncStorageWrapper<
       },
       get cookies() {
         if (!cache.cookies) {
+          // if middleware is setting cookie(s), then include those in
+          // the initial cached cookies so they can be read in render
+          let combinedCookies
+          if (
+            'x-middleware-set-cookie' in req.headers &&
+            typeof req.headers['x-middleware-set-cookie'] === 'string'
+          ) {
+            combinedCookies = `${req.headers.cookie}; ${req.headers['x-middleware-set-cookie']}`
+          }
+
           // Seal the cookies object that'll freeze out any methods that could
           // mutate the underlying data.
-          cache.cookies = getCookies(req.headers)
+          cache.cookies = getCookies(
+            combinedCookies
+              ? {
+                  ...req.headers,
+                  cookie: combinedCookies,
+                }
+              : req.headers
+          )
         }
 
         return cache.cookies

diff --git a/packages/next/src/server/web/spec-extension/response.ts b/packages/next/src/server/web/spec-extension/response.ts
@@ -1,6 +1,7 @@
 import type { I18NConfig } from '../../config-shared'
 import { NextURL } from '../next-url'
 import { toNodeOutgoingHttpHeaders, validateURL } from '../utils'
+import { ReflectAdapter } from './adapters/reflect'
 
 import { ResponseCookies } from './cookies'
 
@@ -41,11 +42,37 @@ export class NextResponse<Body = unknown> extends Response {
   constructor(body?: BodyInit | null, init: ResponseInit = {}) {
     super(body, init)
 
+    const headers = this.headers
+    const cookies = new ResponseCookies(headers)
+
+    const cookiesProxy = new Proxy(cookies, {
+      get(target, prop, receiver) {
+        switch (prop) {
+          case 'delete':
+          case 'set': {
+            return (...args: [string, string]) => {
+              const result = Reflect.apply(target[prop], target, args)
+              const newHeaders = new Headers(headers)
+
+              if (result instanceof ResponseCookies) {
+                headers.set('x-middleware-set-cookie', result.toString())
+              }
+
+              handleMiddlewareField(init, newHeaders)
+              return result
+            }
+          }
+          default:
+            return ReflectAdapter.get(target, prop, receiver)
+        }
+      },
+    })
+
     this[INTERNALS] = {
-      cookies: new ResponseCookies(this.headers),
+      cookies: cookiesProxy,
       url: init.url
         ? new NextURL(init.url, {
-            headers: toNodeOutgoingHttpHeaders(this.headers),
+            headers: toNodeOutgoingHttpHeaders(headers),
             nextConfig: init.nextConfig,
           })
         : undefined,

diff --git a/test/e2e/app-dir/app-middleware/app-middleware.test.ts b/test/e2e/app-dir/app-middleware/app-middleware.test.ts
@@ -1,7 +1,7 @@
 /* eslint-env jest */
 import path from 'path'
 import cheerio from 'cheerio'
-import { check, withQuery } from 'next-test-utils'
+import { check, retry, withQuery } from 'next-test-utils'
 import { createNextDescribe, FileRef } from 'e2e-utils'
 import type { Response } from 'node-fetch'
 
@@ -134,6 +134,41 @@ createNextDescribe(
         expect(bypassCookie).toBeDefined()
       })
     })
+
+    it('should be possible to modify cookies & read them in an RSC in a single request', 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()
+
+      // 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+/)
+
+      await browser.refresh()
+
+      const refreshedRandom1 = await browser.elementById('rsc-cookie-1').text()
+      const refreshedRandom2 = await browser.elementById('rsc-cookie-2').text()
+
+      // the cookies should be refreshed and have new values
+      expect(refreshedRandom1).toMatch(/Cookie 1: \d+\.\d+/)
+      expect(refreshedRandom2).toMatch(/Cookie 2: \d+\.\d+/)
+      expect(refreshedRandom1).not.toBe(initialRandom1)
+      expect(refreshedRandom2).not.toBe(initialRandom2)
+
+      // navigate to delete cookies route
+      await browser.elementByCss('[href="/rsc-cookies-delete"]').click()
+      await retry(async () => {
+        // only the first cookie should be deleted
+        expect(await browser.elementById('rsc-cookie-1').text()).toBe(
+          'Cookie 1:'
+        )
+
+        expect(await browser.elementById('rsc-cookie-2').text()).toMatch(
+          /Cookie 2: \d+\.\d+/
+        )
+      })
+    })
   }
 )
 

diff --git a/test/e2e/app-dir/app-middleware/app/rsc-cookies-delete/page.js b/test/e2e/app-dir/app-middleware/app/rsc-cookies-delete/page.js
@@ -0,0 +1,13 @@
+import { cookies } from 'next/headers'
+
+export default function Page() {
+  const rscCookie1 = cookies().get('rsc-cookie-value-1')?.value
+  const rscCookie2 = cookies().get('rsc-cookie-value-2')?.value
+
+  return (
+    <div>
+      <p id="rsc-cookie-1">Cookie 1: {rscCookie1}</p>
+      <p id="rsc-cookie-2">Cookie 2: {rscCookie2}</p>
+    </div>
+  )
+}
diff --git a/test/e2e/app-dir/app-middleware/app/rsc-cookies/page.js b/test/e2e/app-dir/app-middleware/app/rsc-cookies/page.js
@@ -0,0 +1,15 @@
+import { cookies } from 'next/headers'
+import Link from 'next/link'
+
+export default function Page() {
+  const rscCookie1 = cookies().get('rsc-cookie-value-1')?.value
+  const rscCookie2 = cookies().get('rsc-cookie-value-2')?.value
+
+  return (
+    <div>
+      <p id="rsc-cookie-1">Cookie 1: {rscCookie1}</p>
+      <p id="rsc-cookie-2">Cookie 2: {rscCookie2}</p>
+      <Link href="/rsc-cookies-delete">To Delete Cookies Route</Link>
+    </div>
+  )
+}
diff --git a/test/e2e/app-dir/app-middleware/middleware.js b/test/e2e/app-dir/app-middleware/middleware.js
@@ -44,6 +44,21 @@ export async function middleware(request) {
     return NextResponse.rewrite(request.nextUrl)
   }
 
+  if (request.nextUrl.pathname === '/rsc-cookies') {
+    const res = NextResponse.next()
+    res.cookies.set('rsc-cookie-value-1', `${Math.random()}`)
+    res.cookies.set('rsc-cookie-value-2', `${Math.random()}`)
+
+    return res
+  }
+
+  if (request.nextUrl.pathname === '/rsc-cookies-delete') {
+    const res = NextResponse.next()
+    res.cookies.delete('rsc-cookie-value-1')
+
+    return res
+  }
+
   return NextResponse.next({
     request: {
       headers: headersFromRequest,