From 41ea7d4e5534bc2b034cd65b4c91497a8ff5e1b0 Mon Sep 17 00:00:00 2001 From: Tushar Pandey Date: Tue, 18 Mar 2025 17:13:10 +0530 Subject: [PATCH 1/3] chunked cookies implementation with tests, fixed some warnings in other tests related to await expect --- src/server/auth-client.test.ts | 2 +- src/server/chunked-cookies.test.ts | 361 ++++++++++++++++++ src/server/cookies.ts | 144 ++++++- .../session/stateless-session-store.test.ts | 2 +- src/server/session/stateless-session-store.ts | 39 +- src/server/transaction-store.test.ts | 4 +- 6 files changed, 525 insertions(+), 27 deletions(-) create mode 100644 src/server/chunked-cookies.test.ts diff --git a/src/server/auth-client.test.ts b/src/server/auth-client.test.ts index 83997a96..8c7d551c 100644 --- a/src/server/auth-client.test.ts +++ b/src/server/auth-client.test.ts @@ -875,7 +875,7 @@ ca/T0LLtgmbMmxSv/MmzIg== const response = await authClient.handleLogin(request); expect(response.status).toEqual(500); - expect(await response.text()).toEqual( + expect(await response.text()).toContain( "An error occured while trying to initiate the login request." ); }); diff --git a/src/server/chunked-cookies.test.ts b/src/server/chunked-cookies.test.ts new file mode 100644 index 00000000..99042fee --- /dev/null +++ b/src/server/chunked-cookies.test.ts @@ -0,0 +1,361 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { + CookieOptions, + deleteChunkedCookie, + getChunkedCookie, + RequestCookies, + ResponseCookies, + setChunkedCookie +} from "./cookies"; + +// Create mock implementation for RequestCookies and ResponseCookies +const createMocks = () => { + const cookieStore = new Map(); + + const reqCookies = { + get: vi.fn((...args) => { + const name = typeof args[0] === "string" ? args[0] : args[0].name; + if (cookieStore.has(name)) { + return { name, value: cookieStore.get(name) }; + } + return undefined; + }), + getAll: vi.fn((...args) => { + if (args.length === 0) { + return Array.from(cookieStore.entries()).map(([name, value]) => ({ + name, + value + })); + } + const name = typeof args[0] === "string" ? args[0] : args[0].name; + return cookieStore.has(name) + ? [{ name, value: cookieStore.get(name) }] + : []; + }), + has: vi.fn((name) => cookieStore.has(name)), + set: vi.fn((...args) => { + const name = typeof args[0] === "string" ? args[0] : args[0].name; + const value = typeof args[0] === "string" ? args[1] : args[0].value; + cookieStore.set(name, value); + return reqCookies; + }), + delete: vi.fn((names) => { + if (Array.isArray(names)) { + return names.map((name) => cookieStore.delete(name)); + } + return cookieStore.delete(names); + }), + clear: vi.fn(() => { + cookieStore.clear(); + return reqCookies; + }), + get size() { + return cookieStore.size; + }, + [Symbol.iterator]: vi.fn(() => cookieStore.entries()) + }; + + const resCookies = { + get: vi.fn((...args) => { + const name = typeof args[0] === "string" ? args[0] : args[0].name; + if (cookieStore.has(name)) { + return { name, value: cookieStore.get(name) }; + } + return undefined; + }), + getAll: vi.fn((...args) => { + if (args.length === 0) { + return Array.from(cookieStore.entries()).map(([name, value]) => ({ + name, + value + })); + } + const name = typeof args[0] === "string" ? args[0] : args[0].name; + return cookieStore.has(name) + ? [{ name, value: cookieStore.get(name) }] + : []; + }), + has: vi.fn((name) => cookieStore.has(name)), + set: vi.fn((...args) => { + const name = typeof args[0] === "string" ? args[0] : args[0].name; + const value = typeof args[0] === "string" ? args[1] : args[0].value; + cookieStore.set(name, value); + return resCookies; + }), + delete: vi.fn((...args) => { + const name = typeof args[0] === "string" ? args[0] : args[0].name; + cookieStore.delete(name); + return resCookies; + }), + toString: vi.fn(() => { + return Array.from(cookieStore.entries()) + .map(([name, value]) => `${name}=${value}`) + .join("; "); + }) + }; + + return { reqCookies, resCookies, cookieStore }; +}; + +describe("Chunked Cookie Utils", () => { + let reqCookies: RequestCookies; + let resCookies: ResponseCookies; + let cookieStore: Map; + + beforeEach(() => { + const mocks = createMocks(); + reqCookies = mocks.reqCookies; + resCookies = mocks.resCookies; + cookieStore = mocks.cookieStore; + + // Spy on console.warn + vi.spyOn(console, "warn").mockImplementation(() => {}); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + describe("setChunkedCookie", () => { + it("should set a single cookie when value is small enough", () => { + const name = "testCookie"; + const value = "small value"; + const options = { path: "/" } as CookieOptions; + + setChunkedCookie(name, value, options, reqCookies, resCookies); + + expect(resCookies.set).toHaveBeenCalledTimes(1); + expect(resCookies.set).toHaveBeenCalledWith(name, value, options); + expect(reqCookies.set).toHaveBeenCalledTimes(1); + expect(reqCookies.set).toHaveBeenCalledWith(name, value); + }); + + it("should split cookie into chunks when value exceeds max size", () => { + const name = "largeCookie"; + const options = { path: "/" } as CookieOptions; + + // Create a large string (8000 bytes) + const largeValue = "a".repeat(8000); + + setChunkedCookie(name, largeValue, options, reqCookies, resCookies); + + // Should create 3 chunks (8000 / 3500 ≈ 2.3, rounded up to 3) + expect(resCookies.set).toHaveBeenCalledTimes(3); + expect(reqCookies.set).toHaveBeenCalledTimes(3); + + // Check first chunk + expect(resCookies.set).toHaveBeenCalledWith( + `${name}__0`, + largeValue.slice(0, 3500), + options + ); + + // Check second chunk + expect(resCookies.set).toHaveBeenCalledWith( + `${name}__1`, + largeValue.slice(3500, 7000), + options + ); + + // Check third chunk + expect(resCookies.set).toHaveBeenCalledWith( + `${name}__2`, + largeValue.slice(7000), + options + ); + }); + + it("should log a warning when cookie size exceeds warning threshold", () => { + const name = "warningCookie"; + const options = { path: "/" } as CookieOptions; + + // Create a value that exceeds the warning threshold (4096 bytes) + const value = "a".repeat(4097); + + setChunkedCookie(name, value, options, reqCookies, resCookies); + + expect(console.warn).toHaveBeenCalled(); + }); + + describe("getChunkedCookie", () => { + it("should return undefined when cookie does not exist", () => { + const result = getChunkedCookie("nonexistent", reqCookies); + expect(result).toBeUndefined(); + }); + + it("should return cookie value when it exists as a regular cookie", () => { + const name = "simpleCookie"; + const value = "simple value"; + + // Setup the cookie + cookieStore.set(name, value); + + const result = getChunkedCookie(name, reqCookies); + + expect(result).toBe(value); + expect(reqCookies.get).toHaveBeenCalledWith(name); + }); + + it("should reconstruct value from chunks when cookie is chunked", () => { + const name = "chunkedCookie"; + const chunk0 = "chunk0 value"; + const chunk1 = "chunk1 value"; + const chunk2 = "chunk2 value"; + + // Add the chunks to the store (out of order) + cookieStore.set(`${name}__1`, chunk1); + cookieStore.set(`${name}__0`, chunk0); + cookieStore.set(`${name}__2`, chunk2); + + // Also add some unrelated cookies + cookieStore.set("otherCookie", "other value"); + + const result = getChunkedCookie(name, reqCookies); + + // Should combine chunks in proper order + expect(result).toBe(`${chunk0}${chunk1}${chunk2}`); + }); + + it("should return undefined when chunks are not in a complete sequence", () => { + const name = "incompleteCookie"; + + // Add incomplete chunks (missing chunk1) + cookieStore.set(`${name}__0`, "chunk0"); + cookieStore.set(`${name}__2`, "chunk2"); + + const result = getChunkedCookie(name, reqCookies); + + expect(result).toBeUndefined(); + expect(console.warn).toHaveBeenCalled(); + }); + }); + + describe("deleteChunkedCookie", () => { + it("should delete the regular cookie", () => { + const name = "regularCookie"; + cookieStore.set(name, "regular value"); + + deleteChunkedCookie(name, reqCookies, resCookies); + + expect(resCookies.delete).toHaveBeenCalledWith(name); + }); + + it("should delete all chunks of a cookie", () => { + const name = "chunkedCookie"; + + // Add chunks + cookieStore.set(`${name}__0`, "chunk0"); + cookieStore.set(`${name}__1`, "chunk1"); + cookieStore.set(`${name}__2`, "chunk2"); + + // Add unrelated cookie + cookieStore.set("otherCookie", "other value"); + + deleteChunkedCookie(name, reqCookies, resCookies); + + // Should delete main cookie and 3 chunks + expect(resCookies.delete).toHaveBeenCalledTimes(4); + expect(resCookies.delete).toHaveBeenCalledWith(name); + expect(resCookies.delete).toHaveBeenCalledWith(`${name}__0`); + expect(resCookies.delete).toHaveBeenCalledWith(`${name}__1`); + expect(resCookies.delete).toHaveBeenCalledWith(`${name}__2`); + // Should not delete unrelated cookies + expect(resCookies.delete).not.toHaveBeenCalledWith("otherCookie"); + }); + }); + + describe("Edge Cases", () => { + it("should handle empty values correctly", () => { + const name = "emptyCookie"; + const value = ""; + const options = { path: "/" } as CookieOptions; + + setChunkedCookie(name, value, options, reqCookies, resCookies); + + expect(resCookies.set).toHaveBeenCalledTimes(1); + expect(resCookies.set).toHaveBeenCalledWith(name, value, options); + }); + + it("should handle values at the exact chunk boundary", () => { + const name = "boundaryValueCookie"; + const value = "a".repeat(3500); // Exactly MAX_CHUNK_SIZE + const options = { path: "/" } as CookieOptions; + + setChunkedCookie(name, value, options, reqCookies, resCookies); + + // Should still fit in one cookie + expect(resCookies.set).toHaveBeenCalledTimes(1); + expect(resCookies.set).toHaveBeenCalledWith(name, value, options); + }); + + it("should handle special characters in cookie values", () => { + const name = "specialCharCookie"; + const value = + '{"special":"characters","with":"quotation marks","and":"😀 emoji"}'; + const options = { path: "/" } as CookieOptions; + + setChunkedCookie(name, value, options, reqCookies, resCookies); + + expect(resCookies.set).toHaveBeenCalledWith(name, value, options); + + // Setup for retrieval + cookieStore.set(name, value); + + const result = getChunkedCookie(name, reqCookies); + expect(result).toBe(value); + }); + + it("should handle multi-byte characters correctly", () => { + const name = "multiByteCookie"; + // Create a test string with multi-byte characters (emojis) + const value = "Hello 😀 world 🌍 with emojis 🎉"; + const options = { path: "/" } as CookieOptions; + + // Store the cookie + setChunkedCookie(name, value, options, reqCookies, resCookies); + + // For the retrieval test, manually set up the cookies + // We're testing the retrieval functionality, not the chunking itself + cookieStore.clear(); + cookieStore.set(name, value); + + // Verify retrieval works correctly with multi-byte characters + const result = getChunkedCookie(name, reqCookies); + expect(result).toBe(value); + + // Verify emoji characters were preserved + expect(result).toContain("😀"); + expect(result).toContain("🌍"); + expect(result).toContain("🎉"); + }); + + it("should handle very large cookies properly", () => { + const name = "veryLargeCookie"; + const value = "a".repeat(10000); // Will create multiple chunks + const options = { path: "/" } as CookieOptions; + + setChunkedCookie(name, value, options, reqCookies, resCookies); + + // Get chunks count (10000 / 3500 ≈ 2.86, so we need 3 chunks) + const expectedChunks = Math.ceil(10000 / 3500); + + expect(resCookies.set).toHaveBeenCalledTimes(expectedChunks); + + // Clear and set up cookies for retrieval test + cookieStore.clear(); + + // Setup for getChunkedCookie retrieval + for (let i = 0; i < expectedChunks; i++) { + const start = i * 3500; + const end = Math.min((i + 1) * 3500, 10000); + cookieStore.set(`${name}__${i}`, value.slice(start, end)); + } + + const result = getChunkedCookie(name, reqCookies); + expect(result).toBe(value); + expect(result!.length).toBe(10000); + }); + }); + }); +}); diff --git a/src/server/cookies.ts b/src/server/cookies.ts index eba871aa..ea9eace8 100644 --- a/src/server/cookies.ts +++ b/src/server/cookies.ts @@ -1,4 +1,8 @@ -import { RequestCookies, ResponseCookies } from "@edge-runtime/cookies"; +import { + RequestCookie, + RequestCookies, + ResponseCookies +} from "@edge-runtime/cookies"; import hkdf from "@panva/hkdf"; import * as jose from "jose"; @@ -118,3 +122,141 @@ export type ReadonlyRequestCookies = Omit< Pick; export { ResponseCookies }; export { RequestCookies }; + +// Configuration +const MAX_CHUNK_SIZE = 3500; // Slightly under 4KB +const CHUNK_PREFIX = "__"; +const CHUNK_INDEX_REGEX = new RegExp(`${CHUNK_PREFIX}(\\d+)$`); +const COOKIE_SIZE_WARNING_THRESHOLD = 4096; + +/** + * Retrieves the index of a cookie based on its name. + * + * @param name - The name of the cookie. + * @returns The index of the cookie. Returns undefined if no index is found. + */ +const getChunkedCookieIndex = (name: string): number | undefined => { + const match = CHUNK_INDEX_REGEX.exec(name); + if (!match) { + return undefined; + } + return parseInt(match[1], 10); +}; + +/** + * Retrieves all cookies from the request that have names starting with a specific prefix. + * + * @param reqCookies - The cookies from the request. + * @param name - The base name of the cookies to retrieve. + * @returns An array of cookies that have names starting with the specified prefix. + */ +const getAllChunkedCookies = ( + reqCookies: RequestCookies, + name: string +): RequestCookie[] => { + const PREFIX = `${name}${CHUNK_PREFIX}`; + return reqCookies.getAll().filter((cookie) => cookie.name.startsWith(PREFIX)); +}; + +/** + * Splits cookie value into multiple chunks if needed + */ +export function setChunkedCookie( + name: string, + value: string, + options: CookieOptions, + reqCookies: RequestCookies, + resCookies: ResponseCookies +): void { + const valueBytes = new TextEncoder().encode(value).length; + + if (valueBytes > COOKIE_SIZE_WARNING_THRESHOLD) { + console.warn( + `The cookie size exceeds ${COOKIE_SIZE_WARNING_THRESHOLD} bytes, which may cause issues in some browsers. ` + + "Consider removing any unnecessary custom claims from the access token or the user profile. " + + "Alternatively, you can use a stateful session implementation to store the session data in a data store." + ); + } + + // If value fits in a single cookie, set it directly + if (valueBytes <= MAX_CHUNK_SIZE) { + resCookies.set(name, value, options); + // to enable read-after-write in the same request for middleware + reqCookies.set(name, value); + return; + } + + // Split value into chunks + let position = 0; + let chunkIndex = 0; + + while (position < value.length) { + const chunk = value.slice(position, position + MAX_CHUNK_SIZE); + const chunkName = `${name}${CHUNK_PREFIX}${chunkIndex}`; + + resCookies.set(chunkName, chunk, options); + // to enable read-after-write in the same request for middleware + reqCookies.set(chunkName, chunk); + position += MAX_CHUNK_SIZE; + chunkIndex++; + } +} + +/** + * Reconstructs cookie value from chunks + */ +export function getChunkedCookie( + name: string, + reqCookies: RequestCookies +): string | undefined { + // Check if regular cookie exists + const cookie = reqCookies.get(name); + if (cookie?.value) { + return cookie.value; + } + + const chunks = getAllChunkedCookies(reqCookies, name).sort( + // Extract index from cookie name and sort numerically + (first, second) => { + return ( + getChunkedCookieIndex(first.name)! - getChunkedCookieIndex(second.name)! + ); + } + ); + + if (chunks.length === 0) { + return undefined; + } + + // Validate sequence integrity - check for missing chunks + const highestIndex = getChunkedCookieIndex(chunks[chunks.length - 1].name)!; + if (chunks.length !== highestIndex + 1) { + console.warn( + `Incomplete chunked cookie '${name}': Found ${chunks.length} chunks, expected ${highestIndex + 1}` + ); + + // TODO: Invalid sequence, delete all chunks + // deleteChunkedCookie(name, reqCookies, resCookies); + + return undefined; + } + + // Combine all chunks + return chunks.map((c) => c.value).join(""); +} + +/** + * Deletes all chunks of a cookie + */ +export function deleteChunkedCookie( + name: string, + reqCookies: RequestCookies, + resCookies: ResponseCookies +): void { + // Delete main cookie + resCookies.delete(name); + + getAllChunkedCookies(reqCookies, name).forEach((cookie) => { + resCookies.delete(cookie.name); // Delete each filtered cookie + }); +} diff --git a/src/server/session/stateless-session-store.test.ts b/src/server/session/stateless-session-store.test.ts index 74bfbb10..70c18392 100644 --- a/src/server/session/stateless-session-store.test.ts +++ b/src/server/session/stateless-session-store.test.ts @@ -476,7 +476,7 @@ describe("Stateless Session Store", async () => { secret }); - expect( + await expect( sessionStore.delete(requestCookies, responseCookies) ).resolves.not.toThrow(); }); diff --git a/src/server/session/stateless-session-store.ts b/src/server/session/stateless-session-store.ts index 9e71d3f0..0d597678 100644 --- a/src/server/session/stateless-session-store.ts +++ b/src/server/session/stateless-session-store.ts @@ -1,4 +1,4 @@ -import { SessionData } from "../../types"; +import { CookieOptions, SessionData } from "../../types"; import * as cookies from "../cookies"; import { AbstractSessionStore, @@ -39,8 +39,8 @@ export class StatelessSessionStore extends AbstractSessionStore { async get(reqCookies: cookies.RequestCookies) { const cookieValue = - reqCookies.get(this.sessionCookieName)?.value || - reqCookies.get(LEGACY_COOKIE_NAME)?.value; + cookies.getChunkedCookie(this.sessionCookieName, reqCookies) ?? + cookies.getChunkedCookie(LEGACY_COOKIE_NAME, reqCookies); if (!cookieValue) { return null; @@ -64,33 +64,28 @@ export class StatelessSessionStore extends AbstractSessionStore { const jwe = await cookies.encrypt(session, this.secret); const maxAge = this.calculateMaxAge(session.internal.createdAt); const cookieValue = jwe.toString(); - - resCookies.set(this.sessionCookieName, jwe.toString(), { + const options: CookieOptions = { ...this.cookieConfig, maxAge - }); - // to enable read-after-write in the same request for middleware - reqCookies.set(this.sessionCookieName, cookieValue); + }; - // check if the session cookie size exceeds 4096 bytes, and if so, log a warning - const cookieJarSizeTest = new cookies.ResponseCookies(new Headers()); - cookieJarSizeTest.set(this.sessionCookieName, cookieValue, { - ...this.cookieConfig, - maxAge - }); - if (new TextEncoder().encode(cookieJarSizeTest.toString()).length >= 4096) { - console.warn( - "The session cookie size exceeds 4096 bytes, which may cause issues in some browsers. " + - "Consider removing any unnecessary custom claims from the access token or the user profile. " + - "Alternatively, you can use a stateful session implementation to store the session data in a data store." - ); - } + cookies.setChunkedCookie( + this.sessionCookieName, + cookieValue, + options, + reqCookies, + resCookies + ); } async delete( _reqCookies: cookies.RequestCookies, resCookies: cookies.ResponseCookies ) { - await resCookies.delete(this.sessionCookieName); + cookies.deleteChunkedCookie( + this.sessionCookieName, + _reqCookies, + resCookies + ); } } diff --git a/src/server/transaction-store.test.ts b/src/server/transaction-store.test.ts index 164fd9d9..ad4834e2 100644 --- a/src/server/transaction-store.test.ts +++ b/src/server/transaction-store.test.ts @@ -117,7 +117,7 @@ describe("Transaction Store", async () => { secret }); - expect(() => + await expect(() => transactionStore.save(responseCookies, transactionState) ).rejects.toThrowError(); }); @@ -283,7 +283,7 @@ describe("Transaction Store", async () => { secret }); - expect( + await expect( transactionStore.delete(responseCookies, "non-existent-state") ).resolves.not.toThrow(); }); From aa0c86b3390235c99617c87d55e9aa6590cd8f4b Mon Sep 17 00:00:00 2001 From: Tushar Pandey Date: Tue, 18 Mar 2025 17:46:25 +0530 Subject: [PATCH 2/3] added docstrings --- src/server/cookies.ts | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/src/server/cookies.ts b/src/server/cookies.ts index ea9eace8..9141f25d 100644 --- a/src/server/cookies.ts +++ b/src/server/cookies.ts @@ -123,7 +123,7 @@ export type ReadonlyRequestCookies = Omit< export { ResponseCookies }; export { RequestCookies }; -// Configuration +// Chunked cookies Configuration const MAX_CHUNK_SIZE = 3500; // Slightly under 4KB const CHUNK_PREFIX = "__"; const CHUNK_INDEX_REGEX = new RegExp(`${CHUNK_PREFIX}(\\d+)$`); @@ -159,7 +159,18 @@ const getAllChunkedCookies = ( }; /** - * Splits cookie value into multiple chunks if needed + * Sets a cookie with the given name and value, splitting it into chunks if necessary. + * + * If the value exceeds the maximum chunk size, it will be split into multiple cookies + * with names suffixed by a chunk index. + * + * @param name - The name of the cookie. + * @param value - The value to be stored in the cookie. + * @param options - Options for setting the cookie. + * @param reqCookies - The request cookies object, used to enable read-after-write in the same request for middleware. + * @param resCookies - The response cookies object, used to set the cookies in the response. + * + * @throws {Error} If the cookie size exceeds the warning threshold. */ export function setChunkedCookie( name: string, @@ -203,7 +214,13 @@ export function setChunkedCookie( } /** - * Reconstructs cookie value from chunks + * Retrieves a chunked cookie by its name from the request cookies. + * If a regular cookie with the given name exists, it returns its value. + * Otherwise, it attempts to retrieve and combine all chunks of the cookie. + * + * @param name - The name of the cookie to retrieve. + * @param reqCookies - The request cookies object. + * @returns The combined value of the chunked cookie, or `undefined` if the cookie does not exist or is incomplete. */ export function getChunkedCookie( name: string, @@ -236,6 +253,7 @@ export function getChunkedCookie( ); // TODO: Invalid sequence, delete all chunks + // this cannot be done here rn because we don't have access to the response cookies // deleteChunkedCookie(name, reqCookies, resCookies); return undefined; @@ -246,7 +264,11 @@ export function getChunkedCookie( } /** - * Deletes all chunks of a cookie + * Deletes a chunked cookie and all its associated chunks from the response cookies. + * + * @param name - The name of the main cookie to delete. + * @param reqCookies - The request cookies object containing all cookies from the request. + * @param resCookies - The response cookies object to manipulate the cookies in the response. */ export function deleteChunkedCookie( name: string, From c1f3f406727ba34023a83ebbe187319a53404959 Mon Sep 17 00:00:00 2001 From: Tushar Pandey Date: Wed, 19 Mar 2025 15:21:24 +0530 Subject: [PATCH 3/3] use regex match instead of startsWith --- src/server/cookies.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/server/cookies.ts b/src/server/cookies.ts index 9141f25d..4a5a84b8 100644 --- a/src/server/cookies.ts +++ b/src/server/cookies.ts @@ -154,8 +154,10 @@ const getAllChunkedCookies = ( reqCookies: RequestCookies, name: string ): RequestCookie[] => { - const PREFIX = `${name}${CHUNK_PREFIX}`; - return reqCookies.getAll().filter((cookie) => cookie.name.startsWith(PREFIX)); + const chunkedCookieRegex = new RegExp(`^${name}${CHUNK_PREFIX}\\d+$`); + return reqCookies + .getAll() + .filter((cookie) => chunkedCookieRegex.test(cookie.name)); }; /**