Fix visitor auth logic by testing multiple basePaths (GitbookIO#239)

* WIP

* add search params to target in multi-path

* Simplify

* Format

* Ensure trailing slash

* Add unit test

* fix lint

* use search from target in redirect

* remove unused import

---------

Co-authored-by: Samy Pessé <samypesse@gmail.com>

e2e/pages.spec.ts

@@ -5,6 +5,7 @@ import {
     CustomizationSettings,
 } from '@gitbook/api';
 import { test, Page } from '@playwright/test';
+import jwt from 'jsonwebtoken';
 import rison from 'rison';
 
 import { getContentTestURL } from '../tests/utils';
@@ -282,6 +283,184 @@ const testCases: TestsCase[] = [
             },
         ],
     },
+    {
+        name: 'Visitor Auth - Space',
+        baseUrl: `https://gitbook.gitbook.io/gbo-va-space/`,
+        tests: [
+            {
+                name: 'First',
+                url: (() => {
+                    const privateKey = '70b844d0-c519-4532-8586-5970ce48c537';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return `first?jwt_token=${token}`;
+                })(),
+            },
+            {
+                name: 'Second',
+                url: (() => {
+                    const privateKey = '70b844d0-c519-4532-8586-5970ce48c537';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return `second?jwt_token=${token}`;
+                })(),
+            },
+        ],
+    },
+    {
+        name: 'Visitor Auth - Collection',
+        baseUrl: `https://gitbook.gitbook.io/gbo-va-collection/`,
+        tests: [
+            {
+                name: 'Root',
+                url: (() => {
+                    const privateKey = 'af5688dc-f0b6-4146-9b1d-6d834c62c980';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return `?jwt_token=${token}`;
+                })(),
+            },
+            {
+                name: 'Primary (Space A)',
+                url: (() => {
+                    const privateKey = 'af5688dc-f0b6-4146-9b1d-6d834c62c980';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return `v/spacea?jwt_token=${token}`;
+                })(),
+            },
+            {
+                name: 'Space B',
+                url: (() => {
+                    const privateKey = 'af5688dc-f0b6-4146-9b1d-6d834c62c980';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return `v/spaceb?jwt_token=${token}`;
+                })(),
+            },
+            {
+                name: 'Space C',
+                url: (() => {
+                    const privateKey = 'af5688dc-f0b6-4146-9b1d-6d834c62c980';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return `v/spacec?jwt_token=${token}`;
+                })(),
+            },
+        ],
+    },
+    {
+        name: 'Visitor Auth - Space (custom domain)',
+        baseUrl: `https://test.gitbook.community/`,
+        tests: [
+            {
+                name: 'Root',
+                url: (() => {
+                    const privateKey = '19c8166f-c436-4ed1-a24e-60954b804021';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return `?jwt_token=${token}`;
+                })(),
+            },
+            {
+                name: 'First',
+                url: (() => {
+                    const privateKey = '19c8166f-c436-4ed1-a24e-60954b804021';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return `first?jwt_token=${token}`;
+                })(),
+            },
+            {
+                name: 'Custom page',
+                url: (() => {
+                    const privateKey = '19c8166f-c436-4ed1-a24e-60954b804021';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return `custom-page?jwt_token=${token}`;
+                })(),
+            },
+            {
+                name: 'Inner page',
+                url: (() => {
+                    const privateKey = '19c8166f-c436-4ed1-a24e-60954b804021';
+                    const token = jwt.sign(
+                        {
+                            name: 'gitbook-open-tests',
+                        },
+                        privateKey,
+                        {
+                            expiresIn: '24h',
+                        },
+                    );
+                    return `custom-page/inner-page?jwt_token=${token}`;
+                })(),
+            },
+        ],
+    },
     {
         name: 'Languages',
         baseUrl: 'https://gitbook.gitbook.io/test-1-1/',

package.json

@@ -70,6 +70,7 @@
         "@types/js-cookie": "^3.0.6",
         "@types/js-yaml": "^4.0.9",
         "@types/jsontoxml": "^1.0.5",
+        "@types/jsonwebtoken": "^9.0.6",
         "@types/katex": "^0.16.5",
         "@types/node": "^20",
         "@types/object-hash": "^3.0.6",
@@ -83,6 +84,7 @@
         "eslint": "^8",
         "eslint-config-next": "13.5.6",
         "eslint-plugin-import": "^2.29.0",
+        "jsonwebtoken": "^9.0.2",
         "postcss": "^8",
         "prettier": "^3.0.3",
         "psi": "^4.1.0",

src/lib/visitor-auth.test.ts

@@ -0,0 +1,64 @@
+import { it, describe, expect } from 'bun:test';
+import { NextRequest } from 'next/server';
+
+import {
+    getVisitorAuthCookieName,
+    getVisitorAuthCookieValue,
+    getVisitorAuthToken,
+} from './visitor-auth';
+
+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');
+    });
+
+    it('should return the token from the cookie root basepath', () => {
+        const request = nextRequest('https://example.com', {
+            [getVisitorAuthCookieName('/')]: { value: getVisitorAuthCookieValue('/', '123') },
+        });
+        expect(getVisitorAuthToken(request, request.nextUrl)).toEqual('123');
+    });
+
+    it('should return the token from the cookie root basepath for a sub-path', () => {
+        const request = nextRequest('https://example.com/hello/world', {
+            [getVisitorAuthCookieName('/')]: { value: getVisitorAuthCookieValue('/', '123') },
+        });
+        expect(getVisitorAuthToken(request, request.nextUrl)).toEqual('123');
+    });
+
+    it('should return the closest token from the path', () => {
+        const request = nextRequest('https://example.com/hello/world', {
+            [getVisitorAuthCookieName('/')]: { value: getVisitorAuthCookieValue('/', 'no') },
+            [getVisitorAuthCookieName('/hello/')]: {
+                value: getVisitorAuthCookieValue('/hello/', '123'),
+            },
+        });
+        expect(getVisitorAuthToken(request, request.nextUrl)).toEqual('123');
+    });
+
+    it('should return the token from the cookie in a collection type url', () => {
+        const request = nextRequest('https://example.com/hello/v/space1/cool', {
+            [getVisitorAuthCookieName('/hello/v/space1/')]: {
+                value: getVisitorAuthCookieValue('/hello/v/space1/', '123'),
+            },
+        });
+        expect(getVisitorAuthToken(request, request.nextUrl)).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();
+    });
+});
+
+function nextRequest(url: string, cookies: Record<string, { value: string }> = {}) {
+    const nextUrl = new URL(url);
+    // @ts-ignore
+    return {
+        url: nextUrl.toString(),
+        nextUrl,
+        headers: new Headers(),
+        cookies: Object.entries(cookies),
+    } as NextRequest;
+}

