Merge pull request #2926 from dubinc/fix-saml

Fix SAML SSO Enforcement Implementation

Author: steven-tey

apps/web/app/api/workspaces/[idOrSlug]/route.ts

@@ -6,7 +6,6 @@ import { deleteWorkspace } from "@/lib/api/workspaces/delete-workspace";
 import { prefixWorkspaceId } from "@/lib/api/workspaces/workspace-id";
 import { withWorkspace } from "@/lib/auth";
 import { getFeatureFlags } from "@/lib/edge-config";
-import { isGenericEmail } from "@/lib/is-generic-email";
 import { jackson } from "@/lib/jackson";
 import { storage } from "@/lib/storage";
 import z from "@/lib/zod";
@@ -107,24 +106,14 @@ export const PATCH = withWorkspace(
         )
       : null;
 
-    let ssoEmailDomain: string | null | undefined;
-
     if (enforceSAML) {
       if (workspace.plan !== "enterprise") {
         throw new DubApiError({
           code: "forbidden",
           message: "SAML SSO is only available on enterprise plans.",
         });
       }
-      ssoEmailDomain = session.user.email.split("@")[1];
-      if (isGenericEmail(session.user.email)) {
-        throw new DubApiError({
-          code: "forbidden",
-          message: "SAML SSO is not available for generic emails.",
-        });
-      }
 
-      // Check if SAML is configured before enforcing ssoEmailDomain
       const { apiController } = await jackson();
 
       const connections = await apiController.getConnections({
@@ -138,8 +127,6 @@ export const PATCH = withWorkspace(
           message: "SAML SSO is not configured for this workspace.",
         });
       }
-    } else if (enforceSAML === false) {
-      ssoEmailDomain = null;
     }
 
     try {
@@ -156,7 +143,9 @@ export const PATCH = withWorkspace(
             allowedHostnames: validHostnames,
           }),
           ...(publishableKey !== undefined && { publishableKey }),
