fix(quota): make the claim gate atomic, and stop zeroing the storage cache

Two defects in the §9.3 quota path, both invisible to the hermetic suite:

- Audit #4's TOCTOU was real and documented: QuotaCheck and ClaimServer were
  separate statements, so two concurrent claims by one user for two different
  ownerless servers both read count < max_servers and both won. The gate now
  lives inside ClaimServer, in the SAME transaction as the ownership write,
  under pg_advisory_xact_lock(hashtext(user_id)) — the aggregate read, the
  four-dimension re-check (shared with QuotaCheck via one helper so the two
  cannot drift), and the UPDATE are one serialized decision. The loser gets
  ErrQuotaExceeded, which both claim handlers map to the same 403 the
  sequential path gives; the server row is additionally taken FOR UPDATE so
  same-server races still resolve to exactly one winner.

- The server PATCH path called UpdateServerResources(..., 0) for storage even
  though a resources patch cannot change storage. The cached columns are the
  ONLY input to the quota aggregate, so every resource patch silently dropped
  that server's storage contribution from its owner's cap. The handler now
  reads the current spec and passes storage through.

Red-then-green: the new pgint test drives two real concurrent claims against
max_servers=1 (before: both win; now: exactly one win + one gated 403, and the
DB shows one owned row); the hermetic suite pins the 403 mapping and the
storage-preserving cache write.
This commit is contained in:
Lemon-miaow committed 2026-09-23 03:37:59 +08:00
1 parent d829267f1c
commit bb68fefe04
8 files changed
+274 -50

No files matched your search

