feat: basic token validation for qlty.ai (#1137)

fern-api · Jul 11, 2024 · 028bbf2 · 028bbf2
1 parent 83ccfd4
commit 028bbf2
Show file tree

Hide file tree

Showing 9 changed files with 77 additions and 38 deletions.
diff --git a/packages/ui/app/src/auth/FernJWT.ts b/packages/ui/app/src/auth/FernJWT.ts
@@ -13,9 +13,9 @@ export function signFernJWT(fern: FernUser, user?: any): Promise<string> {
         .sign(getJwtTokenSecret());
 }
 
-export async function verifyFernJWT(token: string): Promise<FernUser> {
-    const verified = await jwtVerify(token, getJwtTokenSecret(), {
-        issuer: "https://buildwithfern.com",
+export async function verifyFernJWT(token: string, secret?: Uint8Array, issuer?: string): Promise<FernUser> {
+    const verified = await jwtVerify(token, secret ?? getJwtTokenSecret(), {
+        issuer: issuer ?? "https://buildwithfern.com",
     });
     return FernUserSchema.parse(verified.payload.fern);
 }

diff --git a/packages/ui/app/src/auth/getAuthEdgeConfig.ts b/packages/ui/app/src/auth/getAuthEdgeConfig.ts
@@ -0,0 +1,16 @@
+import { get } from "@vercel/edge-config";
+import { AuthEdgeConfig, AuthEdgeConfigSchema } from "./types";
+
+const KEY = "authentication";
+
+export async function getAuthEdgeConfig(currentDomain: string): Promise<AuthEdgeConfig | undefined> {
+    const domainToTokenConfigMap = await get<Record<string, any>>(KEY);
+    const toRet = domainToTokenConfigMap?.[currentDomain];
+
+    if (toRet != null) {
+        // if the config is present, it should be valid
+        return AuthEdgeConfigSchema.parse(toRet);
+    }
+
+    return;
+}
diff --git a/packages/ui/app/src/auth/getOAuthEdgeConfig.ts b/packages/ui/app/src/auth/getOAuthEdgeConfig.ts
diff --git a/packages/ui/app/src/auth/getOAuthToken.ts b/packages/ui/app/src/auth/getOAuthToken.ts
@@ -1,8 +1,8 @@
 import urlJoin from "url-join";
-import { OAuthEdgeConfig, OAuthTokenResponse, OAuthTokenResponseSchema } from "./types";
+import { AuthEdgeConfigOAuth2, OAuthTokenResponse, OAuthTokenResponseSchema } from "./types";
 
 export async function getOAuthToken(
-    config: OAuthEdgeConfig,
+    config: AuthEdgeConfigOAuth2,
     code: string,
     redirect_uri: string,
 ): Promise<OAuthTokenResponse> {

diff --git a/packages/ui/app/src/auth/index.ts b/packages/ui/app/src/auth/index.ts
@@ -1,4 +1,4 @@
 export * from "./FernJWT";
-export { getOAuthEdgeConfig } from "./getOAuthEdgeConfig";
+export { getAuthEdgeConfig } from "./getAuthEdgeConfig";
 export { getOAuthToken } from "./getOAuthToken";
 export * from "./types";
diff --git a/packages/ui/app/src/auth/types.ts b/packages/ui/app/src/auth/types.ts
@@ -2,14 +2,14 @@ import { z } from "zod";
 
 export const FernUserSchema = z.object({
     type: z.literal("user"),
-    partner: z.union([z.literal("workos"), z.literal("ory")]),
+    partner: z.union([z.literal("workos"), z.literal("ory"), z.literal("custom")]),
     name: z.string(),
     email: z.string(),
 });
 
 export type FernUser = z.infer<typeof FernUserSchema>;
 
-export const OAuthEdgeConfigSchema = z.object({
+export const AuthEdgeConfigOAuth2Schema = z.object({
     type: z.literal("oauth2"),
     partner: z.literal("ory"),
     environment: z.string(),
@@ -18,7 +18,18 @@ export const OAuthEdgeConfigSchema = z.object({
     "api-key-injection-enabled": z.optional(z.boolean()),
 });
 
-export type OAuthEdgeConfig = z.infer<typeof OAuthEdgeConfigSchema>;
+export const AuthEdgeConfigBasicTokenVerificationSchema = z.object({
+    type: z.literal("basic_token_verification"),
+    secret: z.string(),
+    issuer: z.string(),
+    redirect: z.string(),
+});
+
+export const AuthEdgeConfigSchema = z.union([AuthEdgeConfigOAuth2Schema, AuthEdgeConfigBasicTokenVerificationSchema]);
+
+export type AuthEdgeConfig = z.infer<typeof AuthEdgeConfigSchema>;
+export type AuthEdgeConfigOAuth2 = z.infer<typeof AuthEdgeConfigOAuth2Schema>;
+export type AuthEdgeConfigBasicTokenVerification = z.infer<typeof AuthEdgeConfigBasicTokenVerificationSchema>;
 
 export const OAuthTokenResponseSchema = z.object({
     access_token: z.string(),

diff --git a/packages/ui/docs-bundle/src/pages/api/fern-docs/auth/api-key-injection.ts b/packages/ui/docs-bundle/src/pages/api/fern-docs/auth/api-key-injection.ts
@@ -1,5 +1,5 @@
 // eslint-disable-next-line import/no-internal-modules
-import { getOAuthEdgeConfig } from "@fern-ui/ui/auth";
+import { getAuthEdgeConfig } from "@fern-ui/ui/auth";
 import { NextRequest, NextResponse } from "next/server";
 import urlJoin from "url-join";
 import { getXFernHostEdge } from "../../../../utils/xFernHost";
@@ -8,7 +8,11 @@ export const runtime = "edge";
 
 export default async function handler(req: NextRequest): Promise<NextResponse<string | false>> {
     const domain = getXFernHostEdge(req);
-    const config = await getOAuthEdgeConfig(domain);
+    const config = await getAuthEdgeConfig(domain);
+
+    if (config == null || config.type !== "oauth2") {
+        return NextResponse.json(false);
+    }
 
     // ory is the only partner enabled for api-key-injection (with RightBrain)
     if (config?.["api-key-injection-enabled"] && config.partner === "ory") {

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
@@ -1,5 +1,5 @@
 // eslint-disable-next-line import/no-internal-modules
-import { FernUser, OryAccessTokenSchema, getOAuthEdgeConfig, getOAuthToken, signFernJWT } from "@fern-ui/ui/auth";
+import { FernUser, OryAccessTokenSchema, getAuthEdgeConfig, getOAuthToken, signFernJWT } from "@fern-ui/ui/auth";
 import { decodeJwt } from "jose";
 import { NextRequest, NextResponse } from "next/server";
 import urlJoin from "url-join";
@@ -31,9 +31,9 @@ export default async function GET(req: NextRequest): Promise<NextResponse> {
     }
 
     const domain = getXFernHostEdge(req);
-    const config = await getOAuthEdgeConfig(domain);
+    const config = await getAuthEdgeConfig(domain);
 
-    if (config != null && config.partner === "ory") {
+    if (config != null && config.type === "oauth2" && config.partner === "ory") {
         try {
             const { access_token, refresh_token } = await getOAuthToken(
                 config,

diff --git a/packages/ui/docs-bundle/src/utils/getDocsPageProps.ts b/packages/ui/docs-bundle/src/utils/getDocsPageProps.ts
@@ -13,7 +13,7 @@ import {
     setMdxBundler,
 } from "@fern-ui/ui";
 // eslint-disable-next-line import/no-internal-modules
-import { FernUser, getOAuthEdgeConfig, verifyFernJWT } from "@fern-ui/ui/auth";
+import { FernUser, getAuthEdgeConfig, verifyFernJWT } from "@fern-ui/ui/auth";
 // eslint-disable-next-line import/no-internal-modules
 import { getMdxBundler } from "@fern-ui/ui/bundlers";
 import type { Redirect } from "next";
@@ -60,6 +60,22 @@ export async function getDocsPageProps(
         return { type: "notFound", notFound: true };
     }
 
+    const config = await getAuthEdgeConfig(xFernHost);
+    if (config != null && config.type === "basic_token_verification") {
+        const destination = new URL(config.redirect);
+        destination.searchParams.set(
+            "state",
+            encodeURIComponent(urlJoin(`https://${xFernHost}`, `/${slug.join("/")}`)),
+        );
+        return {
+            type: "redirect",
+            redirect: {
+                destination: destination.toString(),
+                permanent: false,
+            },
+        };
+    }
+
     const pathname = decodeURI(slug != null ? slug.join("/") : "");
     const url = buildUrl({ host: xFernHost, pathname });
     // eslint-disable-next-line no-console
@@ -105,7 +121,18 @@ export async function getDynamicDocsPageProps(
     }
 
     try {
-        const user = await verifyFernJWT(cookies.fern_token);
+        const config = await getAuthEdgeConfig(xFernHost);
+        let user: FernUser | undefined = undefined;
+
+        if (config?.type === "basic_token_verification") {
+            user = await verifyFernJWT(
+                cookies.fern_token,
+                new Uint8Array(Buffer.from(config.secret, "base64")),
+                config.issuer,
+            );
+        } else {
+            user = await verifyFernJWT(cookies.fern_token);
+        }
 
         if (user.partner === "workos") {
             const registryService = getRegistryServiceWithToken(`workos_${cookies.fern_token}`);
@@ -134,15 +161,13 @@ export async function getDynamicDocsPageProps(
             }
 
             return convertDocsToDocsPageProps({ docs: docs.body, slug, url, xFernHost });
-        } else if (user.partner === "ory") {
+        } else if (user.partner === "ory" || user.partner === "custom") {
             const docs = await REGISTRY_SERVICE.docs.v2.read.getDocsForUrl({ url });
 
             if (!docs.ok) {
                 throw new Error("Failed to fetch docs");
             }
 
-            const config = await getOAuthEdgeConfig(xFernHost);
-
             if (config == null) {
                 throw new Error("Failed to fetch OAuth config");
             }
@@ -153,7 +178,8 @@ export async function getDynamicDocsPageProps(
                 url,
                 xFernHost,
                 user,
-                apiKey: config["api-key-injection-enabled"] ? cookies.access_token : undefined,
+                apiKey:
+                    config.type === "oauth2" && config["api-key-injection-enabled"] ? cookies.access_token : undefined,
             });
         }
     } catch (error) {