refactor(auth): Centralize auth file reading with snapshot preference

The logic for reading authentication files, which includes retries and a preference for cookie snapshot files, was previously implemented locally within the `watcher` package. This was done to handle potential file locks during writes.

This change moves this functionality into a shared `ReadAuthFileWithRetry` function in the `util` package to promote code reuse and consistency.

The `watcher` package is updated to use this new centralized function. Additionally, the initial token loading in the `run` command now also uses this logic, making it more resilient to file access issues and consistent with the watcher's behavior.
This commit is contained in:
hkfires
2025-09-20 00:14:26 +08:00
parent 1d7abc95b8
commit aba719f5fe
3 changed files with 52 additions and 44 deletions

View File

@@ -26,6 +26,7 @@ import (
"github.com/luispater/CLIProxyAPI/v5/internal/config" "github.com/luispater/CLIProxyAPI/v5/internal/config"
"github.com/luispater/CLIProxyAPI/v5/internal/interfaces" "github.com/luispater/CLIProxyAPI/v5/internal/interfaces"
"github.com/luispater/CLIProxyAPI/v5/internal/misc" "github.com/luispater/CLIProxyAPI/v5/internal/misc"
"github.com/luispater/CLIProxyAPI/v5/internal/util"
"github.com/luispater/CLIProxyAPI/v5/internal/watcher" "github.com/luispater/CLIProxyAPI/v5/internal/watcher"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson" "github.com/tidwall/gjson"
@@ -76,7 +77,7 @@ func StartService(cfg *config.Config, configPath string) {
if !info.IsDir() && strings.HasSuffix(info.Name(), ".json") { if !info.IsDir() && strings.HasSuffix(info.Name(), ".json") {
misc.LogCredentialSeparator() misc.LogCredentialSeparator()
log.Debugf("Loading token from: %s", path) log.Debugf("Loading token from: %s", path)
data, errReadFile := os.ReadFile(path) data, errReadFile := util.ReadAuthFilePreferSnapshot(path)
if errReadFile != nil { if errReadFile != nil {
return errReadFile return errReadFile
} }

View File

@@ -6,6 +6,7 @@ import (
"os" "os"
"path/filepath" "path/filepath"
"strings" "strings"
"time"
"github.com/luispater/CLIProxyAPI/v5/internal/misc" "github.com/luispater/CLIProxyAPI/v5/internal/misc"
) )
@@ -93,6 +94,52 @@ func WriteCookieSnapshot(mainPath string, v any) error {
return nil return nil
} }
// ReadAuthFilePreferSnapshot returns the first non-empty auth payload preferring snapshots.
func ReadAuthFilePreferSnapshot(path string) ([]byte, error) {
return ReadAuthFileWithRetry(path, 1, 0)
}
// ReadAuthFileWithRetry attempts to read an auth file multiple times and prefers cookie snapshots.
func ReadAuthFileWithRetry(path string, attempts int, delay time.Duration) ([]byte, error) {
if attempts < 1 {
attempts = 1
}
read := func(target string) ([]byte, error) {
var lastErr error
for i := 0; i < attempts; i++ {
data, err := os.ReadFile(target)
if err == nil {
return data, nil
}
lastErr = err
if i < attempts-1 {
time.Sleep(delay)
}
}
return nil, lastErr
}
candidates := []string{
CookieSnapshotPath(path),
path,
}
for idx, candidate := range candidates {
data, err := read(candidate)
if err == nil {
return data, nil
}
if errors.Is(err, os.ErrNotExist) {
if idx < len(candidates)-1 {
continue
}
}
return nil, err
}
return nil, os.ErrNotExist
}
// RemoveCookieSnapshots removes the snapshot file if it exists. // RemoveCookieSnapshots removes the snapshot file if it exists.
func RemoveCookieSnapshots(mainPath string) { func RemoveCookieSnapshots(mainPath string) {
_ = RemoveFile(CookieSnapshotPath(mainPath)) _ = RemoveFile(CookieSnapshotPath(mainPath))

View File

@@ -9,7 +9,6 @@ import (
"crypto/sha256" "crypto/sha256"
"encoding/hex" "encoding/hex"
"encoding/json" "encoding/json"
"errors"
"io/fs" "io/fs"
"net/http" "net/http"
"os" "os"
@@ -324,7 +323,7 @@ func (w *Watcher) reloadClients() {
// Rebuild auth file hash cache for current clients // Rebuild auth file hash cache for current clients
w.lastAuthHashes = make(map[string]string, len(newFileClients)) w.lastAuthHashes = make(map[string]string, len(newFileClients))
for path := range newFileClients { for path := range newFileClients {
if data, err := readAuthFileWithRetry(path, authFileReadMaxAttempts, authFileReadRetryDelay); err == nil && len(data) > 0 { if data, err := util.ReadAuthFileWithRetry(path, authFileReadMaxAttempts, authFileReadRetryDelay); err == nil && len(data) > 0 {
sum := sha256.Sum256(data) sum := sha256.Sum256(data)
w.lastAuthHashes[path] = hex.EncodeToString(sum[:]) w.lastAuthHashes[path] = hex.EncodeToString(sum[:])
} }
@@ -353,7 +352,7 @@ func (w *Watcher) reloadClients() {
// createClientFromFile creates a single client instance from a given token file path. // createClientFromFile creates a single client instance from a given token file path.
func (w *Watcher) createClientFromFile(path string, cfg *config.Config) (interfaces.Client, error) { func (w *Watcher) createClientFromFile(path string, cfg *config.Config) (interfaces.Client, error) {
data, errReadFile := readAuthFileWithRetry(path, authFileReadMaxAttempts, authFileReadRetryDelay) data, errReadFile := util.ReadAuthFileWithRetry(path, authFileReadMaxAttempts, authFileReadRetryDelay)
if errReadFile != nil { if errReadFile != nil {
return nil, errReadFile return nil, errReadFile
} }
@@ -416,48 +415,9 @@ func (w *Watcher) clientsToSlice(clientMap map[string]interfaces.Client) []inter
return s return s
} }
// readAuthFileWithRetry attempts to read the auth file multiple times to work around
// short-lived locks on Windows while token files are being written.
func readAuthFileWithRetry(path string, attempts int, delay time.Duration) ([]byte, error) {
read := func(target string) ([]byte, error) {
var lastErr error
for i := 0; i < attempts; i++ {
data, err := os.ReadFile(target)
if err == nil {
return data, nil
}
lastErr = err
if i < attempts-1 {
time.Sleep(delay)
}
}
return nil, lastErr
}
candidates := []string{
util.CookieSnapshotPath(path),
path,
}
for idx, candidate := range candidates {
data, err := read(candidate)
if err == nil {
return data, nil
}
if errors.Is(err, os.ErrNotExist) {
if idx < len(candidates)-1 {
continue
}
}
return nil, err
}
return nil, os.ErrNotExist
}
// addOrUpdateClient handles the addition or update of a single client. // addOrUpdateClient handles the addition or update of a single client.
func (w *Watcher) addOrUpdateClient(path string) { func (w *Watcher) addOrUpdateClient(path string) {
data, errRead := readAuthFileWithRetry(path, authFileReadMaxAttempts, authFileReadRetryDelay) data, errRead := util.ReadAuthFileWithRetry(path, authFileReadMaxAttempts, authFileReadRetryDelay)
if errRead != nil { if errRead != nil {
log.Errorf("failed to read auth file %s: %v", filepath.Base(path), errRead) log.Errorf("failed to read auth file %s: %v", filepath.Base(path), errRead)
return return