Enhance visitor country handling by prioritizing forwarded header on revalidation requests

.changeset/send-visitor-country.md

@@ -0,0 +1,5 @@
+---
+"gitbook": patch
+---
+
+Send the visitor country to resolvePublishedContentByUrl on dev, preview, staging and selected sites, using the forwarded `x-gitbook-country` header on revalidation requests.

packages/gitbook/src/lib/data/lookup.test.ts

@@ -105,3 +105,30 @@ describe('preview auth redirects', () => {
         }
     );
 });
+
+describe('visitor payload', () => {
+    it('forwards the visitor country to the API', async () => {
+        const calls: { visitor?: unknown }[] = [];
+        const apiClientSpy = spyOn(api, 'apiClient').mockReturnValue({
+            urls: {
+                async resolvePublishedContentByUrl(body: { visitor?: unknown }) {
+                    calls.push(body);
+                    return { data: { target: 'external', redirect: 'https://example.com' } };
+                },
+            },
+        } as unknown as ReturnType<typeof api.apiClient>);
+
+        try {
+            await lookupPublishedContentByUrl({
+                url: 'https://docs.example.com',
+                apiToken: null,
+                redirectOnError: false,
+                visitorPayload: { type: 'human', country: 'FR' },
+            });
+        } finally {
+            apiClientSpy.mockRestore();
+        }
+
+        expect(calls[0]?.visitor).toEqual({ type: 'human', country: 'FR' });
+    });
+});

packages/gitbook/src/lib/data/lookup.ts

@@ -17,7 +17,8 @@ interface LookupPublishedContentByUrlInput {
     url: string;
     redirectOnError: boolean;
     apiToken: string | null;
-    visitorPayload: SiteVisitorPayload;
+    // TODO: remove the country extension once @gitbook/api exposes visitor.country
+    visitorPayload: SiteVisitorPayload & { country?: string };
 }
 
 /**

packages/gitbook/src/lib/env/globals.ts

@@ -14,6 +14,12 @@ export const GITBOOK_RUNTIME = (process.env.GITBOOK_RUNTIME ?? 'unknown') as
     | 'cloudflare'
     | 'unknown';
 
+/**
+ * Deployment stage (set by the Cloudflare wrangler configs), defaulting to `dev` locally.
+ */
+export const GITBOOK_STAGE =
+    process.env.STAGE ?? (process.env.NODE_ENV === 'development' ? 'dev' : undefined);
+
 /**
  * Main host on which GitBook is running.
  */

packages/gitbook/src/lib/visitors.test.ts

@@ -7,11 +7,13 @@ import {
     getVisitorAuthCookieMaxAge,
     getVisitorAuthCookieName,
     getVisitorAuthCookieValue,
+    getVisitorCountry,
     getVisitorToken,
     getVisitorType,
     getVisitorUnsignedClaims,
     isRevalidationRequest,
     normalizeVisitorURL,
+    shouldSendVisitorCountry,
 } from './visitors';
 
 describe('getVisitorAuthToken', () => {
@@ -599,3 +601,76 @@ describe('isRevalidationRequest', () => {
         expect(isRevalidationRequest(new Headers())).toBe(false);
     });
 });
+
+describe('getVisitorCountry', () => {
+    const requestWith = (headers: Record<string, string>) => ({ headers: new Headers(headers) });
+
+    it('should prefer the OpenNext country header', () => {
+        expect(
+            getVisitorCountry(
+                requestWith({ 'x-open-next-country': 'FR', 'x-vercel-ip-country': 'US' })
+            )
+        ).toBe('FR');
+    });
+
+    it('should fall back to the Vercel country header', () => {
+        expect(getVisitorCountry(requestWith({ 'x-vercel-ip-country': 'US' }))).toBe('US');
+    });
+
+    it('should normalize the country code to uppercase', () => {
+        expect(getVisitorCountry(requestWith({ 'x-open-next-country': ' fr ' }))).toBe('FR');
+    });
+
+    it.each(['XX', 'T1', 'FRA', 'F', ''])('should ignore the invalid country code %p', (value) => {
+        expect(getVisitorCountry(requestWith({ 'x-open-next-country': value }))).toBeUndefined();
+    });
+
+    it('should return undefined when no country header is present', () => {
+        expect(getVisitorCountry(requestWith({}))).toBeUndefined();
+    });
+
+    it('should prefer the forwarded country header on revalidation requests', () => {
+        expect(
+            getVisitorCountry(
+                requestWith({
+                    'user-agent': 'gitbook-open-revalidation-worker',
+                    'x-gitbook-country': 'DE',
+                    'x-open-next-country': 'FR',
+                })
+            )
+        ).toBe('DE');
+    });
+
+    it('should ignore the forwarded country header outside revalidation requests', () => {
+        expect(
+            getVisitorCountry(
+                requestWith({ 'x-gitbook-country': 'DE', 'x-open-next-country': 'FR' })
+            )
+        ).toBe('FR');
+        expect(getVisitorCountry(requestWith({ 'x-gitbook-country': 'DE' }))).toBeUndefined();
+    });
+
+    it('should fall back to geolocation headers on revalidation requests without a forwarded country', () => {
+        expect(
+            getVisitorCountry(
+                requestWith({
+                    'user-agent': 'gitbook-open-revalidation-worker',
+                    'x-open-next-country': 'FR',
+                })
+            )
+        ).toBe('FR');
+    });
+});
+
+describe('shouldSendVisitorCountry', () => {
+    it.each(['dev', 'preview', 'staging'])('should be enabled on the %p stage', (stage) => {
+        expect(shouldSendVisitorCountry('docs.example.com', stage)).toBe(true);
+    });
+
+    it.each(['production', undefined])(
+        'should be disabled for other hostnames on the %p stage',
+        (stage) => {
+            expect(shouldSendVisitorCountry('docs.example.com', stage)).toBe(false);
+        }
+    );
+});

