Implement API token compression to handle oversized tokens in URLs

.changeset/compress-large-api-token.md

@@ -0,0 +1,5 @@
+---
+"gitbook": patch
+---
+
+Compress oversized site API tokens in the rewritten URL to stay under URL length limits.

bun.lock

@@ -143,6 +143,7 @@
         "direction": "^2.0.1",
         "event-iterator": "^2.0.0",
         "feed": "^5.1.0",
+        "fflate": "^0.8.3",
         "flexsearch": "^0.8.212",
         "image-size": "^2.0.2",
         "js-cookie": "^3.0.5",
@@ -2117,6 +2118,8 @@
 
     "feed": ["feed@5.1.0", "", { "dependencies": { "xml-js": "^1.6.11" } }, "sha512-qGNhgYygnefSkAHHrNHqC7p3R8J0/xQDS/cYUud8er/qD9EFGWyCdUDfULHTJQN1d3H3WprzVwMc9MfB4J50Wg=="],
 
+    "fflate": ["fflate@0.8.3", "", {}, "sha512-tbZNuJrLwGUp3zshBtdy4W+ORxZuIh8a5ilyIEQDC5rY1f3U20JMry0Ll3WBzU58EZKsEuJFXhb5gwv8CsPvgA=="],
+
     "file-entry-cache": ["file-entry-cache@10.0.7", "", { "dependencies": { "flat-cache": "^6.1.7" } }, "sha512-txsf5fu3anp2ff3+gOJJzRImtrtm/oa9tYLN0iTuINZ++EyVR/nRrg2fKYwvG/pXDofcrvvb0scEbX3NyW/COw=="],
 
     "file-uri-to-path": ["file-uri-to-path@1.0.0", "", {}, "sha512-0Zt+s3L7Vf1biwWZ29aARiVYLx7iMGnEUl9x33fbB/j3jR81u/O2LbqK+Bm1CDSNDKVtJ/YjwY7TUd5SkeLQLw=="],

packages/gitbook/package.json

@@ -34,6 +34,7 @@
         "direction": "^2.0.1",
         "event-iterator": "^2.0.0",
         "feed": "^5.1.0",
+        "fflate": "^0.8.3",
         "flexsearch": "^0.8.212",
         "image-size": "^2.0.2",
         "js-cookie": "^3.0.5",

packages/gitbook/src/app/utils.ts

@@ -5,6 +5,7 @@ import rison from 'rison';
 import type { SiteAPIToken } from '@gitbook/api';
 
 import { getVisitorAuthClaims, getVisitorAuthClaimsFromToken } from '@/lib/adaptive';
+import { decompressAPIToken } from '@/lib/api-token-compression';
 import { type SiteURLData, fetchSiteContextByURLLookup, getBaseContext } from '@/lib/context';
 import { getDynamicCustomizationSettings } from '@/lib/customization';
 
