From 7668113b964fc129f935bed4bc2309e093dfdc5c Mon Sep 17 00:00:00 2001
From: peintnermax <max@caos.ch>
Date: Mon, 11 Mar 2024 17:33:19 +0100
Subject: [PATCH] catch org context

---
 apps/login/app/(login)/headers/page.tsx       | 22 ---------
 .../app/(login)/{template.tsx => layout.tsx}  |  5 +-
 apps/login/app/(login)/login/route.ts         | 48 +++++++++++++------
 3 files changed, 37 insertions(+), 38 deletions(-)
 delete mode 100644 apps/login/app/(login)/headers/page.tsx
 rename apps/login/app/(login)/{template.tsx => layout.tsx} (91%)

diff --git a/apps/login/app/(login)/headers/page.tsx b/apps/login/app/(login)/headers/page.tsx
deleted file mode 100644
index a12ea774dfb..00000000000
--- a/apps/login/app/(login)/headers/page.tsx
+++ /dev/null
@@ -1,22 +0,0 @@
-import { headers } from "next/headers";
-
-export default function Page() {
-  const headersList = headers();
-  const hds = [
-    "x-zitadel-login-client",
-    "forwarded",
-    "x-zitadel-forwarded",
-    "host",
-    "referer",
-  ];
-  return (
-    <div className="space-y-8">
-      <h1 className="text-xl font-medium">Headers</h1>
-      {hds.map((h) => (
-        <p key={h}>
-          {h}:{headersList.get(h)}
-        </p>
-      ))}
-    </div>
-  );
-}
diff --git a/apps/login/app/(login)/template.tsx b/apps/login/app/(login)/layout.tsx
similarity index 91%
rename from apps/login/app/(login)/template.tsx
rename to apps/login/app/(login)/layout.tsx
index e938635fcc1..851b7e75859 100644
--- a/apps/login/app/(login)/template.tsx
+++ b/apps/login/app/(login)/layout.tsx
@@ -3,11 +3,14 @@ import React from "react";
 import { getBrandingSettings, server } from "#/lib/zitadel";
 import { Logo } from "#/ui/Logo";
 
-export default async function Template({
+export default async function Layout({
   children,
+  params,
 }: {
   children: React.ReactNode;
+  params: any;
 }) {
+  console.log(params);
   const branding = await getBrandingSettings(server);
   let partial: Partial<BrandingSettings> | undefined;
   if (branding) {
diff --git a/apps/login/app/(login)/login/route.ts b/apps/login/app/(login)/login/route.ts
index 9cc4eca2d7c..2c6034f8ae8 100644
--- a/apps/login/app/(login)/login/route.ts
+++ b/apps/login/app/(login)/login/route.ts
@@ -5,7 +5,7 @@ import {
   server,
 } from "#/lib/zitadel";
 import { SessionCookie, getAllSessions } from "#/utils/cookies";
-import { Session, AuthRequest, Prompt } from "@zitadel/server";
+import { Session, AuthRequest, Prompt, login } from "@zitadel/server";
 import { NextRequest, NextResponse } from "next/server";
 
 async function loadSessions(ids: string[]): Promise<Session[]> {
@@ -82,25 +82,29 @@ export async function GET(request: NextRequest) {
   if (authRequestId) {
     console.log(`Login with authRequest: ${authRequestId}`);
     const { authRequest } = await getAuthRequest(server, { authRequestId });
+    let organization;
+
+    if (
+      authRequest?.scope &&
+      authRequest.scope.find((s) => ORG_SCOPE_REGEX.test(s))
+    ) {
+      const orgId = authRequest.scope.find((s) => ORG_SCOPE_REGEX.test(s));
+
+      if (orgId) {
+        const matched = orgId.replace("urn:zitadel:iam:org:id:", "");
+        organization = matched;
+      }
+    }
 
     if (authRequest && authRequest.prompt.includes(Prompt.PROMPT_CREATE)) {
       const registerUrl = new URL("/register", request.url);
       if (authRequest?.id) {
         registerUrl.searchParams.set("authRequestId", authRequest?.id);
       }
-
-      if (
-        authRequest.scope &&
-        authRequest.scope.find((s) => ORG_SCOPE_REGEX.test(s))
-      ) {
-        const orgId = authRequest.scope
-          .find((s) => ORG_SCOPE_REGEX.test(s))
-          ?.match(ORG_SCOPE_REGEX)?.[1];
-        console.log(orgId);
-        if (orgId) {
-          registerUrl.searchParams.set("orgId", orgId);
-        }
+      if (organization) {
+        registerUrl.searchParams.set("organization", organization);
       }
+
       return NextResponse.redirect(registerUrl);
     }
 
@@ -112,6 +116,9 @@ export async function GET(request: NextRequest) {
         if (authRequest?.id) {
           accountsUrl.searchParams.set("authRequestId", authRequest?.id);
         }
+        if (organization) {
+          accountsUrl.searchParams.set("organization", organization);
+        }
 
         return NextResponse.redirect(accountsUrl);
       } else if (authRequest.prompt.includes(Prompt.PROMPT_LOGIN)) {
@@ -123,6 +130,9 @@ export async function GET(request: NextRequest) {
         if (authRequest.loginHint) {
           loginNameUrl.searchParams.set("loginName", authRequest.loginHint);
         }
+        if (organization) {
+          loginNameUrl.searchParams.set("organization", organization);
+        }
         return NextResponse.redirect(loginNameUrl);
       } else if (authRequest.prompt.includes(Prompt.PROMPT_NONE)) {
         // NONE prompt - silent authentication
@@ -178,13 +188,17 @@ export async function GET(request: NextRequest) {
           } else {
             const accountsUrl = new URL("/accounts", request.url);
             accountsUrl.searchParams.set("authRequestId", authRequestId);
-
+            if (organization) {
+              accountsUrl.searchParams.set("organization", organization);
+            }
             return NextResponse.redirect(accountsUrl);
           }
         } else {
           const accountsUrl = new URL("/accounts", request.url);
           accountsUrl.searchParams.set("authRequestId", authRequestId);
-
+          if (organization) {
+            accountsUrl.searchParams.set("organization", organization);
+          }
           return NextResponse.redirect(accountsUrl);
         }
       }
@@ -197,6 +211,10 @@ export async function GET(request: NextRequest) {
         loginNameUrl.searchParams.set("submit", "true"); // autosubmit
       }
 
+      if (organization) {
+        loginNameUrl.searchParams.set("organization", organization);
+      }
+
       return NextResponse.redirect(loginNameUrl);
     }
   } else {