From 53c08926144c76bb6e6a503abcbfba4c0fbead60 Mon Sep 17 00:00:00 2001
From: Max Peintner <peintnerm@gmail.com>
Date: Mon, 16 Dec 2024 11:15:01 +0100
Subject: [PATCH] org discovery on idp callback

---
 .../(login)/idp/[provider]/success/page.tsx   | 49 ++++++++++++++++++-
 apps/login/src/app/(login)/loginname/page.tsx |  2 +-
 apps/login/src/lib/zitadel.ts                 | 16 ++----
 3 files changed, 53 insertions(+), 14 deletions(-)

diff --git a/apps/login/src/app/(login)/idp/[provider]/success/page.tsx b/apps/login/src/app/(login)/idp/[provider]/success/page.tsx
index a2bebc3ae8..c81d43de69 100644
--- a/apps/login/src/app/(login)/idp/[provider]/success/page.tsx
+++ b/apps/login/src/app/(login)/idp/[provider]/success/page.tsx
@@ -3,17 +3,27 @@ import { DynamicTheme } from "@/components/dynamic-theme";
 import { IdpSignin } from "@/components/idp-signin";
 import { idpTypeToIdentityProviderType, PROVIDER_MAPPING } from "@/lib/idp";
 import {
+  addHuman,
   addIDPLink,
-  createUser,
   getBrandingSettings,
   getIDPByID,
+  getLoginSettings,
+  getOrgsByDomain,
   listUsers,
   retrieveIDPIntent,
 } from "@/lib/zitadel";
+import { create } from "@zitadel/client";
 import { AutoLinkingOption } from "@zitadel/proto/zitadel/idp/v2/idp_pb";
+import { OrganizationSchema } from "@zitadel/proto/zitadel/object/v2/object_pb";
 import { BrandingSettings } from "@zitadel/proto/zitadel/settings/v2/branding_settings_pb";
+import {
+  AddHumanUserRequest,
+  AddHumanUserRequestSchema,
+} from "@zitadel/proto/zitadel/user/v2/user_service_pb";
 import { getLocale, getTranslations } from "next-intl/server";
 
+const ORG_SUFFIX_REGEX = /(?<=@)(.+)/;
+
 async function loginFailed(branding?: BrandingSettings, error: string = "") {
   const locale = getLocale();
   const t = await getTranslations({ locale, namespace: "idp" });
@@ -192,7 +202,42 @@ export default async function Page(props: {
   }
 
   if (options?.isCreationAllowed && options.isAutoCreation) {
-    const newUser = await createUser(providerType, idpInformation);
+    let orgToRegisterOn: string | undefined = organization;
+
+    let userData: AddHumanUserRequest =
+      PROVIDER_MAPPING[providerType](idpInformation);
+
+    if (
+      !orgToRegisterOn &&
+      userData.username && // username or email?
+      ORG_SUFFIX_REGEX.test(userData.username)
+    ) {
+      const matched = ORG_SUFFIX_REGEX.exec(userData.username);
+      const suffix = matched?.[1] ?? "";
+
+      // this just returns orgs where the suffix is set as primary domain
+      const orgs = await getOrgsByDomain(suffix);
+      const orgToCheckForDiscovery =
+        orgs.result && orgs.result.length === 1 ? orgs.result[0].id : undefined;
+
+      const orgLoginSettings = await getLoginSettings(orgToCheckForDiscovery);
+      if (orgLoginSettings?.allowDomainDiscovery) {
+        orgToRegisterOn = orgToCheckForDiscovery;
+      }
+    }
+
+    if (orgToRegisterOn) {
+      const organizationSchema = create(OrganizationSchema, {
+        org: { case: "orgId", value: orgToRegisterOn },
+      });
+
+      userData = create(AddHumanUserRequestSchema, {
+        ...userData,
+        organization: organizationSchema,
+      });
+    }
+
+    const newUser = await addHuman(userData);
 
     if (newUser) {
       return (
diff --git a/apps/login/src/app/(login)/loginname/page.tsx b/apps/login/src/app/(login)/loginname/page.tsx
index e886de7f87..44601e1845 100644
--- a/apps/login/src/app/(login)/loginname/page.tsx
+++ b/apps/login/src/app/(login)/loginname/page.tsx
@@ -61,7 +61,7 @@ export default async function Page(props: {
             <SignInWithIdp
               identityProviders={identityProviders}
               authRequestId={authRequestId}
-              organization={organization ?? defaultOrganization} // use the organization from the searchParams here otherwise fallback to the default organization
+              organization={organization}
             ></SignInWithIdp>
           )}
         </UsernameForm>
diff --git a/apps/login/src/lib/zitadel.ts b/apps/login/src/lib/zitadel.ts
index 306d168b63..aaf0a7c05b 100644
--- a/apps/login/src/lib/zitadel.ts
+++ b/apps/login/src/lib/zitadel.ts
@@ -10,8 +10,8 @@ import {
 import { createServerTransport } from "@zitadel/node";
 import { RequestChallenges } from "@zitadel/proto/zitadel/session/v2/challenge_pb";
 import { Checks } from "@zitadel/proto/zitadel/session/v2/session_service_pb";
-import { IDPInformation } from "@zitadel/proto/zitadel/user/v2/idp_pb";
 import {
+  AddHumanUserRequest,
   RetrieveIdentityProviderIntentRequest,
   SetPasswordRequest,
   SetPasswordRequestSchema,
@@ -23,7 +23,6 @@ import { create, Duration } from "@zitadel/client";
 import { TextQueryMethod } from "@zitadel/proto/zitadel/object/v2/object_pb";
 import { CreateCallbackRequest } from "@zitadel/proto/zitadel/oidc/v2/oidc_service_pb";
 import { Organization } from "@zitadel/proto/zitadel/org/v2/org_pb";
-import { IdentityProviderType } from "@zitadel/proto/zitadel/settings/v2/login_settings_pb";
 import type { RedirectURLsJson } from "@zitadel/proto/zitadel/user/v2/idp_pb";
 import {
   NotificationType,
@@ -39,7 +38,6 @@ import {
   UserState,
 } from "@zitadel/proto/zitadel/user/v2/user_pb";
 import { unstable_cacheLife as cacheLife } from "next/cache";
-import { PROVIDER_MAPPING } from "./idp";
 
 const transport = createServerTransport(
   process.env.ZITADEL_SERVICE_USER_TOKEN!,
@@ -249,6 +247,10 @@ export async function addHumanUser({
   });
 }
 
+export async function addHuman(request: AddHumanUserRequest) {
+  return userService.addHumanUser(request);
+}
+
 export async function verifyTOTPRegistration(code: string, userId: string) {
   return userService.verifyTOTPRegistration({ code, userId }, {});
 }
@@ -487,14 +489,6 @@ export function addIDPLink(
   );
 }
 
-export function createUser(
-  provider: IdentityProviderType,
-  info: IDPInformation,
-) {
-  const userData = PROVIDER_MAPPING[provider](info);
-  return userService.addHumanUser(userData, {});
-}
-
 /**
  *
  * @param userId the id of the user where the email should be set