Ensure cookies are attached when middleware changes the Response (#8612)

* Ensure cookies are attached when middleware changes the Response

* fix test

7 files changed, 42 insertions, 3 deletions

diff --git a/.changeset/curvy-owls-rest.md b/.changeset/curvy-owls-rest.md
new file mode 100644
index 000000000..41bd9a22e
--- /dev/null
+++ b/.changeset/curvy-owls-rest.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Ensure cookies are attached when middleware changes the Response
diff --git a/packages/astro/src/core/cookies/index.ts b/packages/astro/src/core/cookies/index.ts
index f3c7b6d61..c55d060c4 100644
--- a/packages/astro/src/core/cookies/index.ts
+++ b/packages/astro/src/core/cookies/index.ts
@@ -1,2 +1,2 @@
 export { AstroCookies } from './cookies.js';
-export { attachCookiesToResponse, getSetCookiesFromResponse } from './response.js';
+export { attachCookiesToResponse, responseHasCookies, getSetCookiesFromResponse } from './response.js';
diff --git a/packages/astro/src/core/cookies/response.ts b/packages/astro/src/core/cookies/response.ts
index 8dc35e8c7..013f836bf 100644
--- a/packages/astro/src/core/cookies/response.ts
+++ b/packages/astro/src/core/cookies/response.ts
@@ -6,6 +6,10 @@ export function attachCookiesToResponse(response: Response, cookies: AstroCookie
 	Reflect.set(response, astroCookiesSymbol, cookies);
 }
 
+export function responseHasCookies(response: Response): boolean {
+	return Reflect.has(response, astroCookiesSymbol);
+}
+
 function getFromResponse(response: Response): AstroCookies | undefined {
 	let cookies = Reflect.get(response, astroCookiesSymbol);
 	if (cookies != null) {
diff --git a/packages/astro/src/core/middleware/callMiddleware.ts b/packages/astro/src/core/middleware/callMiddleware.ts
index 1725fd38d..bb0871405 100644
--- a/packages/astro/src/core/middleware/callMiddleware.ts
+++ b/packages/astro/src/core/middleware/callMiddleware.ts
@@ -7,6 +7,7 @@ import type {
 } from '../../@types/astro.js';
 import { AstroError, AstroErrorData } from '../errors/index.js';
 import type { Environment } from '../render/index.js';
+import { attachCookiesToResponse, responseHasCookies } from '../cookies/index.js';
 
 /**
  * Utility function that is in charge of calling the middleware.
@@ -82,7 +83,7 @@ export async function callMiddleware<R>(
 				if (value instanceof Response === false) {
 					throw new AstroError(AstroErrorData.MiddlewareNotAResponse);
 				}
-				return value as R;
+				return ensureCookiesAttached(apiContext, value as Response);
 			} else {
 				/**
 				 * Here we handle the case where `next` was called and returned nothing.
@@ -105,11 +106,18 @@ export async function callMiddleware<R>(
 			throw new AstroError(AstroErrorData.MiddlewareNotAResponse);
 		} else {
 			// Middleware did not call resolve and returned a value
-			return value as R;
+			return ensureCookiesAttached(apiContext, value as Response);
 		}
 	});
 }
 
+function ensureCookiesAttached(apiContext: APIContext, response: Response): Response {
+	if(apiContext.cookies !== undefined && !responseHasCookies(response)) {
+		attachCookiesToResponse(response, apiContext.cookies);
+	}
+	return response;
+}
+
 function isEndpointOutput(endpointResult: any): endpointResult is EndpointOutput {
 	return (
 		!(endpointResult instanceof Response) &&
diff --git a/packages/astro/test/fixtures/middleware space/src/middleware.js b/packages/astro/test/fixtures/middleware space/src/middleware.js
index eeb902fb8..b2e5c7e2d 100644
--- a/packages/astro/test/fixtures/middleware space/src/middleware.js
+++ b/packages/astro/test/fixtures/middleware space/src/middleware.js
@@ -24,6 +24,14 @@ const first = defineMiddleware(async (context, next) => {
 		const /** @type {string} */ html = await newResponse.text();
 		const newhtml = html.replace('<h1>testing</h1>', '<h1>it works</h1>');
 		return new Response(newhtml, { status: 200, headers: response.headers });
+	} else if(context.url.pathname === '/return-response-cookies') {
+		const response = await next();
+    const html = await response.text();
+
+    return new Response(html, {
+        status: 200,
+        headers: response.headers
+    });
 	} else {
 		if (context.url.pathname === '/') {
 			context.cookies.set('foo', 'bar');
diff --git a/packages/astro/test/fixtures/middleware space/src/pages/return-response-cookies.astro b/packages/astro/test/fixtures/middleware space/src/pages/return-response-cookies.astro
new file mode 100644
index 000000000..a3583b590
--- /dev/null
+++ b/packages/astro/test/fixtures/middleware space/src/pages/return-response-cookies.astro
@@ -0,0 +1,8 @@
+---
+Astro.cookies.set("astro", "cookie", {
+	httpOnly: true,
+	path: "/",
+	sameSite: "strict",
+	maxAge: 1704085200,
+});
+---
diff --git a/packages/astro/test/middleware.test.js b/packages/astro/test/middleware.test.js
index 81f167647..d3cb83310 100644
--- a/packages/astro/test/middleware.test.js
+++ b/packages/astro/test/middleware.test.js
@@ -79,6 +79,12 @@ describe('Middleware in DEV mode', () => {
 		let html = await res.text();
 		expect(html).to.contain('<h1>it works</h1>');
 	});
+
+	it('should forward cookies set in a component when the middleware returns a new response', async () => {
+		let res = await fixture.fetch('/return-response-cookies');
+		let headers = res.headers;
+		expect(headers.get('set-cookie')).to.not.equal(null);
+	});
 });
 
 describe('Middleware in PROD mode, SSG', () => {