@@ -124,7 +125,8 @@ function getModeFromParams(mode: string): RouteParamMode {
 export function getSiteURLDataFromParams(params: RouteLayoutParams): SiteURLData {
     try {
         const decoded = decodeURIComponent(params.siteData);
-        return rison.decode(decoded);
+        const data = rison.decode<SiteURLData>(decoded);
+        return { ...data, apiToken: decompressAPIToken(data.apiToken) };
     } catch (error) {
         console.error(
             `Returning 404 after failing to decode site data ${params.siteData}: ${error}`

packages/gitbook/src/lib/api-token-compression.test.ts

@@ -0,0 +1,88 @@
+import { describe, expect, it } from 'bun:test';
+import jwt from 'jsonwebtoken';
+import { inflateRawSync } from 'node:zlib';
+
+import {
+    API_TOKEN_COMPRESSION_THRESHOLD,
+    COMPRESSED_API_TOKEN_PREFIX,
+    compressAPITokenIfNeeded,
+    decompressAPIToken,
+} from './api-token-compression';
+
+const ALPHANUM = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
+
+function randomId() {
+    return Array.from({ length: 20 }, () => ALPHANUM[Math.floor(Math.random() * 62)]).join('');
+}
+
+function signToken(spaces: number, claims: Record<string, unknown> = {}) {
+    return jwt.sign(
+        {
+            kind: 'site',
+            organization: randomId(),
+            site: `site_${randomId()}`,
+            space: randomId(),
+            spaces: Array.from({ length: spaces }, randomId),
+            claims,
+        },
+        'secret'
+    );
+}
+
+describe('compressAPITokenIfNeeded', () => {
+    it('keeps tokens under the threshold unchanged', () => {
+        const token = signToken(10);
+        expect(token.length).toBeLessThanOrEqual(API_TOKEN_COMPRESSION_THRESHOLD);
+        expect(compressAPITokenIfNeeded(token)).toBe(token);
+    });
+
+    it('compresses large tokens and restores them exactly', () => {
+        const token = signToken(600, {
+            groups: Array.from({ length: 100 }, (_, i) => `group-${i}`),
+        });
+        expect(token.length).toBeGreaterThan(API_TOKEN_COMPRESSION_THRESHOLD);
+
+        const compressed = compressAPITokenIfNeeded(token);
+        expect(compressed.startsWith(COMPRESSED_API_TOKEN_PREFIX)).toBe(true);
+        expect(compressed.length).toBeLessThan(token.length * 0.85);
+        expect(decompressAPIToken(compressed)).toBe(token);
+    });
+
+    it('restores tokens with non-ASCII characters and dots in the payload', () => {
+        const token = signToken(600, { name: 'Zoë Ünal', 'a.b': 'é.ü.日本', email: 'z.u@ex.co' });
+        expect(decompressAPIToken(compressAPITokenIfNeeded(token))).toBe(token);
+    });
+
+    it('produces a stable output for the same token', () => {
+        const token = signToken(600);
+        expect(compressAPITokenIfNeeded(token)).toBe(compressAPITokenIfNeeded(token));
+    });
+
+    it('produces raw deflate data', () => {
+        const token = signToken(600);
+        const compressed = compressAPITokenIfNeeded(token);
+        const inflated = inflateRawSync(
+            Buffer.from(compressed.slice(COMPRESSED_API_TOKEN_PREFIX.length), 'base64url')
+        );
+        const [header, payload, signature] = token.split('.');
+        expect(inflated.toString()).toBe(
+            `${header}.${Buffer.from(payload!, 'base64url').toString()}.${signature}`
+        );
+    });
+
+    it('does not compress values that are not a JWT', () => {
+        const value = 'x'.repeat(API_TOKEN_COMPRESSION_THRESHOLD + 1);
+        expect(compressAPITokenIfNeeded(value)).toBe(value);
+    });
+});
+
+describe('decompressAPIToken', () => {
+    it('returns raw tokens unchanged', () => {
+        const token = signToken(10);
+        expect(decompressAPIToken(token)).toBe(token);
+    });
+
+    it('throws on corrupted compressed tokens', () => {
+        expect(() => decompressAPIToken(`${COMPRESSED_API_TOKEN_PREFIX}AAAA`)).toThrow();
+    });
+});

packages/gitbook/src/lib/api-token-compression.ts

@@ -0,0 +1,94 @@
+import { deflateSync, inflateSync } from 'fflate';
+
+/**
+ * Prefix marking a compressed API token. A JWT always starts with `eyJ` and `.` is not a
+ * base64url character, so it can't collide with a raw token.
+ */
+export const COMPRESSED_API_TOKEN_PREFIX = 'gbz1.';
+
+/**
+ * Tokens above this length are compressed before being encoded in the rewritten URL,
+ * which is limited to 16KB. Smaller tokens are kept as-is so their URLs (and cache keys) don't change.
+ */
+export const API_TOKEN_COMPRESSION_THRESHOLD = 14 * 1024;
+
+const DOT = 0x2e;
+
+/**
+ * Compress an API token if it is too large to safely fit in the rewritten URL.
+ * The payload is decoded before compressing, as deflate does much better on JSON than on base64.
+ */
+export function compressAPITokenIfNeeded(token: string): string {
+    if (token.length <= API_TOKEN_COMPRESSION_THRESHOLD) {
+        return token;
+    }
+
+    const parts = token.split('.');
+    if (parts.length !== 3) {
+        return token;
+    }
+    const [header, payload, signature] = parts as [string, string, string];
+
+    const encoder = new TextEncoder();
+    const bytes = concatBytes([
+        encoder.encode(`${header}.`),
+        base64UrlToBytes(payload),
+        encoder.encode(`.${signature}`),
+    ]);
+
+    return `${COMPRESSED_API_TOKEN_PREFIX}${bytesToBase64Url(deflateSync(bytes, { level: 9 }))}`;
+}
+
+/**
+ * Restore the original API token from a value produced by `compressAPITokenIfNeeded`.
+ * Values that are not compressed are returned as-is.
+ */
+export function decompressAPIToken(value: string): string {
+    if (!value.startsWith(COMPRESSED_API_TOKEN_PREFIX)) {
+        return value;
+    }
+
+    const bytes = inflateSync(base64UrlToBytes(value.slice(COMPRESSED_API_TOKEN_PREFIX.length)));
+
+    // Header and signature are base64url so they can't contain dots, but the payload JSON can.
+    const firstDot = bytes.indexOf(DOT);
+    const lastDot = bytes.lastIndexOf(DOT);
+    if (firstDot === -1 || firstDot === lastDot) {
+        throw new Error('Invalid compressed API token');
+    }
+
+    const decoder = new TextDecoder();
+    const header = decoder.decode(bytes.subarray(0, firstDot));
+    const payload = bytesToBase64Url(bytes.subarray(firstDot + 1, lastDot));
+    const signature = decoder.decode(bytes.subarray(lastDot + 1));
+
+    return `${header}.${payload}.${signature}`;
+}
+
+function concatBytes(chunks: Uint8Array[]): Uint8Array {
+    const result = new Uint8Array(chunks.reduce((total, chunk) => total + chunk.length, 0));
+    let offset = 0;
+    for (const chunk of chunks) {
+        result.set(chunk, offset);
+        offset += chunk.length;
+    }
+    return result;
+}
+
+// Implemented with atob/btoa as Buffer is not guaranteed in the edge runtime.
+function bytesToBase64Url(bytes: Uint8Array): string {
+    let binary = '';
+    for (let i = 0; i < bytes.length; i += 0x8000) {
+        binary += String.fromCharCode(...bytes.subarray(i, i + 0x8000));
+    }
+    return btoa(binary).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/, '');
+}
+
+function base64UrlToBytes(value: string): Uint8Array {
+    const binary = atob(value.replace(/-/g, '+').replace(/_/g, '/'));
+    const bytes = new Uint8Array(binary.length);
+    for (let i = 0; i < binary.length; i++) {
+        bytes[i] = binary.charCodeAt(i);
+    }
+    return bytes;
+}

packages/gitbook/src/lib/middleware.ts

@@ -2,6 +2,7 @@ import { headers } from 'next/headers';
 
 import { CustomizationDefaultThemeMode } from '@gitbook/api';
 
+import { decompressAPIToken } from './api-token-compression';
 import type { SiteURLData } from './context';
 
 export enum MiddlewareHeaders {
@@ -70,7 +71,8 @@ export async function getSiteURLDataFromMiddleware(): Promise<SiteURLData> {
         );
     }
 
-    return JSON.parse(siteURLData);
+    const data: SiteURLData = JSON.parse(siteURLData);
+    return { ...data, apiToken: decompressAPIToken(data.apiToken) };
 }
 
 /**

packages/gitbook/src/middleware.ts

@@ -22,6 +22,7 @@ import {
     trackServerInsightsEvents,
 } from './lib/tracking';
 import { AI_CATALOG_PATH, AI_CATALOG_WELL_KNOWN_PATH } from '@/lib/aiCatalog/paths';
+import { compressAPITokenIfNeeded } from '@/lib/api-token-compression';
 import { getAPITokenFromCookies, getAPITokenResponseCookies } from '@/lib/api-token-cookie';
 import { isChatGPTRequest } from '@/lib/chatgpt';
 import { MAX_CHUNKED_COOKIE_LENGTH } from '@/lib/chunked-cookies';
@@ -394,7 +395,7 @@ async function serveSiteRoutes(requestURL: URL, request: NextRequest) {
             revision: siteURLData.revision,
             shareKey: siteURLData.shareKey,
             preview: siteURLData.preview,
-            apiToken: siteURLData.apiToken,
+            apiToken: compressAPITokenIfNeeded(siteURLData.apiToken),
             imagesContextId: imagesContextId,
             contextId: siteURLData.contextId,
             isFallback: requestURL.searchParams.get('fallback') === 'true' ? true : undefined,