packages/gitbook/src/lib/visitors.ts

@@ -10,6 +10,7 @@ import {
     getChunkedCookieValue,
     getChunkedResponseCookies,
 } from './chunked-cookies';
+import { GITBOOK_STAGE } from './env';
 
 const VISITOR_AUTH_PARAM = 'jwt_token';
 const VISITOR_PARAM_PREFIX = 'visitor.';
@@ -106,6 +107,45 @@ export function getVisitorType(request: {
     return detection.detected && detection.method !== 'heuristic' ? 'agent' : 'human';
 }
 
+// Production hostnames for which the visitor country is sent while the feature is rolled out.
+const VISITOR_COUNTRY_HOSTNAMES = new Set<string>([]);
+const VISITOR_COUNTRY_STAGES = new Set(['dev', 'preview', 'staging']);
+
+/**
+ * Whether the visitor country should be sent when resolving the site URL.
+ */
+export function shouldSendVisitorCountry(
+    hostname: string,
+    stage: string | undefined = GITBOOK_STAGE
+): boolean {
+    return (
+        (!!stage && VISITOR_COUNTRY_STAGES.has(stage)) || VISITOR_COUNTRY_HOSTNAMES.has(hostname)
+    );
+}
+
+/**
+ * Get the ISO 3166-1 alpha-2 country code of the visitor from the geolocation headers,
+ * or from `x-gitbook-country` when the request comes from the revalidation worker.
+ */
+export function getVisitorCountry(request: { headers: Headers }): string | undefined {
+    // The revalidation worker forwards the original visitor country; only trust it from there.
+    const country = (
+        (isRevalidationRequest(request.headers) && request.headers.get('x-gitbook-country')) ||
+        request.headers.get('x-open-next-country') ||
+        request.headers.get('x-vercel-ip-country') ||
+        ''
+    )
+        .trim()
+        .toUpperCase();
+
+    // Cloudflare uses XX for unknown and T1 for Tor, which are not ISO codes.
+    if (!/^[A-Z]{2}$/.test(country) || country === 'XX' || country === 'T1') {
+        return undefined;
+    }
+
+    return country;
+}
+
 /**
  * Get the visitor data for the request potentially including:
  *   - a JWT token that may contain signed claims or can be used for VA authentication.

packages/gitbook/src/middleware.ts

@@ -56,11 +56,13 @@ import {
     type ResponseCookies,
     getPathScopedCookieName,
     getResponseCookiesForVisitorAuth,
+    getVisitorCountry,
     getVisitorData,
     getVisitorType,
     isRevalidationRequest,
     normalizeVisitorURL,
     serveVisitorClaimsDataRequest,
+    shouldSendVisitorCountry,
 } from '@/lib/visitors';
 import { waitUntil } from '@/lib/waitUntil';
 import { serveResizedImage } from '@/routes/image';
@@ -220,6 +222,10 @@ async function serveSiteRoutes(requestURL: URL, request: NextRequest) {
     //
     request.headers.delete('x-gitbook-disable-tracking');
 
+    const visitorCountry = shouldSendVisitorCountry(siteRequestURL.hostname)
+        ? getVisitorCountry(request)
+        : undefined;
+
     const withAPIToken = async (apiToken: string | null) => {
         const siteURLData = await throwIfDataError(
             lookupPublishedContentByUrl({
@@ -228,6 +234,7 @@ async function serveSiteRoutes(requestURL: URL, request: NextRequest) {
                     jwtToken: visitorToken?.token ?? undefined,
                     unsignedClaims,
                     type: getVisitorType(request),
+                    ...(visitorCountry ? { country: visitorCountry } : {}),
                 },
                 // When the visitor auth token is pulled from the cookie, set redirectOnError when calling resolvePublishedContentByUrl to allow
                 // redirecting when the token is invalid as we could be dealing with stale token stored in the cookie.