[quick-edit] Partitioned cookies (#5016)

Author: penalosa

.vscode/settings.json

@@ -68,5 +68,8 @@
 	},
 	"[html]": {
 		"editor.defaultFormatter": "esbenp.prettier-vscode"
+	},
+	"[typescript]": {
+		"editor.defaultFormatter": "esbenp.prettier-vscode"
 	}
 }

fixtures/worker-app/package.json

@@ -13,11 +13,12 @@
 		"type:tests": "tsc -p ./tests/tsconfig.json"
 	},
 	"devDependencies": {
-		"wrangler": "workspace:*",
 		"@cloudflare/workers-tsconfig": "workspace:^",
-		"undici": "^5.28.3"
+		"undici": "^5.28.3",
+		"wrangler": "workspace:*"
 	},
 	"dependencies": {
+		"cookie": "^0.6.0",
 		"isomorphic-random-example": "workspace:^"
 	}
 }

fixtures/worker-app/src/index.js

@@ -1,3 +1,4 @@
+import cookie from "cookie";
 import { randomBytes } from "isomorphic-random-example";
 import { now } from "./dep";
 import { logErrors } from "./log";
@@ -15,10 +16,29 @@ export default {
 	async fetch(request) {
 		console.log("request log");
 
-		const { pathname, origin } = new URL(request.url);
+		const { pathname, origin, hostname, host } = new URL(request.url);
 		if (pathname === "/random") return new Response(hexEncode(randomBytes(8)));
 		if (pathname === "/error") throw new Error("Oops!");
 		if (pathname === "/redirect") return Response.redirect(`${origin}/foo`);
+		if (pathname === "/cookie")
+			return new Response("", {
+				headers: [
+					[
+						"Set-Cookie",
+						cookie.serialize("hello", "world", {
+							domain: hostname,
+						}),
+					],
+					[
+						"Set-Cookie",
+						cookie.serialize("hello2", "world2", {
+							domain: host,
+							secure: true,
+						}),
+					],
+				],
+			});
+
 		if (request.headers.get("X-Test-URL") !== null) {
 			return new Response(request.url);
 		}

fixtures/worker-app/tests/index.test.ts

@@ -139,4 +139,15 @@ describe("'wrangler dev' correctly renders pages", () => {
 			`http://${ip}:${port}/foo`
 		);
 	});
+
+	it("rewrites set-cookie headers to the hostname, not host", async ({
+		expect,
+	}) => {
+		const response = await fetch(`http://${ip}:${port}/cookie`);
+
+		expect(response.headers.getSetCookie()).toStrictEqual([
+			`hello=world; Domain=${ip}`,
+			`hello2=world2; Domain=${ip}; Secure`,
+		]);
+	});
 });

packages/edge-preview-authenticated-proxy/package.json

@@ -13,10 +13,10 @@
 	"devDependencies": {
 		"@cloudflare/eslint-config-worker": "*",
 		"@cloudflare/workers-types": "^4.20230321.0",
-		"cookie": "^0.5.0",
+		"@types/cookie": "^0.6.0",
+		"cookie": "^0.6.0",
 		"promjs": "^0.4.2",
 		"toucan-js": "^3.1.0",
-		"wrangler": "workspace:*",
-		"@types/cookie": "^0.5.1"
+		"wrangler": "workspace:*"
 	}
 }

packages/edge-preview-authenticated-proxy/src/index.ts

@@ -297,6 +297,7 @@ async function updatePreviewToken(url: URL, env: Env, ctx: ExecutionContext) {
 				sameSite: "none",
 				httpOnly: true,
 				domain: url.hostname,
+				partitioned: true,
 			}),
 		},
 	});

packages/edge-preview-authenticated-proxy/tests/index.test.ts

