From e5ab83aa840f954353392e121065cc8658062a15 Mon Sep 17 00:00:00 2001
From: dsinghvi <deep@usebirch.com>
Date: Thu, 7 Nov 2024 12:06:05 -0500
Subject: [PATCH] fix authed previews

---
 .../pages/api/fern-docs/org/org-for-url.ts    | 21 +++++++------------
 .../src/server/auth/getAuthState.ts           |  8 ++++---
 .../src/server/auth/metadata-for-url.ts       | 20 ++++++++++++++++++
 .../src/getPreviewUrlAuthConfig.ts            | 14 +++++++------
 .../ui/fern-docs-server/src/DocsLoader.ts     |  9 +++++---
 5 files changed, 47 insertions(+), 25 deletions(-)
 create mode 100644 packages/ui/docs-bundle/src/server/auth/metadata-for-url.ts

diff --git a/packages/ui/docs-bundle/src/pages/api/fern-docs/org/org-for-url.ts b/packages/ui/docs-bundle/src/pages/api/fern-docs/org/org-for-url.ts
index 1497584cfd..5ee043bbef 100644
--- a/packages/ui/docs-bundle/src/pages/api/fern-docs/org/org-for-url.ts
+++ b/packages/ui/docs-bundle/src/pages/api/fern-docs/org/org-for-url.ts
@@ -1,12 +1,12 @@
+import { getOrgMetadataForDomain } from "@/server/auth/meta";
 import { getDocsDomainEdge } from "@/server/xfernhost/edge";
-import { DocsLoader } from "@fern-ui/fern-docs-server";
 import { NextRequest, NextResponse } from "next/server";
 
 export const runtime = "edge";
 
 export interface DocsMetadata {
-    orgId: string | undefined;
-    isPreviewUrl: boolean | undefined;
+    orgId: string;
+    isPreviewUrl: boolean;
 }
 
 export default async function handler(req: NextRequest): Promise<NextResponse<DocsMetadata>> {
@@ -16,16 +16,11 @@ export default async function handler(req: NextRequest): Promise<NextResponse<Do
         return new NextResponse(JSON.stringify({ error: "Invalid domain" }), { status: 400 });
     }
 
-    try {
-        const docsLoader = DocsLoader.create(domain).withEnvironment(process.env.NEXT_PUBLIC_FDR_ORIGIN);
-        const metadata = await docsLoader.getMetadata();
+    const metadata = await getOrgMetadataForDomain(domain);
 
-        if (metadata) {
-            return new NextResponse(JSON.stringify(metadata), { status: 200 });
-        } else {
-            return new NextResponse(JSON.stringify({ error: "Org not found" }), { status: 404 });
-        }
-    } catch (error) {
-        return new NextResponse(JSON.stringify({ error: "Internal server error" }), { status: 500 });
+    if (metadata) {
+        return new NextResponse(JSON.stringify(metadata), { status: 200 });
+    } else {
+        return new NextResponse(JSON.stringify({ error: "Org not found" }), { status: 404 });
     }
 }
diff --git a/packages/ui/docs-bundle/src/server/auth/getAuthState.ts b/packages/ui/docs-bundle/src/server/auth/getAuthState.ts
index c145c7c6e5..e959eb0aaf 100644
--- a/packages/ui/docs-bundle/src/server/auth/getAuthState.ts
+++ b/packages/ui/docs-bundle/src/server/auth/getAuthState.ts
@@ -4,6 +4,7 @@ import { PreviewUrlAuth, getAuthEdgeConfig, getPreviewUrlAuthConfig } from "@fer
 import { removeTrailingSlash } from "next/dist/shared/lib/router/utils/remove-trailing-slash";
 import urlJoin from "url-join";
 import { safeVerifyFernJWTConfig } from "./FernJWT";
+import { getOrgMetadataForDomain } from "./metadata-for-url";
 import { getOryAuthorizationUrl } from "./ory";
 import { getReturnToQueryParam } from "./return-to";
 import { getWebflowAuthorizationUrl } from "./webflow";
@@ -107,7 +108,8 @@ export async function getAuthStateInternal({
                     redirectUri,
                     organization: previewAuthConfig.org,
                 });
-                return { authed: false, ok: true, authorizationUrl, partner: "workos" };
+
+                return { authed: false, ok: false, authorizationUrl, partner: "workos" };
             }
         }
         return { authed: false, ok: true, authorizationUrl: undefined, partner: undefined };
