fix(identity): close the check-then-act race on max_domains

internal/apiserver/apiserver.go

@@ -189,7 +189,7 @@ func BuildDeps(p Params) httpapi.Deps {
 
 		ListDomains:                  p.Store.ListDomainsByUser,
 		SendingRampSnapshot:          rampSnapshot,
-		ClaimDomain:                  p.Store.ClaimOrCreateDomain,
+		ClaimDomain:                  p.Store.ClaimOrCreateDomainWithLimit,
 		EnforceDomainCreate:          p.Enforcer.CheckDomainCreate,
 		DeleteDomain:                 deleteDomainFunc(p),
 		LookupDomainTeardownSnapshot: p.Store.LookupDomainTeardownSnapshot,

internal/e2e/domains_create_race_e2e_test.go

@@ -0,0 +1,96 @@
+//go:build integration
+
+package e2e_test
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"net/http"
+	"sync"
+	"testing"
+
+	"github.com/tokencanopy/e2a/internal/limits"
+	"github.com/tokencanopy/e2a/internal/testutil"
+)
+
+// max_domains must be enforced under concurrency: N simultaneous POST
+// /v1/domains for distinct new domains on one account, capped at 1, must
+// leave exactly one domain claimed (#822). Before the fix, EnforceDomainCreate's
+// count read and ClaimDomain's insert were two independent DB calls with
+// nothing serializing them, so every concurrent request could read the same
+// pre-insert count and all pass the cap.
+func TestRegisterDomainConcurrentRequestsRespectMaxDomainsE2E(t *testing.T) {
+	pool := testutil.TestDB(t)
+	ts := testutil.TestServer(t, pool)
+	ctx := context.Background()
+
+	user, err := ts.Store.CreateOrGetUser(ctx, "race-owner@example.com", "Race Owner", "google-race-owner")
+	if err != nil {
+		t.Fatalf("CreateOrGetUser: %v", err)
+	}
+	apiKey, err := ts.Store.CreateAPIKey(ctx, user.ID, "race-key", nil)
+	if err != nil {
+		t.Fatalf("CreateAPIKey: %v", err)
+	}
+	if err := limits.NewStore(pool).Upsert(ctx, user.ID, limits.Limits{
+		PlanCode: "test", MaxAgents: 100000, MaxDomains: 1,
+		MaxMessagesMonth: 100000, MaxStorageBytes: 1 << 40,
+	}); err != nil {
+		t.Fatalf("Upsert limits: %v", err)
+	}
+
+	const n = 8
+	var wg sync.WaitGroup
+	start := make(chan struct{})
+	codes := make([]int, n)
+	for i := 0; i < n; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			domain := fmt.Sprintf("race-domain-%d.example.com", i)
+			body := []byte(`{"domain":"` + domain + `"}`)
+			req, err := http.NewRequest("POST", ts.HTTPServer.URL+"/v1/domains", bytes.NewReader(body))
+			if err != nil {
+				t.Errorf("build request for %s: %v", domain, err)
+				return
+			}
+			req.Header.Set("Authorization", "Bearer "+apiKey.PlaintextKey)
+			req.Header.Set("Content-Type", "application/json")
+			<-start
+			resp, err := http.DefaultClient.Do(req)
+			if err != nil {
+				t.Errorf("POST %s: %v", domain, err)
+				return
+			}
+			codes[i] = resp.StatusCode
+			resp.Body.Close()
+		}(i)
+	}
+	close(start)
+	wg.Wait()
+
+	var created, rejected int
+	for _, code := range codes {
+		switch code {
+		case http.StatusCreated:
+			created++
+		case http.StatusPaymentRequired:
+			rejected++
+		default:
+			t.Errorf("unexpected status code %d", code)
+		}
+	}
+	if created != 1 || rejected != n-1 {
+		t.Fatalf("want 1 created and %d rejected (max_domains=1), got created=%d rejected=%d (codes=%v)",
+			n-1, created, rejected, codes)
+	}
+
+	var domainCount int
+	if err := pool.QueryRow(ctx, `SELECT count(*) FROM domains WHERE user_id = $1`, user.ID).Scan(&domainCount); err != nil {
+		t.Fatalf("count domains: %v", err)
+	}
+	if domainCount != 1 {
+		t.Fatalf("domains row count = %d, want 1 (max_domains cap was bypassed)", domainCount)
+	}
+}

internal/httpapi/domains.go

