mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-10-19 22:19:12 +08:00
feat(auth): introduce mutation locking and persistence delta merging
- Add `lockAuthMutation` and `lockAuthMutationContext` with context cancellation support for mutation gates and reload barriers - Introduce `persistLocked` to unlock manager state during storage I/O and prevent blocking runtime operations - Implement `mergeAuthSaveDelta` to merge store save changes while preserving concurrent edits Closes: #6329
This commit is contained in:
@@ -2,6 +2,7 @@ package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"net/http"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
@@ -10,6 +11,7 @@ import (
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config"
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor"
|
||||
"github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi"
|
||||
"golang.org/x/sync/semaphore"
|
||||
)
|
||||
|
||||
// ProviderExecutor defines the contract required by Manager to execute provider calls.
|
||||
@@ -208,6 +210,11 @@ type Manager struct {
|
||||
refreshLocks sync.Map
|
||||
// persistLocks serializes disk persistence per auth ID and guards against out-of-order writes.
|
||||
persistLocks sync.Map
|
||||
// authMutationLocks coordinate credential mutations without blocking unrelated readers.
|
||||
authMutationLocks sync.Map
|
||||
// authLoadGate allows concurrent credential transactions, but excludes whole-store reloads.
|
||||
// Mutations acquire one permit; Load acquires all permits before reading the store.
|
||||
authLoadGate *semaphore.Weighted
|
||||
}
|
||||
|
||||
// NewManager constructs a manager with optional custom selector and hook.
|
||||
@@ -220,6 +227,7 @@ func NewManager(store Store, selector Selector, hook Hook) *Manager {
|
||||
}
|
||||
manager := &Manager{
|
||||
store: store,
|
||||
authLoadGate: semaphore.NewWeighted(math.MaxInt64),
|
||||
executors: make(map[string]ProviderExecutor),
|
||||
selector: selector,
|
||||
hook: hook,
|
||||
|
||||
180
sdk/cliproxy/auth/conductor_cancellation_test.go
Normal file
180
sdk/cliproxy/auth/conductor_cancellation_test.go
Normal file
@@ -0,0 +1,180 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"math"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"testing/synctest"
|
||||
)
|
||||
|
||||
type cancellationStore struct {
|
||||
saves atomic.Int32
|
||||
lists atomic.Int32
|
||||
entered chan struct{}
|
||||
release chan struct{}
|
||||
}
|
||||
|
||||
func (s *cancellationStore) Save(_ context.Context, auth *Auth) (string, error) {
|
||||
s.saves.Add(1)
|
||||
if auth.Label == "blocked" {
|
||||
close(s.entered)
|
||||
<-s.release
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
func (s *cancellationStore) List(context.Context) ([]*Auth, error) {
|
||||
s.lists.Add(1)
|
||||
return nil, nil
|
||||
}
|
||||
func (*cancellationStore) Delete(context.Context, string) error { return nil }
|
||||
|
||||
type cancellationHook struct{ calls atomic.Int32 }
|
||||
|
||||
func (h *cancellationHook) OnAuthRegistered(context.Context, *Auth) { h.calls.Add(1) }
|
||||
func (h *cancellationHook) OnAuthUpdated(context.Context, *Auth) { h.calls.Add(1) }
|
||||
func (*cancellationHook) OnResult(context.Context, Result) {}
|
||||
|
||||
func assertCanceledMutation(t *testing.T, done <-chan error) {
|
||||
t.Helper()
|
||||
select {
|
||||
case err := <-done:
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Errorf("expected context cancellation, got %v", err)
|
||||
}
|
||||
default:
|
||||
t.Error("canceled operation still waiting for lock")
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerCanceledLoadUnblocksUnrelatedMutation(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
store := &cancellationStore{entered: make(chan struct{}), release: make(chan struct{})}
|
||||
manager := NewManager(store, nil, nil)
|
||||
saveDone := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := manager.Register(context.Background(), &Auth{ID: "A", Provider: "audit", Label: "blocked", Metadata: map[string]any{"type": "audit"}})
|
||||
saveDone <- err
|
||||
}()
|
||||
<-store.entered
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
loadDone := make(chan error, 1)
|
||||
go func() { loadDone <- manager.Load(ctx) }()
|
||||
synctest.Wait()
|
||||
bDone := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := manager.Register(context.Background(), &Auth{ID: "B", Provider: "audit"})
|
||||
bDone <- err
|
||||
}()
|
||||
synctest.Wait()
|
||||
select {
|
||||
case <-bDone:
|
||||
t.Error("B overtook queued reload")
|
||||
default:
|
||||
}
|
||||
cancel()
|
||||
synctest.Wait()
|
||||
assertCanceledMutation(t, loadDone)
|
||||
select {
|
||||
case err := <-bDone:
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
default:
|
||||
t.Error("canceled reload still blocks unrelated B")
|
||||
}
|
||||
if store.lists.Load() != 0 {
|
||||
t.Error("canceled reload called List")
|
||||
}
|
||||
close(store.release)
|
||||
synctest.Wait()
|
||||
if err := <-saveDone; err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
if !manager.authLoadGate.TryAcquire(math.MaxInt64) {
|
||||
t.Error("reload barrier permit leaked")
|
||||
} else {
|
||||
manager.authLoadGate.Release(math.MaxInt64)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestManagerCanceledMutationWait(t *testing.T) {
|
||||
for _, operation := range []string{"Register", "Update", "Prepare", "Refresh"} {
|
||||
for _, waitOn := range []string{"barrier", "gate"} {
|
||||
t.Run(operation+"/"+waitOn, func(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
store := &cancellationStore{}
|
||||
hook := &cancellationHook{}
|
||||
manager := NewManager(store, nil, hook)
|
||||
base, err := manager.Register(WithSkipPersist(context.Background()), &Auth{ID: "A", Provider: "audit", Metadata: map[string]any{"type": "audit"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
hook.calls.Store(0)
|
||||
var release func()
|
||||
if waitOn == "barrier" {
|
||||
if errAcquire := manager.authLoadGate.Acquire(context.Background(), math.MaxInt64); errAcquire != nil {
|
||||
t.Fatal(errAcquire)
|
||||
}
|
||||
release = func() { manager.authLoadGate.Release(math.MaxInt64) }
|
||||
} else {
|
||||
release = manager.lockAuthMutation(base.ID)
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
var errMutation error
|
||||
switch operation {
|
||||
case "Register":
|
||||
_, errMutation = manager.Register(ctx, base.Clone())
|
||||
case "Update":
|
||||
_, errMutation = manager.Update(ctx, base.Clone())
|
||||
case "Prepare":
|
||||
_, errMutation = manager.UpdatePreparedAuth(ctx, base, base.Clone())
|
||||
case "Refresh":
|
||||
_, errMutation = manager.UpdateRefreshedAuth(ctx, base, base.Clone())
|
||||
}
|
||||
done <- errMutation
|
||||
}()
|
||||
synctest.Wait()
|
||||
cancel()
|
||||
synctest.Wait()
|
||||
assertCanceledMutation(t, done)
|
||||
release()
|
||||
synctest.Wait()
|
||||
if store.saves.Load() != 0 || hook.calls.Load() != 0 {
|
||||
t.Error("canceled mutation invoked Save or hook")
|
||||
}
|
||||
if got, _ := manager.GetByID(base.ID); got.Generation != base.Generation || got.RegistrationEpoch != base.RegistrationEpoch {
|
||||
t.Error("canceled mutation published auth")
|
||||
}
|
||||
if !manager.authLoadGate.TryAcquire(math.MaxInt64) {
|
||||
t.Error("mutation leaked reload barrier permit")
|
||||
} else {
|
||||
manager.authLoadGate.Release(math.MaxInt64)
|
||||
}
|
||||
value, _ := manager.authMutationLocks.Load(base.ID)
|
||||
if len(value.(chan struct{})) != 0 {
|
||||
t.Error("mutation leaked credential gate")
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerMutationNilContext(t *testing.T) {
|
||||
manager := NewManager(&cancellationStore{}, nil, nil)
|
||||
auth, err := manager.Register(nil, &Auth{ID: "A", Provider: "audit", Metadata: map[string]any{"type": "audit"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, errUpdate := manager.Update(nil, auth); errUpdate != nil {
|
||||
t.Fatal(errUpdate)
|
||||
}
|
||||
if errLoad := manager.Load(nil); errLoad != nil {
|
||||
t.Fatal(errLoad)
|
||||
}
|
||||
}
|
||||
@@ -475,6 +475,8 @@ func (m *Manager) ResetQuota(ctx context.Context, authID string) (*Auth, []strin
|
||||
registeredModels := modelsForRegisteredAuth(authID)
|
||||
cooldownStateChanged := false
|
||||
|
||||
releaseMutation := m.lockAuthMutation(authID)
|
||||
defer releaseMutation()
|
||||
m.mu.Lock()
|
||||
auth, ok := m.auths[authID]
|
||||
if !ok || auth == nil {
|
||||
@@ -524,8 +526,9 @@ func (m *Manager) ResetQuota(ctx context.Context, authID string) (*Auth, []strin
|
||||
cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now)
|
||||
cooldownStateChanged = !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter)
|
||||
}
|
||||
errPersist := m.persist(ctx, auth)
|
||||
errPersist := m.persistLocked(ctx, auth)
|
||||
m.mu.Unlock()
|
||||
releaseMutation()
|
||||
|
||||
defer func() {
|
||||
if cooldownStateChanged {
|
||||
@@ -766,6 +769,8 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) {
|
||||
cooldownStateChanged := false
|
||||
now := time.Now()
|
||||
|
||||
releaseMutation := m.lockAuthMutation(result.AuthID)
|
||||
defer releaseMutation()
|
||||
m.mu.Lock()
|
||||
if auth, ok := m.auths[result.AuthID]; ok && auth != nil {
|
||||
if modelKey == "" && strings.TrimSpace(result.RouteModel) != "" {
|
||||
@@ -1016,7 +1021,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) {
|
||||
}
|
||||
}
|
||||
|
||||
_ = m.persist(ctx, auth)
|
||||
_ = m.persistLocked(ctx, auth)
|
||||
authSnapshot = auth.Clone()
|
||||
if trackCooldownState {
|
||||
cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now)
|
||||
@@ -1024,6 +1029,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) {
|
||||
}
|
||||
}
|
||||
m.mu.Unlock()
|
||||
releaseMutation()
|
||||
if m.scheduler != nil && authSnapshot != nil {
|
||||
var targetModels []string
|
||||
if !result.CredentialScope && modelKey != "" {
|
||||
@@ -1096,6 +1102,8 @@ func (m *Manager) recordAvailabilityNeutralResult(ctx context.Context, result Re
|
||||
}
|
||||
|
||||
var authSnapshot *Auth
|
||||
releaseMutation := m.lockAuthMutation(result.AuthID)
|
||||
defer releaseMutation()
|
||||
m.mu.Lock()
|
||||
if auth, ok := m.auths[result.AuthID]; ok && auth != nil {
|
||||
now := time.Now()
|
||||
@@ -1107,10 +1115,11 @@ func (m *Manager) recordAvailabilityNeutralResult(ctx context.Context, result Re
|
||||
}
|
||||
auth.Generation++
|
||||
auth.UpdatedAt = now
|
||||
_ = m.persist(ctx, auth)
|
||||
_ = m.persistLocked(ctx, auth)
|
||||
authSnapshot = auth.Clone()
|
||||
}
|
||||
m.mu.Unlock()
|
||||
releaseMutation()
|
||||
|
||||
m.hook.OnResult(ctx, result)
|
||||
m.publishErrorEvent(result, authSnapshot)
|
||||
|
||||
@@ -3,6 +3,7 @@ package auth
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -88,6 +89,11 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) {
|
||||
if auth.ID == "" {
|
||||
auth.ID = uuid.NewString()
|
||||
}
|
||||
releaseMutation, errMutation := m.lockAuthMutationContext(ctx, auth.ID)
|
||||
if errMutation != nil {
|
||||
return nil, errMutation
|
||||
}
|
||||
defer releaseMutation()
|
||||
now := time.Now()
|
||||
if auth.Generation == 0 {
|
||||
auth.Generation = 1
|
||||
@@ -114,6 +120,16 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) {
|
||||
m.authEpochs[auth.ID]++
|
||||
auth.RegistrationEpoch = m.authEpochs[auth.ID]
|
||||
auth.Generation = 1
|
||||
// Serialize this credential, but release the manager lock during store I/O.
|
||||
// Persist failures stay non-fatal, but must not be silent: a restart would
|
||||
// lose the credential that only exists in memory.
|
||||
if errPersist := m.persistLocked(ctx, auth); errPersist != nil {
|
||||
log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist registered auth %s (%s): %v", auth.Provider, auth.ID, errPersist)
|
||||
}
|
||||
if m.authEpochs[auth.ID] != auth.RegistrationEpoch {
|
||||
m.mu.Unlock()
|
||||
return nil, fmt.Errorf("auth %s changed registration during persistence", auth.ID)
|
||||
}
|
||||
authClone := auth.Clone()
|
||||
m.auths[auth.ID] = authClone
|
||||
m.notifyAuthChangeLocked(auth.ID)
|
||||
@@ -123,6 +139,7 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) {
|
||||
schedulerSnapshot = authClone.Clone()
|
||||
}
|
||||
m.mu.Unlock()
|
||||
releaseMutation()
|
||||
if !shouldDeferAPIKeyModelAliasRebuild(ctx) {
|
||||
m.rebuildAPIKeyModelAliasFromRuntimeConfig()
|
||||
}
|
||||
@@ -131,11 +148,6 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) {
|
||||
}
|
||||
m.structuralEpoch.Add(1)
|
||||
m.queueRefreshReschedule(auth.ID)
|
||||
// Persist failures stay non-fatal, but must not be silent: a restart would
|
||||
// lose the credential that only exists in memory.
|
||||
if errPersist := m.persist(ctx, auth); errPersist != nil {
|
||||
log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist registered auth %s (%s): %v", auth.Provider, auth.ID, errPersist)
|
||||
}
|
||||
m.hook.OnAuthRegistered(ctx, auth.Clone())
|
||||
if cooldownStateChanged {
|
||||
m.persistCooldownStates(context.Background())
|
||||
@@ -176,6 +188,11 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd
|
||||
if errWeight := ValidateAuthWeight(auth); errWeight != nil {
|
||||
return nil, fmt.Errorf("update auth: %w", errWeight)
|
||||
}
|
||||
releaseMutation, errMutation := m.lockAuthMutationContext(ctx, auth.ID)
|
||||
if errMutation != nil {
|
||||
return nil, errMutation
|
||||
}
|
||||
defer releaseMutation()
|
||||
m.mu.Lock()
|
||||
existing, ok := m.auths[auth.ID]
|
||||
if !ok || existing == nil {
|
||||
@@ -260,15 +277,26 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd
|
||||
cooldownStateChanged = clearCooldownStateForAuth(auth, now) || cooldownStateChanged
|
||||
}
|
||||
auth.EnsureIndex()
|
||||
// A minted Meta key must reach the configured store before requests can use it.
|
||||
// Keep the epoch check, save and installation together so a concurrent reload
|
||||
// or removal cannot let an obsolete mint overwrite the credential on disk.
|
||||
// Save before publication, including the transactional Meta mint path.
|
||||
// Runtime-only changes made during I/O are merged below, never overwritten.
|
||||
existingBeforeSave := existing.Clone()
|
||||
persistMetaMint := (mode == updateModePrepare || mode == updateModeRefresh) && strings.EqualFold(strings.TrimSpace(auth.Provider), "meta")
|
||||
if persistMetaMint {
|
||||
if errPersist := m.persist(ctx, auth); errPersist != nil {
|
||||
if errPersist := m.persistLocked(ctx, auth); errPersist != nil {
|
||||
// A minted Meta key must reach the store before requests can use it.
|
||||
if persistMetaMint {
|
||||
m.mu.Unlock()
|
||||
return nil, fmt.Errorf("persist meta auth: %w", errPersist)
|
||||
}
|
||||
// Ordinary persistence failures remain non-fatal, but never silent.
|
||||
log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist updated auth %s (%s): %v", auth.Provider, auth.ID, errPersist)
|
||||
}
|
||||
if m.authEpochs[auth.ID] != auth.RegistrationEpoch {
|
||||
m.mu.Unlock()
|
||||
return nil, fmt.Errorf("auth %s changed registration during persistence", auth.ID)
|
||||
}
|
||||
mergeAuthSaveDelta(auth, existingBeforeSave, existing, false)
|
||||
if existing.Generation >= auth.Generation {
|
||||
auth.Generation = existing.Generation + 1
|
||||
}
|
||||
authClone := auth.Clone()
|
||||
m.auths[auth.ID] = authClone
|
||||
@@ -279,6 +307,7 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd
|
||||
schedulerSnapshot = authClone.Clone()
|
||||
}
|
||||
m.mu.Unlock()
|
||||
releaseMutation()
|
||||
if !shouldDeferAPIKeyModelAliasRebuild(ctx) {
|
||||
m.rebuildAPIKeyModelAliasFromRuntimeConfig()
|
||||
}
|
||||
@@ -287,13 +316,6 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd
|
||||
}
|
||||
m.structuralEpoch.Add(1)
|
||||
m.queueRefreshReschedule(auth.ID)
|
||||
if !persistMetaMint {
|
||||
// Persist failures stay non-fatal, but must not be silent: after a token
|
||||
// refresh the rotated credentials only exist in memory until persisted.
|
||||
if errPersist := m.persist(ctx, auth); errPersist != nil {
|
||||
log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist updated auth %s (%s): %v", auth.Provider, auth.ID, errPersist)
|
||||
}
|
||||
}
|
||||
m.hook.OnAuthUpdated(ctx, auth.Clone())
|
||||
if cooldownStateChanged {
|
||||
m.persistCooldownStates(context.Background())
|
||||
@@ -312,6 +334,8 @@ func (m *Manager) Remove(ctx context.Context, id string) {
|
||||
return
|
||||
}
|
||||
_ = ctx
|
||||
releaseMutation := m.lockAuthMutation(id)
|
||||
defer releaseMutation()
|
||||
|
||||
m.mu.Lock()
|
||||
existing := m.auths[id]
|
||||
@@ -343,6 +367,7 @@ func (m *Manager) Remove(ctx context.Context, id string) {
|
||||
m.authEpochs[id]++
|
||||
tombstoneEpoch := m.authEpochs[id]
|
||||
m.mu.Unlock()
|
||||
releaseMutation()
|
||||
|
||||
if !shouldDeferAPIKeyModelAliasRebuild(ctx) {
|
||||
m.rebuildAPIKeyModelAliasFromRuntimeConfig()
|
||||
@@ -376,16 +401,27 @@ func (m *Manager) invalidateSessionAffinity(authID string) {
|
||||
|
||||
// Load resets manager state from the backing store.
|
||||
func (m *Manager) Load(ctx context.Context) error {
|
||||
m.mu.Lock()
|
||||
if m.store == nil {
|
||||
m.mu.Unlock()
|
||||
// Exclude complete Save/publication transactions, including new registrations
|
||||
// not yet present in auths. Never wait for this barrier while holding m.mu.
|
||||
lockCtx := ctx
|
||||
if lockCtx == nil {
|
||||
lockCtx = context.Background()
|
||||
}
|
||||
if errAcquire := m.authLoadGate.Acquire(lockCtx, math.MaxInt64); errAcquire != nil {
|
||||
return errAcquire
|
||||
}
|
||||
defer m.authLoadGate.Release(math.MaxInt64)
|
||||
m.mu.RLock()
|
||||
store := m.store
|
||||
m.mu.RUnlock()
|
||||
if store == nil {
|
||||
return nil
|
||||
}
|
||||
items, err := m.store.List(ctx)
|
||||
items, err := store.List(ctx)
|
||||
if err != nil {
|
||||
m.mu.Unlock()
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
previousAuths := m.auths
|
||||
m.auths = make(map[string]*Auth, len(items))
|
||||
if m.authEpochs == nil {
|
||||
@@ -448,7 +484,10 @@ type authPersistLock struct {
|
||||
}
|
||||
|
||||
func (m *Manager) persist(ctx context.Context, auth *Auth) error {
|
||||
if m.store == nil || auth == nil {
|
||||
m.mu.RLock()
|
||||
store := m.store
|
||||
m.mu.RUnlock()
|
||||
if store == nil || auth == nil {
|
||||
return nil
|
||||
}
|
||||
if errWeight := ValidateAuthWeight(auth); errWeight != nil {
|
||||
@@ -483,13 +522,13 @@ func (m *Manager) persist(ctx context.Context, auth *Auth) error {
|
||||
if shouldSkipPersist(ctx) {
|
||||
return nil
|
||||
}
|
||||
_, err := m.store.Save(ctx, auth)
|
||||
_, err := store.Save(ctx, auth)
|
||||
return err
|
||||
}
|
||||
|
||||
if shouldSkipPersist(ctx) {
|
||||
return nil
|
||||
}
|
||||
_, err := m.store.Save(ctx, auth)
|
||||
_, err := store.Save(ctx, auth)
|
||||
return err
|
||||
}
|
||||
|
||||
196
sdk/cliproxy/auth/conductor_load_persistence_test.go
Normal file
196
sdk/cliproxy/auth/conductor_load_persistence_test.go
Normal file
@@ -0,0 +1,196 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"reflect"
|
||||
"sync"
|
||||
"testing"
|
||||
"testing/synctest"
|
||||
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor"
|
||||
)
|
||||
|
||||
// reloadAuthStore snapshots records independently of blocked storage I/O.
|
||||
type reloadAuthStore struct {
|
||||
mu sync.Mutex
|
||||
records map[string]*Auth
|
||||
saveEntered chan struct{}
|
||||
saveRelease chan struct{}
|
||||
listEntered chan struct{}
|
||||
listRelease chan struct{}
|
||||
}
|
||||
|
||||
func (*reloadAuthStore) Delete(context.Context, string) error { return nil }
|
||||
func (s *reloadAuthStore) Save(_ context.Context, auth *Auth) (string, error) {
|
||||
if s.saveEntered != nil {
|
||||
close(s.saveEntered)
|
||||
<-s.saveRelease
|
||||
}
|
||||
auth.FileName = auth.ID + ".json"
|
||||
s.mu.Lock()
|
||||
s.records[auth.ID] = auth.Clone()
|
||||
s.mu.Unlock()
|
||||
return auth.FileName, nil
|
||||
}
|
||||
func (s *reloadAuthStore) List(context.Context) ([]*Auth, error) {
|
||||
s.mu.Lock()
|
||||
items := make([]*Auth, 0, len(s.records))
|
||||
for _, auth := range s.records {
|
||||
items = append(items, auth.Clone())
|
||||
}
|
||||
s.mu.Unlock()
|
||||
close(s.listEntered)
|
||||
if s.listRelease != nil {
|
||||
<-s.listRelease
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func assertUnrelatedAuthReadable(t *testing.T, manager *Manager) {
|
||||
t.Helper()
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
if _, ok := manager.GetByID("B"); !ok {
|
||||
done <- errors.New("B disappeared")
|
||||
return
|
||||
}
|
||||
selected, err := manager.SelectAuth(context.Background(), "audit-b", "", cliproxyexecutor.Options{})
|
||||
if err == nil && (selected == nil || selected.ID != "B") {
|
||||
err = errors.New("did not select B")
|
||||
}
|
||||
done <- err
|
||||
}()
|
||||
synctest.Wait()
|
||||
select {
|
||||
case err := <-done:
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
default:
|
||||
t.Error("storage I/O blocked reading/selecting unrelated B")
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerBlockedSaveSerializesLoad(t *testing.T) {
|
||||
for _, operation := range []string{"Register", "ReRegister", "Update", "Prepare", "Refresh", "MarkResult"} {
|
||||
t.Run(operation, func(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
store := &reloadAuthStore{records: make(map[string]*Auth), saveEntered: make(chan struct{}), saveRelease: make(chan struct{}), listEntered: make(chan struct{})}
|
||||
manager := NewManager(store, nil, nil)
|
||||
manager.RegisterExecutor(&aliasRoutingExecutor{id: "audit-b"})
|
||||
b, err := manager.Register(WithSkipPersist(ctx), &Auth{ID: "B", Provider: "audit-b", Status: StatusActive})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
store.records[b.ID] = b.Clone()
|
||||
base := &Auth{ID: "A", Provider: "meta", Status: StatusActive, Metadata: map[string]any{"access_token": "old"}}
|
||||
if operation != "Register" {
|
||||
base, err = manager.Register(WithSkipPersist(ctx), base)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
store.records[base.ID] = base.Clone()
|
||||
}
|
||||
updated := base.Clone()
|
||||
updated.Metadata["access_token"] = "minted"
|
||||
saveDone := make(chan error, 1)
|
||||
go func() {
|
||||
var errSave error
|
||||
switch operation {
|
||||
case "Register", "ReRegister":
|
||||
_, errSave = manager.Register(ctx, updated)
|
||||
case "Update":
|
||||
_, errSave = manager.Update(ctx, updated)
|
||||
case "Prepare":
|
||||
_, errSave = manager.UpdatePreparedAuth(ctx, base, updated)
|
||||
case "Refresh":
|
||||
_, errSave = manager.UpdateRefreshedAuth(ctx, base, updated)
|
||||
case "MarkResult":
|
||||
manager.MarkResult(ctx, Result{AuthID: base.ID, Provider: base.Provider, Success: true})
|
||||
}
|
||||
saveDone <- errSave
|
||||
}()
|
||||
<-store.saveEntered
|
||||
loadDone := make(chan error, 1)
|
||||
go func() { loadDone <- manager.Load(ctx) }()
|
||||
synctest.Wait()
|
||||
select {
|
||||
case <-store.listEntered:
|
||||
t.Error("Load read stale records before Save/publication completed")
|
||||
default:
|
||||
}
|
||||
assertUnrelatedAuthReadable(t, manager)
|
||||
close(store.saveRelease)
|
||||
if errSave := <-saveDone; errSave != nil {
|
||||
t.Errorf("save/publication failed: %v", errSave)
|
||||
}
|
||||
if errLoad := <-loadDone; errLoad != nil {
|
||||
t.Fatal(errLoad)
|
||||
}
|
||||
got, ok := manager.GetByID(base.ID)
|
||||
if !ok {
|
||||
t.Fatal("Load lost the saved credential")
|
||||
}
|
||||
store.mu.Lock()
|
||||
saved := store.records[base.ID].Clone()
|
||||
store.mu.Unlock()
|
||||
if got.RegistrationEpoch <= saved.RegistrationEpoch {
|
||||
t.Error("Load did not run after save/publication")
|
||||
}
|
||||
saved.RegistrationEpoch, saved.Generation = got.RegistrationEpoch, got.Generation
|
||||
if !reflect.DeepEqual(saved, got) {
|
||||
t.Errorf("store and manager diverged: saved=%+v got=%+v", saved, got)
|
||||
}
|
||||
manager.scheduler.mu.Lock()
|
||||
scheduled := manager.scheduler.providers["meta"].auths[base.ID].auth.Clone()
|
||||
manager.scheduler.mu.Unlock()
|
||||
if !reflect.DeepEqual(got, scheduled) {
|
||||
t.Error("scheduler differs from reloaded credential")
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerBlockedLoadSerializesNewRegistration(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
store := &reloadAuthStore{records: make(map[string]*Auth), saveEntered: make(chan struct{}), saveRelease: make(chan struct{}), listEntered: make(chan struct{}), listRelease: make(chan struct{})}
|
||||
manager := NewManager(store, nil, nil)
|
||||
manager.RegisterExecutor(&aliasRoutingExecutor{id: "audit-b"})
|
||||
b, err := manager.Register(WithSkipPersist(ctx), &Auth{ID: "B", Provider: "audit-b", Status: StatusActive})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
store.records[b.ID] = b.Clone()
|
||||
loadDone := make(chan error, 1)
|
||||
go func() { loadDone <- manager.Load(ctx) }()
|
||||
<-store.listEntered
|
||||
saveDone := make(chan error, 1)
|
||||
go func() {
|
||||
_, errRegister := manager.Register(ctx, &Auth{ID: "new", Provider: "audit", Metadata: map[string]any{"type": "audit"}})
|
||||
saveDone <- errRegister
|
||||
}()
|
||||
synctest.Wait()
|
||||
select {
|
||||
case <-store.saveEntered:
|
||||
t.Error("new registration overtook Load's whole-map replacement")
|
||||
default:
|
||||
}
|
||||
assertUnrelatedAuthReadable(t, manager)
|
||||
close(store.listRelease)
|
||||
if errLoad := <-loadDone; errLoad != nil {
|
||||
t.Fatal(errLoad)
|
||||
}
|
||||
<-store.saveEntered
|
||||
close(store.saveRelease)
|
||||
if errSave := <-saveDone; errSave != nil {
|
||||
t.Fatal(errSave)
|
||||
}
|
||||
if got, ok := manager.GetByID("new"); !ok || got.FileName != "new.json" {
|
||||
t.Fatal("Load discarded a new registration")
|
||||
}
|
||||
})
|
||||
}
|
||||
113
sdk/cliproxy/auth/conductor_persistence.go
Normal file
113
sdk/cliproxy/auth/conductor_persistence.go
Normal file
@@ -0,0 +1,113 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"reflect"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// lockAuthMutation acquires the reload barrier, then the credential gate, before
|
||||
// m.mu. Keep both until publication completes. Load never takes credential or
|
||||
// refresh locks. Store callbacks may read the manager, but must not mutate it.
|
||||
func (m *Manager) lockAuthMutation(id string) func() {
|
||||
release, _ := m.lockAuthMutationContext(context.Background(), id)
|
||||
return release
|
||||
}
|
||||
|
||||
// lockAuthMutationContext lets fallible mutations cancel either lock wait.
|
||||
func (m *Manager) lockAuthMutationContext(ctx context.Context, id string) (func(), error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if errAcquire := m.authLoadGate.Acquire(ctx, 1); errAcquire != nil {
|
||||
return nil, errAcquire
|
||||
}
|
||||
value, _ := m.authMutationLocks.LoadOrStore(id, make(chan struct{}, 1))
|
||||
gate := value.(chan struct{})
|
||||
select {
|
||||
case gate <- struct{}{}:
|
||||
// Cancellation may race with the gate becoming available.
|
||||
if errContext := ctx.Err(); errContext != nil {
|
||||
<-gate
|
||||
m.authLoadGate.Release(1)
|
||||
return nil, errContext
|
||||
}
|
||||
case <-ctx.Done():
|
||||
m.authLoadGate.Release(1)
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
var once sync.Once
|
||||
return func() {
|
||||
once.Do(func() {
|
||||
<-gate
|
||||
m.authLoadGate.Release(1)
|
||||
})
|
||||
}, nil
|
||||
}
|
||||
|
||||
// persistLocked requires m.mu and the credential mutation gate. Never pass a
|
||||
// published object to Store.Save: readers and runtime-only writers remain live
|
||||
// during I/O. Only the store's delta is merged back, preserving intervening edits.
|
||||
func (m *Manager) persistLocked(ctx context.Context, auth *Auth) error {
|
||||
before := auth.Clone()
|
||||
saved := auth.Clone()
|
||||
m.mu.Unlock()
|
||||
err := m.persist(ctx, saved)
|
||||
m.mu.Lock()
|
||||
mergeAuthSaveDelta(auth, before, saved, true)
|
||||
return err
|
||||
}
|
||||
|
||||
// mergeAuthSaveDelta handles all exported value fields, not a store-specific
|
||||
// allowlist. Map entries are independent fields; nested values are atomic.
|
||||
// Identity/version fields belong exclusively to the manager. Store-owned shared
|
||||
// objects must be replaced rather than mutated in place (see Store.Save).
|
||||
func mergeAuthSaveDelta(target, before, after *Auth, preserveConcurrent bool) {
|
||||
dst := reflect.ValueOf(target).Elem()
|
||||
base := reflect.ValueOf(before).Elem()
|
||||
next := reflect.ValueOf(after).Elem()
|
||||
for i := 0; i < dst.NumField(); i++ {
|
||||
field := dst.Type().Field(i)
|
||||
if field.PkgPath != "" || field.Name == "ID" || field.Name == "RegistrationEpoch" || field.Name == "Generation" {
|
||||
continue
|
||||
}
|
||||
d, b, n := dst.Field(i), base.Field(i), next.Field(i)
|
||||
if reflect.DeepEqual(b.Interface(), n.Interface()) {
|
||||
continue
|
||||
}
|
||||
if reflect.DeepEqual(d.Interface(), b.Interface()) {
|
||||
d.Set(n)
|
||||
continue
|
||||
}
|
||||
if d.Kind() == reflect.Map {
|
||||
keys := make(map[any]reflect.Value)
|
||||
for _, key := range b.MapKeys() {
|
||||
keys[key.Interface()] = key
|
||||
}
|
||||
for _, key := range n.MapKeys() {
|
||||
keys[key.Interface()] = key
|
||||
}
|
||||
for _, key := range keys {
|
||||
oldValue, newValue := b.MapIndex(key), n.MapIndex(key)
|
||||
if equalSaveValue(oldValue, newValue) || (preserveConcurrent && !equalSaveValue(d.MapIndex(key), oldValue)) {
|
||||
continue
|
||||
}
|
||||
if d.IsNil() {
|
||||
d.Set(reflect.MakeMap(d.Type()))
|
||||
}
|
||||
d.SetMapIndex(key, newValue)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !preserveConcurrent {
|
||||
d.Set(n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func equalSaveValue(a, b reflect.Value) bool {
|
||||
if !a.IsValid() || !b.IsValid() {
|
||||
return a.IsValid() == b.IsValid()
|
||||
}
|
||||
return reflect.DeepEqual(a.Interface(), b.Interface())
|
||||
}
|
||||
@@ -786,7 +786,9 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc
|
||||
if m.shouldRefresh(updated, now) {
|
||||
updated.NextRefreshAfter = now.Add(refreshIneffectiveBackoff)
|
||||
}
|
||||
saved, errUpdate := m.UpdateRefreshedAuth(ctx, base, updated)
|
||||
// A successful refresh may have rotated upstream tokens after cancellation.
|
||||
// Commit its result with context values intact; keep epoch checks and save ordering.
|
||||
saved, errUpdate := m.UpdateRefreshedAuth(context.WithoutCancel(ctx), base, updated)
|
||||
if errUpdate != nil {
|
||||
// Warn, not debug: a restart after a lost persist fails with invalid_grant.
|
||||
log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("persist refreshed auth %s (%s) failed: %v", auth.Provider, auth.ID, errUpdate)
|
||||
|
||||
521
sdk/cliproxy/auth/conductor_saved_fields_test.go
Normal file
521
sdk/cliproxy/auth/conductor_saved_fields_test.go
Normal file
@@ -0,0 +1,521 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"reflect"
|
||||
"sync"
|
||||
"testing"
|
||||
"testing/synctest"
|
||||
"time"
|
||||
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor"
|
||||
)
|
||||
|
||||
type enrichingAuthStore struct{ err error }
|
||||
|
||||
func (*enrichingAuthStore) List(context.Context) ([]*Auth, error) { return nil, nil }
|
||||
func (*enrichingAuthStore) Delete(context.Context, string) error { return nil }
|
||||
func (s *enrichingAuthStore) Save(_ context.Context, auth *Auth) (string, error) {
|
||||
if auth.Attributes == nil {
|
||||
auth.Attributes = make(map[string]string)
|
||||
}
|
||||
auth.Attributes[AttributePath] = "/fixture/" + auth.ID + ".json"
|
||||
auth.Attributes[AttributeSource] = auth.Attributes[AttributePath]
|
||||
auth.Attributes[AttributeSourceBackend] = AuthSourceFile
|
||||
auth.FileName = auth.ID + ".json"
|
||||
auth.Metadata["disabled"] = auth.Disabled
|
||||
return auth.Attributes[AttributePath], s.err
|
||||
}
|
||||
|
||||
type savedFieldsHook struct{ auth *Auth }
|
||||
|
||||
func (h *savedFieldsHook) OnAuthRegistered(_ context.Context, auth *Auth) { h.auth = auth.Clone() }
|
||||
func (h *savedFieldsHook) OnAuthUpdated(_ context.Context, auth *Auth) { h.auth = auth.Clone() }
|
||||
func (*savedFieldsHook) OnResult(context.Context, Result) {}
|
||||
|
||||
func TestManagerSavedFieldsPublished(t *testing.T) {
|
||||
for _, saveErr := range []error{nil, errors.New("save failed after enrichment")} {
|
||||
name := "success"
|
||||
if saveErr != nil {
|
||||
name = "failure"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
testManagerSavedFieldsPublished(t, saveErr)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func testManagerSavedFieldsPublished(t *testing.T, saveErr error) {
|
||||
for _, operation := range []string{"Register", "Update", "Prepare", "Refresh"} {
|
||||
t.Run(operation, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
hook := &savedFieldsHook{}
|
||||
manager := NewManager(&enrichingAuthStore{err: saveErr}, nil, hook)
|
||||
input := &Auth{ID: "saved-fields", Provider: "audit", Status: StatusActive, Metadata: map[string]any{"type": "audit"}}
|
||||
var returned *Auth
|
||||
var err error
|
||||
if operation == "Register" {
|
||||
returned, err = manager.Register(ctx, input)
|
||||
} else {
|
||||
base, errRegister := manager.Register(ctx, input)
|
||||
if errRegister != nil {
|
||||
t.Fatal(errRegister)
|
||||
}
|
||||
input = base.Clone()
|
||||
input.FileName = ""
|
||||
input.Attributes = nil
|
||||
delete(input.Metadata, "disabled")
|
||||
switch operation {
|
||||
case "Update":
|
||||
returned, err = manager.Update(ctx, input)
|
||||
case "Prepare":
|
||||
returned, err = manager.UpdatePreparedAuth(ctx, base, input)
|
||||
case "Refresh":
|
||||
returned, err = manager.UpdateRefreshedAuth(ctx, base, input)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
published, ok := manager.GetByID(input.ID)
|
||||
if !ok {
|
||||
t.Fatal("missing published auth")
|
||||
}
|
||||
for name, snapshot := range map[string]*Auth{"manager": published, "return": returned, "hook": hook.auth} {
|
||||
if snapshot == nil || snapshot.FileName != "saved-fields.json" || snapshot.Attributes[AttributePath] != "/fixture/saved-fields.json" || snapshot.Attributes[AttributeSource] != "/fixture/saved-fields.json" || snapshot.Attributes[AttributeSourceBackend] != AuthSourceFile || snapshot.Metadata["disabled"] != false {
|
||||
t.Errorf("%s snapshot missing Store.Save enrichment: %+v", name, snapshot)
|
||||
}
|
||||
if !reflect.DeepEqual(snapshot, published) {
|
||||
t.Errorf("%s snapshot differs from manager", name)
|
||||
}
|
||||
}
|
||||
manager.scheduler.mu.Lock()
|
||||
scheduled := manager.scheduler.providers["audit"].auths[input.ID].auth.Clone()
|
||||
manager.scheduler.mu.Unlock()
|
||||
if !reflect.DeepEqual(scheduled, published) {
|
||||
t.Error("scheduler snapshot differs from manager")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// blockingEnrichingAuthStore pauses one save before enriching the auth.
|
||||
type blockingEnrichingAuthStore struct {
|
||||
enrichingAuthStore
|
||||
entered chan struct{}
|
||||
release chan struct{}
|
||||
once sync.Once
|
||||
mu sync.Mutex
|
||||
saved []*Auth
|
||||
}
|
||||
|
||||
func (s *blockingEnrichingAuthStore) Save(ctx context.Context, auth *Auth) (string, error) {
|
||||
s.once.Do(func() {
|
||||
close(s.entered)
|
||||
<-s.release
|
||||
})
|
||||
path, err := s.enrichingAuthStore.Save(ctx, auth)
|
||||
s.mu.Lock()
|
||||
s.saved = append(s.saved, auth.Clone())
|
||||
s.mu.Unlock()
|
||||
return path, err
|
||||
}
|
||||
|
||||
func TestManagerSavedFieldsConcurrentLifecycle(t *testing.T) {
|
||||
for _, saveErr := range []error{nil, errors.New("save failed after enrichment")} {
|
||||
name := "success"
|
||||
if saveErr != nil {
|
||||
name = "failure"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) { testManagerSavedFieldsConcurrentLifecycle(t, saveErr) })
|
||||
}
|
||||
}
|
||||
|
||||
func testManagerSavedFieldsConcurrentLifecycle(t *testing.T, saveErr error) {
|
||||
for _, operation := range []string{"Register", "Update"} {
|
||||
for _, concurrent := range []string{"MarkResult", "Update", "Remove", "ReRegister"} {
|
||||
t.Run(operation+"/"+concurrent, func(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
store := &blockingEnrichingAuthStore{enrichingAuthStore: enrichingAuthStore{err: saveErr}, entered: make(chan struct{}), release: make(chan struct{})}
|
||||
var releaseOnce sync.Once
|
||||
release := func() { releaseOnce.Do(func() { close(store.release) }) }
|
||||
defer release()
|
||||
manager := NewManager(store, nil, nil)
|
||||
input := &Auth{ID: "saved-fields", Provider: "audit", Status: StatusActive, Label: "old", Metadata: map[string]any{"type": "audit"}}
|
||||
if operation == "Update" {
|
||||
var errRegister error
|
||||
input, errRegister = manager.Register(WithSkipPersist(ctx), input)
|
||||
if errRegister != nil {
|
||||
t.Fatal(errRegister)
|
||||
}
|
||||
}
|
||||
next := input.Clone()
|
||||
next.Label = "new"
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
var err error
|
||||
if operation == "Register" {
|
||||
_, err = manager.Register(ctx, input)
|
||||
} else {
|
||||
_, err = manager.Update(ctx, input)
|
||||
}
|
||||
done <- err
|
||||
}()
|
||||
<-store.entered
|
||||
started := make(chan struct{})
|
||||
concurrentDone := make(chan error, 1)
|
||||
go func() {
|
||||
close(started)
|
||||
var err error
|
||||
switch concurrent {
|
||||
case "MarkResult":
|
||||
manager.MarkResult(ctx, Result{AuthID: next.ID, Provider: next.Provider, Model: "model", Error: &Error{HTTPStatus: 429, Message: "quota exceeded"}})
|
||||
case "Update":
|
||||
_, err = manager.Update(ctx, next)
|
||||
case "Remove":
|
||||
manager.Remove(ctx, next.ID)
|
||||
case "ReRegister":
|
||||
manager.Remove(ctx, next.ID)
|
||||
_, err = manager.Register(ctx, next)
|
||||
}
|
||||
concurrentDone <- err
|
||||
}()
|
||||
<-started
|
||||
synctest.Wait()
|
||||
select {
|
||||
case <-concurrentDone:
|
||||
t.Fatal("same-credential mutation completed before the active save")
|
||||
default:
|
||||
}
|
||||
release()
|
||||
if errDone := <-done; errDone != nil {
|
||||
t.Fatal(errDone)
|
||||
}
|
||||
if errConcurrent := <-concurrentDone; errConcurrent != nil {
|
||||
t.Fatal(errConcurrent)
|
||||
}
|
||||
published, ok := manager.GetByID(input.ID)
|
||||
manager.scheduler.mu.Lock()
|
||||
var scheduled *Auth
|
||||
if provider := manager.scheduler.providers["audit"]; provider != nil {
|
||||
if meta := provider.auths[input.ID]; meta != nil {
|
||||
scheduled = meta.auth.Clone()
|
||||
}
|
||||
}
|
||||
manager.scheduler.mu.Unlock()
|
||||
if concurrent == "Remove" {
|
||||
if ok || scheduled != nil {
|
||||
t.Fatal("completed save resurrected removed auth")
|
||||
}
|
||||
return
|
||||
}
|
||||
if !ok || published.FileName != "saved-fields.json" || published.Attributes[AttributePath] != "/fixture/saved-fields.json" {
|
||||
t.Fatalf("missing saved fields: %+v", published)
|
||||
}
|
||||
if concurrent == "MarkResult" {
|
||||
state := published.ModelStates["model"]
|
||||
if published.Failed != 1 || published.LastError == nil || state == nil || !state.Quota.Exceeded || !state.Unavailable {
|
||||
t.Fatalf("save overwrote concurrent result: %+v", published)
|
||||
}
|
||||
} else if published.Label != "new" {
|
||||
t.Fatalf("save overwrote newer lifecycle state: %+v", published)
|
||||
}
|
||||
if !reflect.DeepEqual(scheduled, published) {
|
||||
t.Fatal("scheduler retained an obsolete save snapshot")
|
||||
}
|
||||
store.mu.Lock()
|
||||
saved := store.saved[len(store.saved)-1].Clone()
|
||||
store.mu.Unlock()
|
||||
if !reflect.DeepEqual(saved, published) {
|
||||
t.Fatal("store retained an obsolete auth version")
|
||||
}
|
||||
if concurrent == "ReRegister" {
|
||||
if published.RegistrationEpoch <= input.RegistrationEpoch {
|
||||
t.Fatal("re-registration did not advance epoch")
|
||||
}
|
||||
if _, errUpdate := manager.UpdateRefreshedAuth(ctx, input, input.Clone()); errUpdate == nil {
|
||||
t.Fatal("accepted stale registration after re-register")
|
||||
}
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerBlockedSaveAllowsUnrelatedCredential(t *testing.T) {
|
||||
for _, operation := range []string{"Register", "Update"} {
|
||||
t.Run(operation, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
store := &blockingEnrichingAuthStore{entered: make(chan struct{}), release: make(chan struct{})}
|
||||
manager := NewManager(store, nil, nil)
|
||||
manager.RegisterExecutor(&aliasRoutingExecutor{id: "audit"})
|
||||
if _, err := manager.Register(WithSkipPersist(ctx), &Auth{ID: "B", Provider: "audit", Status: StatusActive}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
input := &Auth{ID: "A", Provider: "audit-a", Status: StatusActive, Metadata: map[string]any{"type": "audit-a"}}
|
||||
if operation == "Update" {
|
||||
if _, err := manager.Register(WithSkipPersist(ctx), input); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
if operation == "Register" {
|
||||
_, _ = manager.Register(ctx, input)
|
||||
} else {
|
||||
_, _ = manager.Update(ctx, input)
|
||||
}
|
||||
}()
|
||||
<-store.entered
|
||||
readDone := make(chan error, 1)
|
||||
go func() {
|
||||
if _, ok := manager.GetByID("B"); !ok {
|
||||
readDone <- errors.New("missing B")
|
||||
return
|
||||
}
|
||||
selected, err := manager.SelectAuth(ctx, "audit", "", cliproxyexecutor.Options{})
|
||||
if err == nil && (selected == nil || selected.ID != "B") {
|
||||
err = errors.New("did not select B")
|
||||
}
|
||||
readDone <- err
|
||||
}()
|
||||
select {
|
||||
case err := <-readDone:
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Error("blocked Save(A) prevented reading/selecting unrelated B")
|
||||
}
|
||||
close(store.release)
|
||||
<-done
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type callbackAuthStore struct{ save func(*Auth) error }
|
||||
|
||||
func (*callbackAuthStore) List(context.Context) ([]*Auth, error) { return nil, nil }
|
||||
func (*callbackAuthStore) Delete(context.Context, string) error { return nil }
|
||||
func (s *callbackAuthStore) Save(_ context.Context, auth *Auth) (string, error) {
|
||||
return "", s.save(auth)
|
||||
}
|
||||
|
||||
func TestManagerSaveReadCallbackAndCustomFields(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
store := &callbackAuthStore{}
|
||||
manager := NewManager(store, nil, nil)
|
||||
store.save = func(auth *Auth) error {
|
||||
manager.GetByID(auth.ID)
|
||||
manager.List()
|
||||
auth.Prefix = "custom-prefix"
|
||||
auth.ProxyURL = "http://custom.invalid"
|
||||
auth.Metadata["custom"] = map[string]any{"value": "saved"}
|
||||
delete(auth.Metadata, "remove")
|
||||
return nil
|
||||
}
|
||||
for _, operation := range []string{"Register", "Update"} {
|
||||
input := &Auth{ID: "callback", Provider: "audit", Status: StatusActive, Metadata: map[string]any{"remove": true}}
|
||||
var result *Auth
|
||||
var err error
|
||||
if operation == "Register" {
|
||||
result, err = manager.Register(ctx, input)
|
||||
} else {
|
||||
result, err = manager.Update(ctx, input)
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, _ := manager.GetByID(input.ID)
|
||||
if got.Prefix != "custom-prefix" || got.ProxyURL != "http://custom.invalid" || got.Metadata["custom"] == nil {
|
||||
t.Fatalf("custom store fields missing: %+v", got)
|
||||
}
|
||||
if _, ok := got.Metadata["remove"]; ok {
|
||||
t.Fatal("store deletion lost")
|
||||
}
|
||||
if !reflect.DeepEqual(result, got) {
|
||||
t.Fatal("return differs from manager")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestManagerSavePreservesConcurrentRefreshState(t *testing.T) {
|
||||
for _, operation := range []string{"Update", "MarkResult"} {
|
||||
t.Run(operation, func(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
store := &blockingEnrichingAuthStore{entered: make(chan struct{}), release: make(chan struct{})}
|
||||
manager := NewManager(store, nil, nil)
|
||||
base, err := manager.Register(WithSkipPersist(ctx), &Auth{ID: "runtime", Provider: "audit", Status: StatusActive, Metadata: map[string]any{"type": "audit"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
if operation == "Update" {
|
||||
_, _ = manager.Update(ctx, base.Clone())
|
||||
} else {
|
||||
manager.MarkResult(ctx, Result{AuthID: base.ID, Provider: base.Provider, Success: true})
|
||||
}
|
||||
}()
|
||||
<-store.entered
|
||||
job := manager.markRefreshPending(nil, base.ID, base.RegistrationEpoch, time.Now())
|
||||
if job == nil {
|
||||
t.Fatal("failed to schedule concurrent refresh")
|
||||
}
|
||||
close(store.release)
|
||||
<-done
|
||||
got, _ := manager.GetByID(base.ID)
|
||||
if !got.NextRefreshAfter.Equal(job.pendingUntil) {
|
||||
t.Fatal("save lost concurrent refresh scheduling")
|
||||
}
|
||||
if got.Attributes[AttributePath] == "" {
|
||||
t.Fatal("save enrichment lost")
|
||||
}
|
||||
if operation == "MarkResult" && got.Success != 1 {
|
||||
t.Fatal("save lost result counter")
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerSavedFieldsSkipPersistence(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
store := &callbackAuthStore{save: func(*Auth) error { t.Fatal("unexpected Save"); return nil }}
|
||||
for _, mode := range []string{"context", "nil-metadata", "runtime-only", "config-api-key", "plugin-virtual"} {
|
||||
t.Run(mode, func(t *testing.T) {
|
||||
manager := NewManager(store, nil, nil)
|
||||
input := &Auth{ID: mode, Provider: "audit", Status: StatusActive, Metadata: map[string]any{"type": "audit"}}
|
||||
saveCtx := ctx
|
||||
switch mode {
|
||||
case "context":
|
||||
saveCtx = WithSkipPersist(ctx)
|
||||
case "nil-metadata":
|
||||
input.Metadata = nil
|
||||
case "runtime-only":
|
||||
input.Attributes = map[string]string{"runtime_only": "true"}
|
||||
case "config-api-key":
|
||||
input.Attributes = map[string]string{"source": "config", "api_key": "test"}
|
||||
case "plugin-virtual":
|
||||
MarkPluginVirtualAuth(input, "fixture", 0)
|
||||
}
|
||||
for _, operation := range []string{"Register", "Update"} {
|
||||
var result *Auth
|
||||
var err error
|
||||
if operation == "Register" {
|
||||
result, err = manager.Register(saveCtx, input)
|
||||
} else {
|
||||
result, err = manager.Update(saveCtx, input)
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, _ := manager.GetByID(input.ID)
|
||||
if !reflect.DeepEqual(got, result) {
|
||||
t.Fatal("skip return differs from manager")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeAuthSaveDeltaMapChanges(t *testing.T) {
|
||||
t.Run("empty-map", func(t *testing.T) {
|
||||
before := &Auth{}
|
||||
after := &Auth{Attributes: map[string]string{}, Metadata: map[string]any{}}
|
||||
target := before.Clone()
|
||||
mergeAuthSaveDelta(target, before, after, true)
|
||||
if !reflect.DeepEqual(target, after) {
|
||||
t.Fatal("lost explicit empty map enrichment")
|
||||
}
|
||||
})
|
||||
t.Run("delete-map-preserves-concurrent-key", func(t *testing.T) {
|
||||
before := &Auth{Metadata: map[string]any{"old": true}}
|
||||
after := before.Clone()
|
||||
after.Metadata = nil
|
||||
target := before.Clone()
|
||||
target.Metadata["concurrent"] = true
|
||||
mergeAuthSaveDelta(target, before, after, true)
|
||||
if len(target.Metadata) != 1 || target.Metadata["concurrent"] != true {
|
||||
t.Fatalf("map deletion lost: %v", target.Metadata)
|
||||
}
|
||||
})
|
||||
t.Run("conflicts", func(t *testing.T) {
|
||||
before := &Auth{Label: "old", Metadata: map[string]any{"same": "old", "other": "old"}}
|
||||
after := before.Clone()
|
||||
after.Label = "store"
|
||||
after.Metadata["same"] = "store"
|
||||
after.Metadata["other"] = "store"
|
||||
target := before.Clone()
|
||||
target.Label = "runtime"
|
||||
target.Metadata["same"] = "runtime"
|
||||
mergeAuthSaveDelta(target, before, after, true)
|
||||
if target.Label != "runtime" || target.Metadata["same"] != "runtime" || target.Metadata["other"] != "store" {
|
||||
t.Fatalf("incorrect merge: %+v", target)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestManagerSavedFieldsMetaMintPublication(t *testing.T) {
|
||||
for _, operation := range []string{"Prepare", "Refresh"} {
|
||||
for _, fail := range []bool{false, true} {
|
||||
name := operation + "/success"
|
||||
if fail {
|
||||
name = operation + "/failure"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
store := &blockingEnrichingAuthStore{entered: make(chan struct{}), release: make(chan struct{})}
|
||||
if fail {
|
||||
store.err = errors.New("mint save failed")
|
||||
}
|
||||
hook := &savedFieldsHook{}
|
||||
manager := NewManager(store, nil, hook)
|
||||
base, err := manager.Register(WithSkipPersist(ctx), &Auth{ID: "meta-mint", Provider: "meta", Status: StatusActive, Metadata: map[string]any{"access_token": "old"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
updated := base.Clone()
|
||||
updated.Metadata["access_token"] = "minted"
|
||||
done := make(chan struct{})
|
||||
var returned *Auth
|
||||
var errUpdate error
|
||||
go func() {
|
||||
defer close(done)
|
||||
if operation == "Prepare" {
|
||||
returned, errUpdate = manager.UpdatePreparedAuth(ctx, base, updated)
|
||||
} else {
|
||||
returned, errUpdate = manager.UpdateRefreshedAuth(ctx, base, updated)
|
||||
}
|
||||
}()
|
||||
<-store.entered
|
||||
pending, _ := manager.GetByID(base.ID)
|
||||
if !reflect.DeepEqual(base, pending) {
|
||||
t.Fatal("mint published before persistence completed")
|
||||
}
|
||||
close(store.release)
|
||||
<-done
|
||||
got, _ := manager.GetByID(base.ID)
|
||||
if fail {
|
||||
if errUpdate == nil || returned != nil || !reflect.DeepEqual(base, got) || !reflect.DeepEqual(base, hook.auth) {
|
||||
t.Fatal("failed Meta mint changed published state")
|
||||
}
|
||||
} else {
|
||||
if errUpdate != nil || got.Metadata["access_token"] != "minted" || got.FileName == "" || !reflect.DeepEqual(got, returned) || !reflect.DeepEqual(got, hook.auth) {
|
||||
t.Fatalf("successful mint snapshots differ: %v", errUpdate)
|
||||
}
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -213,6 +213,8 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin
|
||||
cooldownStateChanged bool
|
||||
)
|
||||
|
||||
releaseMutation := m.lockAuthMutation(authID)
|
||||
defer releaseMutation()
|
||||
m.mu.Lock()
|
||||
auth, ok := m.auths[authID]
|
||||
if ok && auth != nil {
|
||||
@@ -347,7 +349,7 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin
|
||||
}
|
||||
auth.Generation++
|
||||
auth.UpdatedAt = now
|
||||
if errPersist := m.persist(context.Background(), auth); errPersist != nil {
|
||||
if errPersist := m.persistLocked(context.Background(), auth); errPersist != nil {
|
||||
logEntryWithRequestID(ctx).WithField("auth_id", auth.ID).Warnf("failed to persist auth changes during model state reconciliation: %v", errPersist)
|
||||
}
|
||||
}
|
||||
@@ -361,6 +363,7 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin
|
||||
}
|
||||
}
|
||||
m.mu.Unlock()
|
||||
releaseMutation()
|
||||
|
||||
if snapshot == nil {
|
||||
return
|
||||
|
||||
@@ -175,7 +175,8 @@ func (e *concurrencyTrackingRefreshExecutor) Refresh(_ context.Context, auth *Au
|
||||
if auth.Metadata == nil {
|
||||
auth.Metadata = make(map[string]any)
|
||||
}
|
||||
auth.Metadata["access_token"] = "refreshed-token"
|
||||
auth.Metadata["access_token"] = "rotated-access-" + auth.ID
|
||||
auth.Metadata["refresh_token"] = "rotated-refresh-" + auth.ID
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
@@ -269,10 +270,34 @@ func TestManager_ForceRefreshAll_CanceledContextSkipsExecutors(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestManager_ForceRefreshAll_DynamicCancellationSkipsRemaining(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
// refreshCommitStore observes the commit context without writing credential files.
|
||||
type refreshCommitStore struct {
|
||||
*memoryAuthTestStore
|
||||
checkContext func(context.Context)
|
||||
}
|
||||
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
func (s *refreshCommitStore) Save(ctx context.Context, auth *Auth) (string, error) {
|
||||
s.checkContext(ctx)
|
||||
return s.memoryAuthTestStore.Save(ctx, auth)
|
||||
}
|
||||
|
||||
func TestManager_ForceRefreshAll_DynamicCancellationSkipsRemaining(t *testing.T) {
|
||||
type contextKey struct{}
|
||||
ctx, cancel := context.WithCancel(context.WithValue(context.Background(), contextKey{}, "refresh-value"))
|
||||
defer cancel()
|
||||
|
||||
store := &refreshCommitStore{
|
||||
memoryAuthTestStore: newMemoryAuthTestStore(),
|
||||
checkContext: func(ctx context.Context) {
|
||||
if ctx.Err() != nil || ctx.Done() != nil {
|
||||
t.Error("successful refresh commit must not inherit cancellation")
|
||||
}
|
||||
if ctx.Value(contextKey{}) != "refresh-value" {
|
||||
t.Error("refresh commit lost context value")
|
||||
}
|
||||
},
|
||||
}
|
||||
manager := NewManager(store, &RoundRobinSelector{}, nil)
|
||||
manager.runtimeConfig.Store(&internalconfig.Config{AuthAutoRefreshWorkers: 2})
|
||||
|
||||
executor := &concurrencyTrackingRefreshExecutor{
|
||||
@@ -286,9 +311,9 @@ func TestManager_ForceRefreshAll_DynamicCancellationSkipsRemaining(t *testing.T)
|
||||
auth := &Auth{
|
||||
ID: fmt.Sprintf("ag-%d", i),
|
||||
Provider: "antigravity",
|
||||
Metadata: map[string]any{"refresh_token": fmt.Sprintf("ref-%d", i)},
|
||||
Metadata: map[string]any{"access_token": "old-access", "refresh_token": "old-refresh"},
|
||||
}
|
||||
if _, err := manager.Register(context.Background(), auth); err != nil {
|
||||
if _, err := manager.Register(WithSkipPersist(context.Background()), auth); err != nil {
|
||||
t.Fatalf("register auth: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -315,7 +340,8 @@ func TestManager_ForceRefreshAll_DynamicCancellationSkipsRemaining(t *testing.T)
|
||||
t.Fatalf("expected 6 results, got %d", len(results))
|
||||
}
|
||||
|
||||
// Exactly 2 workers entered Refresh; remaining 4 were canceled in queue
|
||||
// Exactly 2 workers entered Refresh; remaining 4 were canceled in queue.
|
||||
// Successful rotations must be saved and published despite cancellation.
|
||||
if entered := executor.totalEntered.Load(); entered != 2 {
|
||||
t.Fatalf("expected exactly 2 entered calls, got %d", entered)
|
||||
}
|
||||
@@ -323,10 +349,25 @@ func TestManager_ForceRefreshAll_DynamicCancellationSkipsRemaining(t *testing.T)
|
||||
successCount := 0
|
||||
cancelCount := 0
|
||||
for _, res := range results {
|
||||
current, exists := manager.GetByID(res.ID)
|
||||
if !exists {
|
||||
t.Fatalf("auth %s disappeared", res.ID)
|
||||
}
|
||||
store.mu.Lock()
|
||||
persisted := store.auths[res.ID].Clone()
|
||||
store.mu.Unlock()
|
||||
if res.Success {
|
||||
successCount++
|
||||
for name, auth := range map[string]*Auth{"manager": current, "store": persisted} {
|
||||
if auth == nil || auth.Metadata["access_token"] != "rotated-access-"+res.ID || auth.Metadata["refresh_token"] != "rotated-refresh-"+res.ID {
|
||||
t.Errorf("%s missing rotated tokens for %s", name, res.ID)
|
||||
}
|
||||
}
|
||||
} else if res.Error == context.Canceled.Error() {
|
||||
cancelCount++
|
||||
if persisted != nil || current.Metadata["access_token"] != "old-access" || current.Metadata["refresh_token"] != "old-refresh" {
|
||||
t.Errorf("unstarted auth %s was modified or persisted", res.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
if successCount != 2 {
|
||||
|
||||
@@ -7,6 +7,12 @@ type Store interface {
|
||||
// List returns all auth records stored in the backend.
|
||||
List(ctx context.Context) ([]*Auth, error)
|
||||
// Save persists the provided auth record, replacing any existing one with same ID.
|
||||
// It may enrich exported fields, but must not change ID, RegistrationEpoch or
|
||||
// Generation. The manager merges these changes even on non-fatal save errors.
|
||||
// Do not retain auth or mutate shared nested values in place: replace metadata
|
||||
// values, pointers, Storage and Runtime instead. Auth.Clone shares these objects.
|
||||
// Save may call manager read APIs, but must not synchronously mutate the
|
||||
// manager or call Load (the reload barrier remains held until publication).
|
||||
Save(ctx context.Context, auth *Auth) (string, error)
|
||||
// Delete removes the auth record identified by id.
|
||||
Delete(ctx context.Context, id string) error
|
||||
|
||||
Reference in New Issue
Block a user