fix: allow cookies in middleware to be accessible during the same request

Author: mrstork

edge-runtime/lib/middleware.ts

@@ -1,6 +1,7 @@
 import type { Context } from '@netlify/edge-functions'
 
 import type { ElementHandlers } from '../vendor/deno.land/x/htmlrewriter@v1.0.0/src/index.ts'
+import { getCookies } from '../vendor/deno.land/std@0.175.0/http/cookie.ts'
 
 type NextDataTransform = <T>(data: T) => T
 
@@ -58,3 +59,19 @@ export const addMiddlewareHeaders = async (
   })
   return response
 }
+
+export function mergeMiddlewareCookies(middlewareResponse: Response, request: Request) {
+  let mergedCookies = getCookies(request.headers)
+  const middlewareCookies = middlewareResponse.headers.get('x-middleware-set-cookie') || ''
+  const regex = new RegExp(/,(?!\s)/) // commas that are not followed by whitespace
+
+  middlewareCookies.split(regex).forEach((entry) => {
+    const [cookie] = entry.split(';')
+    const [name, value] = cookie.split('=')
+    mergedCookies[name] = value
+  })
+
+  return Object.entries(mergedCookies)
+    .map(kv => kv.join('='))
+    .join('; ')
+}

edge-runtime/lib/response.ts

@@ -6,7 +6,7 @@ import {
 
 import { updateModifiedHeaders } from './headers.ts'
 import type { StructuredLogger } from './logging.ts'
-import { addMiddlewareHeaders, isMiddlewareRequest, isMiddlewareResponse } from './middleware.ts'
+import { addMiddlewareHeaders, isMiddlewareRequest, isMiddlewareResponse, mergeMiddlewareCookies } from './middleware.ts'
 import { RequestData } from './next-request.ts'
 import {
   addBasePath,
@@ -116,12 +116,13 @@ export const buildResponse = async ({
     }
     return rewriter.transform(response.originResponse)
   }
-  const res = new Response(result.response.body, result.response)
+
+  const edgeResponse = new Response(result.response.body, result.response)
   request.headers.set('x-nf-next-middleware', 'skip')
 
-  let rewrite = res.headers.get('x-middleware-rewrite')
-  let redirect = res.headers.get('location')
-  let nextRedirect = res.headers.get('x-nextjs-redirect')
+  let rewrite = edgeResponse.headers.get('x-middleware-rewrite')
+  let redirect = edgeResponse.headers.get('location')
+  let nextRedirect = edgeResponse.headers.get('x-nextjs-redirect')
 
   // Data requests (i.e. requests for /_next/data ) need special handling
   const isDataReq = request.headers.has('x-nextjs-data')
@@ -152,7 +153,7 @@ export const buildResponse = async ({
       // Data requests might be rewritten to an external URL
       // This header tells the client router the redirect target, and if it's external then it will do a full navigation
 
-      res.headers.set('x-nextjs-rewrite', relativeUrl)
+      edgeResponse.headers.set('x-nextjs-rewrite', relativeUrl)
     }
 
     if (rewriteUrl.origin !== baseUrl.origin) {
@@ -178,7 +179,7 @@ export const buildResponse = async ({
         })
       }
 
-      return addMiddlewareHeaders(fetch(proxyRequest, { redirect: 'manual' }), res)
+      return addMiddlewareHeaders(fetch(proxyRequest, { redirect: 'manual' }), edgeResponse)
     }
 
     if (isDataReq) {
@@ -197,9 +198,9 @@ export const buildResponse = async ({
       logger.withFields({ rewrite_url: rewrite }).debug('Rewrite url is same as original url')
       return
     }
-    res.headers.set('x-middleware-rewrite', relativeUrl)
+    edgeResponse.headers.set('x-middleware-rewrite', relativeUrl)
     request.headers.set('x-middleware-rewrite', target)
-    return addMiddlewareHeaders(context.rewrite(target), res)
+    return addMiddlewareHeaders(context.rewrite(target), edgeResponse)
   }
 
   if (redirect) {
@@ -208,27 +209,32 @@ export const buildResponse = async ({
       logger.withFields({ redirect_url: redirect }).debug('Redirect url is same as original url')
       return
     }
-    res.headers.set('location', redirect)
+    edgeResponse.headers.set('location', redirect)
   }
 
   // Data requests shouldn't automatically redirect in the browser (they might be HTML pages): they're handled by the router
   if (redirect && isDataReq) {
-    res.headers.delete('location')
-    res.headers.set('x-nextjs-redirect', relativizeURL(redirect, request.url))
+    edgeResponse.headers.delete('location')
+    edgeResponse.headers.set('x-nextjs-redirect', relativizeURL(redirect, request.url))
   }
 
-  nextRedirect = res.headers.get('x-nextjs-redirect')
+  nextRedirect = edgeResponse.headers.get('x-nextjs-redirect')
 
   if (nextRedirect && isDataReq) {
-    res.headers.set('x-nextjs-redirect', normalizeDataUrl(nextRedirect))
+    edgeResponse.headers.set('x-nextjs-redirect', normalizeDataUrl(nextRedirect))
   }
 
-  if (res.headers.get('x-middleware-next') === '1') {
-    res.headers.delete('x-middleware-next')
-    return addMiddlewareHeaders(context.next(), res)
+  if (edgeResponse.headers.get('x-middleware-next') === '1') {
+    edgeResponse.headers.delete('x-middleware-next')
+
+    // coookies set in middleware need to be available during the lambda request
+    const newRequest = new Request(request)
+    newRequest.headers.set('Cookie', mergeMiddlewareCookies(edgeResponse, newRequest))
+
+    return addMiddlewareHeaders(context.next(newRequest), edgeResponse)
   }
 
-  return res
+  return edgeResponse
 }
 
 /**