@@ -182,7 +184,7 @@ export async function getAuthState(
     setFernToken?: (token: string) => void,
 ): Promise<AuthState & DomainAndHost> {
     authConfig ??= await getAuthEdgeConfig(domain);
-    const previewAuthConfig = await getPreviewUrlAuthConfig(domain);
+    const orgMetadata = await getOrgMetadataForDomain(domain);
 
     const authState = await getAuthStateInternal({
         host,
@@ -190,7 +192,7 @@ export async function getAuthState(
         pathname,
         authConfig,
         setFernToken,
-        previewAuthConfig,
+        previewAuthConfig: orgMetadata != null ? await getPreviewUrlAuthConfig(orgMetadata) : undefined,
     });
 
     return { ...authState, domain, host };
diff --git a/packages/ui/docs-bundle/src/server/auth/metadata-for-url.ts b/packages/ui/docs-bundle/src/server/auth/metadata-for-url.ts
new file mode 100644
index 0000000000..f8734828cb
--- /dev/null
+++ b/packages/ui/docs-bundle/src/server/auth/metadata-for-url.ts
@@ -0,0 +1,20 @@
+import { DocsLoader } from "@fern-ui/fern-docs-server";
+
+export interface OrgMetadata {
+    orgId: string;
+    isPreviewUrl: boolean;
+}
+
+export async function getOrgMetadataForDomain(domain: string): Promise<OrgMetadata | undefined> {
+    if (!domain || typeof domain !== "string") {
+        return undefined;
+    }
+
+    try {
+        const docsLoader = DocsLoader.create(domain).withEnvironment(process.env.NEXT_PUBLIC_FDR_ORIGIN);
+        const metadata = await docsLoader.getMetadata();
+        return metadata ?? undefined;
+    } catch (error) {
+        return undefined;
+    }
+}
diff --git a/packages/ui/fern-docs-edge-config/src/getPreviewUrlAuthConfig.ts b/packages/ui/fern-docs-edge-config/src/getPreviewUrlAuthConfig.ts
index 16141072b9..5af60b1c76 100644
--- a/packages/ui/fern-docs-edge-config/src/getPreviewUrlAuthConfig.ts
+++ b/packages/ui/fern-docs-edge-config/src/getPreviewUrlAuthConfig.ts
@@ -17,15 +17,17 @@ const PreviewUrlAuthConfigSchema = z.record(PreviewUrlAuthSchema);
 
 type PreviewUrlAuthConfig = z.infer<typeof PreviewUrlAuthConfigSchema>;
 
-export async function getPreviewUrlAuthConfig(currentDomain: string): Promise<PreviewUrlAuth | undefined> {
-    const previewDomain = isPreviewDomain(currentDomain);
-    const org = extractOrgFromPreview(currentDomain);
-    if (!previewDomain || !org) {
+export interface Metadata {
+    isPreviewUrl: boolean;
+    orgId: string;
+}
+
+export async function getPreviewUrlAuthConfig(metadata: Metadata): Promise<PreviewUrlAuth | undefined> {
+    if (!metadata.isPreviewUrl) {
         return undefined;
     }
-
     const config = await get<PreviewUrlAuthConfig>("authed-previews");
-    return config?.[org];
+    return config?.[metadata.orgId];
 }
 
 export function isPreviewDomain(domain: string): boolean {
diff --git a/packages/ui/fern-docs-server/src/DocsLoader.ts b/packages/ui/fern-docs-server/src/DocsLoader.ts
index b3778394cd..63ab215a92 100644
--- a/packages/ui/fern-docs-server/src/DocsLoader.ts
+++ b/packages/ui/fern-docs-server/src/DocsLoader.ts
@@ -2,8 +2,8 @@ import { FdrAPI, FdrClient } from "@fern-api/fdr-sdk";
 import { DocsDomainKVCache } from "./DocsKVCache";
 
 export interface DocsMetadata {
-    orgId: string | undefined;
-    isPreviewUrl: boolean | undefined;
+    orgId: string;
+    isPreviewUrl: boolean;
 }
 
 export class DocsLoader {
@@ -47,7 +47,10 @@ export class DocsLoader {
 
     #getClient = () => new FdrClient({ environment: this.environment });
     #loadMetadataForUrl = async (): Promise<FdrAPI.docs.v2.read.DocsUrlMetadata | undefined> => {
-        const response = await this.#getClient().docs.v2.read.getDocsUrlMetadata({ url: FdrAPI.Url(this.domain) });
+        const response = await this.#getClient().docs.v2.read.getDocsUrlMetadata(
+            { url: FdrAPI.Url(this.domain) },
+            { timeoutInSeconds: 180 },
+        );
         if (!response.ok) {
             return undefined;
         }