@@ -142,7 +142,7 @@ compatibility_date = "2023-01-01"
 		expect(
 			removeUUID(resp.headers.get("set-cookie") ?? "")
 		).toMatchInlineSnapshot(
-			'"token=00000000-0000-0000-0000-000000000000; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; SameSite=None"'
+			'"token=00000000-0000-0000-0000-000000000000; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; Partitioned; SameSite=None"'
 		);
 		tokenId = (resp.headers.get("set-cookie") ?? "")
 			.split(";")[0]
@@ -152,7 +152,7 @@ compatibility_date = "2023-01-01"
 			{
 				method: "GET",
 				headers: {
-					cookie: `token=${tokenId}; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; SameSite=None`,
+					cookie: `token=${tokenId}`,
 				},
 			}
 		);
@@ -183,7 +183,7 @@ compatibility_date = "2023-01-01"
 		expect(
 			removeUUID(resp.headers.get("set-cookie") ?? "")
 		).toMatchInlineSnapshot(
-			'"token=00000000-0000-0000-0000-000000000000; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; SameSite=None"'
+			'"token=00000000-0000-0000-0000-000000000000; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; Partitioned; SameSite=None"'
 		);
 		tokenId = (resp.headers.get("set-cookie") ?? "")
 			.split(";")[0]
@@ -218,7 +218,7 @@ compatibility_date = "2023-01-01"
 			{
 				method: "GET",
 				headers: {
-					cookie: `token=${tokenId}; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; SameSite=None`,
+					cookie: `token=${tokenId}; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; Partitioned; SameSite=None`,
 				},
 			}
 		);
@@ -237,7 +237,7 @@ compatibility_date = "2023-01-01"
 			{
 				method: "GET",
 				headers: {
-					cookie: `token=${tokenId}; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; SameSite=None`,
+					cookie: `token=${tokenId}; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; Partitioned; SameSite=None`,
 				},
 				redirect: "manual",
 			}
@@ -255,7 +255,7 @@ compatibility_date = "2023-01-01"
 			{
 				method: "PUT",
 				headers: {
-					cookie: `token=${tokenId}; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; SameSite=None`,
+					cookie: `token=${tokenId}; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; Partitioned; SameSite=None`,
 				},
 				redirect: "manual",
 			}
@@ -270,7 +270,7 @@ compatibility_date = "2023-01-01"
 				method: "PUT",
 				headers: {
 					"X-Custom-Header": "custom",
-					cookie: `token=${tokenId}; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; SameSite=None`,
+					cookie: `token=${tokenId}; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; Partitioned; SameSite=None`,
 				},
 				redirect: "manual",
 			}
@@ -284,7 +284,7 @@ compatibility_date = "2023-01-01"
 			{
 				method: "PUT",
 				headers: {
-					cookie: `token=${tokenId}; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; SameSite=None`,
+					cookie: `token=${tokenId}; Domain=random-data.preview.devprod.cloudflare.dev; HttpOnly; Secure; Partitioned; SameSite=None`,
 				},
 				redirect: "manual",
 			}

packages/wrangler/templates/startDevWorker/ProxyWorker.ts

@@ -304,6 +304,8 @@ function insertLiveReloadScript(
  * so that this proxy is transparent to the Client Browser and User Worker.
  */
 function rewriteUrlRelatedHeaders(headers: Headers, from: URL, to: URL) {
+	const setCookie = headers.getAll("Set-Cookie");
+	headers.delete("Set-Cookie");
 	headers.forEach((value, key) => {
 		if (typeof value === "string" && value.includes(from.host)) {
 			headers.set(
@@ -312,4 +314,13 @@ function rewriteUrlRelatedHeaders(headers: Headers, from: URL, to: URL) {
 			);
 		}
 	});
+	for (const cookie of setCookie) {
+		headers.append(
+			"Set-Cookie",
+			cookie.replace(
+				new RegExp(`Domain=${from.hostname}($|;|,)`),
+				`Domain=${to.hostname}$1`
+			)
+		);
+	}
 }