Fix the site canonical URL to include the visitor token if necessary

.changeset/dry-knives-hide.md

@@ -0,0 +1,5 @@
+---
+"gitbook-v2": patch
+---
+
+Fix the site canonical URL to include the visitor token if necessary

packages/gitbook-v2/src/lib/data/urls.test.ts

@@ -1,6 +1,6 @@
 import { describe, expect, it } from 'bun:test';
 
-import { getURLLookupAlternatives, normalizeURL } from './urls';
+import { getSiteCanonicalURL, getURLLookupAlternatives, normalizeURL } from './urls';
 
 describe('getURLLookupAlternatives', () => {
     it('should return all URLs up to the root', () => {
@@ -364,6 +364,78 @@ describe('getURLLookupAlternatives', () => {
     });
 });
 
+describe('getSiteCanonicalURL', () => {
+    it('should have the jwt token in canonical url if token was from the source url', () => {
+        expect(
+            getSiteCanonicalURL(
+                {
+                    site: 'site_foo',
+                    siteSpace: 'sitesp_foo',
+                    basePath: '/foo/',
+                    siteBasePath: '/foo/',
+                    organization: 'org_foo',
+                    space: 'space_foo',
+                    pathname: '/hello/world',
+                    complete: false,
+                    apiToken: 'api_token_foo',
+                    canonicalUrl: 'https://example.com/docs/foo/hello/world',
+                },
+                {
+                    source: 'url',
+                    token: 'jwt_foo',
+                }
+            ).toString()
+        ).toEqual('https://example.com/docs/foo/hello/world?jwt_token=jwt_foo');
+    });
+
+    it('should not have the jwt token in canonical url if token was NOT from the source url', () => {
+        // va cookie
+        expect(
+            getSiteCanonicalURL(
+                {
+                    site: 'site_foo',
+                    siteSpace: 'sitesp_foo',
+                    basePath: '/foo/',
+                    siteBasePath: '/foo/',
+                    organization: 'org_foo',
+                    space: 'space_foo',
+                    pathname: '/hello/world',
+                    complete: false,
+                    apiToken: 'api_token_foo',
+                    canonicalUrl: 'https://example.com/docs/foo/hello/world',
+                },
+                {
+                    source: 'visitor-auth-cookie',
+                    basePath: '/foo/',
+                    token: 'jwt_foo',
+                }
+            ).toString()
+        ).toEqual('https://example.com/docs/foo/hello/world');
+
+        // gitbook visitor cookie
+        expect(
+            getSiteCanonicalURL(
+                {
+                    site: 'site_foo',
+                    siteSpace: 'sitesp_foo',
+                    basePath: '/foo/',
+                    siteBasePath: '/foo/',
+                    organization: 'org_foo',
+                    space: 'space_foo',
+                    pathname: '/hello/world',
+                    complete: false,
+                    apiToken: 'api_token_foo',
+                    canonicalUrl: 'https://example.com/docs/foo/hello/world',
+                },
+                {
+                    source: 'gitbook-visitor-cookie',
+                    token: 'jwt_foo',
+                }
+            ).toString()
+        ).toEqual('https://example.com/docs/foo/hello/world');
+    });
+});
+
 describe('normalizeURL', () => {
     it('should remove trailing slashes', () => {
         expect(normalizeURL(new URL('https://docs.mycompany.com/hello/'))).toEqual(

packages/gitbook-v2/src/lib/data/urls.ts

@@ -1,3 +1,6 @@
+import { VISITOR_AUTH_PARAM, type VisitorTokenLookup } from '@/lib/visitor-token';
+import type { PublishedSiteContent } from '@gitbook/api';
+
 /**
  * For a given GitBook URL, return a list of alternative URLs that could be matched against to lookup the content.
  * The approach is optimized to aim at reusing cached lookup results as much as possible.
@@ -110,6 +113,22 @@ export function getURLLookupAlternatives(input: URL) {
     return { urls: alternatives, basePath, changeRequest, revision };
 }
 
+/**
+ * Get the canonical URL for a resolved site,
+ * including the visitor token if available.
+ */
+export function getSiteCanonicalURL(
+    siteURLData: PublishedSiteContent,
+    visitorToken: VisitorTokenLookup
+): URL {
+    const siteCanonicalURL = new URL(siteURLData.canonicalUrl);
+    if (visitorToken?.source === 'url') {
+        siteCanonicalURL.searchParams.set(VISITOR_AUTH_PARAM, visitorToken.token);
+    }
+
+    return siteCanonicalURL;
+}
+
 /**
  * Normalize a URL to remove duplicate slashes and trailing slashes
  * and transform the pathname to lowercase.

packages/gitbook-v2/src/lib/data/visitor.test.ts

@@ -9,8 +9,8 @@ describe('getVisitorAuthBasePath', () => {
                 {
                     site: 'site_foo',
                     siteSpace: 'sitesp_foo',
-                    basePath: '/foo',
-                    siteBasePath: '/foo',
+                    basePath: '/foo/',
+                    siteBasePath: '/foo/',
                     organization: 'org_foo',
                     space: 'space_foo',
                     pathname: '/hello/world',

packages/gitbook-v2/src/middleware.ts

@@ -17,6 +17,7 @@ import { serveResizedImage } from '@/routes/image';
 import {
     DataFetcherError,
     getPublishedContentByURL,
+    getSiteCanonicalURL,
     getVisitorAuthBasePath,
     normalizeURL,
     throwIfDataError,
@@ -146,20 +147,17 @@ async function serveSiteRoutes(requestURL: URL, request: NextRequest) {
 
         // We use the host/origin from the canonical URL to ensure the links are
         // correctly generated when the site is proxied. e.g. https://proxy.gitbook.com/site/siteId/...
-        const siteCanonicalURL = new URL(siteURLData.canonicalUrl);
+        const siteCanonicalURL = getSiteCanonicalURL(siteURLData, visitorToken);
 
         //
         // Make sure the URL is clean of any va token after a successful lookup
         // The token is stored in a cookie that is set on the redirect response
         //
         const incomingURL = mode === 'url' ? requestURL : siteCanonicalURL;
-        const requestURLWithoutToken = normalizeVisitorAuthURL(incomingURL);
-        if (
-            requestURLWithoutToken !== incomingURL &&
-            requestURLWithoutToken.toString() !== incomingURL.toString()
-        ) {
+        const incomingURLWithoutToken = normalizeVisitorAuthURL(incomingURL);
+        if (incomingURLWithoutToken.toString() !== incomingURL.toString()) {
             return writeResponseCookies(
-                NextResponse.redirect(requestURLWithoutToken.toString()),
+                NextResponse.redirect(incomingURLWithoutToken.toString()),
                 cookies
             );
         }

packages/gitbook/src/lib/visitor-token.ts

@@ -2,7 +2,7 @@ import { type JwtPayload, jwtDecode } from 'jwt-decode';
 import type { NextRequest } from 'next/server';
 import hash from 'object-hash';
 
-const VISITOR_AUTH_PARAM = 'jwt_token';
+export const VISITOR_AUTH_PARAM = 'jwt_token';
 export const VISITOR_TOKEN_COOKIE = 'gitbook-visitor-token';
 
 /**