From aba719f5fe10f2a6af246c6a6923952c5b3e0425 Mon Sep 17 00:00:00 2001
From: hkfires <10558748+hkfires@users.noreply.github.com>
Date: Sat, 20 Sep 2025 00:14:26 +0800
Subject: [PATCH] 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.
---
 internal/cmd/run.go              |  3 +-
 internal/util/cookie_snapshot.go | 47 ++++++++++++++++++++++++++++++++
 internal/watcher/watcher.go      | 46 ++-----------------------------
 3 files changed, 52 insertions(+), 44 deletions(-)

diff --git a/internal/cmd/run.go b/internal/cmd/run.go
index f5c5a14c..7e0316be 100644
--- a/internal/cmd/run.go
+++ b/internal/cmd/run.go
@@ -26,6 +26,7 @@ import (
 	"github.com/luispater/CLIProxyAPI/v5/internal/config"
 	"github.com/luispater/CLIProxyAPI/v5/internal/interfaces"
 	"github.com/luispater/CLIProxyAPI/v5/internal/misc"
+	"github.com/luispater/CLIProxyAPI/v5/internal/util"
 	"github.com/luispater/CLIProxyAPI/v5/internal/watcher"
 	log "github.com/sirupsen/logrus"
 	"github.com/tidwall/gjson"
@@ -76,7 +77,7 @@ func StartService(cfg *config.Config, configPath string) {
 		if !info.IsDir() && strings.HasSuffix(info.Name(), ".json") {
 			misc.LogCredentialSeparator()
 			log.Debugf("Loading token from: %s", path)
-			data, errReadFile := os.ReadFile(path)
+			data, errReadFile := util.ReadAuthFilePreferSnapshot(path)
 			if errReadFile != nil {
 				return errReadFile
 			}
diff --git a/internal/util/cookie_snapshot.go b/internal/util/cookie_snapshot.go
index 9a049c59..bed1e03a 100644
--- a/internal/util/cookie_snapshot.go
+++ b/internal/util/cookie_snapshot.go
@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
 
 	"github.com/luispater/CLIProxyAPI/v5/internal/misc"
 )
@@ -93,6 +94,52 @@ func WriteCookieSnapshot(mainPath string, v any) error {
 	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.
 func RemoveCookieSnapshots(mainPath string) {
 	_ = RemoveFile(CookieSnapshotPath(mainPath))
diff --git a/internal/watcher/watcher.go b/internal/watcher/watcher.go
index a6303f31..ecda6478 100644
--- a/internal/watcher/watcher.go
+++ b/internal/watcher/watcher.go
@@ -9,7 +9,6 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"encoding/json"
-	"errors"
 	"io/fs"
 	"net/http"
 	"os"
@@ -324,7 +323,7 @@ func (w *Watcher) reloadClients() {
 	// Rebuild auth file hash cache for current clients
 	w.lastAuthHashes = make(map[string]string, len(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)
 			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.
 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 {
 		return nil, errReadFile
 	}
@@ -416,48 +415,9 @@ func (w *Watcher) clientsToSlice(clientMap map[string]interfaces.Client) []inter
 	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.
 func (w *Watcher) addOrUpdateClient(path string) {
-	data, errRead := readAuthFileWithRetry(path, authFileReadMaxAttempts, authFileReadRetryDelay)
+	data, errRead := util.ReadAuthFileWithRetry(path, authFileReadMaxAttempts, authFileReadRetryDelay)
 	if errRead != nil {
 		log.Errorf("failed to read auth file %s: %v", filepath.Base(path), errRead)
 		return