From 258ef413b1d71801566e03c025c68bbe3d4bbd1f Mon Sep 17 00:00:00 2001
From: Taran Vohra <taranvohra@outlook.com>
Date: Tue, 20 Feb 2024 17:11:23 +0530
Subject: [PATCH] Fix logic to fetch visitor auth token from cookies and
 cleanup URL params (#175)

* Remove va token from query params before using it for response

* sanitize
---
 src/lib/visitor-auth.ts | 12 +++++++++++-
 src/middleware.ts       |  7 +++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/lib/visitor-auth.ts b/src/lib/visitor-auth.ts
index b018b32c8..c9eb46357 100644
--- a/src/lib/visitor-auth.ts
+++ b/src/lib/visitor-auth.ts
@@ -38,13 +38,23 @@ export function getVisitorAuthCookieValue(basePath: string, token: string): stri
     return JSON.stringify(value);
 }
 
+/**
+ * Sanitize the URL by removing the visitor authentication token from the query parameters (if present).
+ */
+export function sanitizeVisitorAuthURL(url: URL): URL {
+    const withoutVAParam = new URL(url);
+    withoutVAParam.searchParams.delete(VISITOR_AUTH_PARAM);
+    return withoutVAParam;
+}
+
 /**
  * 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 urlBasePath = url.pathname.split('/').filter(Boolean)[0] ?? '';
+    const urlPathParts = url.pathname.split('/').filter(Boolean);
+    const urlBasePath = urlPathParts.length === 0 ? `/` : `/${urlPathParts[0]}/`;
 
     return Array.from(request.cookies).reduce<string | undefined>((acc, [name, cookie]) => {
         if (name === getVisitorAuthCookieName(urlBasePath)) {
diff --git a/src/middleware.ts b/src/middleware.ts
index e984eabf6..b7d0ede26 100644
--- a/src/middleware.ts
+++ b/src/middleware.ts
@@ -20,6 +20,7 @@ import {
     getVisitorAuthCookieName,
     getVisitorAuthCookieValue,
     getVisitorAuthToken,
+    sanitizeVisitorAuthURL,
 } from '@/lib/visitor-auth';
 
 export const config = {
@@ -166,8 +167,10 @@ export async function middleware(request: NextRequest) {
         headers.set('x-gitbook-api', apiEndpoint);
     }
 
-    const target = new URL(rewritePathname, request.nextUrl.toString());
-    target.search = url.search;
+    const rewrite = new URL(rewritePathname, request.nextUrl.toString());
+    rewrite.search = url.search;
+    // Make sure the target URL is clean of any va token before we use it for response
+    const target = sanitizeVisitorAuthURL(rewrite);
 
     const response = writeCookies(
         NextResponse.rewrite(target, {