From 1cdf00508752c264542b4abd3a68be0f3ff9e031 Mon Sep 17 00:00:00 2001
From: Bryan Frimin <bryan@getprobo.com>
Date: Sat, 27 Dec 2025 20:15:27 +0100
Subject: [PATCH] WIP

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
---
 pkg/iam/session_service.go                | 81 +++++++++++++++++++++++
 pkg/server/api/connect/v1/saml_handler.go |  4 +-
 pkg/server/api/connect/v1/v1_resolver.go  |  2 +
 3 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/pkg/iam/session_service.go b/pkg/iam/session_service.go
index bb0c31ae5..e9f0959e7 100644
--- a/pkg/iam/session_service.go
+++ b/pkg/iam/session_service.go
@@ -315,6 +315,87 @@ func (s SessionService) GetActiveSessionForMembership(ctx context.Context, rootS
 	return childSession, nil
 }
 
+// OpenSAMLChildSessionForOrganization creates a SAML-authenticated child session for the given
+// organization under the provided root session.
+//
+// This is intended to be used after a successful SAML assertion ("step-up auth") when the user
+// might have an existing PASSWORD root session, but we still want a SAML child session for a
+// SAML-enabled organization.
+func (s SessionService) OpenSAMLChildSessionForOrganization(
+	ctx context.Context,
+	rootSessionID gid.GID,
+	organizationID gid.GID,
+) (*coredata.Session, *coredata.Membership, error) {
+	var (
+		now          = time.Now()
+		rootSession  = &coredata.Session{}
+		identity     = &coredata.Identity{}
+		membership   = &coredata.Membership{}
+		childSession = &coredata.Session{}
+		scope        = coredata.NewScopeFromObjectID(organizationID)
+	)
+
+	err := s.pg.WithTx(
+		ctx,
+		func(tx pg.Conn) 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 = membership.LoadByIdentityInOrganization(ctx, tx, 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.AuthMethodSAML,
+				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/saml_handler.go b/pkg/server/api/connect/v1/saml_handler.go
index 210f3bde9..f52628d64 100644
--- a/pkg/server/api/connect/v1/saml_handler.go
+++ b/pkg/server/api/connect/v1/saml_handler.go
@@ -90,9 +90,9 @@ func (h *SAMLHandler) ConsumeHandler(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	_, _, err = h.iam.SessionService.AssumeOrganizationSession(ctx, rootSession.ID, membership.OrganizationID)
+	_, _, err = h.iam.SessionService.OpenSAMLChildSessionForOrganization(ctx, rootSession.ID, membership.OrganizationID)
 	if err != nil {
-		h.logger.ErrorCtx(ctx, "cannot assume organization session", log.Error(err))
+		h.logger.ErrorCtx(ctx, "cannot open SAML child session", log.Error(err))
 		h.renderInternalServerError(w, r)
 		return
 	}
diff --git a/pkg/server/api/connect/v1/v1_resolver.go b/pkg/server/api/connect/v1/v1_resolver.go
index 97828682f..edff32c91 100644
--- a/pkg/server/api/connect/v1/v1_resolver.go
+++ b/pkg/server/api/connect/v1/v1_resolver.go
@@ -319,6 +319,8 @@ func (r *membershipConnectionResolver) TotalCount(ctx context.Context, obj *type
 
 // SignIn is the resolver for the signIn field.
 func (r *mutationResolver) SignIn(ctx context.Context, input types.SignInInput) (*types.SignInPayload, error) {
+	// TODO: handle existing session to only open child session and chnage root session auth method to PASSWORD
+
 	user, session, err := r.iam.AuthService.OpenSessionWithPassword(ctx, input.Email, input.Password)
 	if err != nil {
 		var ErrInvalidCredentials *iam.ErrInvalidCredentials