From 2c8ae26ea125dd538e3d74c3341be853f22e29f3 Mon Sep 17 00:00:00 2001
From: Bryan Frimin <bryan@probo.com>
Date: Wed, 17 Jun 2026 18:16:47 +0200
Subject: [PATCH] Open OIDC child session when assuming organization

OIDC login dropped organization_id before the provider redirect, so
callbacks with an existing matching root session never created an org
child session. Persist organization_id in OIDC state, open the child
session on callback, and forward the parameter from the sign-in UI.

Signed-off-by: Bryan Frimin <bryan@probo.com>
---
 .../auth/sign-in/_components/OIDCButton.tsx   | 14 ++-
 .../src/pages/auth/_components/OIDCButton.tsx | 14 ++-
 pkg/coredata/migrations/20260617T120000Z.sql  | 15 +++
 pkg/coredata/oidc_state.go                    | 37 ++++----
 pkg/iam/oidc/service.go                       | 44 ++++-----
 pkg/iam/session_service.go                    | 93 +++++++++++++++++++
 pkg/server/api/connect/v1/oidc_handler.go     | 34 ++++++-
 7 files changed, 202 insertions(+), 49 deletions(-)
 create mode 100644 pkg/coredata/migrations/20260617T120000Z.sql

diff --git a/apps/console/src/pages/iam/auth/sign-in/_components/OIDCButton.tsx b/apps/console/src/pages/iam/auth/sign-in/_components/OIDCButton.tsx
index f292fe148..92f6b0c5e 100644
--- a/apps/console/src/pages/iam/auth/sign-in/_components/OIDCButton.tsx
+++ b/apps/console/src/pages/iam/auth/sign-in/_components/OIDCButton.tsx
@@ -16,6 +16,7 @@ import { useTranslate } from "@probo/i18n";
 import { Button, Google, Microsoft } from "@probo/ui";
 import type { ComponentProps } from "react";
 import { useFragment } from "react-relay";
+import { useSearchParams } from "react-router";
 import { graphql } from "relay-runtime";
 
 import type { OIDCButtonFragment$key } from "#/__generated__/iam/OIDCButtonFragment.graphql";
