From 0a4059e5af1af0d077840b0f63a05a3e8a1feadb Mon Sep 17 00:00:00 2001 From: Jonathon Leight Date: Sun, 5 Jul 2026 18:51:09 -0400 Subject: [PATCH] Make the last-admin-guard atomic --- internal/core/admin.go | 6 +- internal/store/orgs.go | 86 ++++++++++----------- internal/store/orgs_admin_race_test.go | 101 +++++++++++++++++++++++++ 3 files changed, 149 insertions(+), 44 deletions(-) create mode 100644 internal/store/orgs_admin_race_test.go diff --git a/internal/core/admin.go b/internal/core/admin.go index 7bdfbb1..30cf506 100644 --- a/internal/core/admin.go +++ b/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, diff --git a/internal/store/orgs.go b/internal/store/orgs.go index 0e64cf3..cacd549 100644 --- a/internal/store/orgs.go +++ b/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 diff --git a/internal/store/orgs_admin_race_test.go b/internal/store/orgs_admin_race_test.go new file mode 100644 index 0000000..1dbd4af --- /dev/null +++ b/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) + } + } +}