From dbe6377012475897efa031c458b0df9c4596d455 Mon Sep 17 00:00:00 2001
From: Andrew Jiang <andrew@buildwithfern.com>
Date: Tue, 8 Oct 2024 12:52:19 -0400
Subject: [PATCH] fix: correct redirect host for proxied domains in auth

---
 packages/ui/docs-bundle/src/middleware.ts             |  9 +--------
 .../src/pages/api/fern-docs/auth/callback.ts          | 11 +++++++++--
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/packages/ui/docs-bundle/src/middleware.ts b/packages/ui/docs-bundle/src/middleware.ts
index ea1a01e00f..aacebab031 100644
--- a/packages/ui/docs-bundle/src/middleware.ts
+++ b/packages/ui/docs-bundle/src/middleware.ts
@@ -4,7 +4,7 @@ import { rewritePosthog } from "@/server/rewritePosthog";
 import { getXFernHostEdge } from "@/server/xfernhost/edge";
 import type { FernUser } from "@fern-ui/fern-docs-auth";
 import { getAuthEdgeConfig } from "@fern-ui/fern-docs-edge-config";
-import { COOKIE_FERN_TOKEN, HEADER_X_FERN_HOST } from "@fern-ui/fern-docs-utils";
+import { COOKIE_FERN_TOKEN } from "@fern-ui/fern-docs-utils";
 import { removeTrailingSlash } from "next/dist/shared/lib/router/utils/remove-trailing-slash";
 import { NextRequest, NextResponse, type NextMiddleware } from "next/server";
 import urlJoin from "url-join";
@@ -29,13 +29,6 @@ export const middleware: NextMiddleware = async (request) => {
         return NextResponse.next();
     }
 
-    /**
-     * Add x-fern-host header to the request
-     */
-    if (!headers.has(HEADER_X_FERN_HOST)) {
-        headers.set(HEADER_X_FERN_HOST, xFernHost);
-    }
-
     /**
      * Rewrite robots.txt
      */
diff --git a/packages/ui/docs-bundle/src/pages/api/fern-docs/auth/callback.ts b/packages/ui/docs-bundle/src/pages/api/fern-docs/auth/callback.ts
index e5eda1a1a0..367481b683 100644
--- a/packages/ui/docs-bundle/src/pages/api/fern-docs/auth/callback.ts
+++ b/packages/ui/docs-bundle/src/pages/api/fern-docs/auth/callback.ts
@@ -4,7 +4,7 @@ import { getWorkOS, getWorkOSClientId } from "@/server/workos";
 import { getXFernHostEdge } from "@/server/xfernhost/edge";
 import { FernUser } from "@fern-ui/fern-docs-auth";
 import { getAuthEdgeConfig } from "@fern-ui/fern-docs-edge-config";
-import { COOKIE_FERN_TOKEN } from "@fern-ui/fern-docs-utils";
+import { COOKIE_FERN_TOKEN, HEADER_X_FERN_HOST } from "@fern-ui/fern-docs-utils";
 import { NextRequest, NextResponse } from "next/server";
 
 export const runtime = "edge";
@@ -44,7 +44,14 @@ export default async function GET(req: NextRequest): Promise<NextResponse> {
             "/api/fern-docs/auth/callback",
             "/api/fern-docs/oauth/ory/callback",
         );
-        // Permanent GET redirect to the Ory callback endpoint
+
+        // Redirect to x-fern-host domain if it exists
+        // this is to ensure proxied origins are used for the redirect
+        if (req.headers.has(HEADER_X_FERN_HOST)) {
+            // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+            nextUrl.host = req.headers.get(HEADER_X_FERN_HOST)!;
+        }
+
         return NextResponse.redirect(nextUrl);
     }