Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions db/queries/model_router_api_keys.sql
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,10 @@ 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 so callers (RotateAPIKey) can tell a 0-row no-op
-- from a real transition — see #817.
-- name: SoftDeleteModelRouterAPIKey :execrows
UPDATE router.model_router_api_keys
SET deleted_at = NOW()
WHERE id = @id::uuid
Expand Down
6 changes: 3 additions & 3 deletions internal/api/admin/keys_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
5 changes: 4 additions & 1 deletion internal/auth/api_key.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,5 +34,8 @@ 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 when it is still active. Returns the
// rows-affected count so callers that must not proceed on a no-op
// (RotateAPIKey) can distinguish a real transition from a lost race.
SoftDelete(ctx context.Context, installationID, id string) (int64, error)
}
289 changes: 289 additions & 0 deletions internal/auth/rotate_race_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,289 @@
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 that models real
// SoftDelete semantics: matching active row → soft-delete; already-deleted
// or missing → silent no-op (0 rows). Used to exercise RotateAPIKey races
// without Postgres (see #817).
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 is the #817 unit repro:
// two RotateAPIKey callers both observe the key as active, then SoftDelete;
// exactly one must mint a successor, the other ErrAPIKeyNotFound.
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 is the #817 Delete-vs-Rotate
// unit repro: Delete wins SoftDelete after Rotate has Listed; Rotate must
// return ErrAPIKeyNotFound and must not issue a zombie replacement.
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")
}
24 changes: 19 additions & 5 deletions internal/auth/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -132,8 +132,16 @@ 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.
// can't be rotated.
//
// SoftDelete's rows-affected count is load-bearing (#817): a 0-row result
// means another caller already transitioned the key (a concurrent
// RotateAPIKey, or a DeleteAPIKey that won the race after this caller
// Listed). In that case we must not mint a successor — doing so produced
// untracked zombie credentials when SoftDelete was :exec and the count was
// discarded. The List→SoftDelete→Issue steps are still not wrapped in a
// single transaction; Postgres row-level locking serializes SoftDelete, and
// reading rows-affected is enough to close both failure modes.
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 {
Expand All @@ -149,9 +157,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
Expand All @@ -162,9 +174,11 @@ 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). SoftDelete's
// rows-affected count is discarded: a 0-row no-op (already deleted) is the
// existing idempotent success behavior and must not surface as an error.
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)
Expand Down
7 changes: 5 additions & 2 deletions internal/auth/service_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading