Make the last-admin-guard atomic

internal/core/admin.go

@@ -67,7 +67,11 @@ func groupCatalog(catalog []*store.Command) []catalogGroup {
 }
 
 func (s *Handlers) pageAdmin(w http.ResponseWriter, r *http.Request) {
-	u, _ := s.Store.GetUserByID(r.Context(), s.Auth.CurrentUserID(r.Context()))
+	u, err := s.Store.GetUserByID(r.Context(), s.Auth.CurrentUserID(r.Context()))
+	if err != nil {
+		s.ServerError(w, r, "could not load account", err)
+		return
+	}
 	s.Render(w, r, "admin.html", map[string]any{
 		"CanCatalog": u.CapManageCatalog,
 		"CanUsers":   u.CapManageUsers,

internal/store/orgs.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"regexp"
+	"slices"
 	"strings"
 	"time"
 
@@ -408,62 +409,61 @@ func (s *Store) ListOrgAdmins(ctx context.Context, orgID int64) ([]OrgMemberInfo
 	})
 }
 
-// countOrgAdmins returns the number of admins in an org.
-func (s *Store) countOrgAdmins(ctx context.Context, orgID int64) (int, error) {
-	var n int
-	err := s.pool.QueryRow(ctx,
-		`SELECT count(*) FROM org_members WHERE org_id = $1 AND role = 'admin'`, orgID).Scan(&n)
-	if err != nil {
-		return 0, fmt.Errorf("count admins: %w", err)
-	}
-	return n, nil
-}
-
 // ErrLastAdmin is returned when an action would leave an org with no admins.
 var ErrLastAdmin = errors.New("store: cannot remove the last org admin")
 
 // SetOrgMemberRole changes a member's role, refusing to demote the last admin.
+// The guard and the update run in one transaction (see guardLastAdminTx) so two
+// concurrent demotions can't both pass the check and leave the org with no admin.
 func (s *Store) SetOrgMemberRole(ctx context.Context, orgID, userID int64, role string) error {
-	if role != "admin" {
-		if err := s.guardLastAdmin(ctx, orgID, userID); err != nil {
+	return s.inTx(ctx, func(tx pgx.Tx) error {
+		if role != "admin" {
+			if err := guardLastAdminTx(ctx, tx, orgID, userID); err != nil {
+				return err
+			}
+		}
+		if _, err := tx.Exec(ctx,
+			`UPDATE org_members SET role = $3 WHERE org_id = $1 AND user_id = $2`, orgID, userID, role); err != nil {
+			return fmt.Errorf("set role: %w", err)
+		}
+		return nil
+	})
+}
+
+// RemoveOrgMember removes a member, refusing to remove the last admin. Guard and
+// delete share one transaction so the last-admin check can't be raced.
+func (s *Store) RemoveOrgMember(ctx context.Context, orgID, userID int64) error {
+	return s.inTx(ctx, func(tx pgx.Tx) error {
+		if err := guardLastAdminTx(ctx, tx, orgID, userID); err != nil {
 			return err
 		}
-	}
-	_, err := s.pool.Exec(ctx,
-		`UPDATE org_members SET role = $3 WHERE org_id = $1 AND user_id = $2`, orgID, userID, role)
-	if err != nil {
-		return fmt.Errorf("set role: %w", err)
-	}
-	return nil
+		if _, err := tx.Exec(ctx,
+			`DELETE FROM org_members WHERE org_id = $1 AND user_id = $2`, orgID, userID); err != nil {
+			return fmt.Errorf("remove member: %w", err)
+		}
+		return nil
+	})
 }
 
-// RemoveOrgMember removes a member, refusing to remove the last admin.
-func (s *Store) RemoveOrgMember(ctx context.Context, orgID, userID int64) error {
-	if err := s.guardLastAdmin(ctx, orgID, userID); err != nil {
-		return err
-	}
-	_, err := s.pool.Exec(ctx,
-		`DELETE FROM org_members WHERE org_id = $1 AND user_id = $2`, orgID, userID)
+// guardLastAdminTx returns ErrLastAdmin if demoting/removing userID would leave
+// the org with no admins. It must run inside a transaction: it takes a row lock
+// (SELECT … FOR UPDATE) on the org's admin rows, so concurrent demotions/removals
+// serialize — the second one blocks until the first commits and then re-reads the
+// now-smaller admin set, instead of both seeing the pre-change count and racing
+// the org down to zero admins.
+func guardLastAdminTx(ctx context.Context, tx pgx.Tx, orgID, userID int64) error {
+	rows, err := tx.Query(ctx,
+		`SELECT user_id FROM org_members WHERE org_id = $1 AND role = 'admin' FOR UPDATE`, orgID)
 	if err != nil {
-		return fmt.Errorf("remove member: %w", err)
+		return fmt.Errorf("lock admins: %w", err)
 	}
-	return nil
-}
-
-// guardLastAdmin returns ErrLastAdmin if userID is the org's only admin.
-func (s *Store) guardLastAdmin(ctx context.Context, orgID, userID int64) error {
-	role, ok, err := s.OrgRole(ctx, orgID, userID)
+	admins, err := collectRows(rows, scanID)
 	if err != nil {
-		return err
+		return fmt.Errorf("lock admins: %w", err)
 	}
-	if !ok || role != "admin" {
-		return nil // not an admin; removing/demoting is fine
-	}
-	n, err := s.countOrgAdmins(ctx, orgID)
-	if err != nil {
-		return err
-	}
-	if n <= 1 {
+	// Only a problem when userID is currently the org's sole admin; demoting a
+	// non-admin, or an admin with peers, is fine.
+	if slices.Contains(admins, userID) && len(admins) <= 1 {
 		return ErrLastAdmin
 	}
 	return nil

internal/store/orgs_admin_race_test.go

@@ -0,0 +1,101 @@
+package store
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+)
+
+// adminCount returns how many admins the org has (test helper; the store no
+// longer exposes a count method).
+func adminCount(t *testing.T, st *Store, ctx context.Context, orgID int64) int {
+	t.Helper()
+	var n int
+	if err := st.pool.QueryRow(ctx,
+		`SELECT count(*) FROM org_members WHERE org_id = $1 AND role = 'admin'`, orgID).Scan(&n); err != nil {
+		t.Fatalf("count admins: %v", err)
+	}
+	return n
+}
+
+// TestLastAdminRace is the regression for the check-then-act race: with exactly
+// two admins, demoting both concurrently must not leave the org with zero admins.
+// The FOR UPDATE lock in guardLastAdminTx serializes the two demotions, so exactly
+// one succeeds and the other gets ErrLastAdmin. Before the fix the two demotions
+// could both read "2 admins" and both proceed, dropping the org to zero admins.
+func TestLastAdminRace(t *testing.T) {
+	t.Parallel()
+	st, ctx := orgTestStore(t)
+
+	a, err := st.CreateUser(ctx, "admin-a", "")
+	if err != nil {
+		t.Fatalf("create a: %v", err)
+	}
+	org, err := st.CreateOrg(ctx, "Org", a.ID) // a is admin
+	if err != nil {
+		t.Fatalf("create org: %v", err)
+	}
+	b, err := st.CreateUser(ctx, "admin-b", "")
+	if err != nil {
+		t.Fatalf("create b: %v", err)
+	}
+	if err := st.AddOrgMember(ctx, org.ID, b.ID, "admin"); err != nil {
+		t.Fatalf("add admin b: %v", err)
+	}
+
+	// Repeat to give the race window many chances to interleave; the invariant must
+	// hold every round.
+	for round := 0; round < 30; round++ {
+		// Restore both to admin (promotion skips the guard, so it never fails).
+		if err := st.SetOrgMemberRole(ctx, org.ID, a.ID, "admin"); err != nil {
+			t.Fatalf("round %d: repromote a: %v", round, err)
+		}
+		if err := st.SetOrgMemberRole(ctx, org.ID, b.ID, "admin"); err != nil {
+			t.Fatalf("round %d: repromote b: %v", round, err)
+		}
+		if got := adminCount(t, st, ctx, org.ID); got != 2 {
+			t.Fatalf("round %d: setup left %d admins, want 2", round, got)
+		}
+
+		// Demote both admins at once.
+		var wg sync.WaitGroup
+		start := make(chan struct{})
+		errs := make([]error, 2)
+		targets := []int64{a.ID, b.ID}
+		for i, uid := range targets {
+			wg.Add(1)
+			go func(i int, uid int64) {
+				defer wg.Done()
+				<-start
+				errs[i] = st.SetOrgMemberRole(ctx, org.ID, uid, "member")
+			}(i, uid)
+		}
+		close(start)
+		wg.Wait()
+
+		// The org must keep at least one admin.
+		if got := adminCount(t, st, ctx, org.ID); got < 1 {
+			t.Fatalf("round %d: org left with %d admins (race lost the last admin)", round, got)
+		}
+
+		// Exactly one demotion succeeds; the other is refused as the last admin.
+		var ok, refused, other int
+		for _, e := range errs {
+			switch {
+			case e == nil:
+				ok++
+			case errors.Is(e, ErrLastAdmin):
+				refused++
+			default:
+				other++
+			}
+		}
+		if other != 0 {
+			t.Fatalf("round %d: unexpected error(s) from demotion: %v", round, errs)
+		}
+		if ok != 1 || refused != 1 {
+			t.Fatalf("round %d: got %d ok / %d refused, want 1/1 (%v)", round, ok, refused, errs)
+		}
+	}
+}