Implement API token cookie handling with chunking support (#4421)

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 }]);
+    });
+});

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<ResponseCookie, 'name' | 'value'>;
+
+/**
+ *
+ * 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<number, string>();
+    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<ResponseCookie['options']>;
+}): 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`);
+    }
+}

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);
 }