From d536110404ed16b2e48fda02b8dc5c02386b80de Mon Sep 17 00:00:00 2001 From: HEUDavid Date: Tue, 10 Feb 2026 08:35:36 +0800 Subject: [PATCH] feat/auth-hook: add post auth hook --- internal/auth/claude/token.go | 19 +++--------- internal/auth/codex/token.go | 19 +++--------- internal/auth/gemini/gemini_token.go | 45 ++++++++++++---------------- internal/auth/iflow/iflow_token.go | 19 +++--------- internal/auth/kimi/token.go | 19 +++--------- internal/auth/qwen/qwen_token.go | 19 +++--------- internal/misc/credentials.go | 35 ++++++++++++++++++++++ 7 files changed, 74 insertions(+), 101 deletions(-) diff --git a/internal/auth/claude/token.go b/internal/auth/claude/token.go index c36f8e76..6ebb0f2f 100644 --- a/internal/auth/claude/token.go +++ b/internal/auth/claude/token.go @@ -75,21 +75,10 @@ func (ts *ClaudeTokenStorage) SaveTokenToFile(authFilePath string) error { _ = f.Close() }() - // Convert struct to map for merging - data := make(map[string]any) - temp, errJson := json.Marshal(ts) - if errJson != nil { - return fmt.Errorf("failed to marshal struct: %w", errJson) - } - if errUnmarshal := json.Unmarshal(temp, &data); errUnmarshal != nil { - return fmt.Errorf("failed to unmarshal struct map: %w", errUnmarshal) - } - - // Merge extra metadata - if ts.Metadata != nil { - for k, v := range ts.Metadata { - data[k] = v - } + // Merge metadata using helper + data, errMerge := misc.MergeMetadata(ts, ts.Metadata) + if errMerge != nil { + return fmt.Errorf("failed to merge metadata: %w", errMerge) } // Encode and write the token data as JSON diff --git a/internal/auth/codex/token.go b/internal/auth/codex/token.go index 1ea84f3a..a3252d1b 100644 --- a/internal/auth/codex/token.go +++ b/internal/auth/codex/token.go @@ -68,21 +68,10 @@ func (ts *CodexTokenStorage) SaveTokenToFile(authFilePath string) error { _ = f.Close() }() - // Convert struct to map for merging - data := make(map[string]any) - temp, errJson := json.Marshal(ts) - if errJson != nil { - return fmt.Errorf("failed to marshal struct: %w", errJson) - } - if errUnmarshal := json.Unmarshal(temp, &data); errUnmarshal != nil { - return fmt.Errorf("failed to unmarshal struct map: %w", errUnmarshal) - } - - // Merge extra metadata - if ts.Metadata != nil { - for k, v := range ts.Metadata { - data[k] = v - } + // Merge metadata using helper + data, errMerge := misc.MergeMetadata(ts, ts.Metadata) + if errMerge != nil { + return fmt.Errorf("failed to merge metadata: %w", errMerge) } if err = json.NewEncoder(f).Encode(data); err != nil { diff --git a/internal/auth/gemini/gemini_token.go b/internal/auth/gemini/gemini_token.go index 24828076..f84564e2 100644 --- a/internal/auth/gemini/gemini_token.go +++ b/internal/auth/gemini/gemini_token.go @@ -11,7 +11,6 @@ import ( "strings" "github.com/router-for-me/CLIProxyAPI/v6/internal/misc" - log "github.com/sirupsen/logrus" ) // GeminiTokenStorage stores OAuth2 token information for Google Gemini API authentication. @@ -58,41 +57,35 @@ func (ts *GeminiTokenStorage) SetMetadata(meta map[string]any) { // - error: An error if the operation fails, nil otherwise func (ts *GeminiTokenStorage) SaveTokenToFile(authFilePath string) error { misc.LogSavingCredentials(authFilePath) - ts.Type = "gemini" - if err := os.MkdirAll(filepath.Dir(authFilePath), 0700); err != nil { - return fmt.Errorf("failed to create directory: %v", err) + ts.Type = "gemini" // Ensure type is set before merging/saving + + // Merge metadata using helper + data, errMerge := misc.MergeMetadata(ts, ts.Metadata) + if errMerge != nil { + return fmt.Errorf("failed to merge metadata: %w", errMerge) } + // Create parent directory + if err := os.MkdirAll(filepath.Dir(authFilePath), os.ModePerm); err != nil { + return fmt.Errorf("failed to create directory: %w", err) + } + + // Create file f, err := os.Create(authFilePath) if err != nil { - return fmt.Errorf("failed to create token file: %w", err) + return fmt.Errorf("failed to create file: %w", err) } defer func() { - if errClose := f.Close(); errClose != nil { - log.Errorf("failed to close file: %v", errClose) - } + _ = f.Close() }() - // Convert struct to map for merging - data := make(map[string]any) - temp, errJson := json.Marshal(ts) - if errJson != nil { - return fmt.Errorf("failed to marshal struct: %w", errJson) - } - if errUnmarshal := json.Unmarshal(temp, &data); errUnmarshal != nil { - return fmt.Errorf("failed to unmarshal struct map: %w", errUnmarshal) + // Write to file + enc := json.NewEncoder(f) + enc.SetIndent("", " ") + if err := enc.Encode(data); err != nil { + return fmt.Errorf("failed to encode token to file: %w", err) } - // Merge extra metadata - if ts.Metadata != nil { - for k, v := range ts.Metadata { - data[k] = v - } - } - - if err = json.NewEncoder(f).Encode(data); err != nil { - return fmt.Errorf("failed to write token to file: %w", err) - } return nil } diff --git a/internal/auth/iflow/iflow_token.go b/internal/auth/iflow/iflow_token.go index 13eb7de1..a515c926 100644 --- a/internal/auth/iflow/iflow_token.go +++ b/internal/auth/iflow/iflow_token.go @@ -46,21 +46,10 @@ func (ts *IFlowTokenStorage) SaveTokenToFile(authFilePath string) error { } defer func() { _ = f.Close() }() - // Convert struct to map for merging - data := make(map[string]any) - temp, errJson := json.Marshal(ts) - if errJson != nil { - return fmt.Errorf("failed to marshal struct: %w", errJson) - } - if errUnmarshal := json.Unmarshal(temp, &data); errUnmarshal != nil { - return fmt.Errorf("failed to unmarshal struct map: %w", errUnmarshal) - } - - // Merge extra metadata - if ts.Metadata != nil { - for k, v := range ts.Metadata { - data[k] = v - } + // Merge metadata using helper + data, errMerge := misc.MergeMetadata(ts, ts.Metadata) + if errMerge != nil { + return fmt.Errorf("failed to merge metadata: %w", errMerge) } if err = json.NewEncoder(f).Encode(data); err != nil { diff --git a/internal/auth/kimi/token.go b/internal/auth/kimi/token.go index 15171d93..7320d760 100644 --- a/internal/auth/kimi/token.go +++ b/internal/auth/kimi/token.go @@ -95,21 +95,10 @@ func (ts *KimiTokenStorage) SaveTokenToFile(authFilePath string) error { _ = f.Close() }() - // Convert struct to map for merging - data := make(map[string]any) - temp, errJson := json.Marshal(ts) - if errJson != nil { - return fmt.Errorf("failed to marshal struct: %w", errJson) - } - if errUnmarshal := json.Unmarshal(temp, &data); errUnmarshal != nil { - return fmt.Errorf("failed to unmarshal struct map: %w", errUnmarshal) - } - - // Merge extra metadata - if ts.Metadata != nil { - for k, v := range ts.Metadata { - data[k] = v - } + // Merge metadata using helper + data, errMerge := misc.MergeMetadata(ts, ts.Metadata) + if errMerge != nil { + return fmt.Errorf("failed to merge metadata: %w", errMerge) } encoder := json.NewEncoder(f) diff --git a/internal/auth/qwen/qwen_token.go b/internal/auth/qwen/qwen_token.go index 8037bdb7..276c8b40 100644 --- a/internal/auth/qwen/qwen_token.go +++ b/internal/auth/qwen/qwen_token.go @@ -66,21 +66,10 @@ func (ts *QwenTokenStorage) SaveTokenToFile(authFilePath string) error { _ = f.Close() }() - // Convert struct to map for merging - data := make(map[string]any) - temp, errJson := json.Marshal(ts) - if errJson != nil { - return fmt.Errorf("failed to marshal struct: %w", errJson) - } - if errUnmarshal := json.Unmarshal(temp, &data); errUnmarshal != nil { - return fmt.Errorf("failed to unmarshal struct map: %w", errUnmarshal) - } - - // Merge extra metadata - if ts.Metadata != nil { - for k, v := range ts.Metadata { - data[k] = v - } + // Merge metadata using helper + data, errMerge := misc.MergeMetadata(ts, ts.Metadata) + if errMerge != nil { + return fmt.Errorf("failed to merge metadata: %w", errMerge) } if err = json.NewEncoder(f).Encode(data); err != nil { diff --git a/internal/misc/credentials.go b/internal/misc/credentials.go index b03cd788..6b4f9ced 100644 --- a/internal/misc/credentials.go +++ b/internal/misc/credentials.go @@ -1,6 +1,7 @@ package misc import ( + "encoding/json" "fmt" "path/filepath" "strings" @@ -24,3 +25,37 @@ func LogSavingCredentials(path string) { func LogCredentialSeparator() { log.Debug(credentialSeparator) } + +// MergeMetadata serializes the source struct into a map and merges the provided metadata into it. +func MergeMetadata(source any, metadata map[string]any) (map[string]any, error) { + var data map[string]any + + // Fast path: if source is already a map, just copy it to avoid mutation of original + if srcMap, ok := source.(map[string]any); ok { + data = make(map[string]any, len(srcMap)+len(metadata)) + for k, v := range srcMap { + data[k] = v + } + } else { + // Slow path: marshal to JSON and back to map to respect JSON tags + temp, err := json.Marshal(source) + if err != nil { + return nil, fmt.Errorf("failed to marshal source: %w", err) + } + if err := json.Unmarshal(temp, &data); err != nil { + return nil, fmt.Errorf("failed to unmarshal to map: %w", err) + } + } + + // Merge extra metadata + if metadata != nil { + if data == nil { + data = make(map[string]any) + } + for k, v := range metadata { + data[k] = v + } + } + + return data, nil +}