fix(identity): keep verification_token stable on domain reclaim

internal/identity/store.go

@@ -262,10 +262,17 @@ func (s *Store) EnsureSharedDomain(ctx context.Context, domain string) error {
 }
 
 // ClaimOrCreateDomain implements the atomic create/claim logic from the design doc.
-// Creates if new, overwrites user_id if unverified, returns if verified+same
-// user, errors if verified+different user. Callers are responsible for
-// rejecting the configured shared domain before invoking this — the store
-// has no concept of a reserved domain.
+// Creates if new, returns the existing row when the same user already owns
+// it (verified or not), and errors if a different user owns it. The
+// verification_token and DKIM keypair are minted on first INSERT and remain
+// stable across re-claims — a caller that has already published the TXT
+// record on DNS (or has mail in flight signed with the DKIM key) isn't
+// silently invalidated by a second call. A different user cannot take
+// over an unverified row; that closes a squatting window where the new
+// owner could verify against a TXT record the original owner already
+// published. Callers are responsible for rejecting the configured shared
+// domain before invoking this — the store has no concept of a reserved
+// domain.
 func (s *Store) ClaimOrCreateDomain(ctx context.Context, domain, userID string) (*Domain, error) {
 	domain = normalizeDomain(domain)
 
@@ -287,19 +294,24 @@ func (s *Store) ClaimOrCreateDomain(ctx context.Context, domain, userID string)
 		log.Printf("[identity] dkim keygen failed for %s: %v", domain, kerr)
 	}
 
-	// Atomic upsert: insert new or overwrite unverified. The DKIM
-	// columns are only written on a true INSERT — the ON CONFLICT
-	// branch leaves any existing key in place so re-claiming an
-	// unverified domain doesn't invalidate signatures on mail already
-	// in flight.
+	// Atomic upsert. The conflict branch only fires for a same-user
+	// re-claim of an unverified row, and runs as a no-op SET so
+	// RETURNING surfaces the existing row. DKIM columns and the
+	// verification_token are only written on a true INSERT, so they
+	// stay stable across re-claims — DKIM stability avoids
+	// invalidating signatures on mail in flight, and token stability
+	// means a caller who already published the TXT record on DNS
+	// isn't silently invalidated. A different-user conflict falls
+	// through to the SELECT below and returns "domain not available",
+	// preventing squatting on an unverified row whose TXT record the
+	// original owner may have already published.
 	d := &Domain{}
 	err := s.pool.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)
 		 ON CONFLICT (domain) DO UPDATE
-		 SET user_id = EXCLUDED.user_id,
-		     verification_token = EXCLUDED.verification_token
-		 WHERE domains.verified = false
+		 SET user_id = domains.user_id
+		 WHERE domains.verified = false AND domains.user_id = $2
 		 RETURNING domain, user_id, verified, verification_token, created_at, verified_at, is_primary, last_checked_at, COALESCE(dkim_selector, ''), COALESCE(dkim_public_key, '')`,
 		domain, userID, verificationToken, nullIfEmpty(dkimSelector), nullIfEmpty(dkimPubKey), nullIfEmptyBytes(dkimPrivKey),
 	).Scan(&d.Domain, &d.UserID, &d.Verified, &d.VerificationToken, &d.CreatedAt, &d.VerifiedAt, &d.IsPrimary, &d.LastCheckedAt, &d.DKIMSelector, &d.DKIMPublicKey)
@@ -308,7 +320,10 @@ func (s *Store) ClaimOrCreateDomain(ctx context.Context, domain, userID string)
 		return d, nil
 	}
 
-	// No row returned — domain exists and is verified. Check ownership.
+	// No row returned — the row exists but the conflict UPDATE was
+	// skipped because either it's already verified or a different user
+	// owns it. Re-read to decide between "verified + same user → return
+	// it" and "different user → not available".
 	existing := &Domain{}
 	err = s.pool.QueryRow(ctx,
 		`SELECT domain, user_id, verified, verification_token, created_at, verified_at, is_primary, last_checked_at, COALESCE(dkim_selector, ''), COALESCE(dkim_public_key, '')

internal/identity/store_test.go

