Adapt code to pull token from customer backend generated custom cookies (#2650)

spastorelli · emmerich · web-flow · commit 98245e5f2026 · 2024-12-19T18:29:32.000+01:00
Co-authored-by: Steven Hall &lt;steven@gitbook.io&gt;
diff --git a/.changeset/blue-baboons-exercise.md b/.changeset/blue-baboons-exercise.md
@@ -0,0 +1,5 @@
+---
+'gitbook': minor
+---
+
+Adapt code to pull token from customer backend generated custom cookies
diff --git a/packages/gitbook/e2e/pages.spec.ts b/packages/gitbook/e2e/pages.spec.ts
@@ -17,7 +17,11 @@ import jwt from 'jsonwebtoken';
 import rison from 'rison';
 import { DeepPartial } from 'ts-essentials';
 
-import { getVisitorAuthCookieName, getVisitorAuthCookieValue } from '@/lib/visitor-auth';
+import {
+    getVisitorAuthCookieName,
+    getVisitorAuthCookieValue,
+    VISITOR_TOKEN_COOKIE,
+} from '@/lib/visitor-token';
 
 import { getContentTestURL } from '../tests/utils';
 
@@ -1081,6 +1085,149 @@ const testCases: TestsCase[] = [
             },
         ],
     },
+    {
+        name: 'Adaptive Content - Public',
+        baseUrl: `https://gitbook-open-e2e-sites.gitbook.io/adaptive-content-public/`,
+        tests: [
+            {
+                name: 'No custom cookie',
+                url: '',
+                run: async (page) => {
+                    const welcomePage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Welcome Page' });
+                    const alphaUserPage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Alpha User' });
+                    const betaUserPage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Beta User' });
+
+                    await expect(welcomePage).toBeVisible();
+                    await expect(alphaUserPage).toHaveCount(0);
+                    await expect(betaUserPage).toHaveCount(0);
+                },
+            },
+            {
+                name: 'Custom cookie with isAlphaUser claim',
+                cookies: (() => {
+                    const privateKey = '4ddd3c2f-e4b7-4e73-840b-526c3be19746';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                            isAlphaUser: true,
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return [
+                        {
+                            name: VISITOR_TOKEN_COOKIE,
+                            value: token,
+                            httpOnly: true,
+                        },
+                    ];
+                })(),
+                url: '',
+                run: async (page) => {
+                    const welcomePage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Welcome Page' });
+                    const alphaUserPage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Alpha User' });
+                    const betaUserPage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Beta User' });
+
+                    await expect(welcomePage).toBeVisible();
+                    await expect(alphaUserPage).toBeVisible();
+                    await expect(betaUserPage).toHaveCount(0);
+                },
+            },
+            {
+                name: 'Custom cookie with isBetaUser claim',
+                cookies: (() => {
+                    const privateKey = '4ddd3c2f-e4b7-4e73-840b-526c3be19746';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                            isBetaUser: true,
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return [
+                        {
+                            name: VISITOR_TOKEN_COOKIE,
+                            value: token,
+                            httpOnly: true,
+                        },
+                    ];
+                })(),
+                url: '',
+                run: async (page) => {
+                    const welcomePage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Welcome Page' });
+                    const alphaUserPage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Alpha User' });
+                    const betaUserPage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Beta User' });
+
+                    await expect(welcomePage).toBeVisible();
+                    await expect(betaUserPage).toBeVisible();
+                    await expect(alphaUserPage).toHaveCount(0);
+                },
+            },
+            {
+                name: 'Custom cookie with isAlphaUser & isBetaUser claims',
+                cookies: (() => {
+                    const privateKey = '4ddd3c2f-e4b7-4e73-840b-526c3be19746';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                            isAlphaUser: true,
+                            isBetaUser: true,
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return [
+                        {
+                            name: VISITOR_TOKEN_COOKIE,
+                            value: token,
+                            httpOnly: true,
+                        },
+                    ];
+                })(),
+                url: '',
+                run: async (page) => {
+                    const welcomePage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Welcome Page' });
+                    const alphaUserPage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Alpha User' });
+                    const betaUserPage = page
+                        .locator('a[class*="group\\/toclink"]')
+                        .filter({ hasText: 'Beta User' });
+
+                    await expect(welcomePage).toBeVisible();
+                    await expect(betaUserPage).toBeVisible();
+                    await expect(alphaUserPage).toBeVisible();
+                },
+            },
+        ],
+    },
     {
         name: 'Tables',
         baseUrl: 'https://gitbook.gitbook.io/test-gitbook-open/',
diff --git a/packages/gitbook/src/lib/api.ts b/packages/gitbook/src/lib/api.ts
@@ -214,7 +214,8 @@ export const getPublishedContentByUrl = cache({
     get: async (
         url: string,
         visitorAuthToken: string | undefined,
-        redirectOnError: boolean | undefined,
+        // Prefer undefined for a better cache key.
+        redirectOnError: true | undefined,
         options: CacheFunctionOptions,
     ) => {
         try {
diff --git a/packages/gitbook/src/lib/visitor-token.test.ts b/packages/gitbook/src/lib/visitor-token.test.ts
@@ -2,23 +2,22 @@ import { it, describe, expect } from 'bun:test';
 import { NextRequest } from 'next/server';
 
 import {
-    VisitorAuthCookieValue,
     getVisitorAuthCookieName,
     getVisitorAuthCookieValue,
-    getVisitorAuthToken,
-} from './visitor-auth';
+    getVisitorToken,
+} from './visitor-token';
 
 describe('getVisitorAuthToken', () => {
     it('should return the token from the query parameters', () => {
         const request = nextRequest('https://example.com?jwt_token=123');
-        expect(getVisitorAuthToken(request, request.nextUrl)).toEqual('123');
+        expect(getVisitorToken(request, request.nextUrl)).toEqual({ source: 'url', token: '123' });
     });
 
     it('should return the token from the cookie root basepath', () => {
         const request = nextRequest('https://example.com', {
             [getVisitorAuthCookieName('/')]: { value: getVisitorAuthCookieValue('/', '123') },
         });
-        const visitorAuth = getVisitorAuthToken(request, request.nextUrl);
+        const visitorAuth = getVisitorToken(request, request.nextUrl);
         assertVisitorAuthCookieValue(visitorAuth);
         expect(visitorAuth.token).toEqual('123');
     });
@@ -27,7 +26,7 @@ describe('getVisitorAuthToken', () => {
         const request = nextRequest('https://example.com/hello/world', {
             [getVisitorAuthCookieName('/')]: { value: getVisitorAuthCookieValue('/', '123') },
         });
-        const visitorAuth = getVisitorAuthToken(request, request.nextUrl);
+        const visitorAuth = getVisitorToken(request, request.nextUrl);
         assertVisitorAuthCookieValue(visitorAuth);
         expect(visitorAuth.token).toEqual('123');
     });
@@ -39,7 +38,7 @@ describe('getVisitorAuthToken', () => {
                 value: getVisitorAuthCookieValue('/hello/', '123'),
             },
         });
-        const visitorAuth = getVisitorAuthToken(request, request.nextUrl);
+        const visitorAuth = getVisitorToken(request, request.nextUrl);
         assertVisitorAuthCookieValue(visitorAuth);
         expect(visitorAuth.token).toEqual('123');
     });
@@ -50,14 +49,14 @@ describe('getVisitorAuthToken', () => {
                 value: getVisitorAuthCookieValue('/hello/v/space1/', '123'),
             },
         });
-        const visitorAuth = getVisitorAuthToken(request, request.nextUrl);
+        const visitorAuth = getVisitorToken(request, request.nextUrl);
         assertVisitorAuthCookieValue(visitorAuth);
         expect(visitorAuth.token).toEqual('123');
     });
 
     it('should return undefined if no cookie and no query param', () => {
         const request = nextRequest('https://example.com');
-        expect(getVisitorAuthToken(request, request.nextUrl)).toBeUndefined();
+        expect(getVisitorToken(request, request.nextUrl)).toBeUndefined();
     });
 
     // For backwards compatibility