@@ -42,19 +43,24 @@ export function OIDCButton({
   providerRef: OIDCButtonFragment$key;
 }) {
   const { __ } = useTranslate();
+  const [searchParams] = useSearchParams();
   const safeContinueUrl = useSafeContinueUrl();
   const provider = useFragment(fragment, providerRef);
   const Icon = providerIcons[provider.name];
+  const organizationId = searchParams.get("organization-id");
 
   return (
     <Button
       variant="secondary"
       className="w-full h-10"
       onClick={() => {
-        window.location.href
-          = provider.loginURL
-            + "?continue="
-            + encodeURIComponent(safeContinueUrl.toString());
+        const loginURL = new URL(provider.loginURL, window.location.origin);
+        loginURL.searchParams.set("continue", safeContinueUrl.toString());
+        if (organizationId) {
+          loginURL.searchParams.set("organization_id", organizationId);
+        }
+
+        window.location.href = loginURL.toString();
       }}
     >
       <span className="flex items-center gap-2">
diff --git a/apps/trust/src/pages/auth/_components/OIDCButton.tsx b/apps/trust/src/pages/auth/_components/OIDCButton.tsx
index 64f7e68dc..d8b534290 100644
--- a/apps/trust/src/pages/auth/_components/OIDCButton.tsx
+++ b/apps/trust/src/pages/auth/_components/OIDCButton.tsx
@@ -16,6 +16,7 @@ import { useTranslate } from "@probo/i18n";
 import { Button, Google, Microsoft } from "@probo/ui";
 import type { ComponentProps } from "react";
 import { useFragment } from "react-relay";
+import { useSearchParams } from "react-router";
 import { graphql } from "relay-runtime";
 
 import { useSafeContinueUrl } from "#/hooks/useSafeContinueUrl";
@@ -43,19 +44,24 @@ export function OIDCButton({
   providerRef: OIDCButtonFragment$key;
 }) {
   const { __ } = useTranslate();
+  const [searchParams] = useSearchParams();
   const safeContinueUrl = useSafeContinueUrl();
   const provider = useFragment(fragment, providerRef);
   const Icon = providerIcons[provider.name];
+  const organizationId = searchParams.get("organization-id");
 
   return (
     <Button
       variant="secondary"
       className="w-full h-10"
       onClick={() => {
-        window.location.href
-          = provider.loginURL
-            + "?continue="
-            + encodeURIComponent(safeContinueUrl.toString());
+        const loginURL = new URL(provider.loginURL, window.location.origin);
+        loginURL.searchParams.set("continue", safeContinueUrl.toString());
+        if (organizationId) {
+          loginURL.searchParams.set("organization_id", organizationId);
+        }
+
+        window.location.href = loginURL.toString();
       }}
     >
       <span className="flex items-center gap-2">
diff --git a/pkg/coredata/migrations/20260617T120000Z.sql b/pkg/coredata/migrations/20260617T120000Z.sql
new file mode 100644
index 000000000..2be8c8f12
--- /dev/null
+++ b/pkg/coredata/migrations/20260617T120000Z.sql
@@ -0,0 +1,15 @@
+-- Copyright (c) 2026 Probo Inc <hello@probo.com>.
+--
+-- Permission to use, copy, modify, and/or distribute this software for any
+-- purpose with or without fee is hereby granted, provided that the above
+-- copyright notice and this permission notice appear in all copies.
+--
+-- THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+-- REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+-- AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+-- INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+-- LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+-- OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+-- PERFORMANCE OF THIS SOFTWARE.
+
+ALTER TABLE iam_oidc_states ADD COLUMN organization_id TEXT;
diff --git a/pkg/coredata/oidc_state.go b/pkg/coredata/oidc_state.go
index 381dd7747..52ff5f703 100644
--- a/pkg/coredata/oidc_state.go
+++ b/pkg/coredata/oidc_state.go
@@ -22,32 +22,35 @@ import (
 
 	"github.com/jackc/pgx/v5"
 	"go.gearno.de/kit/pg"
+	"go.probo.inc/probo/pkg/gid"
 )
 
 type OIDCState struct {
-	ID           string       `db:"id"`
-	Provider     OIDCProvider `db:"provider"`
-	Nonce        string       `db:"nonce"`
-	CodeVerifier string       `db:"code_verifier"`
-	ContinueURL  string       `db:"continue_url"`
-	CreatedAt    time.Time    `db:"created_at"`
-	ExpiresAt    time.Time    `db:"expires_at"`
+	ID             string       `db:"id"`
+	Provider       OIDCProvider `db:"provider"`
+	Nonce          string       `db:"nonce"`
+	CodeVerifier   string       `db:"code_verifier"`
+	ContinueURL    string       `db:"continue_url"`
+	OrganizationID *gid.GID     `db:"organization_id"`
+	CreatedAt      time.Time    `db:"created_at"`
+	ExpiresAt      time.Time    `db:"expires_at"`
 }
 
 func (s *OIDCState) Insert(ctx context.Context, conn pg.Tx) error {
 	query := `
-INSERT INTO iam_oidc_states (id, provider, nonce, code_verifier, continue_url, created_at, expires_at)
-VALUES (@id, @provider, @nonce, @code_verifier, @continue_url, @created_at, @expires_at)
+INSERT INTO iam_oidc_states (id, provider, nonce, code_verifier, continue_url, organization_id, created_at, expires_at)
+VALUES (@id, @provider, @nonce, @code_verifier, @continue_url, @organization_id, @created_at, @expires_at)
 `
 
 	args := pgx.StrictNamedArgs{
-		"id":            s.ID,
-		"provider":      s.Provider,
-		"nonce":         s.Nonce,
-		"code_verifier": s.CodeVerifier,
-		"continue_url":  s.ContinueURL,
-		"created_at":    s.CreatedAt,
-		"expires_at":    s.ExpiresAt,
+		"id":              s.ID,
+		"provider":        s.Provider,
+		"nonce":           s.Nonce,
+		"code_verifier":   s.CodeVerifier,
+		"continue_url":    s.ContinueURL,
+		"organization_id": s.OrganizationID,
+		"created_at":      s.CreatedAt,
+		"expires_at":      s.ExpiresAt,
 	}
 
 	_, err := conn.Exec(ctx, query, args)
@@ -60,7 +63,7 @@ VALUES (@id, @provider, @nonce, @code_verifier, @continue_url, @created_at, @exp
 
 func (s *OIDCState) LoadByIDForUpdate(ctx context.Context, conn pg.Tx, id string) error {
 	query := `
-SELECT id, provider, nonce, code_verifier, continue_url, created_at, expires_at
+SELECT id, provider, nonce, code_verifier, continue_url, organization_id, created_at, expires_at
 FROM iam_oidc_states
 WHERE id = @id
 FOR UPDATE
diff --git a/pkg/iam/oidc/service.go b/pkg/iam/oidc/service.go
index f885ab5d1..dc2fa5961 100644
--- a/pkg/iam/oidc/service.go
+++ b/pkg/iam/oidc/service.go
@@ -274,6 +274,7 @@ func (s *Service) InitiateLogin(
 	ctx context.Context,
 	provider coredata.OIDCProvider,
 	continueURL string,
+	organizationID *gid.GID,
 ) (string, error) {
 	info, ok := s.providers[provider]
 	if !ok {
@@ -297,13 +298,14 @@ func (s *Service) InitiateLogin(
 
 	now := time.Now()
 	oidcState := &coredata.OIDCState{
-		ID:           state,
-		Provider:     provider,
-		Nonce:        nonce,
-		CodeVerifier: codeVerifier,
-		ContinueURL:  continueURL,
-		CreatedAt:    now,
-		ExpiresAt:    now.Add(10 * time.Minute),
+		ID:             state,
+		Provider:       provider,
+		Nonce:          nonce,
+		CodeVerifier:   codeVerifier,
+		ContinueURL:    continueURL,
+		OrganizationID: organizationID,
+		CreatedAt:      now,
+		ExpiresAt:      now.Add(10 * time.Minute),
 	}
 
 	err = s.pg.WithTx(
@@ -337,10 +339,10 @@ func (s *Service) HandleCallback(
 	provider coredata.OIDCProvider,
 	stateParam string,
 	code string,
-) (*coredata.Identity, string, error) {
+) (*coredata.Identity, string, *gid.GID, error) {
 	info, ok := s.providers[provider]
 	if !ok {
-		return nil, "", NewProviderNotEnabledError(provider)
+		return nil, "", nil, NewProviderNotEnabledError(provider)
 	}
 
 	var oidcState coredata.OIDCState
@@ -364,15 +366,15 @@ func (s *Service) HandleCallback(
 		},
 	)
 	if err != nil {
-		return nil, "", err
+		return nil, "", nil, err
 	}
 
 	if time.Now().After(oidcState.ExpiresAt) {
-		return nil, "", NewInvalidStateError()
+		return nil, "", nil, NewInvalidStateError()
 	}
 
 	if oidcState.Provider != provider {
-		return nil, "", NewInvalidStateError()
+		return nil, "", nil, NewInvalidStateError()
 	}
 
 	token, err := info.oauth2Config.Exchange(
@@ -381,34 +383,34 @@ func (s *Service) HandleCallback(
 		oauth2.SetAuthURLParam("code_verifier", oidcState.CodeVerifier),
 	)
 	if err != nil {
-		return nil, "", NewCodeExchangeError(err)
+		return nil, "", nil, NewCodeExchangeError(err)
 	}
 
 	rawIDToken, ok := token.Extra("id_token").(string)
 	if !ok {
-		return nil, "", NewIDTokenMissingError()
+		return nil, "", nil, NewIDTokenMissingError()
 	}
 
 	claims, err := s.verifyAndParseIDToken(ctx, info, rawIDToken, oidcState.Nonce)
 	if err != nil {
-		return nil, "", fmt.Errorf("cannot verify id token: %w", err)
+		return nil, "", nil, fmt.Errorf("cannot verify id token: %w", err)
 	}
 
 	if claims.Email == "" {
-		return nil, "", NewMissingEmailClaimError()
+		return nil, "", nil, NewMissingEmailClaimError()
 	}
 
 	if !info.trustProviderEmail && !claims.isEmailVerified() {
-		return nil, "", NewEmailNotVerifiedError()
+		return nil, "", nil, NewEmailNotVerifiedError()
 	}
 
 	if !info.enterpriseChecker(claims) {
-		return nil, "", NewPersonalAccountNotAllowedError()
+		return nil, "", nil, NewPersonalAccountNotAllowedError()
 	}
 
 	email, err := mail.ParseAddr(claims.Email)
 	if err != nil {
-		return nil, "", fmt.Errorf("cannot parse email from id token: %w", err)
+		return nil, "", nil, fmt.Errorf("cannot parse email from id token: %w", err)
 	}
 
 	var identity *coredata.Identity
@@ -455,10 +457,10 @@ func (s *Service) HandleCallback(
 		},
 	)
 	if err != nil {
-		return nil, "", err
+		return nil, "", nil, err
 	}
 
-	return identity, oidcState.ContinueURL, nil
+	return identity, oidcState.ContinueURL, oidcState.OrganizationID, nil
 }
 
 func (s *Service) verifyAndParseIDToken(ctx context.Context, info *providerInfo, rawIDToken string, expectedNonce string) (*idTokenClaims, error) {
diff --git a/pkg/iam/session_service.go b/pkg/iam/session_service.go
index 973244c27..12d1374f5 100644
--- a/pkg/iam/session_service.go
+++ b/pkg/iam/session_service.go
@@ -505,6 +505,99 @@ func (s SessionService) OpenSAMLChildSessionForOrganization(
 	return childSession, membership, nil
 }
 
+// OpenOIDCChildSessionForOrganization creates an OIDC-authenticated child session for the
+// given organization under the provided root session.
+func (s SessionService) OpenOIDCChildSessionForOrganization(
+	ctx context.Context,
+	rootSessionID gid.GID,
+	organizationID gid.GID,
+) (*coredata.Session, *coredata.Membership, error) {
+	var (
+		now          = time.Now()
+		rootSession  = &coredata.Session{}
+		identity     = &coredata.Identity{}
+		profile      = &coredata.MembershipProfile{}
+		membership   = &coredata.Membership{}
+		childSession = &coredata.Session{}
+		scope        = coredata.NewScopeFromObjectID(organizationID)
+	)
+
+	err := s.pg.WithTx(
+		ctx,
+		func(ctx context.Context, tx pg.Tx) error {
+			err := rootSession.LoadByID(ctx, tx, rootSessionID)
+			if err != nil {
+				if err == coredata.ErrResourceNotFound {
+					return NewSessionNotFoundError(rootSessionID)
+				}
+
+				return fmt.Errorf("cannot load session: %w", err)
+			}
+
+			if !rootSession.IsRootSession() {
+				return fmt.Errorf("session %q is not a root session", rootSessionID)
+			}
+
+			if rootSession.ExpireReason != nil || now.After(rootSession.ExpiredAt) {
+				return NewSessionExpiredError(rootSessionID)
+			}
+
+			err = identity.LoadByID(ctx, tx, rootSession.IdentityID)
+			if err != nil {
+				return fmt.Errorf("cannot load identity: %w", err)
+			}
+
+			err = profile.LoadByIdentityIDAndOrganizationID(ctx, tx, scope, rootSession.IdentityID, organizationID)
+			if err != nil {
+				if err == coredata.ErrResourceNotFound {
+					return NewProfileNotFoundError(gid.Nil)
+				}
+
+				return fmt.Errorf("cannot load profile: %w", err)
+			}
+
+			if profile.State == coredata.ProfileStateInactive {
+				return NewUserInactiveError(profile.ID)
+			}
+
+			err = membership.LoadByIdentityIDAndOrganizationID(ctx, tx, scope, rootSession.IdentityID, organizationID)
+			if err != nil {
+				if err == coredata.ErrResourceNotFound {
+					return NewMembershipNotFoundError(organizationID)
+				}
+
+				return fmt.Errorf("cannot load membership: %w", err)
+			}
+
+			tenantID := scope.GetTenantID()
+			childSession = &coredata.Session{
+				ID:              gid.New(tenantID, coredata.SessionEntityType),
+				IdentityID:      rootSession.IdentityID,
+				TenantID:        &tenantID,
+				MembershipID:    &membership.ID,
+				ParentSessionID: &rootSession.ID,
+				AuthMethod:      coredata.AuthMethodOIDC,
+				AuthenticatedAt: now,
+				ExpiredAt:       rootSession.ExpiredAt,
+				CreatedAt:       now,
+				UpdatedAt:       now,
+			}
+
+			err = childSession.Insert(ctx, tx)
+			if err != nil {
+				return fmt.Errorf("cannot insert child session: %w", err)
+			}
+
+			return nil
+		},
+	)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return childSession, membership, nil
+}
+
 func (s SessionService) AssumeOrganizationSession(
 	ctx context.Context,
 	sessionID gid.GID,
diff --git a/pkg/server/api/connect/v1/oidc_handler.go b/pkg/server/api/connect/v1/oidc_handler.go
index 371bdfb5c..51098ad62 100644
--- a/pkg/server/api/connect/v1/oidc_handler.go
+++ b/pkg/server/api/connect/v1/oidc_handler.go
@@ -25,6 +25,7 @@ import (
 	"go.gearno.de/kit/httpserver"
 	"go.gearno.de/kit/log"
 	"go.probo.inc/probo/pkg/coredata"
+	"go.probo.inc/probo/pkg/gid"
 	"go.probo.inc/probo/pkg/iam"
 	"go.probo.inc/probo/pkg/saferedirect"
 	"go.probo.inc/probo/pkg/securecookie"
@@ -77,7 +78,19 @@ func (h *OIDCHandler) LoginHandler(w http.ResponseWriter, r *http.Request) {
 
 	continueURL := r.URL.Query().Get("continue")
 
-	authURL, err := h.iam.OIDCService.InitiateLogin(ctx, provider, continueURL)
+	var organizationID *gid.GID
+
+	if organizationIDParam := r.URL.Query().Get("organization_id"); organizationIDParam != "" {
+		parsedOrganizationID, err := gid.ParseGID(organizationIDParam)
+		if err != nil {
+			httpserver.RenderError(w, http.StatusBadRequest, errors.New("invalid organization_id parameter"))
+			return
+		}
+
+		organizationID = &parsedOrganizationID
+	}
+
+	authURL, err := h.iam.OIDCService.InitiateLogin(ctx, provider, continueURL, organizationID)
 	if err != nil {
 		h.logger.ErrorCtx(ctx, "cannot initiate OIDC login", log.Error(err))
 		httpserver.RenderError(w, http.StatusInternalServerError, errors.New("internal server error"))
@@ -118,7 +131,7 @@ func (h *OIDCHandler) CallbackHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	identity, continueURL, err := h.iam.OIDCService.HandleCallback(ctx, provider, stateParam, code)
+	identity, continueURL, organizationID, err := h.iam.OIDCService.HandleCallback(ctx, provider, stateParam, code)
 	if err != nil {
 		h.logger.ErrorCtx(ctx, "cannot handle OIDC callback", log.Error(err))
 		httpserver.RenderError(w, http.StatusUnauthorized, errors.New("authentication failed"))
@@ -155,9 +168,24 @@ func (h *OIDCHandler) CallbackHandler(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
+	if organizationID != nil {
+		_, _, err = h.iam.SessionService.OpenOIDCChildSessionForOrganization(ctx, rootSession.ID, *organizationID)
+		if err != nil {
+			h.logger.ErrorCtx(ctx, "cannot open OIDC child session", log.Error(err))
+			httpserver.RenderError(w, http.StatusInternalServerError, errors.New("internal server error"))
+
+			return
+		}
+	}
+
 	h.sessionCookie.Set(w, rootSession)
 
-	redirectURL := h.safeRedirect.GetSafeRedirectURL(ctx, continueURL, "/")
+	defaultRedirect := "/"
+	if organizationID != nil {
+		defaultRedirect = "/organizations/" + organizationID.String()
+	}
+
+	redirectURL := h.safeRedirect.GetSafeRedirectURL(ctx, continueURL, defaultRedirect)
 
 	if transferURL, ok := h.buildSessionTransferURL(ctx, redirectURL, rootSession.ID.String()); ok {
 		http.Redirect(w, r, transferURL, http.StatusFound)