fix: SAML SSO - Stop relying on invitedTo column and use membership table instead (#11331)

apps/web/pages/auth/sso/[provider].tsx

@@ -7,13 +7,8 @@ import { getPremiumMonthlyPlanPriceId } from "@calcom/app-store/stripepayment/li
 import { getServerSession } from "@calcom/features/auth/lib/getServerSession";
 import { orgDomainConfig } from "@calcom/features/ee/organizations/lib/orgDomains";
 import stripe from "@calcom/features/ee/payments/server/stripe";
-import {
-  hostedCal,
-  isSAMLLoginEnabled,
-  samlProductID,
-  samlTenantID,
-  samlTenantProduct,
-} from "@calcom/features/ee/sso/lib/saml";
+import { hostedCal, isSAMLLoginEnabled, samlProductID, samlTenantID } from "@calcom/features/ee/sso/lib/saml";
+import { ssoTenantProduct } from "@calcom/features/ee/sso/lib/sso";
 import { checkUsername } from "@calcom/lib/server/checkUsername";
 import prisma from "@calcom/prisma";
 
@@ -111,7 +106,7 @@ export const getServerSideProps = async (context: GetServerSidePropsContext) =>
       error = "Email not provided";
     } else {
       try {
-        const ret = await samlTenantProduct(prisma, emailParam);
+        const ret = await ssoTenantProduct(prisma, emailParam);
         tenant = ret.tenant;
         product = ret.product;
       } catch (e) {

packages/features/ee/sso/lib/saml.ts

@@ -2,8 +2,6 @@ import type { SAMLSSORecord, OIDCSSORecord } from "@boxyhq/saml-jackson";
 
 import { HOSTED_CAL_FEATURES } from "@calcom/lib/constants";
 import { isTeamAdmin } from "@calcom/lib/server/queries/teams";
-import type { PrismaClient } from "@calcom/prisma";
-import { TRPCError } from "@calcom/trpc/server";
 
 export const samlDatabaseUrl = process.env.SAML_DATABASE_URL || "";
 export const isSAMLLoginEnabled = samlDatabaseUrl.length > 0;
@@ -30,38 +28,6 @@ export const isSAMLAdmin = (email: string) => {
   return false;
 };
 
-export const samlTenantProduct = async (prisma: PrismaClient, email: string) => {
-  const user = await prisma.user.findUnique({
-    where: {
-      email,
-    },
-    select: {
-      id: true,
-      invitedTo: true,
-    },
-  });
-
-  if (!user) {
-    throw new TRPCError({
-      code: "UNAUTHORIZED",
-      message: "no_account_exists",
-    });
-  }
-
-  if (!user.invitedTo) {
-    throw new TRPCError({
-      code: "BAD_REQUEST",
-      message:
-        "Could not find a SAML Identity Provider for your email. Please contact your admin to ensure you have been given access to Cal",
-    });
-  }
-
-  return {
-    tenant: tenantPrefix + user.invitedTo,
-    product: samlProductID,
-  };
-};
-
 export const canAccess = async (user: { id: number; email: string }, teamId: number | null) => {
   const { id: userId, email } = user;
 

packages/features/ee/sso/lib/sso.ts

@@ -0,0 +1,57 @@
+import type { PrismaClient } from "@calcom/prisma";
+import { TRPCError } from "@calcom/trpc/server";
+
+import jackson from "./jackson";
+import { tenantPrefix, samlProductID } from "./saml";
+
+export const ssoTenantProduct = async (prisma: PrismaClient, email: string) => {
+  const { connectionController } = await jackson();
+
+  const memberships = await prisma.membership.findMany({
+    select: {
+      teamId: true,
+    },
+    where: {
+      accepted: true,
+      user: {
+        email,
+      },
+    },
+  });
+
+  if (!memberships || memberships.length === 0) {
+    throw new TRPCError({
+      code: "UNAUTHORIZED",
+      message: "no_account_exists",
+    });
+  }
+
+  // Check SSO connections for each team user is a member of
+  // We'll use the first one we find
+  const promises = memberships.map(({ teamId }) =>
+    connectionController.getConnections({
+      tenant: `${tenantPrefix}${teamId}`,
+      product: samlProductID,
+    })
+  );
+
+  const connectionResults = await Promise.allSettled(promises);
+
+  const connectionsFound = connectionResults
+    .filter((result) => result.status === "fulfilled")
+    .map((result) => (result.status === "fulfilled" ? result.value : []))
+    .filter((connections) => connections.length > 0);
+
+  if (connectionsFound.length === 0) {
+    throw new TRPCError({
+      code: "BAD_REQUEST",
+      message:
+        "Could not find a SSO Identity Provider for your email. Please contact your admin to ensure you have been given access to Cal",
+    });
+  }
+
+  return {
+    tenant: connectionsFound[0][0].tenant,
+    product: samlProductID,
+  };
+};

packages/trpc/server/routers/publicViewer/samlTenantProduct.handler.ts

@@ -1,4 +1,4 @@
-import { samlTenantProduct } from "@calcom/features/ee/sso/lib/saml";
+import { ssoTenantProduct } from "@calcom/features/ee/sso/lib/sso";
 import type { PrismaClient } from "@calcom/prisma";
 
 import type { TSamlTenantProductInputSchema } from "./samlTenantProduct.schema";
@@ -14,7 +14,7 @@ export const samlTenantProductHandler = ({ ctx, input }: SamlTenantProductOption
   const { prisma } = ctx;
   const { email } = input;
 
-  return samlTenantProduct(prisma, email);
+  return ssoTenantProduct(prisma, email);
 };
 
 export default samlTenantProductHandler;