@@ -503,35 +503,48 @@ func (s *Server) handleRegisterDomain(ctx context.Context, in *registerDomainInp
 	// parent/child claim is a genuinely new row that no exact-match lookup
 	// finds. Anything other than a clean hit — dep unwired, lookup failed,
 	// nil row — falls through to enforcing, so a limit is never skipped by
-	// accident. The check is deliberately outside the claim transaction: it
-	// is the same non-transactional shape the cap already had, so two racing
-	// creates could still both pass, exactly as before this change.
+	// accident.
 	alreadyOwned := false
 	if s.deps.LookupDomain != nil {
 		if existing, lookupErr := s.deps.LookupDomain(ctx, normalized, user.ID); lookupErr == nil && existing != nil {
 			alreadyOwned = true
 		}
 	}
-	if !alreadyOwned && s.deps.EnforceDomainCreate != nil {
-		if err := s.deps.EnforceDomainCreate(ctx, user.ID); err != nil {
-			if env, ok := limitEnvelope(err); ok {
-				return nil, env
+	maxDomains := 0
+	if !alreadyOwned {
+		// A richly-detailed pre-check (plan code, upgrade URL); ClaimDomain
+		// below is the race-proof source of truth for the cap itself.
+		if s.deps.EnforceDomainCreate != nil {
+			if err := s.deps.EnforceDomainCreate(ctx, user.ID); err != nil {
+				if env, ok := limitEnvelope(err); ok {
+					return nil, env
+				}
+				return nil, NewError(http.StatusInternalServerError, "internal_error", "limits check failed")
 			}
-			return nil, NewError(http.StatusInternalServerError, "internal_error", "limits check failed")
+		}
+		if s.deps.GetLimits != nil {
+			lim, err := s.deps.GetLimits(ctx, user.ID)
+			if err != nil {
+				return nil, NewError(http.StatusInternalServerError, "internal_error", "limits check failed")
+			}
+			maxDomains = lim.MaxDomains
 		}
 	}
-	d, err := s.deps.ClaimDomain(ctx, normalized, user.ID)
+	d, err := s.deps.ClaimDomain(ctx, normalized, user.ID, maxDomains)
 	if err != nil {
 		if errors.Is(err, identity.ErrReservedDomain) {
 			return nil, NewError(http.StatusBadRequest, "reserved_domain", "reserved domain")
 		}
 		if errors.Is(err, identity.ErrDomainTaken) {
 			return nil, NewError(http.StatusConflict, "domain_taken", "domain is already claimed by another account")
 		}
-		// Any non-taken failure here is a store/lookup error, not a client
-		// error (ClaimOrCreateDomain returns a domain, ErrDomainTaken, or a
-		// wrapped DB error — never nil, nil), so it is a 500 — the former
-		// 400 "domain_unavailable" misclassified it as the caller's fault.
+		var limErr *identity.DomainLimitExceededError
+		if errors.As(err, &limErr) {
+			return nil, NewError(http.StatusPaymentRequired, "limit_exceeded", limErr.Error()).
+				WithDetails(LimitExceededDetails{Resource: "domains", Limit: int64(limErr.Limit), Current: int64(limErr.Current)})
+		}
+		// Any other failure here is a store/lookup error, not a client error,
+		// so it is a 500 (the former 400 "domain_unavailable" misclassified it).
 		return nil, NewError(http.StatusInternalServerError, "internal_error", "failed to register domain")
 	}
 	view, err := s.domainViewWithRamp(ctx, user.ID, d)

internal/httpapi/domains_limit_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/tokencanopy/e2a/internal/identity"
+	"github.com/tokencanopy/e2a/internal/limits"
 )
 
 // POST /v1/domains is idempotent for a domain the caller already owns: the
@@ -75,3 +76,41 @@ func TestRegisterDomainAtCapEnforcesWhenLookupErrors(t *testing.T) {
 		t.Fatalf("want 402 limit_exceeded when the lookup fails, got %d %v", code, body)
 	}
 }
+
+// ClaimDomain is the race-proof source of truth for max_domains (#822): a
+// DomainLimitExceededError from it must still surface as 402 limit_exceeded,
+// even when EnforceDomainCreate's own pre-check already passed.
+func TestRegisterDomainClaimDomainLimitExceededMapsTo402(t *testing.T) {
+	srv := testServer(t, func(d *Deps) {
+		d.ClaimDomain = func(ctx context.Context, domain, userID string, maxDomains int) (*identity.Domain, error) {
+			return nil, &identity.DomainLimitExceededError{Limit: maxDomains, Current: maxDomains}
+		}
+	})
+	code, body := postJSON(t, srv.URL+"/v1/domains", "good", map[string]any{"domain": "raced.com"})
+	if code != 402 || errCode(body) != "limit_exceeded" {
+		t.Fatalf("want 402 limit_exceeded when ClaimDomain reports the cap, got %d %v", code, body)
+	}
+}
+
+// A GetLimits failure must not fall through with maxDomains left at its zero
+// value: that reads as "unlimited" to ClaimDomain and would silently drop
+// the cap for this request instead of failing safe.
+func TestRegisterDomainAtCapEnforcesWhenGetLimitsErrors(t *testing.T) {
+	claimedWith := -1
+	srv := testServer(t, func(d *Deps) {
+		d.GetLimits = func(ctx context.Context, userID string) (limits.Limits, error) {
+			return limits.Limits{}, errors.New("connection refused")
+		}
+		d.ClaimDomain = func(ctx context.Context, domain, userID string, maxDomains int) (*identity.Domain, error) {
+			claimedWith = maxDomains
+			return &identity.Domain{Domain: domain, Verified: false, VerificationToken: "e2a-verify=new"}, nil
+		}
+	})
+	code, body := postJSON(t, srv.URL+"/v1/domains", "good", map[string]any{"domain": "new-domain.com"})
+	if code != 500 || errCode(body) != "internal_error" {
+		t.Fatalf("want 500 internal_error when GetLimits fails, got %d %v", code, body)
+	}
+	if claimedWith != -1 {
+		t.Fatalf("ClaimDomain must not be called after GetLimits fails, but it was called with maxDomains=%d", claimedWith)
+	}
+}

internal/httpapi/domains_test.go

@@ -125,7 +125,7 @@ func TestRegisterDomainReserved(t *testing.T) {
 
 func TestRegisterDomainReservedMailFromSubtree(t *testing.T) {
 	srv := testServer(t, func(d *Deps) {
-		d.ClaimDomain = func(ctx context.Context, domain, userID string) (*identity.Domain, error) {
+		d.ClaimDomain = func(ctx context.Context, domain, userID string, maxDomains int) (*identity.Domain, error) {
 			return nil, identity.ErrReservedDomain
 		}
 	})

internal/httpapi/httpapi.go

@@ -209,7 +209,9 @@ type Deps struct {
 	// page).
 	ListDomains         func(ctx context.Context, userID string, limit int, afterCreatedAt time.Time, afterDomain string) ([]identity.Domain, error)
 	SendingRampSnapshot func(ctx context.Context, userID, domain string, now time.Time) (sendramp.Snapshot, error)
-	ClaimDomain         func(ctx context.Context, domain, userID string) (*identity.Domain, error)
+	// ClaimDomain atomically enforces maxDomains (<=0 means unlimited) and
+	// claims the domain: the race-proof backstop behind EnforceDomainCreate.
+	ClaimDomain         func(ctx context.Context, domain, userID string, maxDomains int) (*identity.Domain, error)
 	EnforceDomainCreate func(ctx context.Context, userID string) error
 	// DeleteDomain atomically either deletes the live incarnation or resolves
 	// the newest historical receipt when the row is already gone. The exact-

internal/httpapi/operations_test.go

@@ -362,7 +362,7 @@ func testServer(t *testing.T, opts ...func(*Deps)) *httptest.Server {
 		ListDomains: func(ctx context.Context, userID string, limit int, afterCreatedAt time.Time, afterDomain string) ([]identity.Domain, error) {
 			return []identity.Domain{{Domain: "acme.com", Verified: true, VerificationToken: "e2a-verify=tok", IsPrimary: true, AgentCount: 2}}, nil
 		},
-		ClaimDomain: func(ctx context.Context, domain, userID string) (*identity.Domain, error) {
+		ClaimDomain: func(ctx context.Context, domain, userID string, maxDomains int) (*identity.Domain, error) {
 			if domain == "taken.com" {
 				return nil, identity.ErrDomainTaken
 			}

internal/identity/store.go

@@ -789,7 +789,21 @@ func (s *Store) EnsureSharedDomain(ctx context.Context, domain string) error {
 	return nil
 }
 
-// ClaimOrCreateDomain atomically claims a DNS namespace. A row reserves its
+// ClaimOrCreateDomain atomically claims a DNS namespace, with no cap on how
+// many domains userID may hold. See ClaimOrCreateDomainWithLimit for the
+// account-facing, cap-enforcing form used outside this unlimited case.
+func (s *Store) ClaimOrCreateDomain(ctx context.Context, domain, userID string) (*Domain, error) {
+	return s.claimOrCreateDomain(ctx, domain, userID, 0)
+}
+
+// ClaimOrCreateDomainWithLimit is ClaimOrCreateDomain plus an atomic
+// max_domains check in the same advisory-locked transaction as the INSERT,
+// closing the pre-insert-count race (#822). maxDomains <= 0 means unlimited.
+func (s *Store) ClaimOrCreateDomainWithLimit(ctx context.Context, domain, userID string, maxDomains int) (*Domain, error) {
+	return s.claimOrCreateDomain(ctx, domain, userID, maxDomains)
+}
+
+// claimOrCreateDomain atomically claims a DNS namespace. A row reserves its
 // exact name plus every ancestor/descendant against other accounts; the same
 // account may explicitly register a child. The
 // verification_token and DKIM keypair are minted on first INSERT and remain
@@ -800,7 +814,7 @@ func (s *Store) EnsureSharedDomain(ctx context.Context, domain string) error {
 // owner could verify against a TXT record the original owner already
 // published. The managed bounce.<domain> subtree is reserved once an account
 // owns the parent because SES custom MAIL FROM uses that namespace.
-func (s *Store) ClaimOrCreateDomain(ctx context.Context, domain, userID string) (*Domain, error) {
+func (s *Store) claimOrCreateDomain(ctx context.Context, domain, userID string, maxDomains int) (*Domain, error) {
 	domain = normalizeDomain(domain)
 
 	verificationToken := "e2a-verify=" + generateID()
@@ -845,6 +859,15 @@ func (s *Store) ClaimOrCreateDomain(ctx context.Context, domain, userID string)
 		}
 	}
 
+	// A cap-enforcing caller also takes a per-user lock (keyspace 1, distinct
+	// from the domain-name locks above), on the SAME connection as the count
+	// check and insert below, serializing this user's creates against a race.
+	if maxDomains > 0 {
+		if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 1))`, userID); err != nil {
+			return nil, err
+		}
+	}
+
 	var crossAccountConflict bool
 	if err := tx.QueryRow(ctx,
 		`SELECT EXISTS(
@@ -891,6 +914,18 @@ func (s *Store) ClaimOrCreateDomain(ctx context.Context, domain, userID string)
 		 FROM domains WHERE domain = $1`, domain,
 	).Scan(&d.Domain, &d.UserID, &d.Verified, &d.VerificationToken, &d.CreatedAt, &d.VerifiedAt, &d.IsPrimary, &d.LastCheckedAt, &d.DKIMSelector, &d.DKIMPublicKey, &d.SendingStatus, &d.SendingError, &d.SendingDNSRecordsJSON, &d.SendingLastCheckedAt, &d.SendingDkimStatus, &d.SendingMailFromStatus, &d.AgentCount)
 	if errors.Is(err, pgx.ErrNoRows) {
+		// Only a genuine new row is chargeable (a re-claim already returned
+		// above). Read under the per-user lock, so this reflects every insert
+		// already committed by a concurrent request, not a stale count (#822).
+		if maxDomains > 0 {
+			var count int
+			if err := tx.QueryRow(ctx, `SELECT count(*) FROM domains WHERE user_id = $1`, userID).Scan(&count); err != nil {
+				return nil, err
+			}
+			if count >= maxDomains {
+				return nil, &DomainLimitExceededError{Limit: maxDomains, Current: count}
+			}
+		}
 		err = tx.QueryRow(ctx,
 			`INSERT INTO domains (domain, user_id, verified, verification_token, dkim_selector, dkim_public_key, dkim_private_key)
 			 VALUES ($1, $2, false, $3, $4, $5, $6)
@@ -1812,6 +1847,17 @@ var ErrDomainTaken = fmt.Errorf("domain not available: already claimed by anothe
 // subtree). The API maps it to reserved_domain.
 var ErrReservedDomain = fmt.Errorf("domain is reserved for managed infrastructure")
 
+// DomainLimitExceededError is returned by ClaimOrCreateDomainWithLimit when
+// userID is already at maxDomains. Identity's own type, so this package does
+// not need to import limits for the shared LimitExceededError shape.
+type DomainLimitExceededError struct {
+	Limit, Current int
+}
+
+func (e *DomainLimitExceededError) Error() string {
+	return fmt.Sprintf("domain limit exceeded: %d/%d domains", e.Current, e.Limit)
+}
+
 // DeleteDomain deletes a domain only if owned by the user.
 // The handler should check for existing agents first.
 func (s *Store) DeleteDomain(ctx context.Context, domain, userID string) error {