From 7cb4b488681b036c982f78f69e88af2879837f1d Mon Sep 17 00:00:00 2001
From: Andrew Jiang <andrew@buildwithfern.com>
Date: Wed, 9 Oct 2024 17:51:06 -0400
Subject: [PATCH] fix: inconsistent deserialization of authprops

---
 .../ui/docs-bundle/src/server/DocsLoader.ts   | 24 +++++++++++--------
 .../src/server/getDynamicDocsPageProps.ts     | 18 +++++++-------
 .../src/server/withInitialProps.ts            |  4 ++--
 3 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/packages/ui/docs-bundle/src/server/DocsLoader.ts b/packages/ui/docs-bundle/src/server/DocsLoader.ts
index 129b0714f8..a222207280 100644
--- a/packages/ui/docs-bundle/src/server/DocsLoader.ts
+++ b/packages/ui/docs-bundle/src/server/DocsLoader.ts
@@ -12,7 +12,7 @@ interface DocsLoaderFlags {
 }
 
 export class DocsLoader {
-    static for(xFernHost: string, fernToken: string | undefined): DocsLoader {
+    static for(xFernHost: string, fernToken?: string): DocsLoader {
         return new DocsLoader(xFernHost, fernToken);
     }
 
@@ -30,25 +30,29 @@ export class DocsLoader {
         return this;
     }
 
-    private auth: AuthEdgeConfig | undefined;
-    public withAuth(auth: AuthEdgeConfig | undefined): DocsLoader {
-        this.auth = auth;
+    private authConfig: AuthEdgeConfig | undefined;
+    private authProps: AuthProps | undefined;
+    public withAuth(authConfig: AuthEdgeConfig | undefined, authProps?: AuthProps): DocsLoader {
+        this.authConfig = authConfig;
+        if (authProps) {
+            this.authProps = authProps;
+        }
         return this;
     }
 
     private async loadAuth(): Promise<[AuthProps | undefined, AuthEdgeConfig | undefined]> {
-        if (!this.auth) {
-            this.auth = await getAuthEdgeConfig(this.xFernHost);
+        if (!this.authConfig) {
+            this.authConfig = await getAuthEdgeConfig(this.xFernHost);
         }
-        if (!this.auth) {
-            return [undefined, undefined];
+        if (!this.authConfig || !this.fernToken || (this.authProps && this.authConfig)) {
+            return [this.authProps, this.authConfig];
         }
         try {
-            return [await withAuthProps(this.auth, this.fernToken), this.auth];
+            return [await withAuthProps(this.authConfig, this.fernToken), this.authConfig];
         } catch (e) {
             // eslint-disable-next-line no-console
             console.error(e);
-            return [undefined, this.auth];
+            return [undefined, this.authConfig];
         }
     }
 
diff --git a/packages/ui/docs-bundle/src/server/getDynamicDocsPageProps.ts b/packages/ui/docs-bundle/src/server/getDynamicDocsPageProps.ts
index 6fb6f6ef9c..8f952839d1 100644
--- a/packages/ui/docs-bundle/src/server/getDynamicDocsPageProps.ts
+++ b/packages/ui/docs-bundle/src/server/getDynamicDocsPageProps.ts
@@ -3,7 +3,7 @@ import { type DocsPage } from "@fern-ui/ui";
 import type { NextApiRequestCookies } from "next/dist/server/api-utils";
 import type { GetServerSidePropsResult } from "next/types";
 import type { ComponentProps } from "react";
-import { withAuthProps } from "./authProps";
+import { AuthProps, withAuthProps } from "./authProps";
 import { getDocsPageProps } from "./getDocsPageProps";
 
 type GetServerSideDocsPagePropsResult = GetServerSidePropsResult<ComponentProps<typeof DocsPage>>;
@@ -13,18 +13,20 @@ export async function getDynamicDocsPageProps(
     slug: string[],
     cookies: NextApiRequestCookies,
 ): Promise<GetServerSideDocsPagePropsResult> {
-    if (cookies[COOKIE_FERN_TOKEN] == null) {
-        /**
-         * this only happens when ?error=true is passed in the URL
-         * Note: custom auth (via edge config) is supported via middleware, so we don't need to handle it here
-         */
-        return getDocsPageProps(xFernHost, slug);
+    let authProps: AuthProps | undefined;
+
+    try {
+        if (cookies[COOKIE_FERN_TOKEN]) {
+            authProps = await withAuthProps(xFernHost, cookies[COOKIE_FERN_TOKEN]);
+        }
+    } catch (e) {
+        // eslint-disable-next-line no-console
+        console.error("Failed to get auth props", e);
     }
 
     /**
      * Authenticated user is guaranteed to have a valid token because the middleware
      * would have redirected them to the login page
      */
-    const authProps = await withAuthProps(xFernHost, cookies[COOKIE_FERN_TOKEN]);
     return getDocsPageProps(xFernHost, slug, authProps);
 }
diff --git a/packages/ui/docs-bundle/src/server/withInitialProps.ts b/packages/ui/docs-bundle/src/server/withInitialProps.ts
index 076dbd625c..fe77eb13f8 100644
--- a/packages/ui/docs-bundle/src/server/withInitialProps.ts
+++ b/packages/ui/docs-bundle/src/server/withInitialProps.ts
@@ -61,9 +61,9 @@ export async function withInitialProps({
     const featureFlags = await getFeatureFlags(xFernHost);
 
     const authConfig = await getAuthEdgeConfig(xFernHost);
-    const loader = DocsLoader.for(xFernHost, auth?.token)
+    const loader = DocsLoader.for(xFernHost)
         .withFeatureFlags(featureFlags)
-        .withAuth(authConfig)
+        .withAuth(authConfig, auth)
         .withLoadDocsForUrlResponse(docs);
 
     const root = await loader.root();