initialize ALS with cookies in middleware

vercel · Apr 25, 2024 · 6494081 · 6494081
1 parent 9c0d792
commit 6494081
Show file tree

Hide file tree

Showing 7 changed files with 118 additions and 3 deletions.
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,6 +101,15 @@ 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
+          if (
+            'x-middleware-set-cookie' in req.headers &&
+            typeof req.headers['x-middleware-set-cookie'] === 'string'
+          ) {
+            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)

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,40 @@ 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 () => {
+        // the cookie text no longer contains a random number as it was cleared in middleware
+        expect(await browser.elementById('rsc-cookie-1').text()).toBe(
+          'Cookie 1:'
+        )
+        expect(await browser.elementById('rsc-cookie-2').text()).toBe(
+          'Cookie 2:'
+        )
+      })
+    })
   }
 )
 

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,22 @@ 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')
+    res.cookies.delete('rsc-cookie-value-2')
+
+    return res
+  }
+
   return NextResponse.next({
     request: {
       headers: headersFromRequest,

diff --git a/test/e2e/app-dir/app-static/app-static.test.ts b/test/e2e/app-dir/app-static/app-static.test.ts
@@ -1511,6 +1511,7 @@ createNextDescribe(
               ],
               "initialHeaders": {
                 "set-cookie": "theme=light; Path=/,my_company=ACME; Path=/",
+                "x-middleware-set-cookie": "theme=light; Path=/,my_company=ACME; Path=/",
                 "x-next-cache-tags": "_N_T_/layout,_N_T_/route-handler/layout,_N_T_/route-handler/static-cookies/layout,_N_T_/route-handler/static-cookies/route,_N_T_/route-handler/static-cookies",
               },
               "initialRevalidateSeconds": false,