Simpliy and ensure tenant for saml consume

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

pkg/auth/saml_mapper.go

@@ -81,7 +81,6 @@ func MapSAMLRoleToSystemRole(samlRole string) string {
 		return samlRole
 	}
 
-	// Default to MEMBER role if SAML role is missing or invalid
 	return "MEMBER"
 }
 

pkg/auth/saml_service.go

@@ -409,7 +409,6 @@ type SAMLUserInfo struct {
 	Role           string
 	SAMLSubject    string
 	OrganizationID gid.GID
-	TenantID       gid.TenantID
 	SAMLConfigID   gid.GID
 }
 
@@ -560,7 +559,6 @@ func (s *SAMLService) HandleSAMLAssertion(
 		Role:           systemRole,
 		SAMLSubject:    samlSubject,
 		OrganizationID: relayState.OrganizationID,
-		TenantID:       org.TenantID,
 		SAMLConfigID:   relayState.SAMLConfigID,
 	}, nil
 }

pkg/auth/service.go

@@ -389,39 +389,24 @@ func (s Service) CreateOrGetSAMLUser(
 	err := s.pg.WithTx(
 		ctx,
 		func(tx pg.Conn) error {
-			// Try to load existing user by email
 			if err := user.LoadByEmail(ctx, tx, emailAddress); err == nil {
-				// User exists - update SAML subject and full name if needed
-				needsUpdate := false
-
-				if user.SAMLSubject == nil || *user.SAMLSubject != samlSubject {
 				user.SAMLSubject = &samlSubject
-					needsUpdate = true
-				}
-				if user.FullName != fullName {
 				user.FullName = fullName
-					needsUpdate = true
-				}
-				if !user.EmailAddressVerified {
 				user.EmailAddressVerified = true
-					needsUpdate = true
-				}
-
-				if needsUpdate {
 				user.UpdatedAt = now
+
 				if err := user.Update(ctx, tx); err != nil {
 					return fmt.Errorf("cannot update user: %w", err)
 				}
-				}
+
 				return nil
 			}
 
-			// No existing user, create new user (all users are global now)
 			user = coredata.User{
 				ID:                   gid.New(gid.NilTenant, coredata.UserEntityType),
 				EmailAddress:         emailAddress,
-				HashedPassword:       nil, // SAML users don't have passwords initially
+				HashedPassword:       nil,
-				EmailAddressVerified: true, // SAML users are verified by IdP
+				EmailAddressVerified: true,
 				FullName:             fullName,
 				SAMLSubject:          &samlSubject,
 				CreatedAt:            now,
@@ -449,6 +434,7 @@ func (s Service) CreateSessionForUser(
 	sessionDuration time.Duration,
 ) (*coredata.Session, error) {
 	now := time.Now()
+
 	session := &coredata.Session{
 		ID:        gid.New(gid.NilTenant, coredata.SessionEntityType),
 		UserID:    userID,

pkg/authz/service.go

@@ -30,9 +30,6 @@ import (
 )
 
 type (
-	// Service handles all authorization logic including organization
-	// membership and permissions. This service is completely independent
-	// of authentication methods.
 	Service struct {
 		pg                      *pg.Client
 		hostname                string
@@ -109,8 +106,6 @@ func (s *Service) GetAllUserOrganizations(
 	return organizations, err
 }
 
-// This method is on Service (not TenantAuthzService) because it operates across tenants
-// and doesn't require tenant-scoped access.
 func (s *Service) GetUserOrganizations(
 	ctx context.Context,
 	userID gid.GID,
@@ -118,18 +113,19 @@ func (s *Service) GetUserOrganizations(
 ) ([]*coredata.Organization, error) {
 	var organizations coredata.Organizations
 
-	err := s.pg.WithConn(ctx, func(conn pg.Conn) error {
+	err := s.pg.WithConn(
+		ctx,
+		func(conn pg.Conn) error {
 			if err := organizations.LoadByUserID(ctx, conn, userID, cursor); err != nil {
 				return fmt.Errorf("cannot load user organizations: %w", err)
 			}
 			return nil
-	})
+		},
+	)
 
 	return organizations, err
 }
 
-// This method is on Service (not TenantAuthzService) because the user accepting
-// the invitation doesn't have tenant access yet.
 func (s *Service) AcceptInvitation(
 	ctx context.Context,
 	token string,
@@ -140,9 +136,11 @@ func (s *Service) AcceptInvitation(
 		TokenTypeOrganizationInvitation,
 		token,
 	)
+
 	if err != nil {
 		return fmt.Errorf("invalid invitation token: %w", err)
 	}
+
 	invitationData := payload.Data
 	scope := coredata.NewScope(invitationData.InvitationID.TenantID())
 
@@ -192,8 +190,6 @@ func (s *Service) AcceptInvitation(
 	)
 }
 
-// This method is on Service (not TenantAuthzService) because the user accepting
-// the invitation doesn't have tenant access yet.
 func (s *Service) AcceptInvitationByID(
 	ctx context.Context,
 	invitationID gid.GID,
@@ -261,64 +257,6 @@ func (s *Service) AcceptInvitationByID(
 	return acceptedInvitation, nil
 }
 
-// EnsureSAMLMembership creates or updates a user's membership in an organization.
-// This is used during SAML authentication to ensure the user has the correct role.
-// This method is on Service (not TenantAuthzService) because SAML authentication
-// happens before the user has tenant access.
-func (s *Service) EnsureSAMLMembership(
-	ctx context.Context,
-	tenantID gid.TenantID,
-	userID gid.GID,
-	organizationID gid.GID,
-	role string,
-) error {
-	scope := coredata.NewScope(tenantID)
-	now := time.Now()
-
-	return s.pg.WithTx(
-		ctx,
-		func(tx pg.Conn) error {
-			var membership coredata.Membership
-
-			err := membership.LoadByUserAndOrg(ctx, tx, scope, userID, organizationID)
-			if err != nil {
-				if _, ok := err.(coredata.ErrMembershipNotFound); !ok {
-					return fmt.Errorf("cannot load membership: %w", err)
-				}
-
-				membershipID := gid.New(tenantID, coredata.MembershipEntityType)
-				membership = coredata.Membership{
-					ID:             membershipID,
-					UserID:         userID,
-					OrganizationID: organizationID,
-					Role:           role,
-					CreatedAt:      now,
-					UpdatedAt:      now,
-				}
-
-				if err := membership.Create(ctx, tx, scope); err != nil {
-					return fmt.Errorf("cannot create membership: %w", err)
-				}
-
-				return nil
-			}
-
-			if membership.Role != role {
-				membership.Role = role
-				membership.UpdatedAt = now
-
-				if err := membership.Update(ctx, tx, scope); err != nil {
-					return fmt.Errorf("cannot update membership role: %w", err)
-				}
-			}
-
-			return nil
-		},
-	)
-}
-
-// This method is on Service (not TenantAuthzService) because the user viewing
-// their invitations doesn't have tenant access yet, and it operates across multiple tenants.
 func (s *Service) GetUserInvitations(
 	ctx context.Context,
 	email string,
@@ -344,8 +282,6 @@ func (s *Service) GetUserInvitations(
 	return page.NewPage(invitations, cursor), nil
 }
 
-// This method is on Service (not TenantAuthzService) because the user viewing
-// their invitations doesn't have tenant access yet, and it operates across multiple tenants.
 func (s *Service) CountUserInvitations(
 	ctx context.Context,
 	email string,
@@ -795,6 +731,58 @@ func (s *TenantAuthzService) InviteUserToOrganization(
 	return invitation, nil
 }
 
+// EnsureSAMLMembership creates or updates a user's membership in an organization.
+// This is used during SAML authentication to ensure the user has the correct role.
+func (s *TenantAuthzService) EnsureSAMLMembership(
+	ctx context.Context,
+	userID gid.GID,
+	organizationID gid.GID,
+	role string,
+) error {
+	now := time.Now()
+
+	return s.pg.WithTx(
+		ctx,
+		func(tx pg.Conn) error {
+			var membership coredata.Membership
+
+			err := membership.LoadByUserAndOrg(ctx, tx, s.scope, userID, organizationID)
+			if err != nil {
+				if _, ok := err.(coredata.ErrMembershipNotFound); !ok {
+					return fmt.Errorf("cannot load membership: %w", err)
+				}
+
+				membershipID := gid.New(s.scope.GetTenantID(), coredata.MembershipEntityType)
+				membership = coredata.Membership{
+					ID:             membershipID,
+					UserID:         userID,
+					OrganizationID: organizationID,
+					Role:           role,
+					CreatedAt:      now,
+					UpdatedAt:      now,
+				}
+
+				if err := membership.Create(ctx, tx, s.scope); err != nil {
+					return fmt.Errorf("cannot create membership: %w", err)
+				}
+
+				return nil
+			}
+
+			if membership.Role != role {
+				membership.Role = role
+				membership.UpdatedAt = now
+
+				if err := membership.Update(ctx, tx, s.scope); err != nil {
+					return fmt.Errorf("cannot update membership role: %w", err)
+				}
+			}
+
+			return nil
+		},
+	)
+}
+
 // This is a placeholder for future permission system
 func (s *TenantAuthzService) HasPermission(
 	ctx context.Context,

pkg/server/auth/saml_acs_handler.go

@@ -75,7 +75,8 @@ func SAMLACSHandler(samlSvc *authsvc.SAMLService, authSvc *authsvc.Service, auth
 			return
 		}
 
-		err = authzSvc.EnsureSAMLMembership(ctx, userInfo.TenantID, user.ID, userInfo.OrganizationID, userInfo.Role)
+		tenantAuthzSvc := authzSvc.WithTenant(userInfo.OrganizationID.TenantID())
+		err = tenantAuthzSvc.EnsureSAMLMembership(ctx, user.ID, userInfo.OrganizationID, userInfo.Role)
 		if err != nil {
 			logger.ErrorCtx(ctx, "cannot ensure membership", log.Error(err), log.String("user_id", user.ID.String()), log.String("org_id", userInfo.OrganizationID.String()))
 			http.Error(w, "cannot create membership", http.StatusInternalServerError)