Support workspace with SSO enforced

front/lib/auth.ts

@@ -395,6 +395,7 @@ export class Authenticator {
             ACTIVATE_ALL_FEATURES_DEV && isDevelopment()
               ? [...WHITELISTABLE_FEATURES]
               : this._flags,
+          ssoEnforced: this._workspace.ssoEnforced,
         }
       : null;
   }

front/lib/iam/session.ts

@@ -1,4 +1,5 @@
 import type { RoleType, UserTypeWithWorkspaces } from "@dust-tt/types";
+import { isEnterpriseConnectionSub } from "@dust-tt/types";
 import type {
   GetServerSidePropsContext,
   GetServerSidePropsResult,
@@ -103,6 +104,22 @@ export type CustomGetServerSideProps<
   session: RequireUserPrivilege extends "none" ? null : SessionWithUser
 ) => Promise<GetServerSidePropsResult<Props>>;
 
+export function statisfiesEnforceEntrepriseConnection(
+  auth: Authenticator,
+  session: SessionWithUser
+) {
+  const owner = auth.workspace();
+  if (!owner) {
+    return true;
+  }
+
+  if (owner.ssoEnforced) {
+    return isEnterpriseConnectionSub(session.user.sub);
+  }
+
+  return true;
+}
+
 async function getAuthenticator(
   context: GetServerSidePropsContext<ParsedUrlQuery, PreviewData>,
   session: SessionWithUser | null,
@@ -156,17 +173,27 @@ export function makeGetServerSidePropsRequirementsWrapper<
         requireUserPrivilege
       );
 
-      if (
-        requireUserPrivilege !== "none" &&
-        (!session || !isValidSession(session))
-      ) {
-        return {
-          redirect: {
-            permanent: false,
-            // TODO(2024-03-04 flav) Add support for `returnTo=`.
-            destination: "/api/auth/login",
-          },
-        };
+      if (requireUserPrivilege !== "none") {
+        if (!session || !isValidSession(session)) {
+          return {
+            redirect: {
+              permanent: false,
+              // TODO(2024-03-04 flav) Add support for `returnTo=`.
+              destination: "/api/auth/login",
+            },
+          };
+        }
+
+        // Validate the user's session to guarantee compliance with the workspace's SSO requirements when SSO is enforced.
+        if (auth && !statisfiesEnforceEntrepriseConnection(auth, session)) {
+          return {
+            redirect: {
+              permanent: false,
+              // TODO(2024-03-04 flav) Add support for `returnTo=`.
+              destination: `/sso-enforced?workspaceId=${auth.workspace()?.sId}`,
+            },
+          };
+        }
       }
 
       const userSession = session as RequireUserPrivilege extends "none"

front/lib/models/workspace.ts

@@ -25,6 +25,7 @@ export class Workspace extends Model<
   declare name: string;
   declare description: string | null;
   declare segmentation: WorkspaceSegmentationType;
+  declare ssoEnforced?: boolean;
   declare subscriptions: NonAttribute<Subscription[]>;
 }
 Workspace.init(
@@ -63,6 +64,10 @@ Workspace.init(
       type: DataTypes.STRING,
       allowNull: true,
     },
+    ssoEnforced: {
+      type: DataTypes.BOOLEAN,
+      defaultValue: false,
+    },
   },
   {
     modelName: "workspace",

types/src/front/user.ts

@@ -15,6 +15,7 @@ export type LightWorkspaceType = {
 
 export type WorkspaceType = LightWorkspaceType & {
   flags: WhitelistableFeature[];
+  ssoEnforced?: boolean;
 };
 
 export type UserProviderType = "github" | "google" | null;

types/src/front/workspace.ts

@@ -8,3 +8,15 @@ export interface WorkspaceEnterpriseConnection {
 }
 
 export type SupportedEnterpriseConnectionStrategies = "okta";
+export const supportedEnterpriseConnectionStrategies: SupportedEnterpriseConnectionStrategies[] =
+  ["okta"];
+
+export function isEnterpriseConnectionSub(
+  sub: string
+): sub is SupportedEnterpriseConnectionStrategies {
+  const [provider] = sub.split("|");
+
+  return supportedEnterpriseConnectionStrategies.includes(
+    provider as SupportedEnterpriseConnectionStrategies
+  );
+}