feat(auth): add support for persisting disabled flag in token storage
- Updated `FileTokenStore` and related stores (`objectstore`, `gitstore`, `postgresstore`) to include the `disabled` flag in metadata for token storage. - Adjusted `Auth` metadata handling to initialize empty maps when absent. - Refined logic in `auto_refresh_loop` and `conductor` to exclude `disabled` tokens from refresh checks. - Added comprehensive unit tests to verify proper handling of the `disabled` flag in storage and retrieval operations.
This commit is contained in:
@@ -287,10 +287,18 @@ func (s *GitTokenStore) Save(_ context.Context, auth *cliproxyauth.Auth) (string
|
|||||||
|
|
||||||
switch {
|
switch {
|
||||||
case auth.Storage != nil:
|
case auth.Storage != nil:
|
||||||
|
if auth.Metadata == nil {
|
||||||
|
auth.Metadata = make(map[string]any)
|
||||||
|
}
|
||||||
|
auth.Metadata["disabled"] = auth.Disabled
|
||||||
|
if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok {
|
||||||
|
setter.SetMetadata(auth.Metadata)
|
||||||
|
}
|
||||||
if err = auth.Storage.SaveTokenToFile(path); err != nil {
|
if err = auth.Storage.SaveTokenToFile(path); err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
case auth.Metadata != nil:
|
case auth.Metadata != nil:
|
||||||
|
auth.Metadata["disabled"] = auth.Disabled
|
||||||
raw, errMarshal := json.Marshal(auth.Metadata)
|
raw, errMarshal := json.Marshal(auth.Metadata)
|
||||||
if errMarshal != nil {
|
if errMarshal != nil {
|
||||||
return "", fmt.Errorf("auth filestore: marshal metadata failed: %w", errMarshal)
|
return "", fmt.Errorf("auth filestore: marshal metadata failed: %w", errMarshal)
|
||||||
|
|||||||
@@ -184,10 +184,18 @@ func (s *ObjectTokenStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (s
|
|||||||
|
|
||||||
switch {
|
switch {
|
||||||
case auth.Storage != nil:
|
case auth.Storage != nil:
|
||||||
|
if auth.Metadata == nil {
|
||||||
|
auth.Metadata = make(map[string]any)
|
||||||
|
}
|
||||||
|
auth.Metadata["disabled"] = auth.Disabled
|
||||||
|
if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok {
|
||||||
|
setter.SetMetadata(auth.Metadata)
|
||||||
|
}
|
||||||
if err = auth.Storage.SaveTokenToFile(path); err != nil {
|
if err = auth.Storage.SaveTokenToFile(path); err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
case auth.Metadata != nil:
|
case auth.Metadata != nil:
|
||||||
|
auth.Metadata["disabled"] = auth.Disabled
|
||||||
raw, errMarshal := json.Marshal(auth.Metadata)
|
raw, errMarshal := json.Marshal(auth.Metadata)
|
||||||
if errMarshal != nil {
|
if errMarshal != nil {
|
||||||
return "", fmt.Errorf("object store: marshal metadata: %w", errMarshal)
|
return "", fmt.Errorf("object store: marshal metadata: %w", errMarshal)
|
||||||
|
|||||||
@@ -214,10 +214,18 @@ func (s *PostgresStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri
|
|||||||
|
|
||||||
switch {
|
switch {
|
||||||
case auth.Storage != nil:
|
case auth.Storage != nil:
|
||||||
|
if auth.Metadata == nil {
|
||||||
|
auth.Metadata = make(map[string]any)
|
||||||
|
}
|
||||||
|
auth.Metadata["disabled"] = auth.Disabled
|
||||||
|
if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok {
|
||||||
|
setter.SetMetadata(auth.Metadata)
|
||||||
|
}
|
||||||
if err = auth.Storage.SaveTokenToFile(path); err != nil {
|
if err = auth.Storage.SaveTokenToFile(path); err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
case auth.Metadata != nil:
|
case auth.Metadata != nil:
|
||||||
|
auth.Metadata["disabled"] = auth.Disabled
|
||||||
raw, errMarshal := json.Marshal(auth.Metadata)
|
raw, errMarshal := json.Marshal(auth.Metadata)
|
||||||
if errMarshal != nil {
|
if errMarshal != nil {
|
||||||
return "", fmt.Errorf("postgres store: marshal metadata: %w", errMarshal)
|
return "", fmt.Errorf("postgres store: marshal metadata: %w", errMarshal)
|
||||||
|
|||||||
@@ -72,6 +72,10 @@ func (s *FileTokenStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (str
|
|||||||
|
|
||||||
switch {
|
switch {
|
||||||
case auth.Storage != nil:
|
case auth.Storage != nil:
|
||||||
|
if auth.Metadata == nil {
|
||||||
|
auth.Metadata = make(map[string]any)
|
||||||
|
}
|
||||||
|
auth.Metadata["disabled"] = auth.Disabled
|
||||||
if setter, ok := auth.Storage.(metadataSetter); ok {
|
if setter, ok := auth.Storage.(metadataSetter); ok {
|
||||||
setter.SetMetadata(auth.Metadata)
|
setter.SetMetadata(auth.Metadata)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,64 @@
|
|||||||
|
package auth
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||||
|
)
|
||||||
|
|
||||||
|
type testTokenStorage struct {
|
||||||
|
meta map[string]any
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *testTokenStorage) SetMetadata(meta map[string]any) { s.meta = meta }
|
||||||
|
|
||||||
|
func (s *testTokenStorage) SaveTokenToFile(authFilePath string) error {
|
||||||
|
raw, err := json.Marshal(s.meta)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return os.WriteFile(authFilePath, raw, 0o600)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileTokenStore_Save_DisabledPersistsFlagForTokenStorage(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
baseDir := t.TempDir()
|
||||||
|
path := filepath.Join(baseDir, "disabled.json")
|
||||||
|
|
||||||
|
if err := os.WriteFile(path, []byte(`{"type":"test","disabled":true}`), 0o600); err != nil {
|
||||||
|
t.Fatalf("seed auth file: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
store := NewFileTokenStore()
|
||||||
|
store.SetBaseDir(baseDir)
|
||||||
|
storage := &testTokenStorage{}
|
||||||
|
|
||||||
|
auth := &cliproxyauth.Auth{
|
||||||
|
ID: "disabled.json",
|
||||||
|
Provider: "test",
|
||||||
|
FileName: "disabled.json",
|
||||||
|
Disabled: true,
|
||||||
|
Storage: storage,
|
||||||
|
Metadata: map[string]any{"type": "test"},
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := store.Save(ctx, auth); err != nil {
|
||||||
|
t.Fatalf("Save() error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
raw, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("read auth file: %v", err)
|
||||||
|
}
|
||||||
|
var meta map[string]any
|
||||||
|
if err := json.Unmarshal(raw, &meta); err != nil {
|
||||||
|
t.Fatalf("unmarshal auth file: %v", err)
|
||||||
|
}
|
||||||
|
if disabled, _ := meta["disabled"].(bool); !disabled {
|
||||||
|
t.Fatalf("disabled=%v, want true (raw=%s)", meta["disabled"], string(raw))
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -336,7 +336,7 @@ func (l *authAutoRefreshLoop) remove(authID string) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func nextRefreshCheckAt(now time.Time, auth *Auth, interval time.Duration) (time.Time, bool) {
|
func nextRefreshCheckAt(now time.Time, auth *Auth, interval time.Duration) (time.Time, bool) {
|
||||||
if auth == nil || auth.Disabled {
|
if auth == nil {
|
||||||
return time.Time{}, false
|
return time.Time{}, false
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -34,9 +34,31 @@ func setRefreshLeadFactory(t *testing.T, provider string, factory func() *time.D
|
|||||||
|
|
||||||
func TestNextRefreshCheckAt_DisabledUnschedule(t *testing.T) {
|
func TestNextRefreshCheckAt_DisabledUnschedule(t *testing.T) {
|
||||||
now := time.Date(2026, 4, 12, 0, 0, 0, 0, time.UTC)
|
now := time.Date(2026, 4, 12, 0, 0, 0, 0, time.UTC)
|
||||||
auth := &Auth{ID: "a1", Provider: "test", Disabled: true}
|
expiry := now.Add(time.Hour)
|
||||||
if _, ok := nextRefreshCheckAt(now, auth, 15*time.Minute); ok {
|
lead := 10 * time.Minute
|
||||||
t.Fatalf("nextRefreshCheckAt() ok = true, want false")
|
setRefreshLeadFactory(t, "disabled-schedule", func() *time.Duration {
|
||||||
|
d := lead
|
||||||
|
return &d
|
||||||
|
})
|
||||||
|
|
||||||
|
auth := &Auth{
|
||||||
|
ID: "a1",
|
||||||
|
Provider: "disabled-schedule",
|
||||||
|
Disabled: true,
|
||||||
|
Status: StatusDisabled,
|
||||||
|
Metadata: map[string]any{
|
||||||
|
"email": "x@example.com",
|
||||||
|
"expires_at": expiry.Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
got, ok := nextRefreshCheckAt(now, auth, 15*time.Minute)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("nextRefreshCheckAt() ok = false, want true")
|
||||||
|
}
|
||||||
|
want := expiry.Add(-lead)
|
||||||
|
if !got.Equal(want) {
|
||||||
|
t.Fatalf("nextRefreshCheckAt() = %s, want %s", got, want)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -3454,7 +3454,7 @@ func (m *Manager) queueRefreshReschedule(authID string) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (m *Manager) shouldRefresh(a *Auth, now time.Time) bool {
|
func (m *Manager) shouldRefresh(a *Auth, now time.Time) bool {
|
||||||
if a == nil || a.Disabled {
|
if a == nil {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
if !a.NextRefreshAfter.IsZero() && now.Before(a.NextRefreshAfter) {
|
if !a.NextRefreshAfter.IsZero() && now.Before(a.NextRefreshAfter) {
|
||||||
@@ -3661,7 +3661,7 @@ func lookupMetadataTime(meta map[string]any, keys ...string) (time.Time, bool) {
|
|||||||
func (m *Manager) markRefreshPending(id string, now time.Time) bool {
|
func (m *Manager) markRefreshPending(id string, now time.Time) bool {
|
||||||
m.mu.Lock()
|
m.mu.Lock()
|
||||||
auth, ok := m.auths[id]
|
auth, ok := m.auths[id]
|
||||||
if !ok || auth == nil || auth.Disabled {
|
if !ok || auth == nil {
|
||||||
m.mu.Unlock()
|
m.mu.Unlock()
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user