+37 -7
View File
@@ -38,8 +38,16 @@ type fakeRepo struct {
owners map[string]string
ownersErr error
claimOK map[string]bool // name -> claim succeeds; absent name -> ErrNotFound
audits []AuditEntry
joins []string
// claimQuotaRefuse simulates ClaimServer's atomic quota gate (audit #4)
// refusing a name whose advisory pre-check already passed.
claimQuotaRefuse map[string]bool
// serverResources / resourceUpdates mirror the cached resource columns:
// ServerResources is what the resize path reads (to preserve storage), and
// UpdateServerResources records the write for assertions.
serverResources map[string]ResourceSpec
resourceUpdates map[string]ResourceSpec
audits []AuditEntry
joins []string
// create-server seeding (spec §15)
seeded map[string]bool // name -> servers row exists
aliases map[string]string // subdomain -> bound server name
@@ -206,8 +214,9 @@ func newFakeRepo() *fakeRepo {
allowlist: map[string]map[string]bool{}, allowUUID: map[string]map[string]bool{},
mine: map[string][]MyServerView{},
owners: map[string]string{},
claimOK: map[string]bool{},
seeded: map[string]bool{}, aliases: map[string]string{},
claimOK: map[string]bool{}, claimQuotaRefuse: map[string]bool{},
serverResources: map[string]ResourceSpec{}, resourceUpdates: map[string]ResourceSpec{},
seeded: map[string]bool{}, aliases: map[string]string{},
linkCodes: map[string]fakeLinkCode{}, links: map[string]string{},
linkAuthSource: map[string]string{},
staff: map[string]*StaffUser{},
@@ -249,10 +258,13 @@ func (f *fakeRepo) QuotaCheck(_ context.Context, userID string, _ string, _ Reso
return f.QuotaAvailable(context.TODO(), userID)
}
func (f *fakeRepo) UpdateServerResources(_ context.Context, _ string, _, _, _ int) error { return nil }
func (f *fakeRepo) UpdateServerResources(_ context.Context, name string, cpu, mem, stor int) error {
f.resourceUpdates[name] = ResourceSpec{CPUMilli: cpu, MemoryMB: mem, StorageMB: stor}
return nil
}
func (f *fakeRepo) ServerResources(_ context.Context, _ string) (ResourceSpec, error) {
return ResourceSpec{}, nil
func (f *fakeRepo) ServerResources(_ context.Context, name string) (ResourceSpec, error) {
return f.serverResources[name], nil
}
func (f *fakeRepo) CreateLinkCode(_ context.Context, code, mcUUID, authSource string, expiresAt time.Time) error {
f.linkCodes[code] = fakeLinkCode{mcUUID: mcUUID, authSource: authSource, expiresAt: expiresAt}
@@ -649,6 +661,9 @@ func (f *fakeRepo) ClaimServer(_ context.Context, n, u string) (bool, error) {
if !present {
return false, ErrNotFound
}
if ok && f.claimQuotaRefuse[n] {
return false, ErrQuotaExceeded // mirrors the atomic gate losing the race
}
return ok, nil
}
func (f *fakeRepo) RecordJoin(_ context.Context, n, uuid string) error {
@@ -1783,6 +1798,21 @@ func TestClaimStateMachine(t *testing.T) {
t.Fatalf("code = %d body %s", w.Code, w.Body.String())
}
})
t.Run("atomic gate refusal -> 403 quota_exceeded", func(t *testing.T) {
// The advisory pre-check passed, but ClaimServer's serialized re-check
// (audit #4) refuses: the caller must see the same 403, not a 500.
repo := newFakeRepo()
repo.linked["u1"] = true
repo.quota["u1"] = true
repo.claimOK["survival"] = true
repo.claimQuotaRefuse["survival"] = true
api := newTestAPI(repo, newFakeCluster())
api.External = staticExternal{p: user}
w := do(api.ExternalHandler(), "POST", "/api/v1/servers/survival/claim", "", nil)
if w.Code != http.StatusForbidden || decodeErr(t, w) != "quota_exceeded" {
t.Fatalf("code = %d body %s, want 403 quota_exceeded", w.Code, w.Body.String())
}
})
t.Run("already claimed -> 409", func(t *testing.T) {
repo := newFakeRepo()
repo.linked["u1"] = true
+7
View File
@@ -15,6 +15,13 @@ var (
ErrNotFound = errors.New("not found")
// ErrConflict means an atomic precondition failed (e.g. claim lost the race).
ErrConflict = errors.New("conflict")
// ErrQuotaExceeded means an ownership write would push the user over a quota
// cap (spec §9.3). ClaimServer — the atomic gate — returns it when a claim
// passes the handler's advisory pre-check but loses the serialized re-check
// (two concurrent claims by one user); handlers map it to a 403
// quota_exceeded, the same answer the pre-check gives, so the CONCURRENT case
// and the SEQUENTIAL case are indistinguishable to the caller.
ErrQuotaExceeded = errors.New("server quota exhausted")
// ErrLinkCodeInvalid means an account-link code is unknown or expired (spec
// §10). It is a client error (the verify endpoint exists; the code is bad), so
// handlers map it to 400, not 404.
+6
View File
@@ -240,6 +240,12 @@ func (a *API) handleInternalClaim(w http.ResponseWriter, r *http.Request) {
// 412 above; a lost race (0 rows) is 409.
claimed, err := a.Repo.ClaimServer(r.Context(), name, userID)
if err != nil {
// Same atomic quota gate as the external face (audit #4): the concurrent
// loser gets the sequential 403, never an over-provisioned tenant.
if errors.Is(err, ErrQuotaExceeded) {
writeError(w, r, newError(http.StatusForbidden, "quota_exceeded", "server quota exhausted"))
return
}
a.writeLookupError(w, r, err)
return
}
+20
View File
@@ -127,6 +127,26 @@ func TestPatchServerMemoryOverride(t *testing.T) {
}
}
// TestPatchServerPreservesStorageCache pins the storage dimension of the quota
// aggregate: a resources patch cannot change storage, so the cached storage
// must survive it — passing 0 would silently zero the owner's aggregate (the
// cached columns are QuotaCheck's only input) from that patch onward.
func TestPatchServerPreservesStorageCache(t *testing.T) {
api, repo, _, _ := newPatchAPI()
repo.byName["survival"].OwnerID = "u1"
repo.quota["u1"] = true
repo.serverResources["survival"] = ResourceSpec{StorageMB: 10240}
w := patchSurvival(api, `{"memory":"2Gi","resources":{"memory":"4Gi","cpu":"2"}}`)
if w.Code != http.StatusOK {
t.Fatalf("code = %d, want 200 (%s)", w.Code, w.Body.String())
}
got := repo.resourceUpdates["survival"]
if got.CPUMilli != 2000 || got.MemoryMB != 4096 || got.StorageMB != 10240 {
t.Fatalf("resource cache = %+v, want cpu 2000 / mem 4096 / storage preserved 10240", got)
}
}
// TestPatchServerRejections is the validation matrix: each malformed request is
// rejected with the right status and stable error code, and (critically) NOTHING
// reaches the cluster on a rejection — the analog of create's "no CRD written".
+20 -1
View File
@@ -142,6 +142,13 @@ func (a *API) handleClaim(w http.ResponseWriter, r *http.Request) {
// ③ atomic claim
claimed, err := a.Repo.ClaimServer(r.Context(), name, p.UserID)
if err != nil {
// The atomic gate re-checks quota under the per-user lock (audit #4): a
// concurrent claim that spent the last slot surfaces here, with the same
// 403 the pre-check gives sequentially.
if errors.Is(err, ErrQuotaExceeded) {
writeError(w, r, newError(http.StatusForbidden, "quota_exceeded", "server quota exhausted"))
return
}
a.writeLookupError(w, r, err)
return
}
@@ -759,7 +766,19 @@ func (a *API) handlePatchServer(w http.ResponseWriter, r *http.Request) {
a.writeLookupError(w, r, err)
return
}
_ = a.Repo.UpdateServerResources(r.Context(), name, newCPU, newMemMB, 0)
// A resource patch cannot change storage, so its cached contribution must
// be preserved: passing 0 would silently zero the storage dimension of the
// owner's four-cap aggregate (the cached columns are its only input).
storMB := 0
if rec != nil {
cur, err := a.Repo.ServerResources(r.Context(), name)
if err != nil {
writeError(w, r, err)
return
}
storMB = cur.StorageMB
}
_ = a.Repo.UpdateServerResources(r.Context(), name, newCPU, newMemMB, storMB)
} else {
if err := a.Cluster.PatchServerSpec(r.Context(), name, patch); err != nil {
a.writeLookupError(w, r, err)
+97 -41
View File
@@ -293,22 +293,10 @@ func (p *PGRepo) CompleteOwnerSetup(ctx context.Context, newUserID, code string,
// QuotaAvailable treats a missing quota row or a NULL max_servers as unlimited;
// otherwise it compares the live owned-server count against the cap (spec §9.3).
//
// KNOWN-LIMITATION (audit #4, quota TOCTOU): this check and ClaimServer are two
// separate statements, not one transaction, so the count read here is not serialized
// against a concurrent claim's UPDATE. Two claims by the same user for two DIFFERENT
// ownerless servers can both read count < max_servers (under READ COMMITTED neither
// sees the other's uncommitted UPDATE) and both succeed, leaving the user one server
// over quota. Severity is low: it over-provisions the quota by a small margin under a
// deliberate concurrent burst — it is NOT an authorization, ownership, or isolation
// break (each server is still claimed atomically via UPDATE ... WHERE owner_id IS
// NULL, so two users never share one server). Closing it needs Postgres transaction
// semantics: wrap the count and a conditional UPDATE (gated on count < max_servers) in
// one tx under pg_advisory_xact_lock(hashtext(user_id)) — or SERIALIZABLE with a retry
// loop — folding the gate out of the two handlers (handleClaim and the internal UUID
// claim) into a single repo method. That is INTEGRATION-dependent: it is verifiable
// only against a real Postgres, not the hermetic fakeRepo suite, so it is documented
// here rather than patched blind.
// It is the single-dimension convenience read; handlers use the four-dimension
// QuotaCheck. The former audit-#4 TOCTOU (check and claim in separate statements)
// is closed inside ClaimServer, which re-runs the four-dimension gate under a
// per-user advisory lock in the SAME transaction as the ownership write.
func (p *PGRepo) QuotaAvailable(ctx context.Context, userID string) (bool, error) {
var maxServers sql.NullInt64
switch err := p.db.QueryRowContext(ctx,
@@ -334,8 +322,11 @@ func (p *PGRepo) QuotaAvailable(ctx context.Context, userID string) (bool, error
// cached resources should be excluded ("" for a fresh claim where the row
// doesn't exist yet). Four dimensions are checked: server count, CPU millicores,
// memory MB, and storage MB. A NULL or missing quota row/column means unlimited
// for that dimension. Like QuotaAvailable, the count check and the write are not
// serialized — see the QuotaAvailable TOCTOU docstring.
// for that dimension. As a standalone read it is advisory — it backs the
// handler's fast-path 403 — while the AUTHORITATIVE gate for claims is the one
// ClaimServer re-runs atomically; the resize path (server PATCH) keeps this
// advisory shape because its write goes through the Kubernetes API, not this
// transaction.
func (p *PGRepo) QuotaCheck(ctx context.Context, userID string, excludeName string, incoming ResourceSpec) (bool, error) {
var maxServers, maxCPU, maxMem, maxStor sql.NullInt64
switch err := p.db.QueryRowContext(ctx,
@@ -348,8 +339,7 @@ func (p *PGRepo) QuotaCheck(ctx context.Context, userID string, excludeName stri
return false, err
}
var count int64
var cpuSum, memSum, storSum sql.NullInt64
var count, cpuSum, memSum, storSum int64
switch err := p.db.QueryRowContext(ctx,
`SELECT COUNT(*), COALESCE(SUM(cached_cpu_milli), 0), COALESCE(SUM(cached_memory_mb), 0), COALESCE(SUM(cached_storage_mb), 0)
FROM servers WHERE owner_id = $1 AND deleted_at IS NULL AND name != $2`,
@@ -358,19 +348,7 @@ func (p *PGRepo) QuotaCheck(ctx context.Context, userID string, excludeName stri
return false, err
}
if maxServers.Valid && count >= maxServers.Int64 {
return false, nil
}
if maxCPU.Valid && cpuSum.Int64+int64(incoming.CPUMilli) > maxCPU.Int64 {
return false, nil
}
if maxMem.Valid && memSum.Int64+int64(incoming.MemoryMB) > maxMem.Int64 {
return false, nil
}
if maxStor.Valid && storSum.Int64+int64(incoming.StorageMB) > maxStor.Int64*1024 {
return false, nil
}
return true, nil
return quotaAllows(maxServers, maxCPU, maxMem, maxStor, count, cpuSum, memSum, storSum, incoming), nil
}
// UpdateServerResources updates the resource cache for a server after a spec
@@ -399,17 +377,67 @@ func (p *PGRepo) ServerResources(ctx context.Context, name string) (ResourceSpec
// ClaimServer performs the atomic ownership transfer (spec §9.3). A missing
// server is ErrNotFound; an existing-but-owned server yields claimed=false so the
// handler can answer 409.
// handler can answer 409; a claim that would push the user over any of the four
// quota caps yields ErrQuotaExceeded (the handler's pre-check is a fast path,
// this gate is the authoritative one). The whole decision — quota read,
// per-owner aggregate, and the ownership UPDATE — runs in ONE transaction under
// pg_advisory_xact_lock(hashtext(user_id)), so two concurrent claims by the same
// user for two DIFFERENT ownerless servers serialize instead of both passing the
// gate (audit #4); the row is additionally taken FOR UPDATE so concurrent claims
// of the SAME server still resolve to exactly one winner.
func (p *PGRepo) ClaimServer(ctx context.Context, name, userID string) (bool, error) {
var exists bool
if err := p.db.QueryRowContext(ctx,
`SELECT EXISTS(SELECT 1 FROM servers WHERE name = $1 AND deleted_at IS NULL)`, name).Scan(&exists); err != nil {
tx, err := p.db.BeginTx(ctx, nil)
if err != nil {
return false, err
}
if !exists {
return false, ErrNotFound
defer tx.Rollback() //nolint:errcheck // no-op after commit
// Serialize this user's claim lane: the aggregate read below and the
// ownership write must observe one consistent quota state. A hashtext
// collision across users merely serializes unrelated claims — never waives a
// cap.
if _, err := tx.ExecContext(ctx, `SELECT pg_advisory_xact_lock(hashtext($1))`, userID); err != nil {
return false, err
}
res, err := p.db.ExecContext(ctx,
var owned sql.NullString
var cpu, mem, stor int
switch err := tx.QueryRowContext(ctx,
`SELECT owner_id, cached_cpu_milli, cached_memory_mb, cached_storage_mb
FROM servers WHERE name = $1 AND deleted_at IS NULL FOR UPDATE`,
name).Scan(&owned, &cpu, &mem, &stor); {
case errors.Is(err, sql.ErrNoRows):
return false, ErrNotFound
case err != nil:
return false, err
}
if owned.Valid {
return false, nil // already claimed → 409 at the handler
}
// The four-dimension gate, re-run inside the transaction. A missing quota
// row leaves every NullInt64 invalid → quotaAllows treats each dimension as
// unlimited, matching QuotaCheck.
var maxServers, maxCPU, maxMem, maxStor sql.NullInt64
if err := tx.QueryRowContext(ctx,
`SELECT max_servers, max_cpu_milli, max_memory_mb, max_storage_gb
FROM quotas WHERE user_id = $1`, userID).Scan(
&maxServers, &maxCPU, &maxMem, &maxStor); err != nil && !errors.Is(err, sql.ErrNoRows) {
return false, err
}
var count, cpuSum, memSum, storSum int64
if err := tx.QueryRowContext(ctx,
`SELECT COUNT(*), COALESCE(SUM(cached_cpu_milli), 0), COALESCE(SUM(cached_memory_mb), 0), COALESCE(SUM(cached_storage_mb), 0)
FROM servers WHERE owner_id = $1 AND deleted_at IS NULL AND name != $2`,
userID, name).Scan(&count, &cpuSum, &memSum, &storSum); err != nil {
return false, err
}
if !quotaAllows(maxServers, maxCPU, maxMem, maxStor, count, cpuSum, memSum, storSum,
ResourceSpec{CPUMilli: cpu, MemoryMB: mem, StorageMB: stor}) {
return false, ErrQuotaExceeded
}
res, err := tx.ExecContext(ctx,
`UPDATE servers SET owner_id = $2, claimed_at = now() WHERE name = $1 AND owner_id IS NULL AND deleted_at IS NULL`,
name, userID)
if err != nil {
@@ -419,7 +447,35 @@ func (p *PGRepo) ClaimServer(ctx context.Context, name, userID string) (bool, er
if err != nil {
return false, err
}
return n == 1, nil
if n != 1 {
return false, nil
}
if err := tx.Commit(); err != nil {
return false, err
}
return true, nil
}
// quotaAllows applies the four spec §9.3 caps to one per-owner aggregate plus
// the incoming spec. Shared by QuotaCheck (the advisory pre-check) and
// ClaimServer (the atomic gate) so the two can never drift. An invalid (NULL or
// missing) cap means unlimited for that dimension; storage is compared in MB
// against max_storage_gb × 1024.
func quotaAllows(maxServers, maxCPU, maxMem, maxStor sql.NullInt64,
count, cpuSum, memSum, storSum int64, incoming ResourceSpec) bool {
if maxServers.Valid && count >= maxServers.Int64 {
return false
}
if maxCPU.Valid && cpuSum+int64(incoming.CPUMilli) > maxCPU.Int64 {
return false
}
if maxMem.Valid && memSum+int64(incoming.MemoryMB) > maxMem.Int64 {
return false
}
if maxStor.Valid && storSum+int64(incoming.StorageMB) > maxStor.Int64*1024 {
return false
}
return true
}
func (p *PGRepo) UserInAllowlist(ctx context.Context, name, userID string) (bool, error) {
+5 -1
View File
@@ -231,7 +231,11 @@ type Repo interface {
QuotaCheck(ctx context.Context, userID string, excludeName string, incoming ResourceSpec) (bool, error)
// ClaimServer atomically sets owner_id where it is currently NULL and returns
// whether a row changed. false means the server was already claimed (spec §9.3:
// 0 rows → 409).
// 0 rows → 409). The claim runs in one transaction that re-checks the four quota
// dimensions under pg_advisory_xact_lock(hashtext(user_id)), so it is the
// authoritative gate: a claim that would exceed a cap → ErrQuotaExceeded (403),
// and two concurrent claims by one user cannot both pass (audit #4).
// QuotaCheck remains the advisory pre-check for the handler's fast-path 403.
ClaimServer(ctx context.Context, name, userID string) (bool, error)
// UserInAllowlist reports whether the user's linked UUID is on the server
// allowlist (spec §9.4).