diff --git a/db/queries/model_router_api_keys.sql b/db/queries/model_router_api_keys.sql index 5d3821aa6..2db352b0d 100644 --- a/db/queries/model_router_api_keys.sql +++ b/db/queries/model_router_api_keys.sql @@ -59,8 +59,9 @@ FROM router.model_router_api_keys WHERE id = @id::uuid AND deleted_at IS NULL; --- Soft-deletes a router API key. Cross-tenant safe via installation_id predicate. --- name: SoftDeleteModelRouterAPIKey :exec +-- Soft-deletes a router API key. Cross-tenant safe via installation_id +-- predicate. :execrows lets callers tell a 0-row no-op from a real transition. +-- name: SoftDeleteModelRouterAPIKey :execrows UPDATE router.model_router_api_keys SET deleted_at = NOW() WHERE id = @id::uuid diff --git a/internal/api/admin/keys_test.go b/internal/api/admin/keys_test.go index 20f5e218b..fb2b43274 100644 --- a/internal/api/admin/keys_test.go +++ b/internal/api/admin/keys_test.go @@ -65,17 +65,17 @@ func (f *fakeAPIKeyRepository) ListForInstallation(_ context.Context, installati func (f *fakeAPIKeyRepository) MarkUsed(context.Context, string) error { return nil } -func (f *fakeAPIKeyRepository) SoftDelete(_ context.Context, installationID, id string) error { +func (f *fakeAPIKeyRepository) SoftDelete(_ context.Context, installationID, id string) (int64, error) { f.mu.Lock() defer f.mu.Unlock() for _, k := range f.keys { if k.ID == id && k.InstallationID == installationID && k.DeletedAt == nil { now := time.Now() k.DeletedAt = &now - return nil + return 1, nil } } - return nil + return 0, nil } func (f *fakeAPIKeyRepository) softDeletedSnapshot() []string { diff --git a/internal/auth/api_key.go b/internal/auth/api_key.go index e5e5e64e6..0da53f667 100644 --- a/internal/auth/api_key.go +++ b/internal/auth/api_key.go @@ -34,5 +34,6 @@ type APIKeyRepository interface { GetActiveByHashWithInstallation(ctx context.Context, keyHash string) (*APIKey, *Installation, error) ListForInstallation(ctx context.Context, installationID string) ([]*APIKey, error) MarkUsed(ctx context.Context, id string) error - SoftDelete(ctx context.Context, installationID, id string) error + // SoftDelete soft-deletes the key and returns the rows-affected count; 0 means the key was already gone. + SoftDelete(ctx context.Context, installationID, id string) (int64, error) } diff --git a/internal/auth/rotate_race_test.go b/internal/auth/rotate_race_test.go new file mode 100644 index 000000000..f2cde4b7c --- /dev/null +++ b/internal/auth/rotate_race_test.go @@ -0,0 +1,285 @@ +package auth_test + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "workweave/router/internal/auth" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// statefulAPIKeyRepo is an in-memory APIKeyRepository whose SoftDelete +// counts rows affected (0 when already gone), matching live Postgres semantics. +type statefulAPIKeyRepo struct { + mu sync.Mutex + keys []*auth.APIKey + nextID int +} + +func newStatefulAPIKeyRepo() *statefulAPIKeyRepo { + return &statefulAPIKeyRepo{} +} + +func (r *statefulAPIKeyRepo) Create(_ context.Context, params auth.CreateAPIKeyParams) (*auth.APIKey, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.nextID++ + key := &auth.APIKey{ + ID: "key-" + itoa(r.nextID), + InstallationID: params.InstallationID, + ExternalID: params.ExternalID, + Name: params.Name, + KeyPrefix: params.KeyPrefix, + KeyHash: params.KeyHash, + KeySuffix: params.KeySuffix, + CreatedBy: params.CreatedBy, + CreatedAt: time.Now(), + } + r.keys = append(r.keys, key) + return key, nil +} + +func (r *statefulAPIKeyRepo) GetActiveByHashWithInstallation(context.Context, string) (*auth.APIKey, *auth.Installation, error) { + return nil, nil, assert.AnError +} + +func (r *statefulAPIKeyRepo) ListForInstallation(_ context.Context, installationID string) ([]*auth.APIKey, error) { + r.mu.Lock() + defer r.mu.Unlock() + out := make([]*auth.APIKey, 0, len(r.keys)) + for _, k := range r.keys { + if k.InstallationID == installationID && k.DeletedAt == nil { + out = append(out, k) + } + } + return out, nil +} + +func (r *statefulAPIKeyRepo) MarkUsed(context.Context, string) error { return nil } + +func (r *statefulAPIKeyRepo) SoftDelete(_ context.Context, installationID, id string) (int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + for _, k := range r.keys { + if k.ID == id && k.InstallationID == installationID && k.DeletedAt == nil { + now := time.Now() + k.DeletedAt = &now + return 1, nil + } + } + return 0, nil +} + +func (r *statefulAPIKeyRepo) activeCount(installationID string) int { + r.mu.Lock() + defer r.mu.Unlock() + n := 0 + for _, k := range r.keys { + if k.InstallationID == installationID && k.DeletedAt == nil { + n++ + } + } + return n +} + +func (r *statefulAPIKeyRepo) createCount() int { + r.mu.Lock() + defer r.mu.Unlock() + return len(r.keys) +} + +func itoa(n int) string { + if n == 0 { + return "0" + } + var b [20]byte + i := len(b) + for n > 0 { + i-- + b[i] = byte('0' + n%10) + n /= 10 + } + return string(b[i:]) +} + +// listHoldRepo parks ListForInstallation after the read so SoftDelete cannot +// run until all waiters arrive — same TOCTOU window as the live #817 repro. +type listHoldRepo struct { + inner *statefulAPIKeyRepo + n int + mu sync.Mutex + arrived int + release chan struct{} +} + +func newListHoldRepo(inner *statefulAPIKeyRepo, n int) *listHoldRepo { + return &listHoldRepo{inner: inner, n: n, release: make(chan struct{})} +} + +func (r *listHoldRepo) Create(ctx context.Context, params auth.CreateAPIKeyParams) (*auth.APIKey, error) { + return r.inner.Create(ctx, params) +} +func (r *listHoldRepo) GetActiveByHashWithInstallation(ctx context.Context, keyHash string) (*auth.APIKey, *auth.Installation, error) { + return r.inner.GetActiveByHashWithInstallation(ctx, keyHash) +} +func (r *listHoldRepo) ListForInstallation(ctx context.Context, installationID string) ([]*auth.APIKey, error) { + keys, err := r.inner.ListForInstallation(ctx, installationID) + if err != nil { + return nil, err + } + r.mu.Lock() + r.arrived++ + release := r.release + if r.arrived == r.n { + close(r.release) + r.arrived = 0 + r.release = make(chan struct{}) + } + r.mu.Unlock() + <-release + return keys, nil +} +func (r *listHoldRepo) MarkUsed(ctx context.Context, id string) error { + return r.inner.MarkUsed(ctx, id) +} +func (r *listHoldRepo) SoftDelete(ctx context.Context, installationID, id string) (int64, error) { + return r.inner.SoftDelete(ctx, installationID, id) +} + +// signalListRepo parks a single List until release closes (Rotate-vs-Delete). +type signalListRepo struct { + inner *statefulAPIKeyRepo + listed chan struct{} + release chan struct{} + once sync.Once +} + +func (r *signalListRepo) Create(ctx context.Context, params auth.CreateAPIKeyParams) (*auth.APIKey, error) { + return r.inner.Create(ctx, params) +} +func (r *signalListRepo) GetActiveByHashWithInstallation(ctx context.Context, keyHash string) (*auth.APIKey, *auth.Installation, error) { + return r.inner.GetActiveByHashWithInstallation(ctx, keyHash) +} +func (r *signalListRepo) ListForInstallation(ctx context.Context, installationID string) ([]*auth.APIKey, error) { + keys, err := r.inner.ListForInstallation(ctx, installationID) + if err != nil { + return nil, err + } + r.once.Do(func() { close(r.listed) }) + <-r.release + return keys, nil +} +func (r *signalListRepo) MarkUsed(ctx context.Context, id string) error { + return r.inner.MarkUsed(ctx, id) +} +func (r *signalListRepo) SoftDelete(ctx context.Context, installationID, id string) (int64, error) { + return r.inner.SoftDelete(ctx, installationID, id) +} + +const raceInstallID = "inst-rotate-race" + +// TestRotateAPIKey_ConcurrentRace_OnlyOneSucceeds verifies the #817 fix: +// the losing racer must return ErrAPIKeyNotFound, not mint an orphan key. +func TestRotateAPIKey_ConcurrentRace_OnlyOneSucceeds(t *testing.T) { + inner := newStatefulAPIKeyRepo() + held := newListHoldRepo(inner, 2) + svc := auth.NewService(nil, held, nil, nil, auth.NoOpAPIKeyCache{}, nil, time.Now) + + name := "race-key" + issued, _, err := svc.IssueAPIKey(context.Background(), raceInstallID, &name, nil) + require.NoError(t, err) + require.Equal(t, 1, inner.activeCount(raceInstallID)) + + var ( + ready sync.WaitGroup + done sync.WaitGroup + start = make(chan struct{}) + okCount atomic.Int64 + errNF atomic.Int64 + ) + ready.Add(2) + done.Add(2) + for i := 0; i < 2; i++ { + go func() { + defer done.Done() + ready.Done() + <-start + _, _, err := svc.RotateAPIKey(context.Background(), raceInstallID, issued.ID, nil) + if err == nil { + okCount.Add(1) + return + } + if errors.Is(err, auth.ErrAPIKeyNotFound) { + errNF.Add(1) + return + } + t.Errorf("unexpected RotateAPIKey error: %v", err) + }() + } + ready.Wait() + close(start) + done.Wait() + + assert.Equal(t, int64(1), okCount.Load(), "exactly one RotateAPIKey must succeed") + assert.Equal(t, int64(1), errNF.Load(), "the loser must return ErrAPIKeyNotFound") + assert.Equal(t, 1, inner.activeCount(raceInstallID), + "exactly one active successor; loser must not mint an orphan key") + assert.Equal(t, 2, inner.createCount(), + "seed + one successor only (not seed + two successors)") +} + +// TestRotateAPIKey_LosingToDelete_DoesNotIssue verifies the #817 Delete-vs-Rotate +// fix: Rotate must return ErrAPIKeyNotFound and not mint a zombie successor. +func TestRotateAPIKey_LosingToDelete_DoesNotIssue(t *testing.T) { + inner := newStatefulAPIKeyRepo() + listed := make(chan struct{}) + release := make(chan struct{}) + wrapped := &signalListRepo{inner: inner, listed: listed, release: release} + svc := auth.NewService(nil, wrapped, nil, nil, auth.NoOpAPIKeyCache{}, nil, time.Now) + + name := "mixed-key" + issued, _, err := svc.IssueAPIKey(context.Background(), raceInstallID, &name, nil) + require.NoError(t, err) + seedCreates := inner.createCount() + + var ( + ready sync.WaitGroup + done sync.WaitGroup + start = make(chan struct{}) + rotateErr error + deleteErr error + ) + ready.Add(2) + done.Add(2) + go func() { + defer done.Done() + ready.Done() + <-start + _, _, rotateErr = svc.RotateAPIKey(context.Background(), raceInstallID, issued.ID, nil) + }() + go func() { + defer done.Done() + ready.Done() + <-start + <-listed + deleteErr = svc.DeleteAPIKey(context.Background(), raceInstallID, issued.ID) + close(release) + }() + ready.Wait() + close(start) + done.Wait() + + require.NoError(t, deleteErr, "DeleteAPIKey must succeed") + assert.ErrorIs(t, rotateErr, auth.ErrAPIKeyNotFound, + "RotateAPIKey must not mint after Delete already soft-deleted the key") + assert.Equal(t, 0, inner.activeCount(raceInstallID), + "no zombie successor after Delete wins") + assert.Equal(t, seedCreates, inner.createCount(), + "Rotate must not IssueAPIKey when SoftDelete matched 0 rows") +} diff --git a/internal/auth/service.go b/internal/auth/service.go index 8ee3d9d1d..0ae9f8059 100644 --- a/internal/auth/service.go +++ b/internal/auth/service.go @@ -131,9 +131,9 @@ func (s *Service) ListAPIKeys(ctx context.Context, installationID string) ([]*AP // RotateAPIKey soft-deletes the named key and issues a replacement under the // same installation, carrying forward its name. Returns ErrAPIKeyNotFound if -// keyID isn't an active key owned by installationID, so a foreign key UUID -// can't be rotated. Not transactional: the brief "no active key" gap is fine -// since rotation's whole point is invalidating the old token anyway. +// keyID isn't an active key owned by installationID, or if SoftDelete matches +// 0 rows: a concurrent rotate or delete already transitioned the key, so +// minting a successor would leave an untracked credential. func (s *Service) RotateAPIKey(ctx context.Context, installationID, keyID string, createdBy *string) (*APIKey, string, error) { existing, err := s.apiKeys.ListForInstallation(ctx, installationID) if err != nil { @@ -149,9 +149,13 @@ func (s *Service) RotateAPIKey(ctx context.Context, installationID, keyID string if target == nil { return nil, "", ErrAPIKeyNotFound } - if err := s.apiKeys.SoftDelete(ctx, installationID, target.ID); err != nil { + n, err := s.apiKeys.SoftDelete(ctx, installationID, target.ID) + if err != nil { return nil, "", err } + if n == 0 { + return nil, "", ErrAPIKeyNotFound + } key, raw, err := s.IssueAPIKey(ctx, installationID, target.Name, createdBy) if err != nil { return nil, "", err @@ -162,9 +166,10 @@ func (s *Service) RotateAPIKey(ctx context.Context, installationID, keyID string // DeleteAPIKey soft-deletes an API key and invalidates the installation's // cache entry on this replica and all peers, so the key doesn't stay usable -// for the remainder of the positive cache TTL (5 min). +// for the remainder of the positive cache TTL (5 min). The rows-affected +// count is intentionally ignored: a 0-row no-op stays idempotent success. func (s *Service) DeleteAPIKey(ctx context.Context, installationID, id string) error { - if err := s.apiKeys.SoftDelete(ctx, installationID, id); err != nil { + if _, err := s.apiKeys.SoftDelete(ctx, installationID, id); err != nil { return err } s.invalidateInstallation(installationID) diff --git a/internal/auth/service_test.go b/internal/auth/service_test.go index 53e23faea..3e84de890 100644 --- a/internal/auth/service_test.go +++ b/internal/auth/service_test.go @@ -71,8 +71,11 @@ func (f *fakeAPIKeyRepository) MarkUsed(ctx context.Context, id string) error { return nil } -func (f *fakeAPIKeyRepository) SoftDelete(ctx context.Context, installationID, id string) error { - return f.override +func (f *fakeAPIKeyRepository) SoftDelete(ctx context.Context, installationID, id string) (int64, error) { + if f.override != nil { + return 0, f.override + } + return 1, nil } type fakeExternalAPIKeyRepo struct { diff --git a/internal/postgres/repository.go b/internal/postgres/repository.go index fbaaa7314..0101e946f 100644 --- a/internal/postgres/repository.go +++ b/internal/postgres/repository.go @@ -257,14 +257,14 @@ func (r *apiKeyRepo) MarkUsed(ctx context.Context, id string) error { return q.MarkModelRouterAPIKeyUsed(ctx, parsed) } -func (r *apiKeyRepo) SoftDelete(ctx context.Context, installationID, id string) error { +func (r *apiKeyRepo) SoftDelete(ctx context.Context, installationID, id string) (int64, error) { installationUUID, err := uuid.Parse(installationID) if err != nil { - return err + return 0, err } parsed, err := uuid.Parse(id) if err != nil { - return err + return 0, err } q := sqlc.New(r.tx) return q.SoftDeleteModelRouterAPIKey(ctx, sqlc.SoftDeleteModelRouterAPIKeyParams{ diff --git a/internal/server/middleware/auth_test.go b/internal/server/middleware/auth_test.go index 8c46e9cd7..a4bf2a34e 100644 --- a/internal/server/middleware/auth_test.go +++ b/internal/server/middleware/auth_test.go @@ -78,8 +78,8 @@ func (f *fakeAPIKeyRepository) MarkUsed(ctx context.Context, id string) error { return nil } -func (f *fakeAPIKeyRepository) SoftDelete(ctx context.Context, installationID, id string) error { - return errors.New("not used") +func (f *fakeAPIKeyRepository) SoftDelete(ctx context.Context, installationID, id string) (int64, error) { + return 0, errors.New("not used") } type fakeInstallationRepository struct{} diff --git a/internal/sqlc/model_router_api_keys.sql.go b/internal/sqlc/model_router_api_keys.sql.go index 60df1f17a..4ba4e20ff 100644 --- a/internal/sqlc/model_router_api_keys.sql.go +++ b/internal/sqlc/model_router_api_keys.sql.go @@ -258,7 +258,7 @@ func (q *Queries) MarkModelRouterAPIKeyUsed(ctx context.Context, id uuid.UUID) e return err } -const softDeleteModelRouterAPIKey = `-- name: SoftDeleteModelRouterAPIKey :exec +const softDeleteModelRouterAPIKey = `-- name: SoftDeleteModelRouterAPIKey :execrows UPDATE router.model_router_api_keys SET deleted_at = NOW() WHERE id = $1::uuid @@ -271,14 +271,18 @@ type SoftDeleteModelRouterAPIKeyParams struct { InstallationID uuid.UUID } -// Soft-deletes a router API key. Cross-tenant safe via installation_id predicate. +// Soft-deletes a router API key. Cross-tenant safe via installation_id +// predicate. :execrows lets callers tell a 0-row no-op from a real transition. // // UPDATE router.model_router_api_keys // SET deleted_at = NOW() // WHERE id = $1::uuid // AND installation_id = $2::uuid // AND deleted_at IS NULL -func (q *Queries) SoftDeleteModelRouterAPIKey(ctx context.Context, arg SoftDeleteModelRouterAPIKeyParams) error { - _, err := q.db.Exec(ctx, softDeleteModelRouterAPIKey, arg.ID, arg.InstallationID) - return err +func (q *Queries) SoftDeleteModelRouterAPIKey(ctx context.Context, arg SoftDeleteModelRouterAPIKeyParams) (int64, error) { + result, err := q.db.Exec(ctx, softDeleteModelRouterAPIKey, arg.ID, arg.InstallationID) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil } diff --git a/scripts/rotate_key_race_check/main.go b/scripts/rotate_key_race_check/main.go new file mode 100644 index 000000000..d196befe9 --- /dev/null +++ b/scripts/rotate_key_race_check/main.go @@ -0,0 +1,236 @@ +// Command rotate_key_race_check reproduces a concurrent RotateAPIKey race +// against a live Postgres: N goroutines soft-delete + re-issue the same key +// with no transaction / rows-affected check, so more than one successor key +// can end up live from a single rotation intent. +// +// It is a separate main package (not a _test.go), so `go test ./...` never +// touches Postgres. It is gated on ROUTER_TEST_DATABASE_URL (a DSN to a +// database with the router migrations applied) and is a no-op without it. +// +// Timing note: a start-only barrier misses the race — List→SoftDelete is +// microseconds, so one goroutine soft-deletes before the others List. A second +// barrier inside ListForInstallation holds every caller at the TOCTOU window. +// +// Usage (from the repo root, against the docker-compose Postgres): +// +// ROUTER_TEST_DATABASE_URL="postgres://router:router@localhost:5433/router?search_path=router" \ +// go run ./scripts/rotate_key_race_check +package main + +import ( + "context" + "fmt" + "os" + "sync" + "sync/atomic" + "time" + + "workweave/router/internal/auth" + "workweave/router/internal/postgres" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" +) + +const concurrentRotations = 5 + +func main() { + dsn := os.Getenv("ROUTER_TEST_DATABASE_URL") + if dsn == "" { + fmt.Println("ROUTER_TEST_DATABASE_URL not set; skipping live-DB rotate-key race check (see file header for usage)") + return + } + ctx := context.Background() + + cfg, err := pgxpool.ParseConfig(dsn) + if err != nil { + fail("parse database url", err) + } + // Enough connections that N concurrent RotateAPIKey calls are not + // artificially serialized by the pool (default is already ≥4). + cfg.MaxConns = int32(concurrentRotations + 2) + cfg.MinConns = int32(concurrentRotations) + + pool, err := pgxpool.NewWithConfig(ctx, cfg) + if err != nil { + fail("connect to database", err) + } + defer pool.Close() + if err := pool.Ping(ctx); err != nil { + fail("ping database", err) + } + + if err := checkRotateKeyRace(ctx, pool); err != nil { + fail("rotate-key race check", err) + } +} + +func fail(step string, err error) { + fmt.Fprintf(os.Stderr, "FAIL: %s: %v\n", step, err) + os.Exit(1) +} + +// checkRotateKeyRace creates a fresh installation + one API key, then fires +// concurrentRotations barrier-synchronized RotateAPIKey calls on that same +// keyID and reports how many succeeded vs how many active keys remain. +func checkRotateKeyRace(ctx context.Context, pool *pgxpool.Pool) error { + repo := postgres.NewRepository(pool, auth.NoOpEncryptor{}) + + // Hold every caller between List and SoftDelete to open the TOCTOU window; + // a start-only barrier misses it because List→SoftDelete is microseconds. + listHold := newListHold(concurrentRotations) + apiKeys := &holdingAPIKeyRepo{inner: repo.APIKeys, hold: listHold} + + svc := auth.NewService( + repo.Installations, + apiKeys, + repo.ExternalAPIKeys, + repo.Users, + auth.NoOpAPIKeyCache{}, + nil, + time.Now, + ) + + install, err := repo.Installations.Create(ctx, auth.CreateInstallationParams{ + ExternalID: "org_rotate_race_" + uuid.NewString()[:8], + Name: "RotateAPIKey race repro", + }) + if err != nil { + return fmt.Errorf("create installation: %w", err) + } + fmt.Printf("installation_id=%s\n", install.ID) + + name := "race-repro-key" + issued, _, err := svc.IssueAPIKey(ctx, install.ID, &name, nil) + if err != nil { + return fmt.Errorf("issue initial key: %w", err) + } + fmt.Printf("initial_key_id=%s\n", issued.ID) + + var ( + ready sync.WaitGroup + done sync.WaitGroup + start = make(chan struct{}) + okCount atomic.Int64 + stopMon = make(chan struct{}) + peak atomic.Int32 + ) + ready.Add(concurrentRotations) + done.Add(concurrentRotations) + + go func() { + for { + select { + case <-stopMon: + return + default: + if c := int32(pool.Stat().AcquiredConns()); c > peak.Load() { + peak.Store(c) + } + time.Sleep(50 * time.Microsecond) + } + } + }() + + for i := 0; i < concurrentRotations; i++ { + go func() { + defer done.Done() + ready.Done() + <-start + _, _, err := svc.RotateAPIKey(ctx, install.ID, issued.ID, nil) + if err == nil { + okCount.Add(1) + } else { + fmt.Printf("RotateAPIKey error: %v\n", err) + } + }() + } + ready.Wait() + close(start) + done.Wait() + close(stopMon) + + // Use the unwrapped repo: holdingAPIKeyRepo would block forever waiting + // for N more List arrivals that never come. + keys, err := repo.APIKeys.ListForInstallation(ctx, install.ID) + if err != nil { + return fmt.Errorf("list API keys: %w", err) + } + + fmt.Printf("concurrent_calls=%d\n", concurrentRotations) + fmt.Printf("rotate_successes=%d\n", okCount.Load()) + fmt.Printf("active_keys=%d\n", len(keys)) + fmt.Printf("pool_max_conns=%d\n", pool.Stat().MaxConns()) + fmt.Printf("pool_peak_acquired_conns=%d\n", peak.Load()) + fmt.Printf("list_hold_releases=%d\n", listHold.releases.Load()) + for _, k := range keys { + fmt.Printf(" active key_id=%s prefix=%s\n", k.ID, k.KeyPrefix) + } + if okCount.Load() > 1 && len(keys) > 1 { + fmt.Println("RACE CONFIRMED: multiple RotateAPIKey calls succeeded and multiple active keys remain") + } else { + fmt.Println("race not observed this run (only one successor / one success)") + } + return nil +} + +// listHold is a reusable N-party barrier: the Nth arrival closes release so +// every waiter proceeds together. +type listHold struct { + n int + mu sync.Mutex + arrived int + release chan struct{} + releases atomic.Int64 +} + +func newListHold(n int) *listHold { + return &listHold{n: n, release: make(chan struct{})} +} + +func (h *listHold) Wait() { + h.mu.Lock() + h.arrived++ + release := h.release + if h.arrived == h.n { + h.releases.Add(1) + close(h.release) + h.arrived = 0 + h.release = make(chan struct{}) + } + h.mu.Unlock() + <-release +} + +// holdingAPIKeyRepo wraps an APIKeyRepository and parks every +// ListForInstallation caller at listHold so SoftDelete cannot run until all +// concurrent rotators have observed the pre-delete key set. +type holdingAPIKeyRepo struct { + inner auth.APIKeyRepository + hold *listHold +} + +func (r *holdingAPIKeyRepo) Create(ctx context.Context, params auth.CreateAPIKeyParams) (*auth.APIKey, error) { + return r.inner.Create(ctx, params) +} + +func (r *holdingAPIKeyRepo) GetActiveByHashWithInstallation(ctx context.Context, keyHash string) (*auth.APIKey, *auth.Installation, error) { + return r.inner.GetActiveByHashWithInstallation(ctx, keyHash) +} + +func (r *holdingAPIKeyRepo) ListForInstallation(ctx context.Context, installationID string) ([]*auth.APIKey, error) { + keys, err := r.inner.ListForInstallation(ctx, installationID) + if err != nil { + return nil, err + } + r.hold.Wait() + return keys, nil +} + +func (r *holdingAPIKeyRepo) MarkUsed(ctx context.Context, id string) error { + return r.inner.MarkUsed(ctx, id) +} + +func (r *holdingAPIKeyRepo) SoftDelete(ctx context.Context, installationID, id string) (int64, error) { + return r.inner.SoftDelete(ctx, installationID, id) +} diff --git a/scripts/rotate_key_race_ext_check/main.go b/scripts/rotate_key_race_ext_check/main.go new file mode 100644 index 000000000..6def94054 --- /dev/null +++ b/scripts/rotate_key_race_ext_check/main.go @@ -0,0 +1,675 @@ +// Command rotate_key_race_ext_check extends the RotateAPIKey race coverage +// confirmed by scripts/rotate_key_race_check. It is a separate main package +// (not a _test.go), gated on ROUTER_TEST_DATABASE_URL, and a no-op without it. +// +// Axes: +// 1. Concurrency sweep — List-held-open barrier at N = 2, 3, 10, 25 +// 2. Mixed RotateAPIKey vs DeleteAPIKey on the same key +// 3. HTTP-level POST /admin/v1/keys/:id/rotate via gin httptest + real +// WithAdminOnly cookie auth (not a live docker listener) +// 4. Timing sensitivity — N=2, 20 runs each: post-List sleep sweep, then +// start-stagger sweep to find the largest stagger with ≥1/20 hit +// +// Usage (from the repo root, against the docker-compose Postgres): +// +// ROUTER_TEST_DATABASE_URL="postgres://router:router@localhost:5433/router?search_path=router" \ +// go run ./scripts/rotate_key_race_ext_check +package main + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "os" + "sync" + "sync/atomic" + "time" + + "workweave/router/internal/api/admin" + "workweave/router/internal/auth" + "workweave/router/internal/postgres" + "workweave/router/internal/server/middleware" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" +) + +func main() { + dsn := os.Getenv("ROUTER_TEST_DATABASE_URL") + if dsn == "" { + fmt.Println("ROUTER_TEST_DATABASE_URL not set; skipping live-DB rotate-key race ext check (see file header for usage)") + return + } + gin.SetMode(gin.TestMode) + ctx := context.Background() + + cfg, err := pgxpool.ParseConfig(dsn) + if err != nil { + fail("parse database url", err) + } + cfg.MaxConns = 32 + cfg.MinConns = 8 + + pool, err := pgxpool.NewWithConfig(ctx, cfg) + if err != nil { + fail("connect to database", err) + } + defer pool.Close() + if err := pool.Ping(ctx); err != nil { + fail("ping database", err) + } + + fmt.Println("=== 1. Concurrency sweep (List-held-open) ===") + if err := checkConcurrencySweep(ctx, pool); err != nil { + fail("concurrency sweep", err) + } + + fmt.Println("\n=== 2. Mixed RotateAPIKey vs DeleteAPIKey ===") + if err := checkRotateVsDelete(ctx, pool); err != nil { + fail("rotate vs delete", err) + } + + fmt.Println("\n=== 3. HTTP-level rotate race (httptest + WithAdminOnly) ===") + if err := checkHTTPRotateRace(ctx, pool); err != nil { + fail("http rotate race", err) + } + + fmt.Println("\n=== 4. Timing sensitivity (N=2, 20 runs each) ===") + if err := checkTimingSensitivity(ctx, pool); err != nil { + fail("timing sensitivity", err) + } +} + +func fail(step string, err error) { + fmt.Fprintf(os.Stderr, "FAIL: %s: %v\n", step, err) + os.Exit(1) +} + +// --------------------------------------------------------------------------- +// 1. Concurrency sweep +// --------------------------------------------------------------------------- + +func checkConcurrencySweep(ctx context.Context, pool *pgxpool.Pool) error { + fmt.Printf("%4s %10s %11s %s\n", "N", "successes", "active_keys", "notes") + for _, n := range []int{2, 3, 10, 25} { + successes, active, errCounts, err := runHeldRotateRace(ctx, pool, n) + if err != nil { + return fmt.Errorf("N=%d: %w", n, err) + } + note := "linear" + if successes != n || active != n { + note = fmt.Sprintf("NON-LINEAR err_counts=%v", errCounts) + } + fmt.Printf("%4d %10d %11d %s\n", n, successes, active, note) + } + return nil +} + +func runHeldRotateRace(ctx context.Context, pool *pgxpool.Pool, n int) (successes, active int, errCounts map[string]int, err error) { + repo := postgres.NewRepository(pool, auth.NoOpEncryptor{}) + hold := newListHold(n) + svc := auth.NewService( + repo.Installations, + &holdingAPIKeyRepo{inner: repo.APIKeys, hold: hold}, + repo.ExternalAPIKeys, + repo.Users, + auth.NoOpAPIKeyCache{}, + nil, + time.Now, + ) + + install, err := repo.Installations.Create(ctx, auth.CreateInstallationParams{ + ExternalID: "org_rot_sweep_" + uuid.NewString()[:8], + Name: fmt.Sprintf("rotate sweep N=%d", n), + }) + if err != nil { + return 0, 0, nil, fmt.Errorf("create installation: %w", err) + } + name := "sweep-key" + issued, _, err := svc.IssueAPIKey(ctx, install.ID, &name, nil) + if err != nil { + return 0, 0, nil, fmt.Errorf("issue key: %w", err) + } + + var ( + ready sync.WaitGroup + done sync.WaitGroup + start = make(chan struct{}) + okCount atomic.Int64 + errMu sync.Mutex + ) + errCounts = map[string]int{} + ready.Add(n) + done.Add(n) + for i := 0; i < n; i++ { + go func() { + defer done.Done() + ready.Done() + <-start + _, _, err := svc.RotateAPIKey(ctx, install.ID, issued.ID, nil) + if err == nil { + okCount.Add(1) + return + } + errMu.Lock() + errCounts[err.Error()]++ + errMu.Unlock() + }() + } + ready.Wait() + close(start) + done.Wait() + + keys, err := repo.APIKeys.ListForInstallation(ctx, install.ID) + if err != nil { + return 0, 0, nil, fmt.Errorf("list keys: %w", err) + } + return int(okCount.Load()), len(keys), errCounts, nil +} + +// --------------------------------------------------------------------------- +// 2. Mixed Rotate vs Delete +// --------------------------------------------------------------------------- + +// checkRotateVsDelete forces DeleteAPIKey to win SoftDelete after Rotate has +// already Listed the key as active. Expected buggy outcome: Rotate still +// IssueAPIKeys a live successor even though the key was deleted (not rotated). +func checkRotateVsDelete(ctx context.Context, pool *pgxpool.Pool) error { + const trials = 10 + var ( + zombieIssued int // rotate succeeded after delete won SoftDelete + rotateFails int + deleteFails int + ) + + for i := 0; i < trials; i++ { + outcome, err := runOneRotateVsDelete(ctx, pool) + if err != nil { + return err + } + if outcome.deleteErr != nil { + deleteFails++ + } + if outcome.rotateErr != nil { + rotateFails++ + continue + } + // Rotate returned a new key. If Delete already soft-deleted the + // original, that new key is a zombie replacement for a deleted key. + if outcome.deleteWonSoftDelete && outcome.activeAfter > 0 { + zombieIssued++ + } + } + + fmt.Printf("trials=%d zombie_replacements=%d rotate_errors=%d delete_errors=%d\n", + trials, zombieIssued, rotateFails, deleteFails) + if zombieIssued > 0 { + fmt.Println("YES: DeleteAPIKey winning SoftDelete still lets RotateAPIKey mint a live replacement") + } else { + fmt.Println("NO: did not observe a zombie replacement this run") + } + return nil +} + +type rotateVsDeleteOutcome struct { + rotateErr error + deleteErr error + deleteWonSoftDelete bool + activeAfter int +} + +func runOneRotateVsDelete(ctx context.Context, pool *pgxpool.Pool) (rotateVsDeleteOutcome, error) { + repo := postgres.NewRepository(pool, auth.NoOpEncryptor{}) + + listed := make(chan struct{}) + releaseRotate := make(chan struct{}) + apiKeys := &signalListRepo{ + inner: repo.APIKeys, + listed: listed, + release: releaseRotate, + } + svc := auth.NewService( + repo.Installations, + apiKeys, + repo.ExternalAPIKeys, + repo.Users, + auth.NoOpAPIKeyCache{}, + nil, + time.Now, + ) + + install, err := repo.Installations.Create(ctx, auth.CreateInstallationParams{ + ExternalID: "org_rot_del_" + uuid.NewString()[:8], + Name: "rotate vs delete", + }) + if err != nil { + return rotateVsDeleteOutcome{}, fmt.Errorf("create installation: %w", err) + } + name := "mixed-key" + issued, _, err := svc.IssueAPIKey(ctx, install.ID, &name, nil) + if err != nil { + return rotateVsDeleteOutcome{}, fmt.Errorf("issue key: %w", err) + } + + var ( + ready sync.WaitGroup + done sync.WaitGroup + start = make(chan struct{}) + rotateErr error + deleteErr error + ) + ready.Add(2) + done.Add(2) + + go func() { + defer done.Done() + ready.Done() + <-start + _, _, rotateErr = svc.RotateAPIKey(ctx, install.ID, issued.ID, nil) + }() + go func() { + defer done.Done() + ready.Done() + <-start + <-listed // Rotate has Listed; key still appears active to it + deleteErr = svc.DeleteAPIKey(ctx, install.ID, issued.ID) + close(releaseRotate) // let Rotate proceed to SoftDelete (no-op) + Issue + }() + + ready.Wait() + close(start) + done.Wait() + + keys, err := repo.APIKeys.ListForInstallation(ctx, install.ID) + if err != nil { + return rotateVsDeleteOutcome{}, fmt.Errorf("list keys: %w", err) + } + return rotateVsDeleteOutcome{ + rotateErr: rotateErr, + deleteErr: deleteErr, + deleteWonSoftDelete: deleteErr == nil, + activeAfter: len(keys), + }, nil +} + +// signalListRepo parks ListForInstallation after the read until release closes, +// and signals listed once the pre-delete snapshot is in hand. +type signalListRepo struct { + inner auth.APIKeyRepository + listed chan struct{} + release chan struct{} + once sync.Once +} + +func (r *signalListRepo) Create(ctx context.Context, params auth.CreateAPIKeyParams) (*auth.APIKey, error) { + return r.inner.Create(ctx, params) +} +func (r *signalListRepo) GetActiveByHashWithInstallation(ctx context.Context, keyHash string) (*auth.APIKey, *auth.Installation, error) { + return r.inner.GetActiveByHashWithInstallation(ctx, keyHash) +} +func (r *signalListRepo) ListForInstallation(ctx context.Context, installationID string) ([]*auth.APIKey, error) { + keys, err := r.inner.ListForInstallation(ctx, installationID) + if err != nil { + return nil, err + } + r.once.Do(func() { close(r.listed) }) + <-r.release + return keys, nil +} +func (r *signalListRepo) MarkUsed(ctx context.Context, id string) error { + return r.inner.MarkUsed(ctx, id) +} +func (r *signalListRepo) SoftDelete(ctx context.Context, installationID, id string) (int64, error) { + return r.inner.SoftDelete(ctx, installationID, id) +} + +// --------------------------------------------------------------------------- +// 3. HTTP-level race +// --------------------------------------------------------------------------- + +func checkHTTPRotateRace(ctx context.Context, pool *pgxpool.Pool) error { + const n = 5 + const adminPassword = "rotate-race-http-check-password" + + repo := postgres.NewRepository(pool, auth.NoOpEncryptor{}) + hold := newListHold(n) + svc := auth.NewService( + repo.Installations, + &holdingAPIKeyRepo{inner: repo.APIKeys, hold: hold}, + repo.ExternalAPIKeys, + repo.Users, + auth.NoOpAPIKeyCache{}, + nil, + time.Now, + ).WithAdminPassword(adminPassword) + + // Admin cookie sessions operate on EnsureAdminInstallation's singleton. + install, err := svc.EnsureAdminInstallation(ctx) + if err != nil { + return fmt.Errorf("ensure admin installation: %w", err) + } + name := "http-race-key-" + uuid.NewString()[:8] + issued, _, err := svc.IssueAPIKey(ctx, install.ID, &name, nil) + if err != nil { + return fmt.Errorf("issue key: %w", err) + } + fmt.Printf("method=httptest+WithAdminOnly installation_id=%s initial_key_id=%s\n", install.ID, issued.ID) + + session, _, err := svc.IssueAdminSession() + if err != nil { + return fmt.Errorf("issue admin session: %w", err) + } + + engine := gin.New() + mgmt := engine.Group("/admin/v1", middleware.WithAdminOnly(svc)) + mgmt.POST("/keys/:id/rotate", admin.RotateAPIKeyHandler(svc)) + + var ( + ready sync.WaitGroup + done sync.WaitGroup + start = make(chan struct{}) + okCount atomic.Int64 + statusMu sync.Mutex + statuses = map[int]int{} + ) + ready.Add(n) + done.Add(n) + for i := 0; i < n; i++ { + go func() { + defer done.Done() + ready.Done() + <-start + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/admin/v1/keys/"+issued.ID+"/rotate", nil) + req.AddCookie(&http.Cookie{Name: auth.AdminSessionCookieName, Value: session}) + engine.ServeHTTP(rec, req) + statusMu.Lock() + statuses[rec.Code]++ + statusMu.Unlock() + if rec.Code == http.StatusCreated { + okCount.Add(1) + } + }() + } + ready.Wait() + close(start) + done.Wait() + + keys, err := repo.APIKeys.ListForInstallation(ctx, install.ID) + if err != nil { + return fmt.Errorf("list keys: %w", err) + } + // Count only keys from this race (name match) — admin install may have leftovers. + activeNamed := 0 + for _, k := range keys { + if k.Name != nil && *k.Name == name { + activeNamed++ + } + } + + fmt.Printf("http_status_counts=%v\n", statuses) + fmt.Printf("rotate_http_201=%d active_keys_with_race_name=%d\n", okCount.Load(), activeNamed) + if okCount.Load() > 1 && activeNamed > 1 { + fmt.Println("YES: race reproduces through gin WithAdminOnly + RotateAPIKeyHandler (httptest)") + } else { + fmt.Println("NO: HTTP path did not show multi-successor race") + } + return nil +} + +// --------------------------------------------------------------------------- +// 4. Timing sensitivity +// --------------------------------------------------------------------------- + +func checkTimingSensitivity(ctx context.Context, pool *pgxpool.Pool) error { + const trials = 20 + + // A) Simultaneous start, optional post-List sleep (widens SoftDelete delay). + fmt.Println("-- A) simultaneous start, post-List sleep --") + fmt.Printf("%12s %8s %s\n", "post_list_sleep", "hits/20", "min_list_gap_on_hit") + delays := []time.Duration{ + 0, + time.Microsecond, + 10 * time.Microsecond, + 100 * time.Microsecond, + time.Millisecond, + } + for _, delay := range delays { + hits, minGap, err := sweepTimedPairs(ctx, pool, trials, delay, 0) + if err != nil { + return err + } + gapStr := "-" + if minGap >= 0 { + gapStr = minGap.String() + } + fmt.Printf("%12s %3d/20 %s\n", delay, hits, gapStr) + } + + // B) No post-List sleep; stagger second caller's start by G after the + // first begins. Largest G with ≥1/20 hit ≈ real-world overlap window + // (e.g. retry-on-timeout re-sending rotate). + fmt.Println("-- B) start-stagger of 2nd caller (no post-List sleep) --") + fmt.Printf("%12s %8s %s\n", "start_stagger", "hits/20", "max_list_gap_on_hit") + staggers := []time.Duration{ + 0, + 10 * time.Microsecond, + 50 * time.Microsecond, + 100 * time.Microsecond, + 250 * time.Microsecond, + 500 * time.Microsecond, + time.Millisecond, + 2 * time.Millisecond, + 5 * time.Millisecond, + 10 * time.Millisecond, + } + var maxHitStagger time.Duration = -1 + for _, stagger := range staggers { + hits, maxGap, err := sweepTimedPairs(ctx, pool, trials, 0, stagger) + if err != nil { + return err + } + gapStr := "-" + if maxGap >= 0 { + gapStr = maxGap.String() + } + fmt.Printf("%12s %3d/20 %s\n", stagger, hits, gapStr) + if hits > 0 { + maxHitStagger = stagger + } + } + if maxHitStagger < 0 { + fmt.Println("timing_threshold: race never hit under start-stagger sweep") + } else { + fmt.Printf("timing_threshold: largest start-stagger with ≥1/20 hit = %s\n", maxHitStagger) + fmt.Println("(N=2 uncoordinated overlap within that stagger is enough; tighter gaps hit more often)") + } + return nil +} + +func sweepTimedPairs(ctx context.Context, pool *pgxpool.Pool, trials int, postListSleep, startStagger time.Duration) (hits int, extremeGap time.Duration, err error) { + extremeGap = -1 + for t := 0; t < trials; t++ { + hit, gap, err := runTimedRotatePair(ctx, pool, postListSleep, startStagger) + if err != nil { + return 0, -1, err + } + if !hit { + continue + } + hits++ + // For sleep sweep report min gap; for stagger sweep report max gap. + if startStagger > 0 { + if extremeGap < 0 || gap > extremeGap { + extremeGap = gap + } + } else if extremeGap < 0 || gap < extremeGap { + extremeGap = gap + } + } + return hits, extremeGap, nil +} + +func runTimedRotatePair(ctx context.Context, pool *pgxpool.Pool, postListSleep, startStagger time.Duration) (hit bool, listGap time.Duration, err error) { + repo := postgres.NewRepository(pool, auth.NoOpEncryptor{}) + timed := &timedListRepo{inner: repo.APIKeys, sleep: postListSleep} + svc := auth.NewService( + repo.Installations, + timed, + repo.ExternalAPIKeys, + repo.Users, + auth.NoOpAPIKeyCache{}, + nil, + time.Now, + ) + + install, err := repo.Installations.Create(ctx, auth.CreateInstallationParams{ + ExternalID: "org_rot_time_" + uuid.NewString()[:8], + Name: "timing", + }) + if err != nil { + return false, 0, fmt.Errorf("create installation: %w", err) + } + name := "timing-key" + issued, _, err := svc.IssueAPIKey(ctx, install.ID, &name, nil) + if err != nil { + return false, 0, fmt.Errorf("issue key: %w", err) + } + + var ( + done sync.WaitGroup + okCount atomic.Int64 + ) + done.Add(2) + go func() { + defer done.Done() + _, _, err := svc.RotateAPIKey(ctx, install.ID, issued.ID, nil) + if err == nil { + okCount.Add(1) + } + }() + go func() { + defer done.Done() + if startStagger > 0 { + time.Sleep(startStagger) + } + _, _, err := svc.RotateAPIKey(ctx, install.ID, issued.ID, nil) + if err == nil { + okCount.Add(1) + } + }() + done.Wait() + + keys, err := repo.APIKeys.ListForInstallation(ctx, install.ID) + if err != nil { + return false, 0, fmt.Errorf("list keys: %w", err) + } + return okCount.Load() == 2 && len(keys) >= 2, timed.listGap(), nil +} + +// timedListRepo optionally sleeps after each List and records List-completion +// timestamps so we can report the observed inter-List gap when a race hits. +type timedListRepo struct { + inner auth.APIKeyRepository + sleep time.Duration + + mu sync.Mutex + listTimes []time.Time +} + +func (r *timedListRepo) Create(ctx context.Context, params auth.CreateAPIKeyParams) (*auth.APIKey, error) { + return r.inner.Create(ctx, params) +} +func (r *timedListRepo) GetActiveByHashWithInstallation(ctx context.Context, keyHash string) (*auth.APIKey, *auth.Installation, error) { + return r.inner.GetActiveByHashWithInstallation(ctx, keyHash) +} +func (r *timedListRepo) ListForInstallation(ctx context.Context, installationID string) ([]*auth.APIKey, error) { + keys, err := r.inner.ListForInstallation(ctx, installationID) + if err != nil { + return nil, err + } + r.mu.Lock() + r.listTimes = append(r.listTimes, time.Now()) + r.mu.Unlock() + if r.sleep > 0 { + time.Sleep(r.sleep) + } + return keys, nil +} +func (r *timedListRepo) MarkUsed(ctx context.Context, id string) error { + return r.inner.MarkUsed(ctx, id) +} +func (r *timedListRepo) SoftDelete(ctx context.Context, installationID, id string) (int64, error) { + return r.inner.SoftDelete(ctx, installationID, id) +} + +func (r *timedListRepo) listGap() time.Duration { + r.mu.Lock() + defer r.mu.Unlock() + if len(r.listTimes) < 2 { + return 0 + } + d := r.listTimes[1].Sub(r.listTimes[0]) + if d < 0 { + return -d + } + return d +} + +// --------------------------------------------------------------------------- +// Shared list-hold helpers (same pattern as rotate_key_race_check) +// --------------------------------------------------------------------------- + +type listHold struct { + n int + mu sync.Mutex + arrived int + release chan struct{} + releases atomic.Int64 +} + +func newListHold(n int) *listHold { + return &listHold{n: n, release: make(chan struct{})} +} + +func (h *listHold) Wait() { + h.mu.Lock() + h.arrived++ + release := h.release + if h.arrived == h.n { + h.releases.Add(1) + close(h.release) + h.arrived = 0 + h.release = make(chan struct{}) + } + h.mu.Unlock() + <-release +} + +type holdingAPIKeyRepo struct { + inner auth.APIKeyRepository + hold *listHold +} + +func (r *holdingAPIKeyRepo) Create(ctx context.Context, params auth.CreateAPIKeyParams) (*auth.APIKey, error) { + return r.inner.Create(ctx, params) +} +func (r *holdingAPIKeyRepo) GetActiveByHashWithInstallation(ctx context.Context, keyHash string) (*auth.APIKey, *auth.Installation, error) { + return r.inner.GetActiveByHashWithInstallation(ctx, keyHash) +} +func (r *holdingAPIKeyRepo) ListForInstallation(ctx context.Context, installationID string) ([]*auth.APIKey, error) { + keys, err := r.inner.ListForInstallation(ctx, installationID) + if err != nil { + return nil, err + } + r.hold.Wait() + return keys, nil +} +func (r *holdingAPIKeyRepo) MarkUsed(ctx context.Context, id string) error { + return r.inner.MarkUsed(ctx, id) +} +func (r *holdingAPIKeyRepo) SoftDelete(ctx context.Context, installationID, id string) (int64, error) { + return r.inner.SoftDelete(ctx, installationID, id) +} diff --git a/scripts/upsert_external_key_race_check/main.go b/scripts/upsert_external_key_race_check/main.go new file mode 100644 index 000000000..f18de7e93 --- /dev/null +++ b/scripts/upsert_external_key_race_check/main.go @@ -0,0 +1,259 @@ +// Command upsert_external_key_race_check races N=2 concurrent +// UpsertExternalAPIKey calls for the same (installation, provider) to see +// whether the UNIQUE (installation_id, provider) WHERE deleted_at IS NULL +// index cleanly fails the loser (vs. silently minting two live BYOK rows). +// +// SoftDeleteExternalAPIKeyByProvider is :exec (same shape as the RotateAPIKey +// bug); the question is whether Create's unique-violation error propagates. +// +// Gated on ROUTER_TEST_DATABASE_URL; no-op when unset. +// +// ROUTER_TEST_DATABASE_URL="postgres://router:router@localhost:5433/router?search_path=router" \ +// go run ./scripts/upsert_external_key_race_check +package main + +import ( + "context" + "errors" + "fmt" + "os" + "sync" + "sync/atomic" + "time" + + "workweave/router/internal/auth" + "workweave/router/internal/postgres" + "workweave/router/internal/providers" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgxpool" +) + +const n = 2 + +func main() { + dsn := os.Getenv("ROUTER_TEST_DATABASE_URL") + if dsn == "" { + fmt.Println("ROUTER_TEST_DATABASE_URL not set; skipping live-DB upsert-external race check (see file header for usage)") + return + } + ctx := context.Background() + + cfg, err := pgxpool.ParseConfig(dsn) + if err != nil { + fail("parse database url", err) + } + cfg.MaxConns = 8 + cfg.MinConns = 2 + + pool, err := pgxpool.NewWithConfig(ctx, cfg) + if err != nil { + fail("connect to database", err) + } + defer pool.Close() + + if err := checkUpsertExternalRace(ctx, pool); err != nil { + fail("upsert external race check", err) + } +} + +func fail(step string, err error) { + fmt.Fprintf(os.Stderr, "FAIL: %s: %v\n", step, err) + os.Exit(1) +} + +func checkUpsertExternalRace(ctx context.Context, pool *pgxpool.Pool) error { + repo := postgres.NewRepository(pool, auth.NoOpEncryptor{}) + + // Seed with the unwrapped repo so SoftDelete hold isn't engaged yet. + seedSvc := auth.NewService( + repo.Installations, + repo.APIKeys, + repo.ExternalAPIKeys, + repo.Users, + auth.NoOpAPIKeyCache{}, + nil, + time.Now, + ) + + install, err := repo.Installations.Create(ctx, auth.CreateInstallationParams{ + ExternalID: "org_upsert_race_" + uuid.NewString()[:8], + Name: "UpsertExternalAPIKey race", + }) + if err != nil { + return fmt.Errorf("create installation: %w", err) + } + fmt.Printf("installation_id=%s\n", install.ID) + fmt.Printf("provider=%s\n", providers.ProviderAnthropic) + + seedName := "seed-byok" + seeded, err := seedSvc.UpsertExternalAPIKey(ctx, install.ID, providers.ProviderAnthropic, "sk-ant-race-seed-key", &seedName, nil) + if err != nil { + return fmt.Errorf("seed upsert: %w", err) + } + fmt.Printf("seeded_key_id=%s\n", seeded.ID) + + // Barrier AFTER SoftDeleteByProvider so both soft-deletes finish before + // either Create — Upsert's SoftDelete→Create window held open. + hold := newBarrier(n) + var softDeleteCalls atomic.Int64 + ext := &holdingExternalRepo{ + inner: repo.ExternalAPIKeys, + hold: hold, + softDeleteCalls: &softDeleteCalls, + } + raceSvc := auth.NewService( + repo.Installations, + repo.APIKeys, + ext, + repo.Users, + auth.NoOpAPIKeyCache{}, + nil, + time.Now, + ) + + type result struct { + key *auth.ExternalAPIKey + err error + } + results := make([]result, n) + + var ready, done sync.WaitGroup + start := make(chan struct{}) + ready.Add(n) + done.Add(n) + for i := 0; i < n; i++ { + i := i + go func() { + defer done.Done() + ready.Done() + <-start + name := fmt.Sprintf("race-byok-%d", i) + key, err := raceSvc.UpsertExternalAPIKey(ctx, install.ID, providers.ProviderAnthropic, + fmt.Sprintf("sk-ant-race-upsert-%d-%s", i, uuid.NewString()[:8]), &name, nil) + results[i] = result{key: key, err: err} + }() + } + ready.Wait() + close(start) + done.Wait() + + var okCount, errCount int + for i, r := range results { + if r.err == nil { + okCount++ + fmt.Printf("caller[%d] OK key_id=%s\n", i, r.key.ID) + continue + } + errCount++ + fmt.Printf("caller[%d] ERR type=%T msg=%v\n", i, r.err, r.err) + var pgErr *pgconn.PgError + if errors.As(r.err, &pgErr) { + fmt.Printf("caller[%d] pg_code=%s pg_constraint=%s pg_detail=%s\n", + i, pgErr.Code, pgErr.ConstraintName, pgErr.Detail) + } else { + fmt.Printf("caller[%d] errors.As(*pgconn.PgError)=false (wrapped?)\n", i) + fmt.Printf("caller[%d] errors.Unwrap=%v\n", i, errors.Unwrap(r.err)) + } + } + fmt.Printf("upsert_successes=%d upsert_errors=%d soft_delete_calls=%d\n", + okCount, errCount, softDeleteCalls.Load()) + + var active, deleted int + if err := pool.QueryRow(ctx, ` + SELECT + COUNT(*) FILTER (WHERE deleted_at IS NULL), + COUNT(*) FILTER (WHERE deleted_at IS NOT NULL) + FROM router.model_router_external_api_keys + WHERE installation_id = $1::uuid AND provider = $2`, + install.ID, providers.ProviderAnthropic, + ).Scan(&active, &deleted); err != nil { + return fmt.Errorf("count rows: %w", err) + } + fmt.Printf("psql active=%d soft_deleted=%d (want active=1)\n", active, deleted) + + rows, err := pool.Query(ctx, ` + SELECT id, key_prefix, deleted_at IS NULL AS active, created_at + FROM router.model_router_external_api_keys + WHERE installation_id = $1::uuid AND provider = $2 + ORDER BY created_at`, + install.ID, providers.ProviderAnthropic) + if err != nil { + return fmt.Errorf("select rows: %w", err) + } + defer rows.Close() + for rows.Next() { + var id, prefix string + var activeRow bool + var created time.Time + if err := rows.Scan(&id, &prefix, &activeRow, &created); err != nil { + return err + } + fmt.Printf(" id=%s prefix=%s active=%v created_at=%s\n", id, prefix, activeRow, created.Format(time.RFC3339Nano)) + } + + switch { + case okCount == 1 && errCount == 1 && active == 1: + fmt.Println("RESULT: unique constraint stops zombie keys; loser error propagates to UpsertExternalAPIKey caller") + case okCount == 2 && active == 2: + fmt.Println("RESULT: ESCALATE — same zombie-key outcome as RotateAPIKey (constraint did not stop it)") + default: + fmt.Printf("RESULT: unexpected (ok=%d err=%d active=%d)\n", okCount, errCount, active) + } + return nil +} + +type barrier struct { + n int + mu sync.Mutex + arrived int + release chan struct{} +} + +func newBarrier(n int) *barrier { + return &barrier{n: n, release: make(chan struct{})} +} + +func (h *barrier) Wait() { + h.mu.Lock() + h.arrived++ + release := h.release + if h.arrived == h.n { + close(h.release) + h.arrived = 0 + h.release = make(chan struct{}) + } + h.mu.Unlock() + <-release +} + +// holdingExternalRepo parks SoftDeleteByProvider after the real UPDATE so both +// racers finish soft-delete before either Create. +type holdingExternalRepo struct { + inner auth.ExternalAPIKeyRepository + hold *barrier + softDeleteCalls *atomic.Int64 +} + +func (r *holdingExternalRepo) Create(ctx context.Context, params auth.CreateExternalAPIKeyParams) (*auth.ExternalAPIKey, error) { + return r.inner.Create(ctx, params) +} +func (r *holdingExternalRepo) GetForInstallation(ctx context.Context, installationID string) ([]*auth.ExternalAPIKey, error) { + return r.inner.GetForInstallation(ctx, installationID) +} +func (r *holdingExternalRepo) SoftDeleteByProvider(ctx context.Context, installationID, provider string) error { + err := r.inner.SoftDeleteByProvider(ctx, installationID, provider) + if err != nil { + return err + } + r.softDeleteCalls.Add(1) + r.hold.Wait() + return nil +} +func (r *holdingExternalRepo) SoftDelete(ctx context.Context, installationID, id string) error { + return r.inner.SoftDelete(ctx, installationID, id) +} +func (r *holdingExternalRepo) MarkUsed(ctx context.Context, id string) error { + return r.inner.MarkUsed(ctx, id) +}