@@ -69,14 +68,21 @@ describe('getVisitorAuthToken', () => {
             },
         });
 
-        const visitorAuth = getVisitorAuthToken(request, request.nextUrl);
+        const visitorAuth = getVisitorToken(request, request.nextUrl);
         assertVisitorAuthCookieValue(visitorAuth);
         expect(visitorAuth.token).toEqual('gotcha');
     });
 });
 
-function assertVisitorAuthCookieValue(value: unknown): asserts value is VisitorAuthCookieValue {
-    if (value && typeof value === 'object' && 'token' in value) {
+function assertVisitorAuthCookieValue(
+    value: unknown,
+): asserts value is { source: 'visitor-auth-cookie'; basePath: string; token: string } {
+    if (
+        value &&
+        typeof value === 'object' &&
+        'source' in value &&
+        value.source === 'visitor-auth-cookie'
+    ) {
         return;
     }
 
diff --git a/packages/gitbook/src/lib/visitor-token.ts b/packages/gitbook/src/lib/visitor-token.ts
@@ -2,35 +2,64 @@ import type { NextRequest } from 'next/server';
 import hash from 'object-hash';
 
 const VISITOR_AUTH_PARAM = 'jwt_token';
-const VISITOR_AUTH_COOKIE_ROOT = 'gitbook-visitor-token';
+const VISITOR_AUTH_COOKIE_ROOT = 'gitbook-visitor-token~';
+export const VISITOR_TOKEN_COOKIE = 'gitbook-visitor-token';
 
 /**
  * The contents of the visitor authentication cookie.
  */
-export type VisitorAuthCookieValue = {
+type VisitorAuthCookieValue = {
     basePath: string;
     token: string;
 };
 
-export function isVisitorAuthTokenFromCookies(
-    visitorAuthToken: NonNullable<ReturnType<typeof getVisitorAuthToken>>,
-) {
-    return (
-        typeof visitorAuthToken !== 'string' &&
-        'basePath' in visitorAuthToken &&
-        'token' in visitorAuthToken
-    );
-}
+/**
+ * The result of a visitor token lookup.
+ */
+export type VisitorTokenLookup =
+    | {
+          /** A visitor token was found in the URL. */
+          source: 'url';
+          token: string;
+      }
+    | {
+          /** A visitor auth token was found in a VA cookie */
+          source: 'visitor-auth-cookie';
+          basePath: string;
+          token: string;
+      }
+    | {
+          /** A visitor token (not auth) was found in a cookie. */
+          source: 'gitbook-visitor-cookie';
+          token: string;
+      }
+    /** Not visitor token was found */
+    | undefined;
 
 /**
- * Get the visitor authentication token for the request. This token can either be in the
+ * Get the visitor token for the request. This token can either be in the
  * query parameters or stored as a cookie.
  */
-export function getVisitorAuthToken(
+export function getVisitorToken(
     request: NextRequest,
     url: URL | NextRequest['nextUrl'],
-): string | VisitorAuthCookieValue | undefined {
-    return url.searchParams.get(VISITOR_AUTH_PARAM) ?? getVisitorAuthTokenFromCookies(request, url);
+): VisitorTokenLookup {
+    const fromUrl = url.searchParams.get(VISITOR_AUTH_PARAM);
+
+    // Allow the empty string to come through
+    if (fromUrl !== null && fromUrl !== undefined) {
+        return { source: 'url', token: fromUrl };
+    }
+
+    const visitorAuthToken = getVisitorAuthTokenFromCookies(request, url);
+    if (visitorAuthToken) {
+        return { source: 'visitor-auth-cookie', ...visitorAuthToken };
+    }
+
+    const visitorCustomToken = getVisitorCustomTokenFromCookies(request);
+    if (visitorCustomToken) {
+        return { source: 'gitbook-visitor-cookie', token: visitorCustomToken };
+    }
 }
 
 /**
@@ -40,7 +69,7 @@ export function getVisitorAuthToken(
  * different content hosted on the same subdomain.
  */
 export function getVisitorAuthCookieName(basePath: string): string {
-    return `${VISITOR_AUTH_COOKIE_ROOT}~${hash(basePath)}`;
+    return `${VISITOR_AUTH_COOKIE_ROOT}${hash(basePath)}`;
 }
 
 /**
@@ -106,6 +135,20 @@ function getVisitorAuthTokenFromCookies(
     return undefined;
 }
 
+/**
+ * Return the value of a custom visitor cookie that can be set by third party backends
+ * when they authenticate their users off flow to relay information in the form of claims
+ * about the visitor.
+ *
+ * The cookie should contain as value a JWT encoded token that contains the claims of the visitor.
+ */
+function getVisitorCustomTokenFromCookies(request: NextRequest): string | undefined {
+    const visitorCustomCookie = Array.from(request.cookies).find(
+        ([, cookie]) => cookie.name === VISITOR_TOKEN_COOKIE,
+    );
+    return visitorCustomCookie ? visitorCustomCookie[1].value : undefined;
+}
+
 /**
  * Loop through all cookies and find the visitor authentication token for a given base path.
  */
diff --git a/packages/gitbook/src/middleware.ts b/packages/gitbook/src/middleware.ts
