Fix missing child session on SAML

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

GNUmakefile

@@ -216,6 +216,7 @@ clean: ## Clean the project (node_modules and build artifacts)
 	$(RM) -rf sbom-docker.json sbom.json
 	$(RM) -rf coverage.out coverage.html coverage-e2e.out coverage-e2e.html coverage-combined.out coverage-combined.html
 	$(RM) -rf coverage/
+	$(RM) -rf compose/keycloak/certs/cert.pem compose/keycloak/certs/private-key.pem compose/keycloak/probo-realm.json
 
 .PHONY: stack-up
 stack-up: compose/pebble/certs/rootCA.pem compose/keycloak/probo-realm.json ## Start the docker stack as a deamon

cfg/dev.yaml

@@ -1,6 +1,6 @@
 unit:
   metrics:
-    addr: "localhost:8081"
+    addr: "localhost:5173"
   tracing:
     addr: "localhost:4317"
     max-batch-size: 512

compose/keycloak/probo-realm.json.tmpl

@@ -29,7 +29,7 @@
   "defaultRoles": ["member"],
   "clients": [
     {
-      "clientId": "http://localhost:5173/connect/saml/metadata",
+      "clientId": "http://localhost:5173/api/connect/v1/saml/2.0/metadata",
       "name": "Probo Console",
       "description": "Probo GRC Platform - SAML Service Provider",
       "enabled": true,
@@ -48,13 +48,13 @@
         "saml_name_id_format": "email",
         "saml.client.signature": "false",
         "saml.authnstatement": "true",
-        "saml_single_logout_service_url_post": "http://localhost:5173/connect/saml/logout",
+        "saml_single_logout_service_url_post": "http://localhost:5173/api/connect/v1/saml/2.0/consume",
-        "saml_single_logout_service_url_redirect": "http://localhost:5173/connect/saml/logout",
+        "saml_single_logout_service_url_redirect": "http://localhost:5173/api/connect/v1/saml/2.0/consume",
         "saml.onetimeuse.condition": "false"
       },
       "rootUrl": "http://localhost:5173",
       "baseUrl": "/",
-      "adminUrl": "http://localhost:5173/connect/saml/consume",
+      "adminUrl": "http://localhost:5173/api/connect/v1/saml/2.0/consume",
       "redirectUris": ["http://localhost:5173/*"],
       "webOrigins": ["http://localhost:5173"],
       "protocolMappers": [

pkg/coredata/migrations/20251224T114827Z.sql

@@ -0,0 +1 @@
+ALTER TABLE iam_saml_requests DROP COLUMN IF EXISTS tenant_id;

pkg/coredata/migrations/20251224T123758Z.sql

@@ -0,0 +1,2 @@
+DROP INDEX IF EXISTS idx_auth_saml_assertions_tenant_id;
+ALTER TABLE iam_saml_assertions DROP COLUMN IF EXISTS tenant_id;

pkg/iam/saml/service.go

@@ -284,7 +284,9 @@ func (s *Service) HandleAssertion(
 				}
 			}
 
-			err = membership.LoadByIdentityAndOrg(ctx, tx, coredata.NewNoScope(), identity.ID, config.OrganizationID)
+			scope := coredata.NewScopeFromObjectID(config.OrganizationID)
+
+			err = membership.LoadByIdentityAndOrg(ctx, tx, scope, identity.ID, config.OrganizationID)
 			if err != nil && err != coredata.ErrResourceNotFound {
 				return fmt.Errorf("cannot load membership: %w", err)
 			}
@@ -300,7 +302,7 @@ func (s *Service) HandleAssertion(
 					UpdatedAt:      now,
 				}
 
-				err = membership.Insert(ctx, tx, coredata.NewNoScope())
+				err = membership.Insert(ctx, tx, scope)
 				if err != nil {
 					return fmt.Errorf("cannot insert membership: %w", err)
 				}
@@ -323,21 +325,21 @@ func (s *Service) HandleAssertion(
 				membership.Role = *role
 				membership.UpdatedAt = now
 
-				err = membership.Update(ctx, tx, coredata.NewNoScope())
+				err = membership.Update(ctx, tx, scope)
 				if err != nil {
 					return fmt.Errorf("cannot update membership: %w", err)
 				}
 			}
 
 			memberProfile := &coredata.MembershipProfile{}
-			err = memberProfile.LoadByMembershipID(ctx, tx, coredata.NewNoScope(), membership.ID)
+			err = memberProfile.LoadByMembershipID(ctx, tx, scope, membership.ID)
 			if err != nil {
 				return fmt.Errorf("cannot load membership profile: %w", err)
 			}
 
 			memberProfile.FullName = fullname
 			memberProfile.UpdatedAt = now
-			err = memberProfile.Update(ctx, tx, coredata.NewNoScope())
+			err = memberProfile.Update(ctx, tx, scope)
 			if err != nil {
 				return fmt.Errorf("cannot update membership profile: %w", err)
 			}

pkg/iam/session_service.go

@@ -374,15 +374,19 @@ func (s SessionService) AssumeOrganizationSession(
 			}
 
 			if err == nil && samlConfig.EnforcementPolicy == coredata.SAMLEnforcementPolicyRequired {
-				redirectURL, err := s.SAMLService.InitiateLogin(ctx, samlConfig.ID)
+				if rootSession.AuthMethod != coredata.AuthMethodSAML {
-				if err != nil {
+					redirectURL, err := s.SAMLService.InitiateLogin(ctx, samlConfig.ID)
-					return fmt.Errorf("cannot initiate SAML login: %w", err)
+					if err != nil {
+						return fmt.Errorf("cannot initiate SAML login: %w", err)
+					}
+
+					return NewSAMLAuthenticationRequiredError("policy_requirement", redirectURL.String())
 				}
-
+			} else if err == nil && samlConfig.EnforcementPolicy == coredata.SAMLEnforcementPolicyOptional {
-				return NewSAMLAuthenticationRequiredError("policy_requirement", redirectURL.String())
+				// SAML is optional: both PASSWORD and SAML root sessions are allowed.
-			}
+			} else if rootSession.AuthMethod != coredata.AuthMethodPassword {
-
+				// No (or non-required) SAML configuration: require a password-authenticated root session
-			if rootSession.AuthMethod != coredata.AuthMethodPassword {
+				// (eg. when switching into a password-based org from a SAML login).
 				return NewPasswordRequiredError("password_authentication_required")
 			}
 
@@ -393,7 +397,7 @@ func (s SessionService) AssumeOrganizationSession(
 				TenantID:        &tenantID,
 				MembershipID:    &membership.ID,
 				ParentSessionID: &rootSession.ID,
-				AuthMethod:      coredata.AuthMethodPassword,
+				AuthMethod:      rootSession.AuthMethod,
 				AuthenticatedAt: now,
 				ExpiredAt:       rootSession.ExpiredAt,
 				CreatedAt:       now,

pkg/server/api/connect/v1/resolver.go

@@ -55,7 +55,7 @@ func NewMux(logger *log.Logger, svc *iam.Service, cookieConfig securecookie.Conf
 
 	sessionMiddleware := NewSessionMiddleware(svc, cookieConfig)
 	graphqlHandler := NewGraphQLHandler(svc, logger, baseURL, cookieConfig)
-	samlHandler := NewSAMLHandler(svc, cookieConfig, baseURL)
+	samlHandler := NewSAMLHandler(svc, cookieConfig, baseURL, logger)
 
 	router := r.With(sessionMiddleware)
 

pkg/server/api/connect/v1/saml_handler.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/go-chi/chi/v5"
 	"go.gearno.de/kit/httpserver"
+	"go.gearno.de/kit/log"
 	"go.probo.inc/probo/pkg/baseurl"
 	"go.probo.inc/probo/pkg/gid"
 	"go.probo.inc/probo/pkg/iam"
@@ -17,10 +18,15 @@ type SAMLHandler struct {
 	iam          *iam.Service
 	cookieConfig securecookie.Config
 	baseURL      *baseurl.BaseURL
+	logger       *log.Logger
 }
 
-func NewSAMLHandler(iam *iam.Service, cookieConfig securecookie.Config, baseURL *baseurl.BaseURL) *SAMLHandler {
+func NewSAMLHandler(iam *iam.Service, cookieConfig securecookie.Config, baseURL *baseurl.BaseURL, logger *log.Logger) *SAMLHandler {
-	return &SAMLHandler{iam: iam, cookieConfig: cookieConfig, baseURL: baseURL}
+	return &SAMLHandler{iam: iam, cookieConfig: cookieConfig, baseURL: baseURL, logger: logger}
+}
+
+func (h *SAMLHandler) renderInternalServerError(w http.ResponseWriter, r *http.Request) {
+	httpserver.RenderError(w, http.StatusInternalServerError, errors.New("internal server error"))
 }
 
 func (h *SAMLHandler) MetadataHandler(w http.ResponseWriter, r *http.Request) {
@@ -58,17 +64,41 @@ func (h *SAMLHandler) ConsumeHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	session := SessionFromContext(ctx)
+	rootSession := SessionFromContext(ctx)
-	if session == nil {
+
-		h.iam.AuthService.OpenSessionWithSAML(ctx, user.ID, membership.OrganizationID)
+	switch {
+	case rootSession == nil:
+		rootSession, err = h.iam.AuthService.OpenSessionWithSAML(ctx, user.ID, membership.OrganizationID)
+		if err != nil {
+			h.logger.ErrorCtx(ctx, "cannot open root session", log.Error(err))
+			h.renderInternalServerError(w, r)
+			return
+		}
+	case rootSession.IdentityID != user.ID:
+		err = h.iam.SessionService.CloseSession(ctx, rootSession.ID)
+		if err != nil {
+			h.logger.ErrorCtx(ctx, "cannot close session", log.Error(err))
+			h.renderInternalServerError(w, r)
+			return
+		}
+
+		rootSession, err = h.iam.AuthService.OpenSessionWithSAML(ctx, user.ID, membership.OrganizationID)
+		if err != nil {
+			h.logger.ErrorCtx(ctx, "cannot open root session", log.Error(err))
+			h.renderInternalServerError(w, r)
+			return
+		}
 	}
 
-	// TODO open or update the organization session
+	_, _, err = h.iam.SessionService.AssumeOrganizationSession(ctx, rootSession.ID, membership.OrganizationID)
-
+	if err != nil {
-	securecookie.Set(w, h.cookieConfig, session.ID.String())
+		h.logger.ErrorCtx(ctx, "cannot assume organization session", log.Error(err))
+		h.renderInternalServerError(w, r)
+		return
+	}
 
+	securecookie.Set(w, h.cookieConfig, rootSession.ID.String())
 	redirectURL := h.baseURL.WithPath("/organizations/" + membership.OrganizationID.String()).MustString()
-
 	http.Redirect(w, r, redirectURL, http.StatusFound)
 }