diff --git a/packages/gitbook/src/lib/api-token-cookie.test.ts b/packages/gitbook/src/lib/api-token-cookie.test.ts new file mode 100644 index 000000000..8d7bc83a1 --- /dev/null +++ b/packages/gitbook/src/lib/api-token-cookie.test.ts @@ -0,0 +1,128 @@ +import { describe, expect, it } from 'bun:test'; + +import { + MAX_API_TOKEN_COOKIE_LENGTH, + getAPITokenFromCookies, + getAPITokenResponseCookies, +} from './api-token-cookie'; + +const cookieName = 'gitbook-api-token~test'; +const options = { + httpOnly: true, + sameSite: 'none' as const, + secure: true, + maxAge: 60 * 60, +}; + +describe('API token cookies', () => { + it('keeps tokens at the cookie limit in one cookie', () => { + const apiToken = 'a'.repeat(4_000); + + expect(getAPITokenResponseCookies({ cookies: [], cookieName, apiToken, options })).toEqual([ + { name: cookieName, value: apiToken, options }, + ]); + }); + + it('splits and reconstructs oversized tokens', () => { + const apiToken = `${'a'.repeat(4_000)}b`; + const cookies = getAPITokenResponseCookies({ cookies: [], cookieName, apiToken, options }); + + expect(cookies).toEqual([ + { name: cookieName, value: 'chunks:2', options }, + { name: `${cookieName}-0`, value: 'a'.repeat(4_000), options }, + { name: `${cookieName}-1`, value: 'b', options }, + ]); + expect(getAPITokenFromCookies(cookies, cookieName)).toBe(apiToken); + }); + + it('rejects tokens that would exceed the response cookie header limit', () => { + expect(() => + getAPITokenResponseCookies({ + cookies: [], + cookieName, + apiToken: 'a'.repeat(MAX_API_TOKEN_COOKIE_LENGTH + 1), + options, + }) + ).toThrow(`API token exceeds the ${MAX_API_TOKEN_COOKIE_LENGTH}-character cookie limit`); + }); + + it('reads legacy single-cookie tokens', () => { + expect( + getAPITokenFromCookies([{ name: cookieName, value: 'legacy-token' }], cookieName) + ).toBe('legacy-token'); + }); + + it('does not return a partial token when a chunk is missing', () => { + expect( + getAPITokenFromCookies( + [ + { name: cookieName, value: 'chunks:2' }, + { name: `${cookieName}-0`, value: 'first' }, + ], + cookieName + ) + ).toBeUndefined(); + }); + + it('does not return a token when the first chunk is missing', () => { + expect( + getAPITokenFromCookies( + [ + { name: cookieName, value: 'chunks:2' }, + { name: `${cookieName}-1`, value: 'second' }, + ], + cookieName + ) + ).toBeUndefined(); + }); + + it('returns the base cookie value when its chunk count is unknown', () => { + expect( + getAPITokenFromCookies( + [ + { name: cookieName, value: '2' }, + { name: `${cookieName}-0`, value: 'first' }, + { name: `${cookieName}-1`, value: 'second' }, + ], + cookieName + ) + ).toBe('2'); + expect(getAPITokenFromCookies([{ name: cookieName, value: 'chunks:02' }], cookieName)).toBe( + 'chunks:02' + ); + }); + + it('expires chunks no longer needed by a replacement token', () => { + const oldCookies = [ + { name: cookieName, value: 'chunks:3' }, + { name: `${cookieName}-0`, value: 'first' }, + { name: `${cookieName}-1`, value: 'second' }, + { name: `${cookieName}-2`, value: 'third' }, + ]; + + expect( + getAPITokenResponseCookies({ + cookies: oldCookies, + cookieName, + apiToken: 'replacement', + options, + }) + ).toEqual([ + { name: cookieName, value: 'replacement', options }, + { name: `${cookieName}-0`, value: '', options: { ...options, maxAge: 0 } }, + { name: `${cookieName}-1`, value: '', options: { ...options, maxAge: 0 } }, + { name: `${cookieName}-2`, value: '', options: { ...options, maxAge: 0 } }, + ]); + }); + + it('does not attempt to expire an unbounded number of forged chunks', () => { + expect( + getAPITokenResponseCookies({ + cookies: [{ name: cookieName, value: 'chunks:999999999' }], + cookieName, + apiToken: 'replacement', + options, + }) + ).toEqual([{ name: cookieName, value: 'replacement', options }]); + }); +}); diff --git a/packages/gitbook/src/lib/api-token-cookie.ts b/packages/gitbook/src/lib/api-token-cookie.ts new file mode 100644 index 000000000..4f161d8f6 --- /dev/null +++ b/packages/gitbook/src/lib/api-token-cookie.ts @@ -0,0 +1,136 @@ +import type { ResponseCookie, ResponseCookies } from './visitors'; + +const COOKIE_CHUNK_SIZE = 4_000; +const MAX_COOKIE_CHUNKS = 3; +export const MAX_API_TOKEN_COOKIE_LENGTH = COOKIE_CHUNK_SIZE * MAX_COOKIE_CHUNKS; + +type RequestCookie = Pick; + +/** + * + * Retrieves the API token from the provided cookies, handling both single-cookie and chunked representations. + * We need to split the token into multiple cookies if it exceeds the size limit of a single cookie (4,000 characters). + * Some sites go over that limit which then cause an infinite redirect loop. + */ +export function getAPITokenFromCookies( + cookies: readonly RequestCookie[], + cookieName: string +): string | undefined { + const cookie = cookies.find(({ name }) => name === cookieName); + if (!cookie) { + return undefined; + } + + const chunkCount = parseChunkCount(cookie.value); + // A base cookie without a recognized chunk marker is always a token value. + if (chunkCount === undefined) { + return cookie.value; + } + + if (chunkCount > MAX_COOKIE_CHUNKS) { + return undefined; + } + + const chunks = new Map(); + const chunkNamePrefix = `${cookieName}-`; + for (const { name, value } of cookies) { + if (!name.startsWith(chunkNamePrefix)) { + continue; + } + + const indexValue = name.slice(chunkNamePrefix.length); + const index = Number(indexValue); + if (/^\d+$/.test(indexValue) && Number.isSafeInteger(index) && index >= 0) { + chunks.set(index, value); + } + } + + // We don't have all the chunks, so we can't reconstruct the token. + // We consider this a missing token + if (chunks.size < chunkCount) { + return undefined; + } + + const tokenChunks: string[] = []; + for (let index = 0; index < chunkCount; index++) { + const chunk = chunks.get(index); + if (chunk === undefined) { + return undefined; + } + tokenChunks.push(chunk); + } + + return tokenChunks.join(''); +} + +export function getAPITokenResponseCookies(input: { + cookies: readonly RequestCookie[]; + cookieName: string; + apiToken: string; + options: NonNullable; +}): ResponseCookies { + const { cookies, cookieName, apiToken, options } = input; + if (apiToken.length > MAX_API_TOKEN_COOKIE_LENGTH) { + throw new APITokenCookieTooLargeError(); + } + + const chunks = splitIntoCookieChunks(apiToken); + const previousChunkCount = getPreviousChunkCount(cookies, cookieName); + const responseCookies: ResponseCookies = + chunks.length === 1 + ? [{ name: cookieName, value: apiToken, options }] + : [ + { name: cookieName, value: `chunks:${chunks.length}`, options }, + ...chunks.map((value, index) => ({ + name: `${cookieName}-${index}`, + value, + options, + })), + ]; + + const firstChunkToExpire = chunks.length === 1 ? 0 : chunks.length; + for (let index = firstChunkToExpire; index < previousChunkCount; index++) { + responseCookies.push({ + name: `${cookieName}-${index}`, + value: '', + options: { ...options, maxAge: 0 }, + }); + } + + return responseCookies; +} + +function splitIntoCookieChunks(value: string): string[] { + if (value.length <= COOKIE_CHUNK_SIZE) { + return [value]; + } + + const chunks: string[] = []; + for (let start = 0; start < value.length; start += COOKIE_CHUNK_SIZE) { + chunks.push(value.slice(start, start + COOKIE_CHUNK_SIZE)); + } + return chunks; +} + +function getPreviousChunkCount(cookies: readonly RequestCookie[], cookieName: string): number { + const cookie = cookies.find(({ name }) => name === cookieName); + const chunkCount = cookie ? parseChunkCount(cookie.value) : undefined; + return chunkCount && chunkCount <= MAX_COOKIE_CHUNKS ? chunkCount : 0; +} + +function parseChunkCount(value: string): number | undefined { + // This regex matches the format "chunks:N" where N is a number between 2 and 9 or any number with two or more digits. + const match = /^chunks:([2-9]|[1-9]\d+)$/.exec(value); + if (!match) { + return undefined; + } + + const count = Number(match[1]); + return Number.isSafeInteger(count) ? count : undefined; +} + +export class APITokenCookieTooLargeError extends Error { + constructor() { + super(`API token exceeds the ${MAX_API_TOKEN_COOKIE_LENGTH}-character cookie limit`); + } +} diff --git a/packages/gitbook/src/middleware.ts b/packages/gitbook/src/middleware.ts index d2edaee43..e4f4d5769 100644 --- a/packages/gitbook/src/middleware.ts +++ b/packages/gitbook/src/middleware.ts @@ -11,6 +11,11 @@ import type { NextRequest } from 'next/server'; import { NextResponse } from 'next/server'; import rison from 'rison'; +import { + MAX_API_TOKEN_COOKIE_LENGTH, + getAPITokenFromCookies, + getAPITokenResponseCookies, +} from '@/lib/api-token-cookie'; import type { SiteURLData } from '@/lib/context'; import { getContentSecurityPolicy } from '@/lib/csp'; import { validateSerializedCustomization } from '@/lib/customization'; @@ -646,22 +651,31 @@ async function serveWithQueryAPIToken(input: { // If found, we redirect to the same URL but with the token in the cookie const queryAPIToken = requestURL.searchParams.get('token'); if (queryAPIToken) { + if (queryAPIToken.length > MAX_API_TOKEN_COOKIE_LENGTH) { + return new Response('API token is too large', { + status: 400, + headers: { 'content-type': 'text/plain' }, + }); + } + requestURL.searchParams.delete('token'); - return writeResponseCookies(NextResponse.redirect(requestURL.toString()), [ - { - name: cookieName, - value: queryAPIToken, + return writeResponseCookies( + NextResponse.redirect(requestURL.toString()), + getAPITokenResponseCookies({ + cookies: requestCookies.getAll(), + cookieName, + apiToken: queryAPIToken, options: { httpOnly: true, sameSite: process.env.NODE_ENV === 'production' ? 'none' : undefined, secure: process.env.NODE_ENV === 'production', maxAge: 60 * 60, // 1 hour }, - }, - ]); + }) + ); } - const apiToken = requestCookies.get(cookieName)?.value; + const apiToken = getAPITokenFromCookies(requestCookies.getAll(), cookieName); return serve(apiToken ?? null); }