feat(auth): centralize token store management and enhance persistence

- Introduced `RegisterTokenStore` and `GetTokenStore` to centralize token store access.
- Replaced direct file operations with a unified token persistence API.
- Updated all components to use the shared token store for consistent behavior.
- Improved logging for token save operations to include file paths.
This commit is contained in:
Luis Pater
2025-09-25 03:17:50 +08:00
parent 19609db13c
commit 8fc73874de
7 changed files with 107 additions and 24 deletions

View File

@@ -21,6 +21,7 @@ import (
// legacy client removed
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
@@ -341,6 +342,18 @@ func (h *Handler) disableAuth(ctx context.Context, id string) {
}
}
func (h *Handler) saveTokenRecord(ctx context.Context, record *sdkAuth.TokenRecord) (string, error) {
if record == nil {
return "", fmt.Errorf("token record is nil")
}
store := h.tokenStore
if store == nil {
store = sdkAuth.GetTokenStore()
h.tokenStore = store
}
return store.Save(ctx, h.cfg, record)
}
func (h *Handler) RequestAnthropicToken(c *gin.Context) {
ctx := context.Background()
@@ -481,15 +494,20 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
// Create token storage
tokenStorage := anthropicAuth.CreateTokenStorage(bundle)
// Persist token to file directly
fileName := filepath.Join(h.cfg.AuthDir, fmt.Sprintf("claude-%s.json", tokenStorage.Email))
if errSave := tokenStorage.SaveTokenToFile(fileName); errSave != nil {
record := &sdkAuth.TokenRecord{
Provider: "claude",
FileName: fmt.Sprintf("claude-%s.json", tokenStorage.Email),
Storage: tokenStorage,
Metadata: map[string]string{"email": tokenStorage.Email},
}
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save authentication tokens: %v", errSave)
oauthStatus[state] = "Failed to save authentication tokens"
return
}
log.Info("Authentication successful!")
log.Infof("Authentication successful! Token saved to %s", savedPath)
if bundle.APIKey != "" {
log.Info("API key obtained and saved")
}
@@ -639,16 +657,24 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
}
log.Info("Authentication successful.")
// Persist token to file directly
fileName := filepath.Join(h.cfg.AuthDir, fmt.Sprintf("gemini-%s.json", ts.Email))
if err = ts.SaveTokenToFile(fileName); err != nil {
log.Fatalf("Failed to save token to file: %v", err)
record := &sdkAuth.TokenRecord{
Provider: "gemini",
FileName: fmt.Sprintf("gemini-%s.json", ts.Email),
Storage: &ts,
Metadata: map[string]string{
"email": ts.Email,
"project_id": ts.ProjectID,
},
}
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save token to file: %v", errSave)
oauthStatus[state] = "Failed to save token to file"
return
}
delete(oauthStatus, state)
log.Info("You can now use Gemini CLI services through this CLI")
log.Infof("You can now use Gemini CLI services through this CLI; token saved to %s", savedPath)
}()
oauthStatus[state] = ""
@@ -783,13 +809,22 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
// Create token storage and persist
tokenStorage := openaiAuth.CreateTokenStorage(bundle)
fileName := filepath.Join(h.cfg.AuthDir, fmt.Sprintf("codex-%s.json", tokenStorage.Email))
if errSave := tokenStorage.SaveTokenToFile(fileName); errSave != nil {
record := &sdkAuth.TokenRecord{
Provider: "codex",
FileName: fmt.Sprintf("codex-%s.json", tokenStorage.Email),
Storage: tokenStorage,
Metadata: map[string]string{
"email": tokenStorage.Email,
"account_id": tokenStorage.AccountID,
},
}
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
oauthStatus[state] = "Failed to save authentication tokens"
log.Fatalf("Failed to save authentication tokens: %v", errSave)
return
}
log.Info("Authentication successful!")
log.Infof("Authentication successful! Token saved to %s", savedPath)
if bundle.APIKey != "" {
log.Info("API key obtained and saved")
}
@@ -831,15 +866,20 @@ func (h *Handler) RequestQwenToken(c *gin.Context) {
tokenStorage := qwenAuth.CreateTokenStorage(tokenData)
tokenStorage.Email = fmt.Sprintf("qwen-%d", time.Now().UnixMilli())
// Save token storage
fileName := filepath.Join(h.cfg.AuthDir, fmt.Sprintf("qwen-%s.json", tokenStorage.Email))
if err = tokenStorage.SaveTokenToFile(fileName); err != nil {
log.Fatalf("Failed to save authentication tokens: %v", err)
record := &sdkAuth.TokenRecord{
Provider: "qwen",
FileName: fmt.Sprintf("qwen-%s.json", tokenStorage.Email),
Storage: tokenStorage,
Metadata: map[string]string{"email": tokenStorage.Email},
}
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save authentication tokens: %v", errSave)
oauthStatus[state] = "Failed to save authentication tokens"
return
}
log.Info("Authentication successful!")
log.Infof("Authentication successful! Token saved to %s", savedPath)
log.Info("You can now use Qwen services through this CLI")
delete(oauthStatus, state)
}()

View File

@@ -12,6 +12,7 @@ import (
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
"golang.org/x/crypto/bcrypt"
)
@@ -31,6 +32,7 @@ type Handler struct {
failedAttempts map[string]*attemptInfo // keyed by client IP
authManager *coreauth.Manager
usageStats *usage.RequestStatistics
tokenStore sdkAuth.TokenStore
}
// NewHandler creates a new management handler instance.
@@ -41,6 +43,7 @@ func NewHandler(cfg *config.Config, configFilePath string, manager *coreauth.Man
failedAttempts: make(map[string]*attemptInfo),
authManager: manager,
usageStats: usage.GetRequestStatistics(),
tokenStore: sdkAuth.GetTokenStore(),
}
}

View File

@@ -11,7 +11,7 @@ import (
// Returns:
// - *sdkAuth.Manager: A configured authentication manager instance
func newAuthManager() *sdkAuth.Manager {
store := sdkAuth.NewFileTokenStore()
store := sdkAuth.GetTokenStore()
manager := sdkAuth.NewManager(store,
sdkAuth.NewGeminiAuthenticator(),
sdkAuth.NewCodexAuthenticator(),

View File

@@ -3,15 +3,16 @@ package cmd
import (
"bufio"
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"os"
"path/filepath"
"strings"
"github.com/router-for-me/CLIProxyAPI/v6/internal/auth/gemini"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
log "github.com/sirupsen/logrus"
)
@@ -48,13 +49,17 @@ func DoGeminiWebAuth(cfg *config.Config) {
hasher.Write([]byte(secure1psid))
hash := hex.EncodeToString(hasher.Sum(nil))
fileName := fmt.Sprintf("gemini-web-%s.json", hash[:16])
filePath := filepath.Join(cfg.AuthDir, fileName)
err := tokenStorage.SaveTokenToFile(filePath)
record := &sdkAuth.TokenRecord{
Provider: "gemini-web",
FileName: fileName,
Storage: tokenStorage,
}
store := sdkAuth.GetTokenStore()
savedPath, err := store.Save(context.Background(), cfg, record)
if err != nil {
log.Fatalf("Failed to save Gemini Web token to file: %v", err)
return
}
log.Infof("Successfully saved Gemini Web token to: %s", filePath)
log.Infof("Successfully saved Gemini Web token to: %s", savedPath)
}