Compare commits

..

3 Commits

Author SHA1 Message Date
Luis Pater
d6bd6f3fb9 feat(vertex, management): enhance token handling and OAuth2 integration
Extend `vertexAccessToken` to support proxy-aware HTTP clients and update calls accordingly for better configurability. Add `deleteTokenRecord` to handle token cleanup, improving management of authentication files.
2025-11-11 23:42:46 +08:00
Luis Pater
8ae8a5c296 Fixed: #233
feat(management): add auth ID normalization and file-based ID resolution

Introduce `authIDForPath` to standardize ID generation from file paths, improving consistency in authentication handling. Update `registerAuthFromFile` and `disableAuth` to utilize normalized IDs, incorporating relative path resolution and file name extraction where applicable.
2025-11-11 19:23:31 +08:00
Luis Pater
dc804e96fb fix(management): improve error handling and normalize YAML comment indentation
Enhance error management for file operations and clean up temporary files. Add `NormalizeCommentIndentation` function to ensure YAML comments maintain consistent formatting.
2025-11-11 08:37:57 +08:00
4 changed files with 133 additions and 31 deletions

View File

@@ -508,6 +508,10 @@ func (h *Handler) DeleteAuthFile(c *gin.Context) {
}
}
if err = os.Remove(full); err == nil {
if errDel := h.deleteTokenRecord(ctx, full); errDel != nil {
c.JSON(500, gin.H{"error": errDel.Error()})
return
}
deleted++
h.disableAuth(ctx, full)
}
@@ -534,10 +538,32 @@ func (h *Handler) DeleteAuthFile(c *gin.Context) {
}
return
}
if err := h.deleteTokenRecord(ctx, full); err != nil {
c.JSON(500, gin.H{"error": err.Error()})
return
}
h.disableAuth(ctx, full)
c.JSON(200, gin.H{"status": "ok"})
}
func (h *Handler) authIDForPath(path string) string {
path = strings.TrimSpace(path)
if path == "" {
return ""
}
if h == nil || h.cfg == nil {
return path
}
authDir := strings.TrimSpace(h.cfg.AuthDir)
if authDir == "" {
return path
}
if rel, err := filepath.Rel(authDir, path); err == nil && rel != "" {
return rel
}
return path
}
func (h *Handler) registerAuthFromFile(ctx context.Context, path string, data []byte) error {
if h.authManager == nil {
return nil
@@ -566,13 +592,18 @@ func (h *Handler) registerAuthFromFile(ctx context.Context, path string, data []
}
lastRefresh, hasLastRefresh := extractLastRefreshTimestamp(metadata)
authID := h.authIDForPath(path)
if authID == "" {
authID = path
}
attr := map[string]string{
"path": path,
"source": path,
}
auth := &coreauth.Auth{
ID: path,
ID: authID,
Provider: provider,
FileName: filepath.Base(path),
Label: label,
Status: coreauth.StatusActive,
Attributes: attr,
@@ -583,7 +614,7 @@ func (h *Handler) registerAuthFromFile(ctx context.Context, path string, data []
if hasLastRefresh {
auth.LastRefreshedAt = lastRefresh
}
if existing, ok := h.authManager.GetByID(path); ok {
if existing, ok := h.authManager.GetByID(authID); ok {
auth.CreatedAt = existing.CreatedAt
if !hasLastRefresh {
auth.LastRefreshedAt = existing.LastRefreshedAt
@@ -598,10 +629,17 @@ func (h *Handler) registerAuthFromFile(ctx context.Context, path string, data []
}
func (h *Handler) disableAuth(ctx context.Context, id string) {
if h.authManager == nil || id == "" {
if h == nil || h.authManager == nil {
return
}
if auth, ok := h.authManager.GetByID(id); ok {
authID := h.authIDForPath(id)
if authID == "" {
authID = strings.TrimSpace(id)
}
if authID == "" {
return
}
if auth, ok := h.authManager.GetByID(authID); ok {
auth.Disabled = true
auth.Status = coreauth.StatusDisabled
auth.StatusMessage = "removed via management API"
@@ -610,9 +648,20 @@ func (h *Handler) disableAuth(ctx context.Context, id string) {
}
}
func (h *Handler) saveTokenRecord(ctx context.Context, record *coreauth.Auth) (string, error) {
if record == nil {
return "", fmt.Errorf("token record is nil")
func (h *Handler) deleteTokenRecord(ctx context.Context, path string) error {
if strings.TrimSpace(path) == "" {
return fmt.Errorf("auth path is empty")
}
store := h.tokenStoreWithBaseDir()
if store == nil {
return fmt.Errorf("token store unavailable")
}
return store.Delete(ctx, path)
}
func (h *Handler) tokenStoreWithBaseDir() coreauth.Store {
if h == nil {
return nil
}
store := h.tokenStore
if store == nil {
@@ -624,6 +673,17 @@ func (h *Handler) saveTokenRecord(ctx context.Context, record *coreauth.Auth) (s
dirSetter.SetBaseDir(h.cfg.AuthDir)
}
}
return store
}
func (h *Handler) saveTokenRecord(ctx context.Context, record *coreauth.Auth) (string, error) {
if record == nil {
return "", fmt.Errorf("token record is nil")
}
store := h.tokenStoreWithBaseDir()
if store == nil {
return "", fmt.Errorf("token store unavailable")
}
return store.Save(ctx, record)
}

View File

@@ -28,7 +28,7 @@ func (h *Handler) GetConfigYAML(c *gin.Context) {
return
}
var node yaml.Node
if err := yaml.Unmarshal(data, &node); err != nil {
if err = yaml.Unmarshal(data, &node); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "parse_failed", "message": err.Error()})
return
}
@@ -41,17 +41,18 @@ func (h *Handler) GetConfigYAML(c *gin.Context) {
}
func WriteConfig(path string, data []byte) error {
data = config.NormalizeCommentIndentation(data)
f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
if err != nil {
return err
}
if _, err := f.Write(data); err != nil {
f.Close()
return err
if _, errWrite := f.Write(data); errWrite != nil {
_ = f.Close()
return errWrite
}
if err := f.Sync(); err != nil {
f.Close()
return err
if errSync := f.Sync(); errSync != nil {
_ = f.Close()
return errSync
}
return f.Close()
}
@@ -63,7 +64,7 @@ func (h *Handler) PutConfigYAML(c *gin.Context) {
return
}
var cfg config.Config
if err := yaml.Unmarshal(body, &cfg); err != nil {
if err = yaml.Unmarshal(body, &cfg); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_yaml", "message": err.Error()})
return
}
@@ -75,18 +76,20 @@ func (h *Handler) PutConfigYAML(c *gin.Context) {
return
}
tempFile := tmpFile.Name()
if _, err := tmpFile.Write(body); err != nil {
tmpFile.Close()
os.Remove(tempFile)
c.JSON(http.StatusInternalServerError, gin.H{"error": "write_failed", "message": err.Error()})
if _, errWrite := tmpFile.Write(body); errWrite != nil {
_ = tmpFile.Close()
_ = os.Remove(tempFile)
c.JSON(http.StatusInternalServerError, gin.H{"error": "write_failed", "message": errWrite.Error()})
return
}
if err := tmpFile.Close(); err != nil {
os.Remove(tempFile)
c.JSON(http.StatusInternalServerError, gin.H{"error": "write_failed", "message": err.Error()})
if errClose := tmpFile.Close(); errClose != nil {
_ = os.Remove(tempFile)
c.JSON(http.StatusInternalServerError, gin.H{"error": "write_failed", "message": errClose.Error()})
return
}
defer os.Remove(tempFile)
defer func() {
_ = os.Remove(tempFile)
}()
_, err = config.LoadConfigOptional(tempFile, false)
if err != nil {
c.JSON(http.StatusUnprocessableEntity, gin.H{"error": "invalid_config", "message": err.Error()})

View File

@@ -5,6 +5,7 @@
package config
import (
"bytes"
"errors"
"fmt"
"os"
@@ -462,13 +463,19 @@ func SaveConfigPreserveComments(configFile string, cfg *Config) error {
return err
}
defer func() { _ = f.Close() }()
enc := yaml.NewEncoder(f)
var buf bytes.Buffer
enc := yaml.NewEncoder(&buf)
enc.SetIndent(2)
if err = enc.Encode(&original); err != nil {
_ = enc.Close()
return err
}
return enc.Close()
if err = enc.Close(); err != nil {
return err
}
data = NormalizeCommentIndentation(buf.Bytes())
_, err = f.Write(data)
return err
}
func sanitizeConfigForPersist(cfg *Config) *Config {
@@ -518,13 +525,40 @@ func SaveConfigPreserveCommentsUpdateNestedScalar(configFile string, path []stri
return err
}
defer func() { _ = f.Close() }()
enc := yaml.NewEncoder(f)
var buf bytes.Buffer
enc := yaml.NewEncoder(&buf)
enc.SetIndent(2)
if err = enc.Encode(&root); err != nil {
_ = enc.Close()
return err
}
return enc.Close()
if err = enc.Close(); err != nil {
return err
}
data = NormalizeCommentIndentation(buf.Bytes())
_, err = f.Write(data)
return err
}
// NormalizeCommentIndentation removes indentation from standalone YAML comment lines to keep them left aligned.
func NormalizeCommentIndentation(data []byte) []byte {
lines := bytes.Split(data, []byte("\n"))
changed := false
for i, line := range lines {
trimmed := bytes.TrimLeft(line, " \t")
if len(trimmed) == 0 || trimmed[0] != '#' {
continue
}
if len(trimmed) == len(line) {
continue
}
lines[i] = append([]byte(nil), trimmed...)
changed = true
}
if !changed {
return data
}
return bytes.Join(lines, []byte("\n"))
}
// getOrCreateMapValue finds the value node for a given key in a mapping node.
@@ -766,6 +800,7 @@ func matchSequenceElement(original []*yaml.Node, used []bool, target *yaml.Node)
}
}
}
default:
}
// Fallback to structural equality to preserve nodes lacking explicit identifiers.
for i := range original {

View File

@@ -22,6 +22,7 @@ import (
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
"golang.org/x/oauth2"
"golang.org/x/oauth2/google"
)
@@ -89,7 +90,7 @@ func (e *GeminiVertexExecutor) Execute(ctx context.Context, auth *cliproxyauth.A
return resp, errNewReq
}
httpReq.Header.Set("Content-Type", "application/json")
if token, errTok := vertexAccessToken(ctx, saJSON); errTok == nil && token != "" {
if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
httpReq.Header.Set("Authorization", "Bearer "+token)
} else if errTok != nil {
log.Errorf("vertex executor: access token error: %v", errTok)
@@ -184,7 +185,7 @@ func (e *GeminiVertexExecutor) ExecuteStream(ctx context.Context, auth *cliproxy
return nil, errNewReq
}
httpReq.Header.Set("Content-Type", "application/json")
if token, errTok := vertexAccessToken(ctx, saJSON); errTok == nil && token != "" {
if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
httpReq.Header.Set("Authorization", "Bearer "+token)
} else if errTok != nil {
log.Errorf("vertex executor: access token error: %v", errTok)
@@ -295,7 +296,7 @@ func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyau
return cliproxyexecutor.Response{}, errNewReq
}
httpReq.Header.Set("Content-Type", "application/json")
if token, errTok := vertexAccessToken(ctx, saJSON); errTok == nil && token != "" {
if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
httpReq.Header.Set("Authorization", "Bearer "+token)
} else if errTok != nil {
log.Errorf("vertex executor: access token error: %v", errTok)
@@ -407,7 +408,10 @@ func vertexBaseURL(location string) string {
return fmt.Sprintf("https://%s-aiplatform.googleapis.com", loc)
}
func vertexAccessToken(ctx context.Context, saJSON []byte) (string, error) {
func vertexAccessToken(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, saJSON []byte) (string, error) {
if httpClient := newProxyAwareHTTPClient(ctx, cfg, auth, 0); httpClient != nil {
ctx = context.WithValue(ctx, oauth2.HTTPClient, httpClient)
}
// Use cloud-platform scope for Vertex AI.
creds, errCreds := google.CredentialsFromJSON(ctx, saJSON, "https://www.googleapis.com/auth/cloud-platform")
if errCreds != nil {