From 9153b5c4743281f3da8c2243c5ac819de310399f Mon Sep 17 00:00:00 2001
From: peintnermax <max@caos.ch>
Date: Wed, 3 Apr 2024 15:16:06 +0200
Subject: [PATCH] org primary domain scope

---
 apps/login/app/login/route.ts        | 23 +++++++++++++++++++----
 apps/login/lib/zitadel.ts            | 10 +++++++++-
 packages/zitadel-server/src/index.ts |  1 +
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/apps/login/app/login/route.ts b/apps/login/app/login/route.ts
index a5531c518d8..61a2c8e5dee 100644
--- a/apps/login/app/login/route.ts
+++ b/apps/login/app/login/route.ts
@@ -1,11 +1,12 @@
 import {
   createCallback,
   getAuthRequest,
+  getOrgByDomain,
   listSessions,
   server,
 } from "#/lib/zitadel";
 import { SessionCookie, getAllSessions } from "#/utils/cookies";
-import { Session, AuthRequest, Prompt, login } from "@zitadel/server";
+import { Session, AuthRequest, Prompt } from "@zitadel/server";
 import { NextRequest, NextResponse } from "next/server";
 
 async function loadSessions(ids: string[]): Promise<Session[]> {
@@ -18,6 +19,7 @@ async function loadSessions(ids: string[]): Promise<Session[]> {
 }
 
 const ORG_SCOPE_REGEX = /urn:zitadel:iam:org:id:([0-9]+)/;
+const ORG_DOMAIN_SCOPE_REGEX = /urn:zitadel:iam:org:domain:primary:(.+)/; // TODO: check regex for all domain character options
 
 function findSession(
   sessions: Session[],
@@ -91,13 +93,26 @@ export async function GET(request: NextRequest) {
       authRequest?.scope &&
       authRequest.scope.find((s: string) => ORG_SCOPE_REGEX.test(s))
     ) {
-      const orgId = authRequest.scope.find((s: string) =>
+      const orgScope = authRequest.scope.find((s: string) =>
         ORG_SCOPE_REGEX.test(s)
       );
 
-      if (orgId) {
-        const matched = ORG_SCOPE_REGEX.exec(orgId);
+      if (orgScope) {
+        const matched = ORG_SCOPE_REGEX.exec(orgScope);
         organization = matched?.[1] ?? "";
+      } else {
+        const orgDomainScope = authRequest.scope.find((s: string) =>
+          ORG_DOMAIN_SCOPE_REGEX.test(s)
+        );
+
+        if (orgDomainScope) {
+          const matched = ORG_DOMAIN_SCOPE_REGEX.exec(orgDomainScope);
+          const orgDomain = matched?.[1] ?? "";
+          if (orgDomain) {
+            const org = await getOrgByDomain(orgDomain);
+            organization = org?.org?.id ?? "";
+          }
+        }
       }
     }
 
diff --git a/apps/login/lib/zitadel.ts b/apps/login/lib/zitadel.ts
index bb5e9190ab1..92739fd1f67 100644
--- a/apps/login/lib/zitadel.ts
+++ b/apps/login/lib/zitadel.ts
@@ -22,10 +22,11 @@ import {
   SetSessionResponse,
   SetSessionRequest,
   ListUsersResponse,
-  ListUsersRequest,
+  management,
   DeleteSessionResponse,
   VerifyPasskeyRegistrationResponse,
   LoginSettings,
+  GetOrgByDomainGlobalResponse,
   GetLoginSettingsResponse,
   ListAuthenticationMethodTypesResponse,
   StartIdentityProviderIntentRequest,
@@ -331,6 +332,13 @@ export async function listUsers(
   );
 }
 
+export async function getOrgByDomain(
+  domain: string
+): Promise<GetOrgByDomainGlobalResponse> {
+  const mgmtService = management.getManagement(server);
+  return mgmtService.getOrgByDomainGlobal({ domain }, {});
+}
+
 export async function startIdentityProviderFlow(
   server: ZitadelServer,
   { idpId, urls }: StartIdentityProviderIntentRequest
diff --git a/packages/zitadel-server/src/index.ts b/packages/zitadel-server/src/index.ts
index 2c8f16a5000..6ec282fe75a 100644
--- a/packages/zitadel-server/src/index.ts
+++ b/packages/zitadel-server/src/index.ts
@@ -87,6 +87,7 @@ export {
 export {
   SetHumanPasswordResponse,
   SetHumanPasswordRequest,
+  GetOrgByDomainGlobalResponse,
 } from "./proto/server/zitadel/management";
 export * from "./proto/server/zitadel/idp";
 export { type LegalAndSupportSettings } from "./proto/server/zitadel/settings/v2beta/legal_settings";