src/lib/visitor-auth.ts

@@ -47,22 +47,41 @@ export function normalizeVisitorAuthURL(url: URL): URL {
     return withoutVAParam;
 }
 
+/**
+ * Get all possible basePaths for a given URL. This is used to find the visitor
+ * authentication cookie token.
+ * It returns the longest one first, and the shortest one last.
+ */
+function getUrlBasePathCombinations(url: URL): string[] {
+    const parts = url.pathname.split('/').filter(Boolean);
+    const baseNames = ['/'];
+
+    for (let index = 0; index < parts.length; index++) {
+        baseNames.push('/' + parts.slice(0, index + 1).join('/') + '/');
+    }
+
+    return baseNames.reverse();
+}
+
 /**
  * Find the visitor authentication token from the request cookies. This is done by
  * checking all cookies for a matching "visitor authentication cookie" and returning the
  * best possible match for the current URL.
  */
 function getVisitorAuthTokenFromCookies(request: NextRequest, url: URL): string | undefined {
-    const urlPathParts = url.pathname.split('/').filter(Boolean);
-    const urlBasePath = urlPathParts.length === 0 ? null : `/${urlPathParts[0]}/`;
-    /**
-     * First, try to find a visitor authentication token for the current URL. The request could be
-     * something like example.gitbook.io/foo/bar, and we want to find the token for the `/foo/` base path.
-     * If we can't find a token for the current URL, we'll try to find a token for the `/` base path. These
-     * are the only two possible base paths for a given URL for which we try to find a token.
-     */
-    const found = urlBasePath ? findVisitorAuthCookieForBasePath(request, urlBasePath) : null;
-    return found ?? findVisitorAuthCookieForBasePath(request, '/');
+    const urlBasePaths = getUrlBasePathCombinations(url);
+    // Try to find a visitor authentication token for the current URL. The request
+    // for the content could be hosted on a base path like `/foo/v/bar` or `/foo` or just `/`
+    // We keep trying to find with each of these base paths until we find a token.
+    for (const basePath of urlBasePaths) {
+        const found = findVisitorAuthCookieForBasePath(request, basePath);
+        if (found) {
+            return found;
+        }
+    }
+
+    // couldn't find any token for the current URL
+    return undefined;
 }
 
 /**

src/middleware.ts

@@ -484,6 +484,7 @@ async function lookupSpaceInMultiPathMode(request: NextRequest, url: URL): Promi
         if (lookup.target === 'content') {
             // Redirect to the content URL in the same application
             const redirect = new URL(lookup.redirect);
+            redirect.search = url.search;
 
             return {
                 target: 'content',