fix: authed previews in production redirect to the CDN URI (#1803)

fern-api · Nov 8, 2024 · db2586c · db2586c
1 parent ad08da0
commit db2586c
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 56 deletions.
diff --git a/packages/ui/docs-bundle/src/server/auth/getAuthState.ts b/packages/ui/docs-bundle/src/server/auth/getAuthState.ts
@@ -9,8 +9,7 @@ import { getOryAuthorizationUrl } from "./ory";
 import { getReturnToQueryParam } from "./return-to";
 import { getWebflowAuthorizationUrl } from "./webflow";
 import { getWorkosSSOAuthorizationUrl } from "./workos";
-import { encryptSession, getSessionFromToken, refreshSession, toSessionUserInfo } from "./workos-session";
-import { toFernUser } from "./workos-user-to-fern-user";
+import { handleWorkosAuth } from "./workos-handler";
 
 export type AuthPartner = "workos" | "ory" | "webflow" | "custom";
 
@@ -87,31 +86,13 @@ export async function getAuthStateInternal({
     if (!authConfig) {
         if (previewAuthConfig != null) {
             if (previewAuthConfig.type === "workos") {
-                const state = urlJoin(removeTrailingSlash(withDefaultProtocol(host)), pathname ?? "");
-                const session = fernToken != null ? await getSessionFromToken(fernToken) : undefined;
-                const workosUserInfo = await toSessionUserInfo(session);
-                if (workosUserInfo.user) {
-                    // TODO: should this be stored in the session itself?
-                    const roles = await getWorkosRbacRoles(previewAuthConfig.org, workosUserInfo.user.email);
-                    return {
-                        authed: true,
-                        ok: true,
-                        user: toFernUser(workosUserInfo, roles),
-                        partner: "workos",
-                    };
-                }
-
-                const redirectUri = urlJoin(
-                    removeTrailingSlash(withDefaultProtocol(host)),
-                    "/api/fern-docs/auth/sso/callback",
-                );
-                const authorizationUrl = getWorkosSSOAuthorizationUrl({
-                    redirectUri,
+                return handleWorkosAuth({
+                    fernToken,
                     organization: previewAuthConfig.org,
-                    state,
+                    host,
+                    pathname,
+                    setFernToken,
                 });
-
-                return { authed: false, ok: false, authorizationUrl, partner: "workos" };
             }
         }
         return { authed: false, ok: true, authorizationUrl: undefined, partner: undefined };
@@ -132,37 +113,19 @@ export async function getAuthStateInternal({
 
     // check if the user is logged in via WorkOS
     if (authConfig.type === "sso" && authConfig.partner === "workos") {
-        const session = fernToken != null ? await getSessionFromToken(fernToken) : undefined;
-        const workosUserInfo = await toSessionUserInfo(session);
-        if (workosUserInfo.user) {
-            // TODO: should this be stored in the session itself?
-            const roles = await getWorkosRbacRoles(authConfig.organization, workosUserInfo.user.email);
-            return {
-                authed: true,
-                ok: true,
-                user: toFernUser(workosUserInfo, roles),
-                partner: authConfig.partner,
-            };
-        }
-
-        if (session?.refreshToken) {
-            const updatedSession = await refreshSession(session);
-            if (updatedSession) {
-                if (setFernToken) {
-                    setFernToken(await encryptSession(updatedSession));
-                }
-                // TODO: should this be stored in the session itself?
-                const roles = await getWorkosRbacRoles(authConfig.organization, updatedSession.user.email);
-                return {
-                    authed: true,
-                    ok: true,
-                    user: toFernUser(await toSessionUserInfo(updatedSession), roles),
-                    partner: authConfig.partner,
-                };
-            }
-        }
-
-        return { authed: false, ok: false, authorizationUrl, partner: authConfig.partner };
+        return handleWorkosAuth({
+            fernToken,
+            organization: authConfig.organization,
+            host,
+            pathname,
+            setFernToken,
+            authorizationUrl: {
+                connection: authConfig.connection,
+                provider: authConfig.provider,
+                domainHint: authConfig.domainHint,
+                loginHint: authConfig.loginHint,
+            },
+        });
     }
 
     return { authed: false, ok: false, authorizationUrl: undefined, partner: undefined };

diff --git a/packages/ui/docs-bundle/src/server/auth/workos-handler.ts b/packages/ui/docs-bundle/src/server/auth/workos-handler.ts
@@ -0,0 +1,72 @@
+import { withDefaultProtocol } from "@fern-api/ui-core-utils";
+import { removeTrailingSlash } from "next/dist/shared/lib/router/utils/remove-trailing-slash";
+import urlJoin from "url-join";
+import { AuthState, getWorkosRbacRoles } from "./getAuthState";
+import { getWorkosSSOAuthorizationUrl } from "./workos";
+import { encryptSession, getSessionFromToken, refreshSession, toSessionUserInfo } from "./workos-session";
+import { toFernUser } from "./workos-user-to-fern-user";
+
+interface WorkosAuthParams {
+    fernToken: string | undefined;
+    organization: string;
+    host: string;
+    pathname?: string;
+    setFernToken?: (token: string) => void;
+    authorizationUrl?: {
+        connection?: string;
+        provider?: string;
+        domainHint?: string;
+        loginHint?: string;
+    };
+}
+
+export async function handleWorkosAuth({
+    fernToken,
+    organization,
+    host,
+    pathname,
+    setFernToken,
+    authorizationUrl,
+}: WorkosAuthParams): Promise<AuthState> {
+    const state = urlJoin(removeTrailingSlash(withDefaultProtocol(host)), pathname ?? "");
+    const session = fernToken != null ? await getSessionFromToken(fernToken) : undefined;
+    const workosUserInfo = await toSessionUserInfo(session);
+
+    if (workosUserInfo.user) {
+        const roles = await getWorkosRbacRoles(organization, workosUserInfo.user.email);
+        return {
+            authed: true,
+            ok: true,
+            user: toFernUser(workosUserInfo, roles),
+            partner: "workos",
+        };
+    }
+
+    if (session?.refreshToken) {
+        const updatedSession = await refreshSession(session);
+        if (updatedSession) {
+            if (setFernToken) {
+                setFernToken(await encryptSession(updatedSession));
+            }
+            const roles = await getWorkosRbacRoles(organization, updatedSession.user.email);
+            return {
+                authed: true,
+                ok: true,
+                user: toFernUser(await toSessionUserInfo(updatedSession), roles),
+                partner: "workos",
+            };
+        }
+    }
+
+    const redirectUri = String(
+        new URL("/api/fern-docs/auth/sso/callback", withDefaultProtocol(process.env.NEXT_PUBLIC_CDN_URI ?? host)),
+    );
+    const authorizationUrlParams = getWorkosSSOAuthorizationUrl({
+        redirectUri,
+        organization,
+        state,
+        ...authorizationUrl,
+    });
+
+    return { authed: false, ok: false, authorizationUrl: authorizationUrlParams, partner: "workos" };
+}