@@ -59,6 +59,82 @@ func TestCreateAgentDuplicate(t *testing.T) {
 	}
 }
 
+// TestClaimOrCreateDomain_StableOnReclaim asserts that re-claiming an
+// unverified domain returns the row unchanged: the verification_token
+// and DKIM public key minted on the first call must survive the second.
+// A caller that has already published the TXT record on DNS would
+// otherwise be silently invalidated by a benign second register call
+// (e.g. an agent re-fetching the records to show the user).
+func TestClaimOrCreateDomain_StableOnReclaim(t *testing.T) {
+	pool := testutil.TestDB(t)
+	store := identity.NewStore(pool)
+	ctx := context.Background()
+
+	user, _ := store.CreateOrGetUser(ctx, "owner-stable@example.com", "Owner", "google-stable-token")
+
+	first, err := store.ClaimOrCreateDomain(ctx, "stable.example.com", user.ID)
+	if err != nil {
+		t.Fatalf("first ClaimOrCreateDomain: %v", err)
+	}
+	if first.VerificationToken == "" {
+		t.Fatal("first call returned empty VerificationToken")
+	}
+
+	second, err := store.ClaimOrCreateDomain(ctx, "stable.example.com", user.ID)
+	if err != nil {
+		t.Fatalf("second ClaimOrCreateDomain: %v", err)
+	}
+
+	if second.VerificationToken != first.VerificationToken {
+		t.Errorf("VerificationToken rotated on reclaim: first=%q second=%q", first.VerificationToken, second.VerificationToken)
+	}
+	if second.DKIMPublicKey != first.DKIMPublicKey {
+		t.Errorf("DKIMPublicKey rotated on reclaim: first=%q second=%q", first.DKIMPublicKey, second.DKIMPublicKey)
+	}
+	if !second.CreatedAt.Equal(first.CreatedAt) {
+		t.Errorf("CreatedAt changed on reclaim: first=%v second=%v", first.CreatedAt, second.CreatedAt)
+	}
+}
+
+// TestClaimOrCreateDomain_CrossUserReclaimRejected asserts that a second
+// user cannot take over an unverified domain that another user has
+// already claimed. Combined with the stable verification_token, this
+// closes the squatting window where the takeover user could verify
+// against a TXT record the original owner had already published.
+func TestClaimOrCreateDomain_CrossUserReclaimRejected(t *testing.T) {
+	pool := testutil.TestDB(t)
+	store := identity.NewStore(pool)
+	ctx := context.Background()
+
+	userA, _ := store.CreateOrGetUser(ctx, "owner-a@example.com", "Owner A", "google-a")
+	userB, _ := store.CreateOrGetUser(ctx, "owner-b@example.com", "Owner B", "google-b")
+
+	first, err := store.ClaimOrCreateDomain(ctx, "squat.example.com", userA.ID)
+	if err != nil {
+		t.Fatalf("userA ClaimOrCreateDomain: %v", err)
+	}
+
+	if _, err := store.ClaimOrCreateDomain(ctx, "squat.example.com", userB.ID); err == nil {
+		t.Fatal("userB reclaim should fail when userA already owns the unverified row")
+	}
+
+	// userB cannot read the row either; userA still owns it and the
+	// verification_token is unchanged.
+	if _, err := store.LookupDomain(ctx, "squat.example.com", userB.ID); err == nil {
+		t.Error("userB LookupDomain should not see squat.example.com")
+	}
+	after, err := store.LookupDomain(ctx, "squat.example.com", userA.ID)
+	if err != nil {
+		t.Fatalf("userA LookupDomain: %v", err)
+	}
+	if after.UserID == nil || *after.UserID != userA.ID {
+		t.Errorf("ownership changed: got user_id=%v, want %s", after.UserID, userA.ID)
+	}
+	if after.VerificationToken != first.VerificationToken {
+		t.Errorf("verification_token rotated under cross-user reclaim: first=%q after=%q", first.VerificationToken, after.VerificationToken)
+	}
+}
+
 func TestGetAgentByID(t *testing.T) {
 	pool := testutil.TestDB(t)
 	store := identity.NewStore(pool)