-          ...(ssoEmailDomain !== undefined && { ssoEmailDomain }),
+          ...(enforceSAML !== undefined && {
+            ssoEnforcedAt: enforceSAML ? new Date() : null,
+          }),
         },
         include: {
           domains: {
@@ -224,7 +213,7 @@ export const PATCH = withWorkspace(
       if (error.code === "P2002") {
         throw new DubApiError({
           code: "conflict",
-          message: `The ${ssoEmailDomain ? "email domain" : "slug"} "${ssoEmailDomain || slug}" is already in use.`,
+          message: `The slug "${slug}" is already in use.`,
         });
       } else {
         throw new DubApiError({

apps/web/app/api/workspaces/[idOrSlug]/saml/route.ts

@@ -1,4 +1,6 @@
+import { DubApiError } from "@/lib/api/errors";
 import { withWorkspace } from "@/lib/auth";
+import { isGenericEmail } from "@/lib/is-generic-email";
 import { jackson, samlAudience } from "@/lib/jackson";
 import z from "@/lib/zod";
 import { prisma } from "@dub/prisma";
@@ -52,10 +54,20 @@ export const GET = withWorkspace(
 
 // POST /api/workspaces/[idOrSlug]/saml – create a new SAML connection
 export const POST = withWorkspace(
-  async ({ req, workspace }) => {
+  async ({ req, workspace, session }) => {
     const { metadataUrl, encodedRawMetadata } =
       createSAMLConnectionSchema.parse(await req.json());
 
+    const ssoEmailDomain = session.user.email.split("@")[1].toLocaleLowerCase();
+
+    if (isGenericEmail(ssoEmailDomain)) {
+      throw new DubApiError({
+        code: "bad_request",
+        message:
+          "SAML configuration requires you to be logged in with your organization’s work email.",
+      });
+    }
+
     const { apiController } = await jackson();
 
     const data = await apiController.createSAMLConnection({
@@ -67,6 +79,15 @@ export const POST = withWorkspace(
       product: "Dub",
     });
 
+    await prisma.project.update({
+      where: {
+        id: workspace.id,
+      },
+      data: {
+        ssoEmailDomain,
+      },
+    });
+
     return NextResponse.json(data);
   },
   {
@@ -94,6 +115,7 @@ export const DELETE = withWorkspace(
       },
       data: {
         ssoEmailDomain: null,
+        ssoEnforcedAt: null,
       },
     });
 

apps/web/app/app.dub.co/(dashboard)/[slug]/(ee)/settings/security/saml.tsx

@@ -17,12 +17,10 @@ import {
 } from "@dub/ui";
 import { SAML_PROVIDERS } from "@dub/utils";
 import { Lock, ShieldOff } from "lucide-react";
-import { useSession } from "next-auth/react";
 import { useMemo, useState } from "react";
 
 export function SAML() {
-  const { data: session } = useSession();
-  const { id: workspaceId, plan } = useWorkspace();
+  const { id: workspaceId, plan, ssoEmailDomain } = useWorkspace();
   const { SAMLModal, setShowSAMLModal } = useSAMLModal();
   const { RemoveSAMLModal, setShowRemoveSAMLModal } = useRemoveSAMLModal();
   const { provider, configured, loading } = useSAML();
@@ -33,7 +31,7 @@ export function SAML() {
     isLoading,
     update,
   } = useOptimisticUpdate<{
-    ssoEmailDomain: string | null;
+    ssoEnforcedAt: string | null;
   }>(`/api/workspaces/${workspaceId}`, {
     loading: "Saving SAML SSO login setting...",
     success: "SAML SSO login setting has been updated successfully.",
@@ -96,17 +94,16 @@ export function SAML() {
         const { error } = await response.json();
         throw new Error(error.message || "Failed to update workspace.");
       }
+
       const data = await response.json();
 
       return {
-        ssoEmailDomain: data.ssoEmailDomain,
+        ssoEnforcedAt: data.ssoEnforcedAt,
       };
     };
 
     await update(updateWorkspace, {
-      ssoEmailDomain: workspaceData?.ssoEmailDomain
-        ? null
-        : session?.user?.email?.split("@")[1] || "domain.com", // fallback to dummy domain for optimistic update
+      ssoEnforcedAt: enforceSAML ? new Date().toISOString() : null,
     });
   };
 
@@ -208,18 +205,18 @@ export function SAML() {
                 Require workspace members to login with SAML to access this
                 workspace
               </label>
-              {workspaceData?.ssoEmailDomain && (
+              {workspaceData?.ssoEnforcedAt && (
                 <Badge
                   variant="blueGradient"
                   className="flex items-center gap-1"
                 >
                   <Globe2 className="size-3" />
-                  {workspaceData.ssoEmailDomain}
+                  {ssoEmailDomain}
                 </Badge>
               )}
             </div>
             <Switch
-              checked={!!workspaceData?.ssoEmailDomain}
+              checked={workspaceData?.ssoEnforcedAt !== null}
               loading={isLoading}
               disabled={plan !== "enterprise"}
               fn={handleSSOEnforcementChange}

apps/web/lib/api/workspaces/is-saml-enforced-for-email-domain.ts

@@ -6,7 +6,7 @@ import { isGenericEmail } from "../../is-generic-email";
 // Checks if SAML SSO is enforced for a given email domain
 export const isSamlEnforcedForEmailDomain = async (email: string) => {
   const hostname = (await headers()).get("host");
-  const emailDomain = email.split("@")[1].toLowerCase();
+  const emailDomain = email.split("@")[1].toLocaleLowerCase();
 
   if (
     !hostname ||
@@ -17,11 +17,18 @@ export const isSamlEnforcedForEmailDomain = async (email: string) => {
     return false;
   }
 
-  const workspace = await prisma.project.count({
+  const workspace = await prisma.project.findUnique({
     where: {
       ssoEmailDomain: emailDomain,
     },
+    select: {
+      ssoEnforcedAt: true,
+    },
   });
 
-  return workspace > 0;
+  if (workspace?.ssoEnforcedAt) {
+    return true;
+  }
+
+  return false;
 };

apps/web/lib/auth/options.ts

@@ -427,16 +427,19 @@ export const authOptions: NextAuthOptions = {
 
         if (workspace) {
           const { ssoEmailDomain } = workspace;
+          const emailDomain = user.email.split("@")[1];
 
-          if (ssoEmailDomain) {
-            const emailDomain = user.email.split("@")[1];
+          // ssoEmailDomain should be required for all SAML enabled workspace
+          // this should not happen
+          if (!ssoEmailDomain) {
+            return false;
+          }
 
-            if (
-              emailDomain.toLocaleLowerCase() !==
-              ssoEmailDomain.toLocaleLowerCase()
-            ) {
-              return false;
-            }
+          if (
+            emailDomain.toLocaleLowerCase() !==
+            ssoEmailDomain.toLocaleLowerCase()
+          ) {
+            return false;
           }
 
           await Promise.allSettled([

apps/web/lib/zod/schemas/workspaces.ts

@@ -122,6 +122,8 @@ export const WorkspaceSchema = z
       .nullable()
       .describe("Specifies hostnames permitted for client-side click tracking.")
       .openapi({ example: ["dub.sh"] }),
+    ssoEmailDomain: z.string().nullable(),
+    ssoEnforcedAt: z.date().nullable(),
   })
   .openapi({
     title: "Workspace",
@@ -167,7 +169,6 @@ export const WorkspaceSchemaExtended = WorkspaceSchema.extend({
     }),
   ),
   publishableKey: z.string().nullable(),
-  ssoEmailDomain: z.string().nullable(),
 });
 
 export const OnboardingUsageSchema = z.object({

apps/web/scripts/migrations/backfill-saml-sso.ts

@@ -0,0 +1,67 @@
+import { prisma } from "@dub/prisma";
+import "dotenv-flow/config";
+
+async function main() {
+  const workspaces = await prisma.project.findMany({
+    where: {
+      ssoEmailDomain: null,
+      plan: "enterprise",
+    },
+    select: {
+      id: true,
+      name: true,
+      ssoEmailDomain: true,
+      users: {
+        where: {
+          role: "owner",
+        },
+        orderBy: {
+          createdAt: "asc",
+        },
+        take: 1,
+        select: {
+          user: {
+            select: {
+              email: true,
+            },
+          },
+        },
+      },
+    },
+  });
+
+  const workspacesToUpdate: any[] = [];
+
+  for (const workspace of workspaces) {
+    const email = workspace.users[0]?.user?.email;
+    const emailDomain = email ? email.split("@")[1]?.toLowerCase() : undefined;
+
+    if (!emailDomain) {
+      console.log(`Workspace ${workspace.name} has no email domain`);
+      continue;
+    }
+
+    workspacesToUpdate.push({
+      id: workspace.id,
+      name: workspace.name,
+      ssoEmailDomain: emailDomain,
+    });
+  }
+
+  console.table(workspacesToUpdate);
+
+  // await Promise.allSettled(
+  //   workspacesToUpdate.map((workspace) =>
+  //     prisma.project.update({
+  //       where: {
+  //         id: workspace.id,
+  //       },
+  //       data: {
+  //         ssoEmailDomain: workspace.ssoEmailDomain,
+  //       },
+  //     }),
+  //   ),
+  // );
+}
+
+main();

packages/prisma/schema/workspace.prisma

@@ -41,15 +41,15 @@ model Project {
   allowedHostnames Json?
   publishableKey   String? @unique // for the client-side publishable key
 
-  conversionEnabled Boolean @default(false) // Whether to enable conversion tracking for links by default
-  webhookEnabled    Boolean @default(false)
-  ssoEnabled        Boolean @default(false) // TODO: this is not used
-  dotLinkClaimed    Boolean @default(false)
-  ssoEmailDomain    String? @unique
-
-  createdAt        DateTime @default(now())
-  updatedAt        DateTime @updatedAt
-  usageLastChecked DateTime @default(now())
+  conversionEnabled Boolean   @default(false) // Whether to enable conversion tracking for links by default
+  webhookEnabled    Boolean   @default(false)
+  ssoEnabled        Boolean   @default(false) // TODO: this is not used
+  dotLinkClaimed    Boolean   @default(false)
+  ssoEmailDomain    String?   @unique
+  ssoEnforcedAt     DateTime?
+  createdAt         DateTime  @default(now())
+  updatedAt         DateTime  @updatedAt
+  usageLastChecked  DateTime  @default(now())
 
   users                 ProjectUsers[]
   invites               ProjectInvite[]