Accept/Expire invitations on SAML & SCIM operations

Signed-off-by: Émile Ré <emile@getprobo.com>

pkg/coredata/invitation.go

@@ -164,6 +164,72 @@ WHERE
 	return nil
 }
 
+func (i *Invitations) AcceptByEmailAndOrganization(
+	ctx context.Context,
+	conn pg.Conn,
+	scope Scoper,
+	email mail.Addr,
+	organizationID gid.GID,
+	filter *InvitationFilter,
+) error {
+	q := `
+UPDATE iam_invitations SET accepted_at = NOW()
+WHERE
+    email = @email
+    AND organization_id = @organization_id
+    AND %s
+    AND %s
+`
+
+	q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment())
+
+	args := pgx.StrictNamedArgs{
+		"email":           email,
+		"organization_id": organizationID,
+	}
+	maps.Copy(args, scope.SQLArguments())
+	maps.Copy(args, filter.SQLArguments())
+
+	if _, err := conn.Exec(ctx, q, args); err != nil {
+		return fmt.Errorf("cannot accept invitations: %w", err)
+	}
+
+	return nil
+}
+
+func (i *Invitations) ExpireByEmailAndOrganization(
+	ctx context.Context,
+	conn pg.Conn,
+	scope Scoper,
+	email mail.Addr,
+	organizationID gid.GID,
+	filter *InvitationFilter,
+) error {
+	q := `
+UPDATE iam_invitations SET expires_at = NOW()
+WHERE
+    email = @email
+    AND organization_id = @organization_id
+    AND %s
+    AND %s
+`
+
+	q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment())
+
+	args := pgx.StrictNamedArgs{
+		"email":           email,
+		"organization_id": organizationID,
+	}
+	maps.Copy(args, scope.SQLArguments())
+	maps.Copy(args, filter.SQLArguments())
+
+	if _, err := conn.Exec(ctx, q, args); err != nil {
+		return fmt.Errorf("cannot expire invitations: %w", err)
+	}
+
+	return nil
+}
+
 // AuthorizationAttributes loads the minimal authorization attributes for policy condition evaluation.
 // It is intentionally lightweight and does not populate the Invitation struct.
 func (i *Invitation) AuthorizationAttributes(ctx context.Context, conn pg.Conn) (map[string]string, error) {
@@ -230,7 +296,7 @@ DELETE FROM
     iam_invitations
 WHERE
     %s
-	AND id = @invitation_id
+    AND id = @invitation_id
 `
 
 	query = fmt.Sprintf(query, scope.SQLFragment())
@@ -369,11 +435,11 @@ func (i *Invitations) CountByOrganizationID(
 ) (int, error) {
 	q := `
 SELECT
-	COUNT(*)
+    COUNT(*)
 FROM
-	iam_invitations
+    iam_invitations
 WHERE
-	organization_id = @organization_id AND %s AND %s
+    organization_id = @organization_id AND %s AND %s
 `
 
 	q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment())
@@ -405,12 +471,12 @@ func (i *Invitations) CountByEmail(
 ) (int, error) {
 	q := `
 SELECT
-	COUNT(*)
+    COUNT(*)
 FROM
-	iam_invitations
+    iam_invitations
 WHERE
-	email = @email
+    email = @email
-	AND %s
+    AND %s
 `
 
 	q = fmt.Sprintf(q, filter.SQLFragment())

pkg/coredata/membership.go

@@ -339,78 +339,6 @@ LEFT JOIN
 	return nil
 }
 
-func (m *Membership) LoadByEmailAndOrganization(
-	ctx context.Context,
-	conn pg.Conn,
-	scope Scoper,
-	email mail.Addr,
-	organizationID gid.GID,
-) error {
-	q := `
-WITH mbr AS (
-    SELECT
-        am.id,
-        am.identity_id,
-        am.organization_id,
-        am.role,
-        am.source,
-        am.state,
-        am.created_at,
-        am.updated_at
-    FROM
-        iam_memberships am
-    JOIN
-        identities i ON am.identity_id = i.id
-    WHERE
-        i.email_address = @email
-        AND am.organization_id = @organization_id
-        AND %s
-)
-SELECT
-    mbr.id,
-    mbr.identity_id,
-    mbr.organization_id,
-    mbr.role,
-    mbr.source,
-    mbr.state,
-    COALESCE(mp.full_name, i.full_name, '') as full_name,
-    i.email_address,
-    mbr.created_at,
-    mbr.updated_at
-FROM
-    mbr
-JOIN
-    identities i ON mbr.identity_id = i.id
-LEFT JOIN
-    iam_membership_profiles mp ON mp.membership_id = mbr.id
-`
-
-	q = fmt.Sprintf(q, scope.SQLFragment())
-
-	args := pgx.StrictNamedArgs{
-		"email":           email,
-		"organization_id": organizationID,
-	}
-	maps.Copy(args, scope.SQLArguments())
-
-	rows, err := conn.Query(ctx, q, args)
-	if err != nil {
-		return fmt.Errorf("cannot query membership by email: %w", err)
-	}
-
-	membership, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Membership])
-	if err != nil {
-		if errors.Is(err, pgx.ErrNoRows) {
-			return ErrResourceNotFound
-		}
-
-		return fmt.Errorf("cannot collect membership: %w", err)
-	}
-
-	*m = membership
-	return nil
-}
-
 func (m *Membership) Update(ctx context.Context, conn pg.Conn, scope Scoper) error {
 	query := `
 UPDATE

pkg/iam/account_service.go

@@ -210,8 +210,7 @@ func (s *AccountService) AcceptInvitation(
 		func(tx pg.Conn) error {
 			identity := coredata.Identity{}
 
-			err := identity.LoadByID(ctx, tx, identityID)
+			if err := identity.LoadByID(ctx, tx, identityID); err != nil {
-			if err != nil {
 				if err == coredata.ErrResourceNotFound {
 					return NewIdentityNotFoundError(identityID)
 				}
@@ -219,8 +218,7 @@ func (s *AccountService) AcceptInvitation(
 				return fmt.Errorf("cannot load identity: %w", err)
 			}
 
-			err = invitation.LoadByID(ctx, tx, coredata.NewNoScope(), invitationID)
+			if err := invitation.LoadByID(ctx, tx, coredata.NewNoScope(), invitationID); err != nil {
-			if err != nil {
 				if err == coredata.ErrResourceNotFound {
 					return NewInvitationNotFoundError(invitationID)
 				}
@@ -244,8 +242,13 @@ func (s *AccountService) AcceptInvitation(
 			scope := coredata.NewScope(invitation.OrganizationID.TenantID())
 
 			existingMembership := &coredata.Membership{}
-			err = existingMembership.LoadByIdentityAndOrg(ctx, tx, scope, identityID, invitation.OrganizationID)
+			if err := existingMembership.LoadByIdentityAndOrg(
-			if err != nil && err != coredata.ErrResourceNotFound {
+				ctx,
+				tx,
+				scope,
+				identityID,
+				invitation.OrganizationID,
+			); err != nil && err != coredata.ErrResourceNotFound {
 				return fmt.Errorf("cannot load existing membership: %w", err)
 			}
 
@@ -254,8 +257,7 @@ func (s *AccountService) AcceptInvitation(
 				existingMembership.Role = invitation.Role
 				existingMembership.UpdatedAt = now
 
-				err = existingMembership.Update(ctx, tx, scope)
+				if err := existingMembership.Update(ctx, tx, scope); err != nil {
-				if err != nil {
 					return fmt.Errorf("cannot reactivate membership: %w", err)
 				}
 
@@ -272,8 +274,7 @@ func (s *AccountService) AcceptInvitation(
 					UpdatedAt:      now,
 				}
 
-				err = membership.Insert(ctx, tx, scope)
+				if err := membership.Insert(ctx, tx, scope); err != nil {
-				if err != nil {
 					return fmt.Errorf("cannot create membership: %w", err)
 				}
 
@@ -285,15 +286,13 @@ func (s *AccountService) AcceptInvitation(
 					UpdatedAt:    now,
 				}
 
-				err = profile.Insert(ctx, tx)
+				if err := profile.Insert(ctx, tx); err != nil {
-				if err != nil {
 					return fmt.Errorf("cannot insert profile: %w", err)
 				}
 			}
 
 			invitation.AcceptedAt = &now
-			err = invitation.Update(ctx, tx, scope)
+			if err := invitation.Update(ctx, tx, scope); err != nil {
-			if err != nil {
 				if err == coredata.ErrResourceNotFound {
 					return NewInvitationNotFoundError(invitationID)
 				}
@@ -301,6 +300,20 @@ func (s *AccountService) AcceptInvitation(
 				return fmt.Errorf("cannot update invitation: %w", err)
 			}
 
+			// Accept other pending invitations for email in organization
+			invitations := &coredata.Invitations{}
+			onlyPending := coredata.NewInvitationFilter([]coredata.InvitationStatus{coredata.InvitationStatusPending})
+			if err := invitations.AcceptByEmailAndOrganization(
+				ctx,
+				tx,
+				coredata.NewScopeFromObjectID(invitation.OrganizationID),
+				invitation.Email,
+				invitation.OrganizationID,
+				onlyPending,
+			); err != nil {
+				return fmt.Errorf("cannot accept pending invitations by email: %w", err)
+			}
+
 			return nil
 		},
 	)

pkg/iam/saml/service.go

@@ -325,6 +325,20 @@ func (s *Service) HandleAssertion(
 				if err != nil {
 					return fmt.Errorf("cannot insert membership profile: %w", err)
 				}
+
+				// Accept all pending invitations for email in organization
+				invitations := &coredata.Invitations{}
+				onlyPending := coredata.NewInvitationFilter([]coredata.InvitationStatus{coredata.InvitationStatusPending})
+				if err := invitations.AcceptByEmailAndOrganization(
+					ctx,
+					tx,
+					coredata.NewScopeFromObjectID(config.OrganizationID),
+					email,
+					config.OrganizationID,
+					onlyPending,
+				); err != nil {
+					return fmt.Errorf("cannot accept pending invitations by email: %w", err)
+				}
 			}
 
 			if membership.Source != coredata.MembershipSourceSCIM {

pkg/iam/scim/service.go

@@ -172,6 +172,22 @@ func (s *Service) CreateUser(
 			if err != nil {
 				return fmt.Errorf("cannot insert membership profile: %w", err)
 			}
+
+			// Accept all pending invitations for email in organization
+			invitations := &coredata.Invitations{}
+			onlyPending := coredata.NewInvitationFilter([]coredata.InvitationStatus{coredata.InvitationStatusPending})
+			err := invitations.AcceptByEmailAndOrganization(
+				ctx,
+				tx,
+				coredata.NewScopeFromObjectID(config.OrganizationID),
+				emailAddr,
+				config.OrganizationID,
+				onlyPending,
+			)
+
+			if err != nil {
+				return fmt.Errorf("cannot accept pending invitations by email")
+			}
 		} else if err != nil {
 			return fmt.Errorf("cannot load membership: %w", err)
 		} else {
@@ -335,8 +351,7 @@ func (s *Service) updateUser(
 		ctx,
 		func(tx pg.Conn) error {
 			membership = &coredata.Membership{}
-			err := membership.LoadByID(ctx, tx, scope, membershipID)
+			if err := membership.LoadByID(ctx, tx, scope, membershipID); err != nil {
-			if err != nil {
 				if err == coredata.ErrResourceNotFound {
 					return scimerrors.ScimErrorResourceNotFound(membershipID.String())
 				}
@@ -350,13 +365,46 @@ func (s *Service) updateUser(
 			needsUpdate := false
 
 			if active != nil {
+				identity := &coredata.Identity{}
+				if err := identity.LoadByID(ctx, tx, membership.IdentityID); err != nil {
+					return fmt.Errorf("cannot load identity: %w", err)
+				}
+
 				if *active && membership.State == coredata.MembershipStateInactive {
 					membership.State = coredata.MembershipStateActive
 					membership.Role = coredata.MembershipRoleEmployee
 					needsUpdate = true
+
+					// Accept all pending invitations for email in organization
+					invitations := &coredata.Invitations{}
+					onlyPending := coredata.NewInvitationFilter([]coredata.InvitationStatus{coredata.InvitationStatusPending})
+					if err := invitations.AcceptByEmailAndOrganization(
+						ctx,
+						tx,
+						coredata.NewScopeFromObjectID(config.OrganizationID),
+						identity.EmailAddress,
+						config.OrganizationID,
+						onlyPending,
+					); err != nil {
+						return fmt.Errorf("cannot accept pending invitations by email: %w", err)
+					}
 				} else if !*active && membership.State == coredata.MembershipStateActive {
 					membership.State = coredata.MembershipStateInactive
 					needsUpdate = true
+
+					// Expire all pending invitations for email in organization
+					invitations := &coredata.Invitations{}
+					onlyPending := coredata.NewInvitationFilter([]coredata.InvitationStatus{coredata.InvitationStatusPending})
+					if err := invitations.ExpireByEmailAndOrganization(
+						ctx,
+						tx,
+						coredata.NewScopeFromObjectID(config.OrganizationID),
+						identity.EmailAddress,
+						config.OrganizationID,
+						onlyPending,
+					); err != nil {
+						return fmt.Errorf("cannot expire pending invitations: %w", err)
+					}
 				}
 			}
 
@@ -367,21 +415,18 @@ func (s *Service) updateUser(
 
 			if needsUpdate {
 				membership.UpdatedAt = now
-				err = membership.Update(ctx, tx, scope)
+				if err := membership.Update(ctx, tx, scope); err != nil {
-				if err != nil {
 					return fmt.Errorf("cannot update membership: %w", err)
 				}
 			}
 
 			profile := &coredata.MembershipProfile{}
-			err = profile.LoadByMembershipID(ctx, tx, scope, membershipID)
+			if err := profile.LoadByMembershipID(ctx, tx, scope, membershipID); err == nil {
-			if err == nil {
 				if fullName != "" {
 					profile.FullName = fullName
 					profile.UpdatedAt = now
 
-					err = profile.Update(ctx, tx, scope)
+					if err := profile.Update(ctx, tx, scope); err != nil {
-					if err != nil {
 						return fmt.Errorf("cannot update membership profile: %w", err)
 					}
 				}