From 1cc31343c13edb4e078f435f1658e964b66a2171 Mon Sep 17 00:00:00 2001
From: Andrew Jiang <andrew@buildwithfern.com>
Date: Tue, 29 Oct 2024 19:49:08 -0400
Subject: [PATCH] fix: redirect auth callback for app-staging (#1746)

---
 .../pages/api/fern-docs/auth/sso/callback.ts  | 23 +++++++++++-------
 .../src/server/auth/getAuthState.ts           | 24 +------------------
 .../docs-bundle/src/server/xfernhost/edge.ts  |  2 +-
 3 files changed, 16 insertions(+), 33 deletions(-)

diff --git a/packages/ui/docs-bundle/src/pages/api/fern-docs/auth/sso/callback.ts b/packages/ui/docs-bundle/src/pages/api/fern-docs/auth/sso/callback.ts
index a836a4b287..82a541857d 100644
--- a/packages/ui/docs-bundle/src/pages/api/fern-docs/auth/sso/callback.ts
+++ b/packages/ui/docs-bundle/src/pages/api/fern-docs/auth/sso/callback.ts
@@ -2,7 +2,7 @@ import { withSecureCookie } from "@/server/auth/with-secure-cookie";
 import { getWorkOSClientId, workos } from "@/server/auth/workos";
 import { encryptSession } from "@/server/auth/workos-session";
 import { safeUrl } from "@/server/safeUrl";
-import { getHostEdge } from "@/server/xfernhost/edge";
+import { getDocsDomainEdge } from "@/server/xfernhost/edge";
 import { COOKIE_FERN_TOKEN } from "@fern-ui/fern-docs-utils";
 import { NextRequest, NextResponse } from "next/server";
 
@@ -17,12 +17,6 @@ export default async function handler(req: NextRequest): Promise<NextResponse> {
         return new NextResponse(null, { status: 405 });
     }
 
-    if (req.nextUrl.searchParams.get(FORWARDED_HOST_QUERY) === req.nextUrl.host) {
-        // eslint-disable-next-line no-console
-        console.error(FORWARDED_HOST_QUERY, "is the same as the host");
-        return new NextResponse(null, { status: 400 });
-    }
-
     const state = req.nextUrl.searchParams.get(STATE_QUERY);
 
     if (state == null) {
@@ -41,10 +35,21 @@ export default async function handler(req: NextRequest): Promise<NextResponse> {
     }
 
     // TODO: this is a security risk (open redirect)! We need to verify that the target host is one of ours.
-    if (getHostEdge(req) !== url.host) {
+    // if the current url is app.buildwithfern.com, we should redirect to ***.docs.buildwithfern.com
+    if (req.nextUrl.host !== url.host && getDocsDomainEdge(req) !== url.host) {
+        if (req.nextUrl.searchParams.get(FORWARDED_HOST_QUERY) === req.nextUrl.host) {
+            // eslint-disable-next-line no-console
+            console.error(
+                FORWARDED_HOST_QUERY,
+                "is the same as the host:",
+                String(req.nextUrl.searchParams.get(FORWARDED_HOST_QUERY)),
+            );
+            return new NextResponse(null, { status: 400 });
+        }
+
+        // TODO: need to support docs instances with subpaths (forward-proxied from the origin).
         const destination = new URL(`${req.nextUrl.pathname}${req.nextUrl.search}`, url.origin);
         destination.searchParams.set(FORWARDED_HOST_QUERY, req.nextUrl.host);
-
         return NextResponse.redirect(destination);
     }
 
diff --git a/packages/ui/docs-bundle/src/server/auth/getAuthState.ts b/packages/ui/docs-bundle/src/server/auth/getAuthState.ts
index 4d0963a34c..2974440436 100644
--- a/packages/ui/docs-bundle/src/server/auth/getAuthState.ts
+++ b/packages/ui/docs-bundle/src/server/auth/getAuthState.ts
@@ -147,31 +147,9 @@ function getAuthorizationUrl(authConfig: AuthEdgeConfig, host: string, pathname?
         destination.searchParams.set("state", state);
         return destination.toString();
     } else if (authConfig.type === "sso" && authConfig.partner === "workos") {
-        const redirectUri = urlJoin(
-            removeTrailingSlash(withDefaultProtocol(getRedirectUri())),
-            "/api/fern-docs/auth/sso/callback",
-        );
+        const redirectUri = urlJoin(removeTrailingSlash(withDefaultProtocol(host)), "/api/fern-docs/auth/sso/callback");
         return getWorkOSAuthorizationUrl({ state, redirectUri, organization: authConfig.organization });
     }
 
     return undefined;
 }
-
-/*
- * Note: our WorkOS prod/staging is not 1:1 with FDR (app/app-dev2) so instead, we:
- * - use the WorkOS production url for ONLY the production docs deployments
- * - use the WorkOS staging instance for all other deployments (prod-preview, dev2, local dev, etc.)
- *
- * This is so that we can test workos using open redirects, and not have to worry about the authkit redirect uri changing:
- */
-function getRedirectUri(): string {
-    if (process.env.NODE_ENV !== "production" || process.env.VERCEL_ENV === "development") {
-        return `http://localhost:${process.env.PORT ?? 3000}`;
-    }
-    return (
-        process.env.NEXT_PUBLIC_CDN_URI ??
-        process.env.VERCEL_BRANCH_URL ??
-        process.env.VERCEL_DEPLOYMENT_URL ??
-        "https://app.buildwithfern.com"
-    );
-}
diff --git a/packages/ui/docs-bundle/src/server/xfernhost/edge.ts b/packages/ui/docs-bundle/src/server/xfernhost/edge.ts
index ce3949fc08..eb1ca03386 100644
--- a/packages/ui/docs-bundle/src/server/xfernhost/edge.ts
+++ b/packages/ui/docs-bundle/src/server/xfernhost/edge.ts
@@ -30,7 +30,7 @@ export function getDocsDomainEdge(req: NextRequest): string {
 export function getHostEdge(req: NextRequest): string {
     if (
         process.env.NODE_ENV === "development" ||
-        (process.env.VERCEL_ENV === "preview" && req.cookies.get(COOKIE_FERN_DOCS_PREVIEW)?.value != null) ||
+        (process.env.VERCEL_ENV === "preview" && req.cookies.has(COOKIE_FERN_DOCS_PREVIEW)) ||
         process.env.VERCEL_ENV === "development"
     ) {
         return req.nextUrl.host;