Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
1802 changed files with 503346 additions and 2 deletions
32
backend/internal/api/buffered_conn.go
Normal file
32
backend/internal/api/buffered_conn.go
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"crypto/tls"
|
||||
"net"
|
||||
)
|
||||
|
||||
type bufferedConn struct {
|
||||
net.Conn
|
||||
reader *bufio.Reader
|
||||
}
|
||||
|
||||
func (c *bufferedConn) Read(p []byte) (int, error) {
|
||||
if c == nil {
|
||||
return 0, net.ErrClosed
|
||||
}
|
||||
if c.reader == nil {
|
||||
return c.Conn.Read(p)
|
||||
}
|
||||
return c.reader.Read(p)
|
||||
}
|
||||
|
||||
func (c *bufferedConn) ConnectionState() tls.ConnectionState {
|
||||
if c == nil || c.Conn == nil {
|
||||
return tls.ConnectionState{}
|
||||
}
|
||||
if stater, ok := c.Conn.(interface{ ConnectionState() tls.ConnectionState }); ok {
|
||||
return stater.ConnectionState()
|
||||
}
|
||||
return tls.ConnectionState{}
|
||||
}
|
||||
117
backend/internal/api/handlers/management/api_key_usage.go
Normal file
117
backend/internal/api/handlers/management/api_key_usage.go
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
type apiKeyUsageEntry struct {
|
||||
Success int64 `json:"success"`
|
||||
Failed int64 `json:"failed"`
|
||||
RecentRequests []coreauth.RecentRequestBucket `json:"recent_requests"`
|
||||
}
|
||||
|
||||
func mergeRecentRequestBuckets(dst, src []coreauth.RecentRequestBucket) []coreauth.RecentRequestBucket {
|
||||
if len(dst) == 0 {
|
||||
return src
|
||||
}
|
||||
if len(src) == 0 {
|
||||
return dst
|
||||
}
|
||||
if len(dst) != len(src) {
|
||||
n := len(dst)
|
||||
if len(src) < n {
|
||||
n = len(src)
|
||||
}
|
||||
for i := 0; i < n; i++ {
|
||||
dst[i].Success += src[i].Success
|
||||
dst[i].Failed += src[i].Failed
|
||||
}
|
||||
return dst
|
||||
}
|
||||
for i := range dst {
|
||||
dst[i].Success += src[i].Success
|
||||
dst[i].Failed += src[i].Failed
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
func apiKeyUsageProviderKey(auth *coreauth.Auth) string {
|
||||
provider := strings.ToLower(strings.TrimSpace(auth.Provider))
|
||||
if auth.Attributes != nil {
|
||||
if compatName := strings.TrimSpace(auth.Attributes["compat_name"]); compatName != "" {
|
||||
provider = strings.ToLower(compatName)
|
||||
}
|
||||
}
|
||||
if provider == "" {
|
||||
return "unknown"
|
||||
}
|
||||
return provider
|
||||
}
|
||||
|
||||
// GetAPIKeyUsage returns recent request buckets for all in-memory api_key auths,
|
||||
// grouped by provider and keyed by "base_url|api_key".
|
||||
func (h *Handler) GetAPIKeyUsage(c *gin.Context) {
|
||||
if h == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "handler not initialized"})
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
manager := h.authManager
|
||||
h.mu.Unlock()
|
||||
if manager == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"})
|
||||
return
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
out := make(map[string]map[string]apiKeyUsageEntry)
|
||||
for _, auth := range manager.List() {
|
||||
if auth == nil {
|
||||
continue
|
||||
}
|
||||
kind, apiKey := auth.AccountInfo()
|
||||
if !strings.EqualFold(strings.TrimSpace(kind), "api_key") {
|
||||
continue
|
||||
}
|
||||
apiKey = strings.TrimSpace(apiKey)
|
||||
if apiKey == "" {
|
||||
continue
|
||||
}
|
||||
baseURL := ""
|
||||
if auth.Attributes != nil {
|
||||
baseURL = strings.TrimSpace(auth.Attributes["base_url"])
|
||||
if baseURL == "" {
|
||||
baseURL = strings.TrimSpace(auth.Attributes["base-url"])
|
||||
}
|
||||
}
|
||||
compositeKey := baseURL + "|" + apiKey
|
||||
provider := apiKeyUsageProviderKey(auth)
|
||||
|
||||
recent := auth.RecentRequestsSnapshot(now)
|
||||
providerBucket, ok := out[provider]
|
||||
if !ok {
|
||||
providerBucket = make(map[string]apiKeyUsageEntry)
|
||||
out[provider] = providerBucket
|
||||
}
|
||||
if existing, exists := providerBucket[compositeKey]; exists {
|
||||
existing.Success += auth.Success
|
||||
existing.Failed += auth.Failed
|
||||
existing.RecentRequests = mergeRecentRequestBuckets(existing.RecentRequests, recent)
|
||||
providerBucket[compositeKey] = existing
|
||||
continue
|
||||
}
|
||||
providerBucket[compositeKey] = apiKeyUsageEntry{
|
||||
Success: auth.Success,
|
||||
Failed: auth.Failed,
|
||||
RecentRequests: recent,
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, out)
|
||||
}
|
||||
142
backend/internal/api/handlers/management/api_key_usage_test.go
Normal file
142
backend/internal/api/handlers/management/api_key_usage_test.go
Normal file
|
|
@ -0,0 +1,142 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func sumRecentRequestBuckets(buckets []coreauth.RecentRequestBucket) (int64, int64) {
|
||||
var success int64
|
||||
var failed int64
|
||||
for _, bucket := range buckets {
|
||||
success += bucket.Success
|
||||
failed += bucket.Failed
|
||||
}
|
||||
return success, failed
|
||||
}
|
||||
|
||||
func TestGetAPIKeyUsage_GroupsByProviderAndAPIKey(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
if _, err := manager.Register(context.Background(), &coreauth.Auth{
|
||||
ID: "codex-auth",
|
||||
Provider: "codex",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "codex-key",
|
||||
"base_url": "https://codex.example.com",
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("register codex auth: %v", err)
|
||||
}
|
||||
if _, err := manager.Register(context.Background(), &coreauth.Auth{
|
||||
ID: "claude-auth",
|
||||
Provider: "claude",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "claude-key",
|
||||
"base_url": "https://claude.example.com",
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("register claude auth: %v", err)
|
||||
}
|
||||
|
||||
manager.MarkResult(context.Background(), coreauth.Result{AuthID: "codex-auth", Provider: "codex", Model: "gpt-5", Success: true})
|
||||
manager.MarkResult(context.Background(), coreauth.Result{AuthID: "codex-auth", Provider: "codex", Model: "gpt-5", Success: false})
|
||||
manager.MarkResult(context.Background(), coreauth.Result{AuthID: "claude-auth", Provider: "claude", Model: "claude-4", Success: true})
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ginCtx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v0/management/api-key-usage", nil)
|
||||
ginCtx.Request = req
|
||||
h.GetAPIKeyUsage(ginCtx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
|
||||
var payload map[string]map[string]apiKeyUsageEntry
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
|
||||
t.Fatalf("decode payload: %v", err)
|
||||
}
|
||||
|
||||
codexEntry := payload["codex"]["https://codex.example.com|codex-key"]
|
||||
if codexEntry.Success != 1 || codexEntry.Failed != 1 {
|
||||
t.Fatalf("codex totals = %d/%d, want 1/1", codexEntry.Success, codexEntry.Failed)
|
||||
}
|
||||
if len(codexEntry.RecentRequests) != 20 {
|
||||
t.Fatalf("codex buckets len = %d, want 20", len(codexEntry.RecentRequests))
|
||||
}
|
||||
codexSuccess, codexFailed := sumRecentRequestBuckets(codexEntry.RecentRequests)
|
||||
if codexSuccess != 1 || codexFailed != 1 {
|
||||
t.Fatalf("codex totals = %d/%d, want 1/1", codexSuccess, codexFailed)
|
||||
}
|
||||
|
||||
claudeEntry := payload["claude"]["https://claude.example.com|claude-key"]
|
||||
if claudeEntry.Success != 1 || claudeEntry.Failed != 0 {
|
||||
t.Fatalf("claude totals = %d/%d, want 1/0", claudeEntry.Success, claudeEntry.Failed)
|
||||
}
|
||||
if len(claudeEntry.RecentRequests) != 20 {
|
||||
t.Fatalf("claude buckets len = %d, want 20", len(claudeEntry.RecentRequests))
|
||||
}
|
||||
claudeSuccess, claudeFailed := sumRecentRequestBuckets(claudeEntry.RecentRequests)
|
||||
if claudeSuccess != 1 || claudeFailed != 0 {
|
||||
t.Fatalf("claude totals = %d/%d, want 1/0", claudeSuccess, claudeFailed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetAPIKeyUsage_GroupsOpenAICompatibleByCompatName(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
if _, err := manager.Register(context.Background(), &coreauth.Auth{
|
||||
ID: "vast-auth",
|
||||
Provider: "openai-compatible-vast",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "vast-key",
|
||||
"base_url": "https://www.vastnum.com/v1",
|
||||
"compat_name": "VAST",
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("register vast auth: %v", err)
|
||||
}
|
||||
|
||||
manager.MarkResult(context.Background(), coreauth.Result{AuthID: "vast-auth", Provider: "openai-compatible-vast", Model: "gpt-5", Success: true})
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ginCtx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v0/management/api-key-usage", nil)
|
||||
ginCtx.Request = req
|
||||
h.GetAPIKeyUsage(ginCtx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
|
||||
var payload map[string]map[string]apiKeyUsageEntry
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
|
||||
t.Fatalf("decode payload: %v", err)
|
||||
}
|
||||
|
||||
if _, exists := payload["openai-compatible-vast"]; exists {
|
||||
t.Fatalf("unexpected namespaced provider bucket in payload: %#v", payload)
|
||||
}
|
||||
vastBucket, exists := payload["vast"]
|
||||
if !exists {
|
||||
t.Fatalf("missing compat provider bucket in payload: %#v", payload)
|
||||
}
|
||||
vastEntry := vastBucket["https://www.vastnum.com/v1|vast-key"]
|
||||
if vastEntry.Success != 1 || vastEntry.Failed != 0 {
|
||||
t.Fatalf("vast totals = %d/%d, want 1/0", vastEntry.Success, vastEntry.Failed)
|
||||
}
|
||||
}
|
||||
663
backend/internal/api/handlers/management/api_tools.go
Normal file
663
backend/internal/api/handlers/management/api_tools.go
Normal file
|
|
@ -0,0 +1,663 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/proxyutil"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const defaultAPICallTimeout = 60 * time.Second
|
||||
|
||||
const (
|
||||
antigravityOAuthClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
|
||||
antigravityOAuthClientSecret = "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"
|
||||
)
|
||||
|
||||
var antigravityOAuthTokenURL = "https://oauth2.googleapis.com/token"
|
||||
|
||||
type apiCallRequest struct {
|
||||
AuthIndexSnake *string `json:"auth_index"`
|
||||
AuthIndexCamel *string `json:"authIndex"`
|
||||
AuthIndexPascal *string `json:"AuthIndex"`
|
||||
Method string `json:"method"`
|
||||
URL string `json:"url"`
|
||||
ProxyURL string `json:"proxy_url"`
|
||||
Header map[string]string `json:"header"`
|
||||
Data string `json:"data"`
|
||||
}
|
||||
|
||||
type apiCallResponse struct {
|
||||
StatusCode int `json:"status_code"`
|
||||
Header map[string][]string `json:"header"`
|
||||
Body string `json:"body"`
|
||||
}
|
||||
|
||||
// APICall makes a generic HTTP request on behalf of the management API caller.
|
||||
// It is protected by the management middleware.
|
||||
//
|
||||
// Endpoint:
|
||||
//
|
||||
// POST /v0/management/api-call
|
||||
//
|
||||
// Authentication:
|
||||
//
|
||||
// Same as other management APIs (requires a management key and remote-management rules).
|
||||
// You can provide the key via:
|
||||
// - Authorization: Bearer <key>
|
||||
// - X-Management-Key: <key>
|
||||
//
|
||||
// Request JSON:
|
||||
// - auth_index / authIndex / AuthIndex (optional):
|
||||
// The credential "auth_index" from GET /v0/management/auth-files (or other endpoints returning it).
|
||||
// If omitted or not found, credential-specific proxy/token substitution is skipped.
|
||||
// - method (required): HTTP method, e.g. GET, POST, PUT, PATCH, DELETE.
|
||||
// - url (required): Absolute URL including scheme and host, e.g. "https://api.example.com/v1/ping".
|
||||
// - proxy_url (optional): Proxy used for this request. Supports HTTP, HTTPS, SOCKS5, SOCKS5H,
|
||||
// and "direct"/"none" to explicitly bypass proxies. When set, credential and global proxies are ignored.
|
||||
// - header (optional): Request headers map.
|
||||
// Supports magic variable "$TOKEN$" which is replaced using the selected credential:
|
||||
// 1) metadata.access_token
|
||||
// 2) attributes.api_key
|
||||
// 3) metadata.token / metadata.id_token / metadata.cookie
|
||||
// Example: {"Authorization":"Bearer $TOKEN$"}.
|
||||
// Note: if you need to override the HTTP Host header, set header["Host"].
|
||||
// - data (optional): Raw request body as string (useful for POST/PUT/PATCH).
|
||||
//
|
||||
// Proxy selection (highest priority first):
|
||||
// 1. Request proxy_url (when set, lower-priority proxy settings are ignored)
|
||||
// 2. Selected credential proxy_url
|
||||
// 3. Global config proxy-url
|
||||
// 4. Direct connect (environment proxies are not used)
|
||||
//
|
||||
// Response JSON (returned with HTTP 200 when the APICall itself succeeds):
|
||||
// - status_code: Upstream HTTP status code.
|
||||
// - header: Upstream response headers.
|
||||
// - body: Upstream response body as string.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// curl -sS -X POST "http://127.0.0.1:8317/v0/management/api-call" \
|
||||
// -H "Authorization: Bearer <MANAGEMENT_KEY>" \
|
||||
// -H "Content-Type: application/json" \
|
||||
// -d '{"auth_index":"<AUTH_INDEX>","method":"GET","url":"https://api.example.com/v1/ping","header":{"Authorization":"Bearer $TOKEN$"}}'
|
||||
//
|
||||
// curl -sS -X POST "http://127.0.0.1:8317/v0/management/api-call" \
|
||||
// -H "Authorization: Bearer 831227" \
|
||||
// -H "Content-Type: application/json" \
|
||||
// -d '{"auth_index":"<AUTH_INDEX>","method":"POST","url":"https://api.example.com/v1/fetchAvailableModels","header":{"Authorization":"Bearer $TOKEN$","Content-Type":"application/json","User-Agent":"cliproxyapi"},"data":"{}"}'
|
||||
func (h *Handler) APICall(c *gin.Context) {
|
||||
var body apiCallRequest
|
||||
if errBindJSON := c.ShouldBindJSON(&body); errBindJSON != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
|
||||
method := strings.ToUpper(strings.TrimSpace(body.Method))
|
||||
if method == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "missing method"})
|
||||
return
|
||||
}
|
||||
|
||||
urlStr := strings.TrimSpace(body.URL)
|
||||
if urlStr == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "missing url"})
|
||||
return
|
||||
}
|
||||
parsedURL, errParseURL := url.Parse(urlStr)
|
||||
if errParseURL != nil || parsedURL.Scheme == "" || parsedURL.Host == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
||||
return
|
||||
}
|
||||
|
||||
requestProxyURL := strings.TrimSpace(body.ProxyURL)
|
||||
if requestProxyURL != "" {
|
||||
if _, errParseProxy := proxyutil.Parse(requestProxyURL); errParseProxy != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid proxy_url"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
authIndex := firstNonEmptyString(body.AuthIndexSnake, body.AuthIndexCamel, body.AuthIndexPascal)
|
||||
auth := h.authByIndex(authIndex)
|
||||
|
||||
reqHeaders := body.Header
|
||||
if reqHeaders == nil {
|
||||
reqHeaders = map[string]string{}
|
||||
}
|
||||
|
||||
var hostOverride string
|
||||
var token string
|
||||
var tokenResolved bool
|
||||
var tokenErr error
|
||||
for key, value := range reqHeaders {
|
||||
if !strings.Contains(value, "$TOKEN$") {
|
||||
continue
|
||||
}
|
||||
if !tokenResolved {
|
||||
token, tokenErr = h.resolveTokenForAuth(c.Request.Context(), auth, requestProxyURL)
|
||||
tokenResolved = true
|
||||
}
|
||||
if auth != nil && token == "" {
|
||||
if tokenErr != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "auth token refresh failed"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "auth token not found"})
|
||||
return
|
||||
}
|
||||
if token == "" {
|
||||
continue
|
||||
}
|
||||
reqHeaders[key] = strings.ReplaceAll(value, "$TOKEN$", token)
|
||||
}
|
||||
|
||||
var requestBody io.Reader
|
||||
if body.Data != "" {
|
||||
requestBody = strings.NewReader(body.Data)
|
||||
}
|
||||
|
||||
req, errNewRequest := http.NewRequestWithContext(c.Request.Context(), method, urlStr, requestBody)
|
||||
if errNewRequest != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "failed to build request"})
|
||||
return
|
||||
}
|
||||
|
||||
for key, value := range reqHeaders {
|
||||
if strings.EqualFold(key, "host") {
|
||||
hostOverride = strings.TrimSpace(value)
|
||||
continue
|
||||
}
|
||||
req.Header.Set(key, value)
|
||||
}
|
||||
if hostOverride != "" {
|
||||
req.Host = hostOverride
|
||||
}
|
||||
|
||||
httpClient := &http.Client{
|
||||
Timeout: defaultAPICallTimeout,
|
||||
}
|
||||
httpClient.Transport = h.apiCallTransport(auth, requestProxyURL)
|
||||
|
||||
resp, errDo := httpClient.Do(req)
|
||||
if errDo != nil {
|
||||
log.WithError(errDo).Debug("management APICall request failed")
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "request failed"})
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if errClose := resp.Body.Close(); errClose != nil {
|
||||
log.Errorf("response body close error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
|
||||
respBody, errReadAll := io.ReadAll(resp.Body)
|
||||
if errReadAll != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "failed to read response"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, apiCallResponse{
|
||||
StatusCode: resp.StatusCode,
|
||||
Header: resp.Header,
|
||||
Body: string(respBody),
|
||||
})
|
||||
}
|
||||
|
||||
func firstNonEmptyString(values ...*string) string {
|
||||
for _, v := range values {
|
||||
if v == nil {
|
||||
continue
|
||||
}
|
||||
if out := strings.TrimSpace(*v); out != "" {
|
||||
return out
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func tokenValueForAuth(auth *coreauth.Auth) string {
|
||||
if auth == nil {
|
||||
return ""
|
||||
}
|
||||
if v := tokenValueFromMetadata(auth.Metadata); v != "" {
|
||||
return v
|
||||
}
|
||||
if auth.Attributes != nil {
|
||||
if v := strings.TrimSpace(auth.Attributes["api_key"]); v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (h *Handler) resolveTokenForAuth(ctx context.Context, auth *coreauth.Auth, requestProxyURL string) (string, error) {
|
||||
if auth == nil {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
if strings.EqualFold(strings.TrimSpace(auth.Provider), "antigravity") {
|
||||
token, errToken := h.refreshAntigravityOAuthAccessToken(ctx, auth, requestProxyURL)
|
||||
return token, errToken
|
||||
}
|
||||
|
||||
return tokenValueForAuth(auth), nil
|
||||
}
|
||||
|
||||
func (h *Handler) refreshAntigravityOAuthAccessToken(ctx context.Context, auth *coreauth.Auth, requestProxyURL string) (string, error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if auth == nil {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
metadata := auth.Metadata
|
||||
if len(metadata) == 0 {
|
||||
return "", fmt.Errorf("antigravity oauth metadata missing")
|
||||
}
|
||||
|
||||
current := strings.TrimSpace(tokenValueFromMetadata(metadata))
|
||||
if current != "" && !antigravityTokenNeedsRefresh(metadata) {
|
||||
return current, nil
|
||||
}
|
||||
|
||||
refreshToken := stringValue(metadata, "refresh_token")
|
||||
if refreshToken == "" {
|
||||
return "", fmt.Errorf("antigravity refresh token missing")
|
||||
}
|
||||
|
||||
tokenURL := strings.TrimSpace(antigravityOAuthTokenURL)
|
||||
if tokenURL == "" {
|
||||
tokenURL = "https://oauth2.googleapis.com/token"
|
||||
}
|
||||
form := url.Values{}
|
||||
form.Set("client_id", antigravityOAuthClientID)
|
||||
form.Set("client_secret", antigravityOAuthClientSecret)
|
||||
form.Set("grant_type", "refresh_token")
|
||||
form.Set("refresh_token", refreshToken)
|
||||
|
||||
req, errReq := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(form.Encode()))
|
||||
if errReq != nil {
|
||||
return "", errReq
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
|
||||
httpClient := &http.Client{
|
||||
Timeout: defaultAPICallTimeout,
|
||||
Transport: h.apiCallTransport(auth, requestProxyURL),
|
||||
}
|
||||
resp, errDo := httpClient.Do(req)
|
||||
if errDo != nil {
|
||||
return "", errDo
|
||||
}
|
||||
defer func() {
|
||||
if errClose := resp.Body.Close(); errClose != nil {
|
||||
log.Errorf("response body close error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
|
||||
bodyBytes, errRead := io.ReadAll(resp.Body)
|
||||
if errRead != nil {
|
||||
return "", errRead
|
||||
}
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
return "", fmt.Errorf("antigravity oauth token refresh failed: status %d: %s", resp.StatusCode, strings.TrimSpace(string(bodyBytes)))
|
||||
}
|
||||
|
||||
var tokenResp struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
RefreshToken string `json:"refresh_token"`
|
||||
ExpiresIn int64 `json:"expires_in"`
|
||||
TokenType string `json:"token_type"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(bodyBytes, &tokenResp); errUnmarshal != nil {
|
||||
return "", errUnmarshal
|
||||
}
|
||||
|
||||
if strings.TrimSpace(tokenResp.AccessToken) == "" {
|
||||
return "", fmt.Errorf("antigravity oauth token refresh returned empty access_token")
|
||||
}
|
||||
|
||||
if auth.Metadata == nil {
|
||||
auth.Metadata = make(map[string]any)
|
||||
}
|
||||
now := time.Now()
|
||||
auth.Metadata["access_token"] = strings.TrimSpace(tokenResp.AccessToken)
|
||||
if strings.TrimSpace(tokenResp.RefreshToken) != "" {
|
||||
auth.Metadata["refresh_token"] = strings.TrimSpace(tokenResp.RefreshToken)
|
||||
}
|
||||
if tokenResp.ExpiresIn > 0 {
|
||||
auth.Metadata["expires_in"] = tokenResp.ExpiresIn
|
||||
auth.Metadata["timestamp"] = now.UnixMilli()
|
||||
auth.Metadata["expired"] = now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339)
|
||||
}
|
||||
auth.Metadata["type"] = "antigravity"
|
||||
|
||||
if h != nil && h.authManager != nil {
|
||||
auth.LastRefreshedAt = now
|
||||
auth.UpdatedAt = now
|
||||
_, _ = h.authManager.Update(ctx, auth)
|
||||
}
|
||||
|
||||
return strings.TrimSpace(tokenResp.AccessToken), nil
|
||||
}
|
||||
|
||||
func antigravityTokenNeedsRefresh(metadata map[string]any) bool {
|
||||
// Refresh a bit early to avoid requests racing token expiry.
|
||||
const skew = 30 * time.Second
|
||||
|
||||
if metadata == nil {
|
||||
return true
|
||||
}
|
||||
if expStr, ok := metadata["expired"].(string); ok {
|
||||
if ts, errParse := time.Parse(time.RFC3339, strings.TrimSpace(expStr)); errParse == nil {
|
||||
return !ts.After(time.Now().Add(skew))
|
||||
}
|
||||
}
|
||||
expiresIn := int64Value(metadata["expires_in"])
|
||||
timestampMs := int64Value(metadata["timestamp"])
|
||||
if expiresIn > 0 && timestampMs > 0 {
|
||||
exp := time.UnixMilli(timestampMs).Add(time.Duration(expiresIn) * time.Second)
|
||||
return !exp.After(time.Now().Add(skew))
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func int64Value(raw any) int64 {
|
||||
switch typed := raw.(type) {
|
||||
case int:
|
||||
return int64(typed)
|
||||
case int32:
|
||||
return int64(typed)
|
||||
case int64:
|
||||
return typed
|
||||
case uint:
|
||||
return int64(typed)
|
||||
case uint32:
|
||||
return int64(typed)
|
||||
case uint64:
|
||||
if typed > uint64(^uint64(0)>>1) {
|
||||
return 0
|
||||
}
|
||||
return int64(typed)
|
||||
case float32:
|
||||
return int64(typed)
|
||||
case float64:
|
||||
return int64(typed)
|
||||
case json.Number:
|
||||
if i, errParse := typed.Int64(); errParse == nil {
|
||||
return i
|
||||
}
|
||||
case string:
|
||||
if s := strings.TrimSpace(typed); s != "" {
|
||||
if i, errParse := json.Number(s).Int64(); errParse == nil {
|
||||
return i
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func stringValue(metadata map[string]any, key string) string {
|
||||
if len(metadata) == 0 || key == "" {
|
||||
return ""
|
||||
}
|
||||
if v, ok := metadata[key].(string); ok {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func tokenValueFromMetadata(metadata map[string]any) string {
|
||||
if len(metadata) == 0 {
|
||||
return ""
|
||||
}
|
||||
if v, ok := metadata["accessToken"].(string); ok && strings.TrimSpace(v) != "" {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
if v, ok := metadata["access_token"].(string); ok && strings.TrimSpace(v) != "" {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
if tokenRaw, ok := metadata["token"]; ok && tokenRaw != nil {
|
||||
switch typed := tokenRaw.(type) {
|
||||
case string:
|
||||
if v := strings.TrimSpace(typed); v != "" {
|
||||
return v
|
||||
}
|
||||
case map[string]any:
|
||||
if v, ok := typed["access_token"].(string); ok && strings.TrimSpace(v) != "" {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
if v, ok := typed["accessToken"].(string); ok && strings.TrimSpace(v) != "" {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
case map[string]string:
|
||||
if v := strings.TrimSpace(typed["access_token"]); v != "" {
|
||||
return v
|
||||
}
|
||||
if v := strings.TrimSpace(typed["accessToken"]); v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
}
|
||||
if v, ok := metadata["token"].(string); ok && strings.TrimSpace(v) != "" {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
if v, ok := metadata["id_token"].(string); ok && strings.TrimSpace(v) != "" {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
if v, ok := metadata["cookie"].(string); ok && strings.TrimSpace(v) != "" {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (h *Handler) authByIndex(authIndex string) *coreauth.Auth {
|
||||
authIndex = strings.TrimSpace(authIndex)
|
||||
if authIndex == "" || h == nil || h.authManager == nil {
|
||||
return nil
|
||||
}
|
||||
auths := h.authManager.List()
|
||||
for _, auth := range auths {
|
||||
if auth == nil {
|
||||
continue
|
||||
}
|
||||
auth.EnsureIndex()
|
||||
if auth.Index == authIndex {
|
||||
return auth
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *Handler) apiCallTransport(auth *coreauth.Auth, requestProxyURL string) http.RoundTripper {
|
||||
if proxyStr := strings.TrimSpace(requestProxyURL); proxyStr != "" {
|
||||
if transport := buildProxyTransport(proxyStr); transport != nil {
|
||||
return transport
|
||||
}
|
||||
return directAPICallTransport()
|
||||
}
|
||||
|
||||
var proxyCandidates []string
|
||||
if auth != nil {
|
||||
if proxyStr := strings.TrimSpace(auth.ProxyURL); proxyStr != "" {
|
||||
proxyCandidates = append(proxyCandidates, proxyStr)
|
||||
}
|
||||
if h != nil && h.cfg != nil {
|
||||
if proxyStr := strings.TrimSpace(proxyURLFromAPIKeyConfig(h.cfg, auth)); proxyStr != "" {
|
||||
proxyCandidates = append(proxyCandidates, proxyStr)
|
||||
}
|
||||
}
|
||||
}
|
||||
if h != nil && h.cfg != nil {
|
||||
if proxyStr := strings.TrimSpace(h.cfg.ProxyURL); proxyStr != "" {
|
||||
proxyCandidates = append(proxyCandidates, proxyStr)
|
||||
}
|
||||
}
|
||||
|
||||
for _, proxyStr := range proxyCandidates {
|
||||
if transport := buildProxyTransport(proxyStr); transport != nil {
|
||||
return transport
|
||||
}
|
||||
}
|
||||
|
||||
return directAPICallTransport()
|
||||
}
|
||||
|
||||
func directAPICallTransport() http.RoundTripper {
|
||||
transport, ok := http.DefaultTransport.(*http.Transport)
|
||||
if !ok || transport == nil {
|
||||
return &http.Transport{Proxy: nil}
|
||||
}
|
||||
clone := transport.Clone()
|
||||
clone.Proxy = nil
|
||||
return clone
|
||||
}
|
||||
|
||||
type apiKeyConfigEntry interface {
|
||||
GetAPIKey() string
|
||||
GetBaseURL() string
|
||||
}
|
||||
|
||||
func resolveAPIKeyConfig[T apiKeyConfigEntry](entries []T, auth *coreauth.Auth) *T {
|
||||
if auth == nil || len(entries) == 0 {
|
||||
return nil
|
||||
}
|
||||
attrKey, attrBase := "", ""
|
||||
if auth.Attributes != nil {
|
||||
attrKey = strings.TrimSpace(auth.Attributes["api_key"])
|
||||
attrBase = strings.TrimSpace(auth.Attributes["base_url"])
|
||||
}
|
||||
for i := range entries {
|
||||
entry := &entries[i]
|
||||
cfgKey := strings.TrimSpace((*entry).GetAPIKey())
|
||||
cfgBase := strings.TrimSpace((*entry).GetBaseURL())
|
||||
if attrKey != "" && attrBase != "" {
|
||||
if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
continue
|
||||
}
|
||||
if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
|
||||
if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
if attrKey != "" {
|
||||
for i := range entries {
|
||||
entry := &entries[i]
|
||||
if strings.EqualFold(strings.TrimSpace((*entry).GetAPIKey()), attrKey) {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func proxyURLFromAPIKeyConfig(cfg *config.Config, auth *coreauth.Auth) string {
|
||||
if cfg == nil || auth == nil {
|
||||
return ""
|
||||
}
|
||||
authKind, authAccount := auth.AccountInfo()
|
||||
if !strings.EqualFold(strings.TrimSpace(authKind), "api_key") {
|
||||
return ""
|
||||
}
|
||||
|
||||
attrs := auth.Attributes
|
||||
compatName := ""
|
||||
providerKey := ""
|
||||
if len(attrs) > 0 {
|
||||
compatName = strings.TrimSpace(attrs["compat_name"])
|
||||
providerKey = strings.TrimSpace(attrs["provider_key"])
|
||||
}
|
||||
if compatName != "" || strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") {
|
||||
return resolveOpenAICompatAPIKeyProxyURL(cfg, auth, strings.TrimSpace(authAccount), providerKey, compatName)
|
||||
}
|
||||
|
||||
switch strings.ToLower(strings.TrimSpace(auth.Provider)) {
|
||||
case "gemini":
|
||||
if entry := resolveAPIKeyConfig(cfg.GeminiKey, auth); entry != nil {
|
||||
return strings.TrimSpace(entry.ProxyURL)
|
||||
}
|
||||
case "gemini-interactions":
|
||||
if entry := resolveAPIKeyConfig(cfg.InteractionsKey, auth); entry != nil {
|
||||
return strings.TrimSpace(entry.ProxyURL)
|
||||
}
|
||||
case "claude":
|
||||
if entry := resolveAPIKeyConfig(cfg.ClaudeKey, auth); entry != nil {
|
||||
return strings.TrimSpace(entry.ProxyURL)
|
||||
}
|
||||
case "codex":
|
||||
if entry := resolveAPIKeyConfig(cfg.CodexKey, auth); entry != nil {
|
||||
return strings.TrimSpace(entry.ProxyURL)
|
||||
}
|
||||
case "xai":
|
||||
if entry := resolveAPIKeyConfig(cfg.XAIKey, auth); entry != nil {
|
||||
return strings.TrimSpace(entry.ProxyURL)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func resolveOpenAICompatAPIKeyProxyURL(cfg *config.Config, auth *coreauth.Auth, apiKey, providerKey, compatName string) string {
|
||||
if cfg == nil || auth == nil {
|
||||
return ""
|
||||
}
|
||||
apiKey = strings.TrimSpace(apiKey)
|
||||
if apiKey == "" {
|
||||
return ""
|
||||
}
|
||||
candidates := make([]string, 0, 3)
|
||||
if v := strings.TrimSpace(compatName); v != "" {
|
||||
candidates = append(candidates, v)
|
||||
}
|
||||
if v := strings.TrimSpace(providerKey); v != "" {
|
||||
candidates = append(candidates, v)
|
||||
}
|
||||
if v := strings.TrimSpace(auth.Provider); v != "" {
|
||||
candidates = append(candidates, v)
|
||||
}
|
||||
|
||||
for i := range cfg.OpenAICompatibility {
|
||||
compat := &cfg.OpenAICompatibility[i]
|
||||
if compat.Disabled {
|
||||
continue
|
||||
}
|
||||
for _, candidate := range candidates {
|
||||
if candidate != "" && strings.EqualFold(strings.TrimSpace(candidate), compat.Name) {
|
||||
for j := range compat.APIKeyEntries {
|
||||
entry := &compat.APIKeyEntries[j]
|
||||
if strings.EqualFold(strings.TrimSpace(entry.APIKey), apiKey) {
|
||||
return strings.TrimSpace(entry.ProxyURL)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func buildProxyTransport(proxyStr string) *http.Transport {
|
||||
transport, _, errBuild := proxyutil.BuildHTTPTransport(proxyStr)
|
||||
if errBuild != nil {
|
||||
log.WithError(errBuild).Debug("build proxy transport failed")
|
||||
return nil
|
||||
}
|
||||
return transport
|
||||
}
|
||||
317
backend/internal/api/handlers/management/api_tools_test.go
Normal file
317
backend/internal/api/handlers/management/api_tools_test.go
Normal file
|
|
@ -0,0 +1,317 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
)
|
||||
|
||||
func TestAPICallUsesRequestProxyURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
proxyServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
_, _ = w.Write([]byte("proxied"))
|
||||
}))
|
||||
defer proxyServer.Close()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
SDKConfig: sdkconfig.SDKConfig{ProxyURL: "http://127.0.0.1:1"},
|
||||
},
|
||||
}
|
||||
router := gin.New()
|
||||
router.POST("/", h.APICall)
|
||||
|
||||
body := `{"method":"GET","url":"http://upstream.invalid/test","proxy_url":"` + proxyServer.URL + `"}`
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Fatalf("status code = %d, want %d; body = %s", recorder.Code, http.StatusOK, recorder.Body.String())
|
||||
}
|
||||
|
||||
var response apiCallResponse
|
||||
if errDecode := json.NewDecoder(recorder.Body).Decode(&response); errDecode != nil {
|
||||
t.Fatalf("decode response: %v", errDecode)
|
||||
}
|
||||
if response.StatusCode != http.StatusCreated {
|
||||
t.Fatalf("upstream status code = %d, want %d", response.StatusCode, http.StatusCreated)
|
||||
}
|
||||
if response.Body != "proxied" {
|
||||
t.Fatalf("upstream body = %q, want %q", response.Body, "proxied")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPICallTransportDirectBypassesGlobalProxy(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
SDKConfig: sdkconfig.SDKConfig{ProxyURL: "http://global-proxy.example.com:8080"},
|
||||
},
|
||||
}
|
||||
|
||||
transport := h.apiCallTransport(&coreauth.Auth{ProxyURL: "direct"}, "")
|
||||
httpTransport, ok := transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport type = %T, want *http.Transport", transport)
|
||||
}
|
||||
if httpTransport.Proxy != nil {
|
||||
t.Fatal("expected direct transport to disable proxy function")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPICallTransportInvalidAuthFallsBackToGlobalProxy(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
SDKConfig: sdkconfig.SDKConfig{ProxyURL: "http://global-proxy.example.com:8080"},
|
||||
},
|
||||
}
|
||||
|
||||
transport := h.apiCallTransport(&coreauth.Auth{ProxyURL: "bad-value"}, "")
|
||||
httpTransport, ok := transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport type = %T, want *http.Transport", transport)
|
||||
}
|
||||
|
||||
req, errRequest := http.NewRequest(http.MethodGet, "https://example.com", nil)
|
||||
if errRequest != nil {
|
||||
t.Fatalf("http.NewRequest returned error: %v", errRequest)
|
||||
}
|
||||
|
||||
proxyURL, errProxy := httpTransport.Proxy(req)
|
||||
if errProxy != nil {
|
||||
t.Fatalf("httpTransport.Proxy returned error: %v", errProxy)
|
||||
}
|
||||
if proxyURL == nil || proxyURL.String() != "http://global-proxy.example.com:8080" {
|
||||
t.Fatalf("proxy URL = %v, want http://global-proxy.example.com:8080", proxyURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPICallTransportRequestProxyOverridesCredentialAndGlobalProxy(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
SDKConfig: sdkconfig.SDKConfig{ProxyURL: "http://global-proxy.example.com:8080"},
|
||||
},
|
||||
}
|
||||
auth := &coreauth.Auth{ProxyURL: "http://credential-proxy.example.com:8080"}
|
||||
|
||||
transport := h.apiCallTransport(auth, " http://request-proxy.example.com:8080 ")
|
||||
httpTransport, ok := transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport type = %T, want *http.Transport", transport)
|
||||
}
|
||||
|
||||
req, errRequest := http.NewRequest(http.MethodGet, "https://example.com", nil)
|
||||
if errRequest != nil {
|
||||
t.Fatalf("http.NewRequest returned error: %v", errRequest)
|
||||
}
|
||||
|
||||
proxyURL, errProxy := httpTransport.Proxy(req)
|
||||
if errProxy != nil {
|
||||
t.Fatalf("httpTransport.Proxy returned error: %v", errProxy)
|
||||
}
|
||||
if proxyURL == nil || proxyURL.String() != "http://request-proxy.example.com:8080" {
|
||||
t.Fatalf("proxy URL = %v, want http://request-proxy.example.com:8080", proxyURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPICallTransportInvalidRequestProxyDoesNotFallBack(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
SDKConfig: sdkconfig.SDKConfig{ProxyURL: "http://global-proxy.example.com:8080"},
|
||||
},
|
||||
}
|
||||
auth := &coreauth.Auth{ProxyURL: "http://credential-proxy.example.com:8080"}
|
||||
|
||||
transport := h.apiCallTransport(auth, "bad-value")
|
||||
httpTransport, ok := transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport type = %T, want *http.Transport", transport)
|
||||
}
|
||||
if httpTransport.Proxy != nil {
|
||||
t.Fatal("expected invalid request proxy to avoid lower-priority proxy settings")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPICallTransportAPIKeyAuthFallsBackToConfigProxyURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
SDKConfig: sdkconfig.SDKConfig{ProxyURL: "http://global-proxy.example.com:8080"},
|
||||
GeminiKey: []config.GeminiKey{{
|
||||
APIKey: "gemini-key",
|
||||
ProxyURL: "http://gemini-proxy.example.com:8080",
|
||||
}},
|
||||
ClaudeKey: []config.ClaudeKey{{
|
||||
APIKey: "claude-key",
|
||||
ProxyURL: "http://claude-proxy.example.com:8080",
|
||||
}},
|
||||
CodexKey: []config.CodexKey{{
|
||||
APIKey: "codex-key",
|
||||
ProxyURL: "http://codex-proxy.example.com:8080",
|
||||
}},
|
||||
XAIKey: []config.XAIKey{{
|
||||
APIKey: "xai-key",
|
||||
ProxyURL: "http://xai-proxy.example.com:8080",
|
||||
}},
|
||||
OpenAICompatibility: []config.OpenAICompatibility{{
|
||||
Name: "bohe",
|
||||
BaseURL: "https://bohe.example.com",
|
||||
APIKeyEntries: []config.OpenAICompatibilityAPIKey{{
|
||||
APIKey: "compat-key",
|
||||
ProxyURL: "http://compat-proxy.example.com:8080",
|
||||
}},
|
||||
}},
|
||||
},
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
auth *coreauth.Auth
|
||||
wantProxy string
|
||||
}{
|
||||
{
|
||||
name: "gemini",
|
||||
auth: &coreauth.Auth{
|
||||
Provider: "gemini",
|
||||
Attributes: map[string]string{"api_key": "gemini-key"},
|
||||
},
|
||||
wantProxy: "http://gemini-proxy.example.com:8080",
|
||||
},
|
||||
{
|
||||
name: "claude",
|
||||
auth: &coreauth.Auth{
|
||||
Provider: "claude",
|
||||
Attributes: map[string]string{"api_key": "claude-key"},
|
||||
},
|
||||
wantProxy: "http://claude-proxy.example.com:8080",
|
||||
},
|
||||
{
|
||||
name: "codex",
|
||||
auth: &coreauth.Auth{
|
||||
Provider: "codex",
|
||||
Attributes: map[string]string{"api_key": "codex-key"},
|
||||
},
|
||||
wantProxy: "http://codex-proxy.example.com:8080",
|
||||
},
|
||||
{
|
||||
name: "xai",
|
||||
auth: &coreauth.Auth{
|
||||
Provider: "xai",
|
||||
Attributes: map[string]string{"api_key": "xai-key"},
|
||||
},
|
||||
wantProxy: "http://xai-proxy.example.com:8080",
|
||||
},
|
||||
{
|
||||
name: "openai-compatibility",
|
||||
auth: &coreauth.Auth{
|
||||
Provider: "bohe",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "compat-key",
|
||||
"compat_name": "bohe",
|
||||
"provider_key": "bohe",
|
||||
},
|
||||
},
|
||||
wantProxy: "http://compat-proxy.example.com:8080",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
transport := h.apiCallTransport(tc.auth, "")
|
||||
httpTransport, ok := transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport type = %T, want *http.Transport", transport)
|
||||
}
|
||||
|
||||
req, errRequest := http.NewRequest(http.MethodGet, "https://example.com", nil)
|
||||
if errRequest != nil {
|
||||
t.Fatalf("http.NewRequest returned error: %v", errRequest)
|
||||
}
|
||||
|
||||
proxyURL, errProxy := httpTransport.Proxy(req)
|
||||
if errProxy != nil {
|
||||
t.Fatalf("httpTransport.Proxy returned error: %v", errProxy)
|
||||
}
|
||||
if proxyURL == nil || proxyURL.String() != tc.wantProxy {
|
||||
t.Fatalf("proxy URL = %v, want %s", proxyURL, tc.wantProxy)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthByIndexDistinguishesSharedAPIKeysAcrossProviders(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
geminiAuth := &coreauth.Auth{
|
||||
ID: "gemini:apikey:123",
|
||||
Provider: "gemini",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "shared-key",
|
||||
},
|
||||
}
|
||||
compatAuth := &coreauth.Auth{
|
||||
ID: "openai-compatibility:bohe:456",
|
||||
Provider: "bohe",
|
||||
Label: "bohe",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "shared-key",
|
||||
"compat_name": "bohe",
|
||||
"provider_key": "bohe",
|
||||
},
|
||||
}
|
||||
|
||||
if _, errRegister := manager.Register(context.Background(), geminiAuth); errRegister != nil {
|
||||
t.Fatalf("register gemini auth: %v", errRegister)
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), compatAuth); errRegister != nil {
|
||||
t.Fatalf("register compat auth: %v", errRegister)
|
||||
}
|
||||
|
||||
geminiIndex := geminiAuth.EnsureIndex()
|
||||
compatIndex := compatAuth.EnsureIndex()
|
||||
if geminiIndex == compatIndex {
|
||||
t.Fatalf("shared api key produced duplicate auth_index %q", geminiIndex)
|
||||
}
|
||||
|
||||
h := &Handler{authManager: manager}
|
||||
|
||||
gotGemini := h.authByIndex(geminiIndex)
|
||||
if gotGemini == nil {
|
||||
t.Fatal("expected gemini auth by index")
|
||||
}
|
||||
if gotGemini.ID != geminiAuth.ID {
|
||||
t.Fatalf("authByIndex(gemini) returned %q, want %q", gotGemini.ID, geminiAuth.ID)
|
||||
}
|
||||
|
||||
gotCompat := h.authByIndex(compatIndex)
|
||||
if gotCompat == nil {
|
||||
t.Fatal("expected compat auth by index")
|
||||
}
|
||||
if gotCompat.ID != compatAuth.ID {
|
||||
t.Fatalf("authByIndex(compat) returned %q, want %q", gotCompat.ID, compatAuth.ID)
|
||||
}
|
||||
}
|
||||
598
backend/internal/api/handlers/management/auth_files.go
Normal file
598
backend/internal/api/handlers/management/auth_files.go
Normal file
|
|
@ -0,0 +1,598 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/auth/codex"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/credentialweight"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
var lastRefreshKeys = []string{"last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"}
|
||||
|
||||
var (
|
||||
callbackForwardersMu sync.Mutex
|
||||
callbackForwarders = make(map[int]*callbackForwarder)
|
||||
authFileEntryMu sync.Mutex
|
||||
errAuthFileMustBeJSON = errors.New("auth file must be .json")
|
||||
errAuthFileNotFound = errors.New("auth file not found")
|
||||
errPluginVirtualAuth = errors.New("plugin virtual auth cannot be modified directly; edit or delete the source auth file")
|
||||
newCodexOAuthService = func(cfg *config.Config) codexOAuthService { return codex.NewCodexAuth(cfg) }
|
||||
)
|
||||
|
||||
func extractLastRefreshTimestamp(meta map[string]any) (time.Time, bool) {
|
||||
if len(meta) == 0 {
|
||||
return time.Time{}, false
|
||||
}
|
||||
for _, key := range lastRefreshKeys {
|
||||
if val, ok := meta[key]; ok {
|
||||
if ts, ok1 := parseLastRefreshValue(val); ok1 {
|
||||
return ts, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return time.Time{}, false
|
||||
}
|
||||
|
||||
func parseLastRefreshValue(v any) (time.Time, bool) {
|
||||
switch val := v.(type) {
|
||||
case string:
|
||||
s := strings.TrimSpace(val)
|
||||
if s == "" {
|
||||
return time.Time{}, false
|
||||
}
|
||||
layouts := []string{time.RFC3339, time.RFC3339Nano, "2006-01-02 15:04:05", "2006-01-02T15:04:05Z07:00"}
|
||||
for _, layout := range layouts {
|
||||
if ts, err := time.Parse(layout, s); err == nil {
|
||||
return ts.UTC(), true
|
||||
}
|
||||
}
|
||||
if unix, err := strconv.ParseInt(s, 10, 64); err == nil && unix > 0 {
|
||||
return time.Unix(unix, 0).UTC(), true
|
||||
}
|
||||
case float64:
|
||||
if val <= 0 {
|
||||
return time.Time{}, false
|
||||
}
|
||||
return time.Unix(int64(val), 0).UTC(), true
|
||||
case int64:
|
||||
if val <= 0 {
|
||||
return time.Time{}, false
|
||||
}
|
||||
return time.Unix(val, 0).UTC(), true
|
||||
case int:
|
||||
if val <= 0 {
|
||||
return time.Time{}, false
|
||||
}
|
||||
return time.Unix(int64(val), 0).UTC(), true
|
||||
case json.Number:
|
||||
if i, err := val.Int64(); err == nil && i > 0 {
|
||||
return time.Unix(i, 0).UTC(), true
|
||||
}
|
||||
}
|
||||
return time.Time{}, false
|
||||
}
|
||||
|
||||
func (h *Handler) ListAuthFiles(c *gin.Context) {
|
||||
if h == nil {
|
||||
c.JSON(500, gin.H{"error": "handler not initialized"})
|
||||
return
|
||||
}
|
||||
if h.authManager == nil {
|
||||
h.listAuthFilesFromDisk(c)
|
||||
return
|
||||
}
|
||||
nameFilter := strings.TrimSpace(c.Query("name"))
|
||||
authIndexFilter := strings.TrimSpace(c.Query("auth_index"))
|
||||
auths := h.authManager.List()
|
||||
files := make([]gin.H, 0, len(auths))
|
||||
for _, auth := range auths {
|
||||
if !matchesAuthFileLookup(auth, nameFilter, authIndexFilter) {
|
||||
continue
|
||||
}
|
||||
if entry := h.buildAuthFileEntry(auth); entry != nil {
|
||||
files = append(files, entry)
|
||||
}
|
||||
}
|
||||
sort.Slice(files, func(i, j int) bool {
|
||||
nameI, _ := files[i]["name"].(string)
|
||||
nameJ, _ := files[j]["name"].(string)
|
||||
return strings.ToLower(nameI) < strings.ToLower(nameJ)
|
||||
})
|
||||
c.JSON(200, gin.H{"files": files})
|
||||
}
|
||||
|
||||
func lockedAuthIndex(auth *coreauth.Auth) string {
|
||||
if auth == nil {
|
||||
return ""
|
||||
}
|
||||
authFileEntryMu.Lock()
|
||||
defer authFileEntryMu.Unlock()
|
||||
return strings.TrimSpace(auth.EnsureIndex())
|
||||
}
|
||||
|
||||
func matchesAuthFileLookup(auth *coreauth.Auth, name string, authIndex string) bool {
|
||||
if auth == nil {
|
||||
return false
|
||||
}
|
||||
if name != "" && strings.TrimSpace(auth.ID) != name && strings.TrimSpace(auth.FileName) != name {
|
||||
return false
|
||||
}
|
||||
if authIndex != "" && lockedAuthIndex(auth) != authIndex {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *Handler) lookupAuthFile(name string, authIndex string) (*coreauth.Auth, bool) {
|
||||
name = strings.TrimSpace(name)
|
||||
authIndex = strings.TrimSpace(authIndex)
|
||||
if h == nil || h.authManager == nil || name == "" {
|
||||
return nil, false
|
||||
}
|
||||
if authIndex == "" {
|
||||
if auth, ok := h.authManager.GetByID(name); ok {
|
||||
return auth, true
|
||||
}
|
||||
auths := h.authManager.List()
|
||||
for _, auth := range auths {
|
||||
if auth != nil && strings.TrimSpace(auth.FileName) == name {
|
||||
return auth, true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
auths := h.authManager.List()
|
||||
for _, auth := range auths {
|
||||
if matchesAuthFileLookup(auth, name, authIndex) {
|
||||
return auth, true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// GetAuthFileModels returns the models supported by a specific auth file
|
||||
func (h *Handler) GetAuthFileModels(c *gin.Context) {
|
||||
name := c.Query("name")
|
||||
if name == "" {
|
||||
c.JSON(400, gin.H{"error": "name is required"})
|
||||
return
|
||||
}
|
||||
|
||||
// Try to find auth ID via authManager
|
||||
var authID string
|
||||
if h.authManager != nil {
|
||||
auths := h.authManager.List()
|
||||
for _, auth := range auths {
|
||||
if auth.FileName == name || auth.ID == name {
|
||||
authID = auth.ID
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if authID == "" {
|
||||
authID = name // fallback to filename as ID
|
||||
}
|
||||
|
||||
// Get models from registry
|
||||
reg := registry.GetGlobalRegistry()
|
||||
models := reg.GetModelsForClient(authID)
|
||||
|
||||
result := make([]gin.H, 0, len(models))
|
||||
for _, m := range models {
|
||||
entry := gin.H{
|
||||
"id": m.ID,
|
||||
}
|
||||
if m.DisplayName != "" {
|
||||
entry["display_name"] = m.DisplayName
|
||||
}
|
||||
if m.Type != "" {
|
||||
entry["type"] = m.Type
|
||||
}
|
||||
if m.OwnedBy != "" {
|
||||
entry["owned_by"] = m.OwnedBy
|
||||
}
|
||||
result = append(result, entry)
|
||||
}
|
||||
|
||||
c.JSON(200, gin.H{"models": result})
|
||||
}
|
||||
|
||||
// List auth files from disk when the auth manager is unavailable.
|
||||
func (h *Handler) listAuthFilesFromDisk(c *gin.Context) {
|
||||
nameFilter := strings.TrimSpace(c.Query("name"))
|
||||
authIndexFilter := strings.TrimSpace(c.Query("auth_index"))
|
||||
entries, err := os.ReadDir(h.cfg.AuthDir)
|
||||
if err != nil {
|
||||
c.JSON(500, gin.H{"error": fmt.Sprintf("failed to read auth dir: %v", err)})
|
||||
return
|
||||
}
|
||||
files := make([]gin.H, 0)
|
||||
if authIndexFilter != "" {
|
||||
c.JSON(200, gin.H{"files": files})
|
||||
return
|
||||
}
|
||||
for _, e := range entries {
|
||||
if e.IsDir() {
|
||||
continue
|
||||
}
|
||||
name := e.Name()
|
||||
if nameFilter != "" && name != nameFilter {
|
||||
continue
|
||||
}
|
||||
if !strings.HasSuffix(strings.ToLower(name), ".json") {
|
||||
continue
|
||||
}
|
||||
if info, errInfo := e.Info(); errInfo == nil {
|
||||
fileData := gin.H{"name": name, "size": info.Size(), "modtime": info.ModTime()}
|
||||
|
||||
// Read file to get type field
|
||||
full := filepath.Join(h.cfg.AuthDir, name)
|
||||
if data, errRead := os.ReadFile(full); errRead == nil {
|
||||
typeValue := gjson.GetBytes(data, "type").String()
|
||||
emailValue := gjson.GetBytes(data, "email").String()
|
||||
fileData["type"] = typeValue
|
||||
fileData["email"] = emailValue
|
||||
if projectID := strings.TrimSpace(gjson.GetBytes(data, "project_id").String()); projectID != "" {
|
||||
fileData["project_id"] = projectID
|
||||
}
|
||||
if pv := gjson.GetBytes(data, "priority"); pv.Exists() {
|
||||
switch pv.Type {
|
||||
case gjson.Number:
|
||||
fileData["priority"] = int(pv.Int())
|
||||
case gjson.String:
|
||||
if parsed, errAtoi := strconv.Atoi(strings.TrimSpace(pv.String())); errAtoi == nil {
|
||||
fileData["priority"] = parsed
|
||||
}
|
||||
}
|
||||
}
|
||||
if wv := gjson.GetBytes(data, coreauth.AttributeWeight); wv.Exists() {
|
||||
var rawWeight string
|
||||
switch wv.Type {
|
||||
case gjson.Number:
|
||||
rawWeight = wv.Raw
|
||||
case gjson.String:
|
||||
rawWeight = wv.String()
|
||||
}
|
||||
if rawWeight != "" {
|
||||
if weight, errWeight := credentialweight.ParseString(rawWeight); errWeight == nil {
|
||||
fileData[coreauth.AttributeWeight] = weight
|
||||
}
|
||||
}
|
||||
}
|
||||
if nv := gjson.GetBytes(data, "note"); nv.Exists() && nv.Type == gjson.String {
|
||||
if trimmed := strings.TrimSpace(nv.String()); trimmed != "" {
|
||||
fileData["note"] = trimmed
|
||||
}
|
||||
}
|
||||
if wv := gjson.GetBytes(data, "websockets"); wv.Exists() {
|
||||
switch wv.Type {
|
||||
case gjson.True:
|
||||
fileData["websockets"] = true
|
||||
case gjson.False:
|
||||
fileData["websockets"] = false
|
||||
case gjson.String:
|
||||
if parsed, errParse := strconv.ParseBool(strings.TrimSpace(wv.String())); errParse == nil {
|
||||
fileData["websockets"] = parsed
|
||||
}
|
||||
}
|
||||
}
|
||||
if requestRetry, okRetry := authFileRequestRetryFromJSON(data); okRetry {
|
||||
fileData["request_retry"] = requestRetry
|
||||
}
|
||||
}
|
||||
|
||||
files = append(files, fileData)
|
||||
}
|
||||
}
|
||||
c.JSON(200, gin.H{"files": files})
|
||||
}
|
||||
|
||||
func (h *Handler) buildAuthFileEntry(auth *coreauth.Auth) gin.H {
|
||||
authFileEntryMu.Lock()
|
||||
defer authFileEntryMu.Unlock()
|
||||
return h.buildAuthFileEntryLocked(auth)
|
||||
}
|
||||
|
||||
func (h *Handler) buildAuthFileEntryLocked(auth *coreauth.Auth) gin.H {
|
||||
if auth == nil {
|
||||
return nil
|
||||
}
|
||||
auth.EnsureIndex()
|
||||
runtimeOnly := isRuntimeOnlyAuth(auth)
|
||||
if runtimeOnly && (auth.Disabled || auth.Status == coreauth.StatusDisabled) {
|
||||
return nil
|
||||
}
|
||||
path := strings.TrimSpace(authAttribute(auth, "path"))
|
||||
if path == "" && !runtimeOnly {
|
||||
return nil
|
||||
}
|
||||
name := strings.TrimSpace(auth.FileName)
|
||||
if name == "" {
|
||||
name = auth.ID
|
||||
}
|
||||
entry := gin.H{
|
||||
"id": auth.ID,
|
||||
"auth_index": auth.Index,
|
||||
"name": name,
|
||||
"type": strings.TrimSpace(auth.Provider),
|
||||
"provider": strings.TrimSpace(auth.Provider),
|
||||
"label": auth.Label,
|
||||
"status": auth.Status,
|
||||
"status_message": auth.StatusMessage,
|
||||
"disabled": auth.Disabled,
|
||||
"unavailable": auth.Unavailable,
|
||||
"runtime_only": runtimeOnly,
|
||||
"source": "memory",
|
||||
"size": int64(0),
|
||||
}
|
||||
entry["success"] = auth.Success
|
||||
entry["failed"] = auth.Failed
|
||||
entry["recent_requests"] = auth.RecentRequestsSnapshot(time.Now())
|
||||
if email := authEmail(auth); email != "" {
|
||||
entry["email"] = email
|
||||
}
|
||||
if projectID := authProjectID(auth); projectID != "" {
|
||||
entry["project_id"] = projectID
|
||||
}
|
||||
if accountType, account := auth.AccountInfo(); accountType != "" || account != "" {
|
||||
if accountType != "" {
|
||||
entry["account_type"] = accountType
|
||||
}
|
||||
if account != "" {
|
||||
entry["account"] = account
|
||||
}
|
||||
}
|
||||
if !auth.CreatedAt.IsZero() {
|
||||
entry["created_at"] = auth.CreatedAt
|
||||
}
|
||||
if !auth.UpdatedAt.IsZero() {
|
||||
entry["modtime"] = auth.UpdatedAt
|
||||
entry["updated_at"] = auth.UpdatedAt
|
||||
}
|
||||
if !auth.LastRefreshedAt.IsZero() {
|
||||
entry["last_refresh"] = auth.LastRefreshedAt
|
||||
}
|
||||
if !auth.NextRetryAfter.IsZero() {
|
||||
entry["next_retry_after"] = auth.NextRetryAfter
|
||||
}
|
||||
if path != "" {
|
||||
entry["path"] = path
|
||||
entry["source"] = "file"
|
||||
if info, err := os.Stat(path); err == nil {
|
||||
entry["size"] = info.Size()
|
||||
entry["modtime"] = info.ModTime()
|
||||
} else if os.IsNotExist(err) {
|
||||
// Hide credentials removed from disk but still lingering in memory.
|
||||
if !runtimeOnly && (auth.Disabled || auth.Status == coreauth.StatusDisabled || strings.EqualFold(strings.TrimSpace(auth.StatusMessage), "removed via management api")) {
|
||||
return nil
|
||||
}
|
||||
entry["source"] = "memory"
|
||||
} else {
|
||||
log.WithError(err).Warnf("failed to stat auth file %s", path)
|
||||
}
|
||||
}
|
||||
if claims := extractCodexIDTokenClaims(auth); claims != nil {
|
||||
entry["id_token"] = claims
|
||||
}
|
||||
// Expose priority from Attributes (set by synthesizer from JSON "priority" field).
|
||||
// Fall back to Metadata for auths registered via UploadAuthFile (no synthesizer).
|
||||
if p := strings.TrimSpace(authAttribute(auth, "priority")); p != "" {
|
||||
if parsed, err := strconv.Atoi(p); err == nil {
|
||||
entry["priority"] = parsed
|
||||
}
|
||||
} else if auth.Metadata != nil {
|
||||
if rawPriority, ok := auth.Metadata["priority"]; ok {
|
||||
switch v := rawPriority.(type) {
|
||||
case float64:
|
||||
entry["priority"] = int(v)
|
||||
case int:
|
||||
entry["priority"] = v
|
||||
case string:
|
||||
if parsed, err := strconv.Atoi(strings.TrimSpace(v)); err == nil {
|
||||
entry["priority"] = parsed
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Expose note from Attributes (set by synthesizer from JSON "note" field).
|
||||
// Fall back to Metadata for auths registered via UploadAuthFile (no synthesizer).
|
||||
if note := strings.TrimSpace(authAttribute(auth, "note")); note != "" {
|
||||
entry["note"] = note
|
||||
} else if auth.Metadata != nil {
|
||||
if rawNote, ok := auth.Metadata["note"].(string); ok {
|
||||
if trimmed := strings.TrimSpace(rawNote); trimmed != "" {
|
||||
entry["note"] = trimmed
|
||||
}
|
||||
}
|
||||
}
|
||||
if weight, ok := authWeightValue(auth); ok {
|
||||
entry[coreauth.AttributeWeight] = weight
|
||||
}
|
||||
if websockets, ok := authWebsocketsValue(auth); ok {
|
||||
entry["websockets"] = websockets
|
||||
}
|
||||
if requestRetry, ok := auth.RequestRetryOverride(); ok {
|
||||
entry["request_retry"] = requestRetry
|
||||
}
|
||||
return entry
|
||||
}
|
||||
|
||||
func authFileRequestRetryFromJSON(data []byte) (int, bool) {
|
||||
var metadata map[string]any
|
||||
if errUnmarshal := json.Unmarshal(data, &metadata); errUnmarshal != nil {
|
||||
return 0, false
|
||||
}
|
||||
return (&coreauth.Auth{Metadata: metadata}).RequestRetryOverride()
|
||||
}
|
||||
|
||||
func authWeightValue(auth *coreauth.Auth) (int64, bool) {
|
||||
if auth == nil {
|
||||
return 0, false
|
||||
}
|
||||
if rawWeight := strings.TrimSpace(authAttribute(auth, coreauth.AttributeWeight)); rawWeight != "" {
|
||||
weight, errWeight := credentialweight.ParseString(rawWeight)
|
||||
return weight, errWeight == nil
|
||||
}
|
||||
if auth.Metadata == nil {
|
||||
return 0, false
|
||||
}
|
||||
rawWeight, ok := auth.Metadata[coreauth.AttributeWeight]
|
||||
if !ok || rawWeight == nil {
|
||||
return 0, false
|
||||
}
|
||||
weight, errWeight := credentialweight.ParseValue(rawWeight)
|
||||
return weight, errWeight == nil
|
||||
}
|
||||
|
||||
func authWebsocketsValue(auth *coreauth.Auth) (bool, bool) {
|
||||
if auth == nil {
|
||||
return false, false
|
||||
}
|
||||
if auth.Attributes != nil {
|
||||
if raw := strings.TrimSpace(auth.Attributes["websockets"]); raw != "" {
|
||||
parsed, errParse := strconv.ParseBool(raw)
|
||||
if errParse == nil {
|
||||
return parsed, true
|
||||
}
|
||||
}
|
||||
}
|
||||
if auth.Metadata == nil {
|
||||
return false, false
|
||||
}
|
||||
raw, ok := auth.Metadata["websockets"]
|
||||
if !ok || raw == nil {
|
||||
return false, false
|
||||
}
|
||||
switch v := raw.(type) {
|
||||
case bool:
|
||||
return v, true
|
||||
case string:
|
||||
parsed, errParse := strconv.ParseBool(strings.TrimSpace(v))
|
||||
if errParse == nil {
|
||||
return parsed, true
|
||||
}
|
||||
}
|
||||
return false, false
|
||||
}
|
||||
|
||||
func authProjectID(auth *coreauth.Auth) string {
|
||||
if auth == nil {
|
||||
return ""
|
||||
}
|
||||
if auth.Metadata != nil {
|
||||
if v, ok := auth.Metadata["project_id"].(string); ok {
|
||||
if projectID := strings.TrimSpace(v); projectID != "" {
|
||||
return projectID
|
||||
}
|
||||
}
|
||||
}
|
||||
if auth.Attributes != nil {
|
||||
if projectID := strings.TrimSpace(auth.Attributes["project_id"]); projectID != "" {
|
||||
return projectID
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func extractCodexIDTokenClaims(auth *coreauth.Auth) gin.H {
|
||||
if auth == nil || auth.Metadata == nil {
|
||||
return nil
|
||||
}
|
||||
if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") {
|
||||
return nil
|
||||
}
|
||||
idTokenRaw, ok := auth.Metadata["id_token"].(string)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
idToken := strings.TrimSpace(idTokenRaw)
|
||||
if idToken == "" {
|
||||
return nil
|
||||
}
|
||||
claims, err := codex.ParseJWTToken(idToken)
|
||||
if err != nil || claims == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
result := gin.H{}
|
||||
if v := strings.TrimSpace(claims.CodexAuthInfo.ChatgptAccountID); v != "" {
|
||||
result["chatgpt_account_id"] = v
|
||||
}
|
||||
if v := strings.TrimSpace(claims.CodexAuthInfo.ChatgptPlanType); v != "" {
|
||||
result["plan_type"] = v
|
||||
}
|
||||
if v := claims.CodexAuthInfo.ChatgptSubscriptionActiveStart; v != nil {
|
||||
result["chatgpt_subscription_active_start"] = v
|
||||
}
|
||||
if v := claims.CodexAuthInfo.ChatgptSubscriptionActiveUntil; v != nil {
|
||||
result["chatgpt_subscription_active_until"] = v
|
||||
}
|
||||
|
||||
if len(result) == 0 {
|
||||
return nil
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func authEmail(auth *coreauth.Auth) string {
|
||||
if auth == nil {
|
||||
return ""
|
||||
}
|
||||
if auth.Metadata != nil {
|
||||
if v, ok := auth.Metadata["email"].(string); ok {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
}
|
||||
if auth.Attributes != nil {
|
||||
if v := strings.TrimSpace(auth.Attributes["email"]); v != "" {
|
||||
return v
|
||||
}
|
||||
if v := strings.TrimSpace(auth.Attributes["account_email"]); v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func authAttribute(auth *coreauth.Auth, key string) string {
|
||||
if auth == nil || len(auth.Attributes) == 0 {
|
||||
return ""
|
||||
}
|
||||
return auth.Attributes[key]
|
||||
}
|
||||
|
||||
func isRuntimeOnlyAuth(auth *coreauth.Auth) bool {
|
||||
if auth == nil || len(auth.Attributes) == 0 {
|
||||
return false
|
||||
}
|
||||
return strings.EqualFold(strings.TrimSpace(auth.Attributes["runtime_only"]), "true")
|
||||
}
|
||||
|
||||
func isUnsafeAuthFileName(name string) bool {
|
||||
if strings.TrimSpace(name) == "" {
|
||||
return true
|
||||
}
|
||||
if strings.ContainsAny(name, "/\\") {
|
||||
return true
|
||||
}
|
||||
if filepath.VolumeName(name) != "" {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
|
@ -0,0 +1,194 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestUploadAuthFile_BatchMultipart(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
|
||||
files := []struct {
|
||||
name string
|
||||
content string
|
||||
}{
|
||||
{name: "alpha.json", content: `{"type":"codex","email":"alpha@example.com"}`},
|
||||
{name: "beta.json", content: `{"type":"claude","email":"beta@example.com"}`},
|
||||
}
|
||||
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
for _, file := range files {
|
||||
part, err := writer.CreateFormFile("file", file.name)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create multipart file: %v", err)
|
||||
}
|
||||
if _, err = part.Write([]byte(file.content)); err != nil {
|
||||
t.Fatalf("failed to write multipart content: %v", err)
|
||||
}
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
t.Fatalf("failed to close multipart writer: %v", err)
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v0/management/auth-files", &body)
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
ctx.Request = req
|
||||
|
||||
h.UploadAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected upload status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
|
||||
t.Fatalf("failed to decode response: %v", err)
|
||||
}
|
||||
if got, ok := payload["uploaded"].(float64); !ok || int(got) != len(files) {
|
||||
t.Fatalf("expected uploaded=%d, got %#v", len(files), payload["uploaded"])
|
||||
}
|
||||
|
||||
for _, file := range files {
|
||||
fullPath := filepath.Join(authDir, file.name)
|
||||
data, err := os.ReadFile(fullPath)
|
||||
if err != nil {
|
||||
t.Fatalf("expected uploaded file %s to exist: %v", file.name, err)
|
||||
}
|
||||
if string(data) != file.content {
|
||||
t.Fatalf("expected file %s content %q, got %q", file.name, file.content, string(data))
|
||||
}
|
||||
}
|
||||
|
||||
auths := manager.List()
|
||||
if len(auths) != len(files) {
|
||||
t.Fatalf("expected %d auth entries, got %d", len(files), len(auths))
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadAuthFile_BatchMultipart_InvalidJSONDoesNotOverwriteExistingFile(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
|
||||
existingName := "alpha.json"
|
||||
existingContent := `{"type":"codex","email":"alpha@example.com"}`
|
||||
if err := os.WriteFile(filepath.Join(authDir, existingName), []byte(existingContent), 0o600); err != nil {
|
||||
t.Fatalf("failed to seed existing auth file: %v", err)
|
||||
}
|
||||
|
||||
files := []struct {
|
||||
name string
|
||||
content string
|
||||
}{
|
||||
{name: existingName, content: `{"type":"codex"`},
|
||||
{name: "beta.json", content: `{"type":"claude","email":"beta@example.com"}`},
|
||||
}
|
||||
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
for _, file := range files {
|
||||
part, err := writer.CreateFormFile("file", file.name)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create multipart file: %v", err)
|
||||
}
|
||||
if _, err = part.Write([]byte(file.content)); err != nil {
|
||||
t.Fatalf("failed to write multipart content: %v", err)
|
||||
}
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
t.Fatalf("failed to close multipart writer: %v", err)
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v0/management/auth-files", &body)
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
ctx.Request = req
|
||||
|
||||
h.UploadAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusMultiStatus {
|
||||
t.Fatalf("expected upload status %d, got %d with body %s", http.StatusMultiStatus, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(filepath.Join(authDir, existingName))
|
||||
if err != nil {
|
||||
t.Fatalf("expected existing auth file to remain readable: %v", err)
|
||||
}
|
||||
if string(data) != existingContent {
|
||||
t.Fatalf("expected existing auth file to remain %q, got %q", existingContent, string(data))
|
||||
}
|
||||
|
||||
betaData, err := os.ReadFile(filepath.Join(authDir, "beta.json"))
|
||||
if err != nil {
|
||||
t.Fatalf("expected valid auth file to be created: %v", err)
|
||||
}
|
||||
if string(betaData) != files[1].content {
|
||||
t.Fatalf("expected beta auth file content %q, got %q", files[1].content, string(betaData))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteAuthFile_BatchQuery(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
files := []string{"alpha.json", "beta.json"}
|
||||
for _, name := range files {
|
||||
if err := os.WriteFile(filepath.Join(authDir, name), []byte(`{"type":"codex"}`), 0o600); err != nil {
|
||||
t.Fatalf("failed to write auth file %s: %v", name, err)
|
||||
}
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
h.tokenStore = &memoryAuthStore{}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(
|
||||
http.MethodDelete,
|
||||
"/v0/management/auth-files?name="+url.QueryEscape(files[0])+"&name="+url.QueryEscape(files[1]),
|
||||
nil,
|
||||
)
|
||||
ctx.Request = req
|
||||
|
||||
h.DeleteAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected delete status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
|
||||
t.Fatalf("failed to decode response: %v", err)
|
||||
}
|
||||
if got, ok := payload["deleted"].(float64); !ok || int(got) != len(files) {
|
||||
t.Fatalf("expected deleted=%d, got %#v", len(files), payload["deleted"])
|
||||
}
|
||||
|
||||
for _, name := range files {
|
||||
if _, err := os.Stat(filepath.Join(authDir, name)); !os.IsNotExist(err) {
|
||||
t.Fatalf("expected auth file %s to be removed, stat err: %v", name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
555
backend/internal/api/handlers/management/auth_files_crud.go
Normal file
555
backend/internal/api/handlers/management/auth_files_crud.go
Normal file
|
|
@ -0,0 +1,555 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/synthesizer"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
// Download single auth file by name
|
||||
func (h *Handler) DownloadAuthFile(c *gin.Context) {
|
||||
name := strings.TrimSpace(c.Query("name"))
|
||||
if isUnsafeAuthFileName(name) {
|
||||
c.JSON(400, gin.H{"error": "invalid name"})
|
||||
return
|
||||
}
|
||||
if !strings.HasSuffix(strings.ToLower(name), ".json") {
|
||||
c.JSON(400, gin.H{"error": "name must end with .json"})
|
||||
return
|
||||
}
|
||||
full := filepath.Join(h.cfg.AuthDir, name)
|
||||
data, err := os.ReadFile(full)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
c.JSON(404, gin.H{"error": "file not found"})
|
||||
} else {
|
||||
c.JSON(500, gin.H{"error": fmt.Sprintf("failed to read file: %v", err)})
|
||||
}
|
||||
return
|
||||
}
|
||||
c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s\"", name))
|
||||
c.Data(200, "application/json", data)
|
||||
}
|
||||
|
||||
// Upload auth file: multipart or raw JSON with ?name=
|
||||
func (h *Handler) UploadAuthFile(c *gin.Context) {
|
||||
if h.authManager == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"})
|
||||
return
|
||||
}
|
||||
ctx := c.Request.Context()
|
||||
|
||||
fileHeaders, errMultipart := h.multipartAuthFileHeaders(c)
|
||||
if errMultipart != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("invalid multipart form: %v", errMultipart)})
|
||||
return
|
||||
}
|
||||
if len(fileHeaders) == 1 {
|
||||
if _, errUpload := h.storeUploadedAuthFile(ctx, fileHeaders[0]); errUpload != nil {
|
||||
if errors.Is(errUpload, errAuthFileMustBeJSON) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "file must be .json"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": errUpload.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||||
return
|
||||
}
|
||||
if len(fileHeaders) > 1 {
|
||||
uploaded := make([]string, 0, len(fileHeaders))
|
||||
failed := make([]gin.H, 0)
|
||||
for _, file := range fileHeaders {
|
||||
name, errUpload := h.storeUploadedAuthFile(ctx, file)
|
||||
if errUpload != nil {
|
||||
failureName := ""
|
||||
if file != nil {
|
||||
failureName = filepath.Base(file.Filename)
|
||||
}
|
||||
msg := errUpload.Error()
|
||||
if errors.Is(errUpload, errAuthFileMustBeJSON) {
|
||||
msg = "file must be .json"
|
||||
}
|
||||
failed = append(failed, gin.H{"name": failureName, "error": msg})
|
||||
continue
|
||||
}
|
||||
uploaded = append(uploaded, name)
|
||||
}
|
||||
if len(failed) > 0 {
|
||||
c.JSON(http.StatusMultiStatus, gin.H{
|
||||
"status": "partial",
|
||||
"uploaded": len(uploaded),
|
||||
"files": uploaded,
|
||||
"failed": failed,
|
||||
})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok", "uploaded": len(uploaded), "files": uploaded})
|
||||
return
|
||||
}
|
||||
if c.ContentType() == "multipart/form-data" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "no files uploaded"})
|
||||
return
|
||||
}
|
||||
name := strings.TrimSpace(c.Query("name"))
|
||||
if isUnsafeAuthFileName(name) {
|
||||
c.JSON(400, gin.H{"error": "invalid name"})
|
||||
return
|
||||
}
|
||||
if !strings.HasSuffix(strings.ToLower(name), ".json") {
|
||||
c.JSON(400, gin.H{"error": "name must end with .json"})
|
||||
return
|
||||
}
|
||||
data, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
c.JSON(400, gin.H{"error": "failed to read body"})
|
||||
return
|
||||
}
|
||||
if err = h.writeAuthFile(ctx, filepath.Base(name), data); err != nil {
|
||||
c.JSON(500, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(200, gin.H{"status": "ok"})
|
||||
}
|
||||
|
||||
// Delete auth files: single by name or all
|
||||
func (h *Handler) DeleteAuthFile(c *gin.Context) {
|
||||
if h.authManager == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"})
|
||||
return
|
||||
}
|
||||
ctx := c.Request.Context()
|
||||
if all := c.Query("all"); all == "true" || all == "1" || all == "*" {
|
||||
entries, err := os.ReadDir(h.cfg.AuthDir)
|
||||
if err != nil {
|
||||
c.JSON(500, gin.H{"error": fmt.Sprintf("failed to read auth dir: %v", err)})
|
||||
return
|
||||
}
|
||||
deleted := 0
|
||||
for _, e := range entries {
|
||||
if e.IsDir() {
|
||||
continue
|
||||
}
|
||||
name := e.Name()
|
||||
if !strings.HasSuffix(strings.ToLower(name), ".json") {
|
||||
continue
|
||||
}
|
||||
full := filepath.Join(h.cfg.AuthDir, name)
|
||||
if !filepath.IsAbs(full) {
|
||||
if abs, errAbs := filepath.Abs(full); errAbs == nil {
|
||||
full = abs
|
||||
}
|
||||
}
|
||||
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.removeAuth(ctx, full)
|
||||
}
|
||||
}
|
||||
c.JSON(200, gin.H{"status": "ok", "deleted": deleted})
|
||||
return
|
||||
}
|
||||
|
||||
names, errNames := requestedAuthFileNamesForDelete(c)
|
||||
if errNames != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": errNames.Error()})
|
||||
return
|
||||
}
|
||||
if len(names) == 0 {
|
||||
c.JSON(400, gin.H{"error": "invalid name"})
|
||||
return
|
||||
}
|
||||
if len(names) == 1 {
|
||||
if _, status, errDelete := h.deleteAuthFileByName(ctx, names[0]); errDelete != nil {
|
||||
c.JSON(status, gin.H{"error": errDelete.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||||
return
|
||||
}
|
||||
|
||||
deletedFiles := make([]string, 0, len(names))
|
||||
failed := make([]gin.H, 0)
|
||||
for _, name := range names {
|
||||
deletedName, _, errDelete := h.deleteAuthFileByName(ctx, name)
|
||||
if errDelete != nil {
|
||||
failed = append(failed, gin.H{"name": name, "error": errDelete.Error()})
|
||||
continue
|
||||
}
|
||||
deletedFiles = append(deletedFiles, deletedName)
|
||||
}
|
||||
if len(failed) > 0 {
|
||||
c.JSON(http.StatusMultiStatus, gin.H{
|
||||
"status": "partial",
|
||||
"deleted": len(deletedFiles),
|
||||
"files": deletedFiles,
|
||||
"failed": failed,
|
||||
})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok", "deleted": len(deletedFiles), "files": deletedFiles})
|
||||
}
|
||||
|
||||
func (h *Handler) multipartAuthFileHeaders(c *gin.Context) ([]*multipart.FileHeader, error) {
|
||||
if h == nil || c == nil || c.ContentType() != "multipart/form-data" {
|
||||
return nil, nil
|
||||
}
|
||||
form, err := c.MultipartForm()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if form == nil || len(form.File) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
keys := make([]string, 0, len(form.File))
|
||||
for key := range form.File {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
|
||||
headers := make([]*multipart.FileHeader, 0)
|
||||
for _, key := range keys {
|
||||
headers = append(headers, form.File[key]...)
|
||||
}
|
||||
return headers, nil
|
||||
}
|
||||
|
||||
func (h *Handler) storeUploadedAuthFile(ctx context.Context, file *multipart.FileHeader) (string, error) {
|
||||
if file == nil {
|
||||
return "", fmt.Errorf("no file uploaded")
|
||||
}
|
||||
name := filepath.Base(strings.TrimSpace(file.Filename))
|
||||
if !strings.HasSuffix(strings.ToLower(name), ".json") {
|
||||
return "", errAuthFileMustBeJSON
|
||||
}
|
||||
src, err := file.Open()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to open uploaded file: %w", err)
|
||||
}
|
||||
defer src.Close()
|
||||
|
||||
data, err := io.ReadAll(src)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to read uploaded file: %w", err)
|
||||
}
|
||||
if err := h.writeAuthFile(ctx, name, data); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return name, nil
|
||||
}
|
||||
|
||||
func (h *Handler) writeAuthFile(ctx context.Context, name string, data []byte) error {
|
||||
dst := filepath.Join(h.cfg.AuthDir, filepath.Base(name))
|
||||
if !filepath.IsAbs(dst) {
|
||||
if abs, errAbs := filepath.Abs(dst); errAbs == nil {
|
||||
dst = abs
|
||||
}
|
||||
}
|
||||
auth, err := h.buildAuthFromFileData(dst, data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if errWrite := os.WriteFile(dst, data, 0o600); errWrite != nil {
|
||||
return fmt.Errorf("failed to write file: %w", errWrite)
|
||||
}
|
||||
if err := h.upsertAuthRecord(ctx, auth); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func requestedAuthFileNamesForDelete(c *gin.Context) ([]string, error) {
|
||||
if c == nil {
|
||||
return nil, nil
|
||||
}
|
||||
names := uniqueAuthFileNames(c.QueryArray("name"))
|
||||
if len(names) > 0 {
|
||||
return names, nil
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read body")
|
||||
}
|
||||
body = bytes.TrimSpace(body)
|
||||
if len(body) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var objectBody struct {
|
||||
Name string `json:"name"`
|
||||
Names []string `json:"names"`
|
||||
}
|
||||
if body[0] == '[' {
|
||||
var arrayBody []string
|
||||
if err := json.Unmarshal(body, &arrayBody); err != nil {
|
||||
return nil, fmt.Errorf("invalid request body")
|
||||
}
|
||||
return uniqueAuthFileNames(arrayBody), nil
|
||||
}
|
||||
if err := json.Unmarshal(body, &objectBody); err != nil {
|
||||
return nil, fmt.Errorf("invalid request body")
|
||||
}
|
||||
|
||||
out := make([]string, 0, len(objectBody.Names)+1)
|
||||
if strings.TrimSpace(objectBody.Name) != "" {
|
||||
out = append(out, objectBody.Name)
|
||||
}
|
||||
out = append(out, objectBody.Names...)
|
||||
return uniqueAuthFileNames(out), nil
|
||||
}
|
||||
|
||||
func uniqueAuthFileNames(names []string) []string {
|
||||
if len(names) == 0 {
|
||||
return nil
|
||||
}
|
||||
seen := make(map[string]struct{}, len(names))
|
||||
out := make([]string, 0, len(names))
|
||||
for _, name := range names {
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[name]; ok {
|
||||
continue
|
||||
}
|
||||
seen[name] = struct{}{}
|
||||
out = append(out, name)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *Handler) deleteAuthFileByName(ctx context.Context, name string) (string, int, error) {
|
||||
name = strings.TrimSpace(name)
|
||||
if isUnsafeAuthFileName(name) {
|
||||
return "", http.StatusBadRequest, fmt.Errorf("invalid name")
|
||||
}
|
||||
|
||||
targetPath := filepath.Join(h.cfg.AuthDir, filepath.Base(name))
|
||||
targetID := ""
|
||||
if targetAuth := h.findAuthForDelete(name); targetAuth != nil {
|
||||
if !isPluginVirtualSourceDelete(name, targetAuth) {
|
||||
return filepath.Base(name), http.StatusConflict, errPluginVirtualAuth
|
||||
}
|
||||
targetID = strings.TrimSpace(targetAuth.ID)
|
||||
if path := strings.TrimSpace(authAttribute(targetAuth, "path")); path != "" {
|
||||
targetPath = path
|
||||
}
|
||||
}
|
||||
if !filepath.IsAbs(targetPath) {
|
||||
if abs, errAbs := filepath.Abs(targetPath); errAbs == nil {
|
||||
targetPath = abs
|
||||
}
|
||||
}
|
||||
if errRemove := os.Remove(targetPath); errRemove != nil {
|
||||
if os.IsNotExist(errRemove) {
|
||||
return filepath.Base(name), http.StatusNotFound, errAuthFileNotFound
|
||||
}
|
||||
return filepath.Base(name), http.StatusInternalServerError, fmt.Errorf("failed to remove file: %w", errRemove)
|
||||
}
|
||||
if errDeleteRecord := h.deleteTokenRecord(ctx, targetPath); errDeleteRecord != nil {
|
||||
return filepath.Base(name), http.StatusInternalServerError, errDeleteRecord
|
||||
}
|
||||
h.removeAuthsForPath(ctx, targetPath, targetID)
|
||||
return filepath.Base(name), http.StatusOK, nil
|
||||
}
|
||||
|
||||
func isPluginVirtualSourceDelete(name string, auth *coreauth.Auth) bool {
|
||||
if !coreauth.IsPluginVirtualAuth(auth) {
|
||||
return true
|
||||
}
|
||||
sourcePath := strings.TrimSpace(authAttribute(auth, coreauth.AttributeVirtualSource))
|
||||
if sourcePath == "" {
|
||||
sourcePath = strings.TrimSpace(authAttribute(auth, "path"))
|
||||
}
|
||||
if sourcePath == "" {
|
||||
return false
|
||||
}
|
||||
return strings.EqualFold(filepath.Base(strings.TrimSpace(name)), filepath.Base(sourcePath))
|
||||
}
|
||||
|
||||
func (h *Handler) findAuthForDelete(name string) *coreauth.Auth {
|
||||
if h == nil || h.authManager == nil {
|
||||
return nil
|
||||
}
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
if auth, ok := h.authManager.GetByID(name); ok {
|
||||
return auth
|
||||
}
|
||||
auths := h.authManager.List()
|
||||
for _, auth := range auths {
|
||||
if auth == nil {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(auth.FileName) == name {
|
||||
return auth
|
||||
}
|
||||
if filepath.Base(strings.TrimSpace(authAttribute(auth, "path"))) == name {
|
||||
return auth
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *Handler) authIDForPath(path string) string {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return ""
|
||||
}
|
||||
path = filepath.Clean(path)
|
||||
if !filepath.IsAbs(path) {
|
||||
if abs, errAbs := filepath.Abs(path); errAbs == nil {
|
||||
path = abs
|
||||
}
|
||||
}
|
||||
id := path
|
||||
if h != nil && h.cfg != nil {
|
||||
authDir := strings.TrimSpace(h.cfg.AuthDir)
|
||||
if resolvedAuthDir, errResolve := util.ResolveAuthDir(authDir); errResolve == nil && resolvedAuthDir != "" {
|
||||
authDir = resolvedAuthDir
|
||||
}
|
||||
if authDir != "" {
|
||||
authDir = filepath.Clean(authDir)
|
||||
if !filepath.IsAbs(authDir) {
|
||||
if abs, errAbs := filepath.Abs(authDir); errAbs == nil {
|
||||
authDir = abs
|
||||
}
|
||||
}
|
||||
if rel, errRel := filepath.Rel(authDir, path); errRel == nil && rel != "" {
|
||||
id = rel
|
||||
}
|
||||
}
|
||||
}
|
||||
// On Windows, normalize ID casing to avoid duplicate auth entries caused by case-insensitive paths.
|
||||
if runtime.GOOS == "windows" {
|
||||
id = strings.ToLower(id)
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
func (h *Handler) registerAuthFromFile(ctx context.Context, path string, data []byte) error {
|
||||
if h.authManager == nil {
|
||||
return nil
|
||||
}
|
||||
auth, err := h.buildAuthFromFileData(path, data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return h.upsertAuthRecord(ctx, auth)
|
||||
}
|
||||
|
||||
func (h *Handler) buildAuthFromFileData(path string, data []byte) (*coreauth.Auth, error) {
|
||||
if path == "" {
|
||||
return nil, fmt.Errorf("auth path is empty")
|
||||
}
|
||||
if data == nil {
|
||||
var err error
|
||||
data, err = os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read auth file: %w", err)
|
||||
}
|
||||
}
|
||||
metadata := make(map[string]any)
|
||||
if err := json.Unmarshal(data, &metadata); err != nil {
|
||||
return nil, fmt.Errorf("invalid auth file: %w", err)
|
||||
}
|
||||
coreauth.NormalizeCredentialMetadata(metadata)
|
||||
provider, _ := metadata["type"].(string)
|
||||
if provider == "" {
|
||||
provider = "unknown"
|
||||
}
|
||||
label := provider
|
||||
if email, ok := metadata["email"].(string); ok && email != "" {
|
||||
label = email
|
||||
}
|
||||
lastRefresh, hasLastRefresh := extractLastRefreshTimestamp(metadata)
|
||||
|
||||
authID := h.authIDForPath(path)
|
||||
if authID == "" {
|
||||
authID = path
|
||||
}
|
||||
auth := (*coreauth.Auth)(nil)
|
||||
if h != nil && h.cfg != nil {
|
||||
sctx := &synthesizer.SynthesisContext{
|
||||
Config: h.cfg,
|
||||
AuthDir: h.cfg.AuthDir,
|
||||
Now: time.Now(),
|
||||
IDGenerator: synthesizer.NewStableIDGenerator(),
|
||||
}
|
||||
generated, errSynthesize := synthesizer.SynthesizeAuthFile(sctx, path, data)
|
||||
if errSynthesize != nil {
|
||||
return nil, fmt.Errorf("invalid auth file: %w", errSynthesize)
|
||||
}
|
||||
if len(generated) > 0 && generated[0] != nil {
|
||||
auth = generated[0].Clone()
|
||||
}
|
||||
}
|
||||
if auth == nil {
|
||||
auth = &coreauth.Auth{
|
||||
ID: authID,
|
||||
Provider: provider,
|
||||
Label: label,
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"path": path,
|
||||
"source": path,
|
||||
},
|
||||
Metadata: metadata,
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
}
|
||||
}
|
||||
auth.ID = authID
|
||||
auth.FileName = filepath.Base(path)
|
||||
if hasLastRefresh {
|
||||
auth.LastRefreshedAt = lastRefresh
|
||||
}
|
||||
if h != nil && h.authManager != nil {
|
||||
if existing, ok := h.authManager.GetByID(authID); ok {
|
||||
auth.CreatedAt = existing.CreatedAt
|
||||
if !hasLastRefresh {
|
||||
auth.LastRefreshedAt = existing.LastRefreshedAt
|
||||
}
|
||||
auth.NextRefreshAfter = existing.NextRefreshAfter
|
||||
auth.Runtime = existing.Runtime
|
||||
}
|
||||
}
|
||||
coreauth.ApplyCustomHeadersFromMetadata(auth)
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
func (h *Handler) upsertAuthRecord(ctx context.Context, auth *coreauth.Auth) error {
|
||||
if h == nil || h.authManager == nil || auth == nil {
|
||||
return nil
|
||||
}
|
||||
if existing, ok := h.authManager.GetByID(auth.ID); ok {
|
||||
auth.CreatedAt = existing.CreatedAt
|
||||
_, err := h.authManager.Update(ctx, auth)
|
||||
return err
|
||||
}
|
||||
_, err := h.authManager.Register(ctx, auth)
|
||||
return err
|
||||
}
|
||||
|
|
@ -0,0 +1,172 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestDeleteAuthFile_UsesAuthPathFromManager(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
tempDir := t.TempDir()
|
||||
authDir := filepath.Join(tempDir, "auth")
|
||||
externalDir := filepath.Join(tempDir, "external")
|
||||
if errMkdirAuth := os.MkdirAll(authDir, 0o700); errMkdirAuth != nil {
|
||||
t.Fatalf("failed to create auth dir: %v", errMkdirAuth)
|
||||
}
|
||||
if errMkdirExternal := os.MkdirAll(externalDir, 0o700); errMkdirExternal != nil {
|
||||
t.Fatalf("failed to create external dir: %v", errMkdirExternal)
|
||||
}
|
||||
|
||||
fileName := "codex-user@example.com-plus.json"
|
||||
shadowPath := filepath.Join(authDir, fileName)
|
||||
realPath := filepath.Join(externalDir, fileName)
|
||||
if errWriteShadow := os.WriteFile(shadowPath, []byte(`{"type":"codex","email":"shadow@example.com"}`), 0o600); errWriteShadow != nil {
|
||||
t.Fatalf("failed to write shadow file: %v", errWriteShadow)
|
||||
}
|
||||
if errWriteReal := os.WriteFile(realPath, []byte(`{"type":"codex","email":"real@example.com"}`), 0o600); errWriteReal != nil {
|
||||
t.Fatalf("failed to write real file: %v", errWriteReal)
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: "legacy/" + fileName,
|
||||
FileName: fileName,
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusError,
|
||||
Unavailable: true,
|
||||
Attributes: map[string]string{
|
||||
"path": realPath,
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "codex",
|
||||
"email": "real@example.com",
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
h.tokenStore = &memoryAuthStore{}
|
||||
|
||||
deleteRec := httptest.NewRecorder()
|
||||
deleteCtx, _ := gin.CreateTestContext(deleteRec)
|
||||
deleteReq := httptest.NewRequest(http.MethodDelete, "/v0/management/auth-files?name="+url.QueryEscape(fileName), nil)
|
||||
deleteCtx.Request = deleteReq
|
||||
h.DeleteAuthFile(deleteCtx)
|
||||
|
||||
if deleteRec.Code != http.StatusOK {
|
||||
t.Fatalf("expected delete status %d, got %d with body %s", http.StatusOK, deleteRec.Code, deleteRec.Body.String())
|
||||
}
|
||||
if _, errStatReal := os.Stat(realPath); !os.IsNotExist(errStatReal) {
|
||||
t.Fatalf("expected managed auth file to be removed, stat err: %v", errStatReal)
|
||||
}
|
||||
if _, errStatShadow := os.Stat(shadowPath); errStatShadow != nil {
|
||||
t.Fatalf("expected shadow auth file to remain, stat err: %v", errStatShadow)
|
||||
}
|
||||
|
||||
listRec := httptest.NewRecorder()
|
||||
listCtx, _ := gin.CreateTestContext(listRec)
|
||||
listReq := httptest.NewRequest(http.MethodGet, "/v0/management/auth-files", nil)
|
||||
listCtx.Request = listReq
|
||||
h.ListAuthFiles(listCtx)
|
||||
|
||||
if listRec.Code != http.StatusOK {
|
||||
t.Fatalf("expected list status %d, got %d with body %s", http.StatusOK, listRec.Code, listRec.Body.String())
|
||||
}
|
||||
var listPayload map[string]any
|
||||
if errUnmarshal := json.Unmarshal(listRec.Body.Bytes(), &listPayload); errUnmarshal != nil {
|
||||
t.Fatalf("failed to decode list payload: %v", errUnmarshal)
|
||||
}
|
||||
filesRaw, ok := listPayload["files"].([]any)
|
||||
if !ok {
|
||||
t.Fatalf("expected files array, payload: %#v", listPayload)
|
||||
}
|
||||
if len(filesRaw) != 0 {
|
||||
t.Fatalf("expected removed auth to be hidden from list, got %d entries", len(filesRaw))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteAuthFile_FallbackToAuthDirPath(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "fallback-user.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex"}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
h.tokenStore = &memoryAuthStore{}
|
||||
|
||||
deleteRec := httptest.NewRecorder()
|
||||
deleteCtx, _ := gin.CreateTestContext(deleteRec)
|
||||
deleteReq := httptest.NewRequest(http.MethodDelete, "/v0/management/auth-files?name="+url.QueryEscape(fileName), nil)
|
||||
deleteCtx.Request = deleteReq
|
||||
h.DeleteAuthFile(deleteCtx)
|
||||
|
||||
if deleteRec.Code != http.StatusOK {
|
||||
t.Fatalf("expected delete status %d, got %d with body %s", http.StatusOK, deleteRec.Code, deleteRec.Body.String())
|
||||
}
|
||||
if _, errStat := os.Stat(filePath); !os.IsNotExist(errStat) {
|
||||
t.Fatalf("expected auth file to be removed from auth dir, stat err: %v", errStat)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteAuthFile_RemovesRuntimeAuth(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "runtime-remove-user.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex","email":"runtime@example.com"}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: "runtime-remove-auth",
|
||||
FileName: fileName,
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"path": filePath,
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "codex",
|
||||
"email": "runtime@example.com",
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
h.tokenStore = &memoryAuthStore{}
|
||||
|
||||
deleteRec := httptest.NewRecorder()
|
||||
deleteCtx, _ := gin.CreateTestContext(deleteRec)
|
||||
deleteReq := httptest.NewRequest(http.MethodDelete, "/v0/management/auth-files?name="+url.QueryEscape(fileName), nil)
|
||||
deleteCtx.Request = deleteReq
|
||||
h.DeleteAuthFile(deleteCtx)
|
||||
|
||||
if deleteRec.Code != http.StatusOK {
|
||||
t.Fatalf("expected delete status %d, got %d with body %s", http.StatusOK, deleteRec.Code, deleteRec.Body.String())
|
||||
}
|
||||
if _, ok := manager.GetByID(record.ID); ok {
|
||||
t.Fatalf("expected runtime auth %q to be removed", record.ID)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,60 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestDownloadAuthFile_ReturnsFile(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "download-user.json"
|
||||
expected := []byte(`{"type":"codex"}`)
|
||||
if err := os.WriteFile(filepath.Join(authDir, fileName), expected, 0o600); err != nil {
|
||||
t.Fatalf("failed to write auth file: %v", err)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files/download?name="+url.QueryEscape(fileName), nil)
|
||||
h.DownloadAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected download status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
if got := rec.Body.Bytes(); string(got) != string(expected) {
|
||||
t.Fatalf("unexpected download content: %q", string(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadAuthFile_RejectsPathSeparators(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, nil)
|
||||
|
||||
for _, name := range []string{
|
||||
"../external/secret.json",
|
||||
`..\\external\\secret.json`,
|
||||
"nested/secret.json",
|
||||
`nested\\secret.json`,
|
||||
} {
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files/download?name="+url.QueryEscape(name), nil)
|
||||
h.DownloadAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected %d for name %q, got %d with body %s", http.StatusBadRequest, name, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
//go:build windows
|
||||
|
||||
package management
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestDownloadAuthFile_PreventsWindowsSlashTraversal(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
tempDir := t.TempDir()
|
||||
authDir := filepath.Join(tempDir, "auth")
|
||||
externalDir := filepath.Join(tempDir, "external")
|
||||
if err := os.MkdirAll(authDir, 0o700); err != nil {
|
||||
t.Fatalf("failed to create auth dir: %v", err)
|
||||
}
|
||||
if err := os.MkdirAll(externalDir, 0o700); err != nil {
|
||||
t.Fatalf("failed to create external dir: %v", err)
|
||||
}
|
||||
|
||||
secretName := "secret.json"
|
||||
secretPath := filepath.Join(externalDir, secretName)
|
||||
if err := os.WriteFile(secretPath, []byte(`{"secret":true}`), 0o600); err != nil {
|
||||
t.Fatalf("failed to write external file: %v", err)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(
|
||||
http.MethodGet,
|
||||
"/v0/management/auth-files/download?name="+url.QueryEscape("../external/"+secretName),
|
||||
nil,
|
||||
)
|
||||
h.DownloadAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusBadRequest, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
866
backend/internal/api/handlers/management/auth_files_fields.go
Normal file
866
backend/internal/api/handlers/management/auth_files_fields.go
Normal file
|
|
@ -0,0 +1,866 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/credentialweight"
|
||||
sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
// PatchAuthFileStatus toggles the disabled state of an auth file
|
||||
func (h *Handler) PatchAuthFileStatus(c *gin.Context) {
|
||||
if h.authManager == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"})
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
Name string `json:"name"`
|
||||
AuthIndex string `json:"auth_index"`
|
||||
Disabled *bool `json:"disabled"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
|
||||
return
|
||||
}
|
||||
|
||||
name := strings.TrimSpace(req.Name)
|
||||
authIndex := strings.TrimSpace(req.AuthIndex)
|
||||
if name == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"})
|
||||
return
|
||||
}
|
||||
if req.Disabled == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "disabled is required"})
|
||||
return
|
||||
}
|
||||
|
||||
ctx := c.Request.Context()
|
||||
|
||||
targetAuth, _ := h.lookupAuthFile(name, authIndex)
|
||||
if targetAuth == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "auth file not found"})
|
||||
return
|
||||
}
|
||||
if coreauth.IsPluginVirtualAuth(targetAuth) {
|
||||
// Allow status changes only when targeting the source auth file name, matching delete semantics.
|
||||
// Expanded virtual project auths still cannot be modified independently.
|
||||
if !isPluginVirtualSourceDelete(name, targetAuth) {
|
||||
c.JSON(http.StatusConflict, gin.H{"error": errPluginVirtualAuth.Error()})
|
||||
return
|
||||
}
|
||||
if errPatch := h.patchPluginVirtualSourceStatus(ctx, targetAuth, *req.Disabled); errPatch != nil {
|
||||
status := http.StatusInternalServerError
|
||||
if errors.Is(errPatch, errAuthFileNotFound) || os.IsNotExist(errPatch) {
|
||||
status = http.StatusNotFound
|
||||
}
|
||||
c.JSON(status, gin.H{"error": errPatch.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok", "disabled": *req.Disabled})
|
||||
return
|
||||
}
|
||||
|
||||
if coreauth.IsConfigAPIKeyAuth(targetAuth) {
|
||||
h.mu.Lock()
|
||||
handled, errToggle := toggleConfigAPIKeyExcludedAll(h.cfg, targetAuth, *req.Disabled)
|
||||
if errToggle != nil {
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to update config api key: %v", errToggle)})
|
||||
return
|
||||
}
|
||||
if !handled {
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "config api key entry not found"})
|
||||
return
|
||||
}
|
||||
cfgSnapshot, okSnapshot := h.saveConfigAndSnapshotLocked(c)
|
||||
h.mu.Unlock()
|
||||
if !okSnapshot {
|
||||
return
|
||||
}
|
||||
h.reloadConfigAfterManagementSave(ctx, cfgSnapshot)
|
||||
if h.tokenStore != nil {
|
||||
_ = h.tokenStore.Delete(ctx, targetAuth.ID)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"status": "ok",
|
||||
"disabled": *req.Disabled,
|
||||
"via": "config:excluded-models",
|
||||
"excluded_pattern": configAPIKeyDisablePattern,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
applyAuthDisabledState(targetAuth, *req.Disabled)
|
||||
if _, err := h.authManager.Update(ctx, targetAuth); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to update auth: %v", err)})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok", "disabled": *req.Disabled})
|
||||
}
|
||||
|
||||
// patchPluginVirtualSourceStatus toggles disabled on a plugin multi-auth source file and all
|
||||
// runtime auths expanded from it. Virtual project children cannot be toggled independently.
|
||||
func (h *Handler) patchPluginVirtualSourceStatus(ctx context.Context, targetAuth *coreauth.Auth, disabled bool) error {
|
||||
if h == nil || h.authManager == nil || targetAuth == nil {
|
||||
return fmt.Errorf("core auth manager unavailable")
|
||||
}
|
||||
sourcePath := strings.TrimSpace(authAttribute(targetAuth, coreauth.AttributeVirtualSource))
|
||||
if sourcePath == "" {
|
||||
sourcePath = strings.TrimSpace(authAttribute(targetAuth, "path"))
|
||||
}
|
||||
if sourcePath == "" {
|
||||
return errPluginVirtualAuth
|
||||
}
|
||||
if errWrite := setSourceAuthFileDisabled(sourcePath, disabled); errWrite != nil {
|
||||
if os.IsNotExist(errWrite) {
|
||||
return errAuthFileNotFound
|
||||
}
|
||||
return fmt.Errorf("failed to update source auth file: %w", errWrite)
|
||||
}
|
||||
now := time.Now()
|
||||
for _, auth := range h.authManager.List() {
|
||||
if auth == nil {
|
||||
continue
|
||||
}
|
||||
if !sameAuthFilePath(authAttribute(auth, "path"), sourcePath) &&
|
||||
!sameAuthFilePath(authAttribute(auth, coreauth.AttributeVirtualSource), sourcePath) {
|
||||
continue
|
||||
}
|
||||
applyAuthDisabledState(auth, disabled)
|
||||
auth.UpdatedAt = now
|
||||
if _, errUpdate := h.authManager.Update(ctx, auth); errUpdate != nil {
|
||||
return fmt.Errorf("failed to update auth %s: %w", auth.ID, errUpdate)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func setSourceAuthFileDisabled(path string, disabled bool) error {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return fmt.Errorf("source auth path is empty")
|
||||
}
|
||||
data, errRead := os.ReadFile(path)
|
||||
if errRead != nil {
|
||||
return errRead
|
||||
}
|
||||
metadata := make(map[string]any)
|
||||
if len(bytes.TrimSpace(data)) > 0 {
|
||||
if errUnmarshal := json.Unmarshal(data, &metadata); errUnmarshal != nil {
|
||||
return fmt.Errorf("invalid auth file: %w", errUnmarshal)
|
||||
}
|
||||
}
|
||||
if metadata == nil {
|
||||
metadata = make(map[string]any)
|
||||
}
|
||||
coreauth.NormalizeCredentialMetadata(metadata)
|
||||
metadata["disabled"] = disabled
|
||||
raw, errMarshal := json.Marshal(metadata)
|
||||
if errMarshal != nil {
|
||||
return fmt.Errorf("marshal auth file: %w", errMarshal)
|
||||
}
|
||||
if errWrite := os.WriteFile(path, raw, 0o600); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func applyAuthDisabledState(auth *coreauth.Auth, disabled bool) {
|
||||
if auth == nil {
|
||||
return
|
||||
}
|
||||
auth.Disabled = disabled
|
||||
if disabled {
|
||||
auth.Status = coreauth.StatusDisabled
|
||||
auth.StatusMessage = "disabled via management API"
|
||||
} else {
|
||||
auth.Status = coreauth.StatusActive
|
||||
auth.StatusMessage = ""
|
||||
}
|
||||
auth.UpdatedAt = time.Now()
|
||||
if auth.Metadata == nil {
|
||||
auth.Metadata = make(map[string]any)
|
||||
}
|
||||
auth.Metadata["disabled"] = disabled
|
||||
}
|
||||
|
||||
// PatchAuthFileFields updates arbitrary metadata fields of an auth file.
|
||||
func (h *Handler) PatchAuthFileFields(c *gin.Context) {
|
||||
if h.authManager == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"})
|
||||
return
|
||||
}
|
||||
|
||||
var req map[string]json.RawMessage
|
||||
decoder := json.NewDecoder(c.Request.Body)
|
||||
decoder.UseNumber()
|
||||
if err := decoder.Decode(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
|
||||
return
|
||||
}
|
||||
|
||||
nameRaw, ok := req["name"]
|
||||
if !ok {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"})
|
||||
return
|
||||
}
|
||||
var nameValue string
|
||||
if err := json.Unmarshal(nameRaw, &nameValue); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"})
|
||||
return
|
||||
}
|
||||
name := strings.TrimSpace(nameValue)
|
||||
if name == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"})
|
||||
return
|
||||
}
|
||||
delete(req, "name")
|
||||
var errNormalize error
|
||||
req, errNormalize = normalizeAuthFilePatchFields(req)
|
||||
if errNormalize != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": errNormalize.Error()})
|
||||
return
|
||||
}
|
||||
requestRetryPatch, errRequestRetry := decodeAuthFileRequestRetryPatch(req)
|
||||
if errRequestRetry != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": errRequestRetry.Error()})
|
||||
return
|
||||
}
|
||||
for key := range req {
|
||||
if strings.TrimSpace(key) == "request_retry" {
|
||||
delete(req, key)
|
||||
}
|
||||
}
|
||||
|
||||
ctx := c.Request.Context()
|
||||
|
||||
// Find auth by name or ID
|
||||
var targetAuth *coreauth.Auth
|
||||
if auth, ok := h.authManager.GetByID(name); ok {
|
||||
targetAuth = auth
|
||||
} else {
|
||||
auths := h.authManager.List()
|
||||
for _, auth := range auths {
|
||||
if auth.FileName == name {
|
||||
targetAuth = auth
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if targetAuth == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "auth file not found"})
|
||||
return
|
||||
}
|
||||
if coreauth.IsPluginVirtualAuth(targetAuth) {
|
||||
c.JSON(http.StatusConflict, gin.H{"error": errPluginVirtualAuth.Error()})
|
||||
return
|
||||
}
|
||||
coreauth.NormalizeCredentialMetadata(targetAuth.Metadata)
|
||||
|
||||
changed := false
|
||||
touchedRoots := make(map[string]struct{}, len(req))
|
||||
for key, rawValue := range req {
|
||||
fieldPath := strings.TrimSpace(key)
|
||||
if fieldPath == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "field name is required"})
|
||||
return
|
||||
}
|
||||
value, errDecode := decodeAuthFileFieldValue(rawValue)
|
||||
if errDecode != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("invalid field %s", fieldPath)})
|
||||
return
|
||||
}
|
||||
if targetAuth.Metadata == nil {
|
||||
targetAuth.Metadata = make(map[string]any)
|
||||
}
|
||||
|
||||
if fieldPath == coreauth.AttributeWeight {
|
||||
if value == nil {
|
||||
delete(targetAuth.Metadata, coreauth.AttributeWeight)
|
||||
} else {
|
||||
if _, okNumber := value.(json.Number); !okNumber {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "weight must be an integer"})
|
||||
return
|
||||
}
|
||||
weight, errWeight := credentialweight.ParseValue(value)
|
||||
if errWeight != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": errWeight.Error()})
|
||||
return
|
||||
}
|
||||
targetAuth.Metadata[coreauth.AttributeWeight] = weight
|
||||
}
|
||||
} else if rootAuthFileField(fieldPath) == coreauth.AttributeWeight {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "weight does not support nested fields"})
|
||||
return
|
||||
} else if fieldPath == "headers" {
|
||||
applyAuthFileHeadersPatch(targetAuth, value)
|
||||
} else if errSet := setAuthFileMetadataValue(targetAuth.Metadata, fieldPath, value); errSet != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": errSet.Error()})
|
||||
return
|
||||
}
|
||||
if root := rootAuthFileField(fieldPath); root != "" {
|
||||
touchedRoots[root] = struct{}{}
|
||||
}
|
||||
changed = true
|
||||
}
|
||||
if requestRetryPatch.Set {
|
||||
if targetAuth.Metadata == nil {
|
||||
targetAuth.Metadata = make(map[string]any)
|
||||
}
|
||||
if requestRetryPatch.Value == nil {
|
||||
delete(targetAuth.Metadata, "request_retry")
|
||||
} else {
|
||||
targetAuth.Metadata["request_retry"] = *requestRetryPatch.Value
|
||||
}
|
||||
changed = true
|
||||
}
|
||||
if changed {
|
||||
syncAuthFileMetadataFields(targetAuth, touchedRoots)
|
||||
}
|
||||
|
||||
if !changed {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "no fields to update"})
|
||||
return
|
||||
}
|
||||
|
||||
targetAuth.UpdatedAt = time.Now()
|
||||
|
||||
if _, err := h.authManager.Update(ctx, targetAuth); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to update auth: %v", err)})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||||
}
|
||||
|
||||
func decodeAuthFileFieldValue(raw json.RawMessage) (any, error) {
|
||||
decoder := json.NewDecoder(bytes.NewReader(raw))
|
||||
decoder.UseNumber()
|
||||
var value any
|
||||
if err := decoder.Decode(&value); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
type authFileRequestRetryPatch struct {
|
||||
Set bool
|
||||
Value *int
|
||||
}
|
||||
|
||||
func normalizeAuthFilePatchFields(fields map[string]json.RawMessage) (map[string]json.RawMessage, error) {
|
||||
normalized := make(map[string]json.RawMessage, len(fields))
|
||||
originalNames := make(map[string]string, len(fields))
|
||||
canonicalNames := make(map[string]bool, len(fields))
|
||||
for key, value := range fields {
|
||||
parts := strings.Split(strings.TrimSpace(key), ".")
|
||||
for index := range parts {
|
||||
parts[index] = strings.TrimSpace(parts[index])
|
||||
}
|
||||
originalRoot := parts[0]
|
||||
parts[0] = coreauth.CanonicalCredentialMetadataKey(originalRoot)
|
||||
canonicalPath := strings.Join(parts, ".")
|
||||
if original, exists := originalNames[canonicalPath]; exists {
|
||||
currentCanonical := originalRoot == parts[0]
|
||||
if canonicalNames[canonicalPath] != currentCanonical {
|
||||
if currentCanonical {
|
||||
normalized[canonicalPath] = value
|
||||
originalNames[canonicalPath] = key
|
||||
canonicalNames[canonicalPath] = true
|
||||
}
|
||||
continue
|
||||
}
|
||||
return nil, fmt.Errorf("auth file fields %q and %q refer to the same field", original, key)
|
||||
}
|
||||
normalized[canonicalPath] = value
|
||||
originalNames[canonicalPath] = key
|
||||
canonicalNames[canonicalPath] = originalRoot == parts[0]
|
||||
}
|
||||
return normalized, nil
|
||||
}
|
||||
|
||||
func decodeAuthFileRequestRetryPatch(fields map[string]json.RawMessage) (authFileRequestRetryPatch, error) {
|
||||
var raw json.RawMessage
|
||||
found := false
|
||||
for key, value := range fields {
|
||||
fieldPath := strings.TrimSpace(key)
|
||||
fieldRoot := rootAuthFileField(fieldPath)
|
||||
if fieldRoot == "request_retry" && fieldPath != fieldRoot {
|
||||
return authFileRequestRetryPatch{}, fmt.Errorf("request_retry does not support nested fields")
|
||||
}
|
||||
if fieldPath == "request_retry" {
|
||||
found = true
|
||||
raw = value
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
return authFileRequestRetryPatch{}, nil
|
||||
}
|
||||
value, errDecode := decodeAuthFileFieldValue(raw)
|
||||
if errDecode != nil {
|
||||
return authFileRequestRetryPatch{}, fmt.Errorf("request_retry must be an integer or null")
|
||||
}
|
||||
if value == nil {
|
||||
return authFileRequestRetryPatch{Set: true}, nil
|
||||
}
|
||||
number, okNumber := value.(json.Number)
|
||||
if !okNumber {
|
||||
return authFileRequestRetryPatch{}, fmt.Errorf("request_retry must be an integer or null")
|
||||
}
|
||||
parsed, errInt := number.Int64()
|
||||
if errInt != nil {
|
||||
return authFileRequestRetryPatch{}, fmt.Errorf("request_retry must be an integer or null")
|
||||
}
|
||||
normalized := int(parsed)
|
||||
if int64(normalized) != parsed {
|
||||
return authFileRequestRetryPatch{}, fmt.Errorf("request_retry must be an integer or null")
|
||||
}
|
||||
if normalized < 0 {
|
||||
return authFileRequestRetryPatch{Set: true}, nil
|
||||
}
|
||||
return authFileRequestRetryPatch{Set: true, Value: &normalized}, nil
|
||||
}
|
||||
|
||||
func rootAuthFileField(path string) string {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return ""
|
||||
}
|
||||
if idx := strings.Index(path, "."); idx >= 0 {
|
||||
return strings.TrimSpace(path[:idx])
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
func setAuthFileMetadataValue(metadata map[string]any, path string, value any) error {
|
||||
if metadata == nil {
|
||||
return fmt.Errorf("metadata is nil")
|
||||
}
|
||||
parts := strings.Split(path, ".")
|
||||
current := metadata
|
||||
for i, rawPart := range parts {
|
||||
part := strings.TrimSpace(rawPart)
|
||||
if part == "" {
|
||||
return fmt.Errorf("invalid field path: %s", path)
|
||||
}
|
||||
if i == len(parts)-1 {
|
||||
current[part] = value
|
||||
return nil
|
||||
}
|
||||
next, ok := current[part].(map[string]any)
|
||||
if !ok {
|
||||
next = make(map[string]any)
|
||||
current[part] = next
|
||||
}
|
||||
current = next
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func applyAuthFileHeadersPatch(auth *coreauth.Auth, value any) {
|
||||
if auth == nil {
|
||||
return
|
||||
}
|
||||
if auth.Metadata == nil {
|
||||
auth.Metadata = make(map[string]any)
|
||||
}
|
||||
headersPatch, ok := authFileHeadersStringMap(value)
|
||||
if !ok {
|
||||
auth.Metadata["headers"] = value
|
||||
return
|
||||
}
|
||||
|
||||
existingHeaders := coreauth.ExtractCustomHeadersFromMetadata(auth.Metadata)
|
||||
nextHeaders := make(map[string]string, len(existingHeaders))
|
||||
for key, val := range existingHeaders {
|
||||
nextHeaders[key] = val
|
||||
}
|
||||
for key, value := range headersPatch {
|
||||
name := strings.TrimSpace(key)
|
||||
if name == "" {
|
||||
continue
|
||||
}
|
||||
val := strings.TrimSpace(value)
|
||||
if val == "" {
|
||||
delete(nextHeaders, name)
|
||||
continue
|
||||
}
|
||||
nextHeaders[name] = val
|
||||
}
|
||||
|
||||
if len(nextHeaders) == 0 {
|
||||
delete(auth.Metadata, "headers")
|
||||
return
|
||||
}
|
||||
metaHeaders := make(map[string]any, len(nextHeaders))
|
||||
for key, value := range nextHeaders {
|
||||
metaHeaders[key] = value
|
||||
}
|
||||
auth.Metadata["headers"] = metaHeaders
|
||||
}
|
||||
|
||||
func authFileHeadersStringMap(value any) (map[string]string, bool) {
|
||||
switch typed := value.(type) {
|
||||
case map[string]string:
|
||||
return typed, true
|
||||
case map[string]any:
|
||||
out := make(map[string]string, len(typed))
|
||||
for key, rawValue := range typed {
|
||||
value, ok := rawValue.(string)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
out[key] = value
|
||||
}
|
||||
return out, true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
func syncAuthFileMetadataFields(auth *coreauth.Auth, touchedRoots map[string]struct{}) {
|
||||
if auth == nil || len(touchedRoots) == 0 {
|
||||
return
|
||||
}
|
||||
if _, ok := touchedRoots["prefix"]; ok {
|
||||
if prefix, okString := auth.Metadata["prefix"].(string); okString {
|
||||
auth.Prefix = strings.TrimSpace(prefix)
|
||||
}
|
||||
}
|
||||
if _, ok := touchedRoots["proxy_url"]; ok {
|
||||
if proxyURL, okString := auth.Metadata["proxy_url"].(string); okString {
|
||||
auth.ProxyURL = strings.TrimSpace(proxyURL)
|
||||
}
|
||||
}
|
||||
if _, ok := touchedRoots["headers"]; ok {
|
||||
syncAuthFileHeaderAttributes(auth)
|
||||
}
|
||||
if _, ok := touchedRoots["priority"]; ok {
|
||||
syncAuthFilePriorityAttribute(auth)
|
||||
}
|
||||
if _, ok := touchedRoots[coreauth.AttributeWeight]; ok {
|
||||
syncAuthFileWeightAttribute(auth)
|
||||
}
|
||||
if _, ok := touchedRoots["note"]; ok {
|
||||
syncAuthFileNoteAttribute(auth)
|
||||
}
|
||||
if _, ok := touchedRoots["websockets"]; ok {
|
||||
syncAuthFileWebsocketsAttribute(auth)
|
||||
}
|
||||
if _, ok := touchedRoots["disabled"]; ok {
|
||||
syncAuthFileDisabledState(auth)
|
||||
}
|
||||
}
|
||||
|
||||
func syncAuthFileHeaderAttributes(auth *coreauth.Auth) {
|
||||
if auth == nil {
|
||||
return
|
||||
}
|
||||
if auth.Attributes == nil {
|
||||
auth.Attributes = make(map[string]string)
|
||||
}
|
||||
for key := range auth.Attributes {
|
||||
if strings.HasPrefix(key, "header:") {
|
||||
delete(auth.Attributes, key)
|
||||
}
|
||||
}
|
||||
for name, value := range coreauth.ExtractCustomHeadersFromMetadata(auth.Metadata) {
|
||||
auth.Attributes["header:"+name] = value
|
||||
}
|
||||
}
|
||||
|
||||
func syncAuthFilePriorityAttribute(auth *coreauth.Auth) {
|
||||
if auth == nil {
|
||||
return
|
||||
}
|
||||
if auth.Attributes == nil {
|
||||
auth.Attributes = make(map[string]string)
|
||||
}
|
||||
priority, ok := authFileIntValue(auth.Metadata["priority"])
|
||||
if !ok {
|
||||
delete(auth.Attributes, "priority")
|
||||
return
|
||||
}
|
||||
if priority == 0 {
|
||||
delete(auth.Attributes, "priority")
|
||||
return
|
||||
}
|
||||
auth.Attributes["priority"] = strconv.Itoa(priority)
|
||||
}
|
||||
|
||||
func syncAuthFileWeightAttribute(auth *coreauth.Auth) {
|
||||
if auth == nil {
|
||||
return
|
||||
}
|
||||
if auth.Attributes == nil {
|
||||
auth.Attributes = make(map[string]string)
|
||||
}
|
||||
weight, errWeight := credentialweight.ParseValue(auth.Metadata[coreauth.AttributeWeight])
|
||||
if errWeight != nil {
|
||||
delete(auth.Attributes, coreauth.AttributeWeight)
|
||||
return
|
||||
}
|
||||
auth.Attributes[coreauth.AttributeWeight] = strconv.FormatInt(weight, 10)
|
||||
}
|
||||
|
||||
func authFileIntValue(value any) (int, bool) {
|
||||
switch typed := value.(type) {
|
||||
case int:
|
||||
return typed, true
|
||||
case int64:
|
||||
return int(typed), true
|
||||
case float64:
|
||||
return int(typed), true
|
||||
case json.Number:
|
||||
if i, err := typed.Int64(); err == nil {
|
||||
return int(i), true
|
||||
}
|
||||
case string:
|
||||
if i, err := strconv.Atoi(strings.TrimSpace(typed)); err == nil {
|
||||
return i, true
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func syncAuthFileNoteAttribute(auth *coreauth.Auth) {
|
||||
if auth == nil {
|
||||
return
|
||||
}
|
||||
if auth.Attributes == nil {
|
||||
auth.Attributes = make(map[string]string)
|
||||
}
|
||||
note, ok := auth.Metadata["note"].(string)
|
||||
if !ok {
|
||||
delete(auth.Attributes, "note")
|
||||
return
|
||||
}
|
||||
note = strings.TrimSpace(note)
|
||||
if note == "" {
|
||||
delete(auth.Attributes, "note")
|
||||
return
|
||||
}
|
||||
auth.Attributes["note"] = note
|
||||
}
|
||||
|
||||
func syncAuthFileWebsocketsAttribute(auth *coreauth.Auth) {
|
||||
if auth == nil {
|
||||
return
|
||||
}
|
||||
if auth.Attributes == nil {
|
||||
auth.Attributes = make(map[string]string)
|
||||
}
|
||||
websockets, ok := authFileBoolValue(auth.Metadata["websockets"])
|
||||
if !ok {
|
||||
delete(auth.Attributes, "websockets")
|
||||
return
|
||||
}
|
||||
auth.Attributes["websockets"] = strconv.FormatBool(websockets)
|
||||
}
|
||||
|
||||
func authFileBoolValue(value any) (bool, bool) {
|
||||
switch typed := value.(type) {
|
||||
case bool:
|
||||
return typed, true
|
||||
case string:
|
||||
parsed, errParse := strconv.ParseBool(strings.TrimSpace(typed))
|
||||
if errParse == nil {
|
||||
return parsed, true
|
||||
}
|
||||
}
|
||||
return false, false
|
||||
}
|
||||
|
||||
func syncAuthFileDisabledState(auth *coreauth.Auth) {
|
||||
if auth == nil {
|
||||
return
|
||||
}
|
||||
disabled, ok := authFileBoolValue(auth.Metadata["disabled"])
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
auth.Disabled = disabled
|
||||
if disabled {
|
||||
auth.Status = coreauth.StatusDisabled
|
||||
if strings.TrimSpace(auth.StatusMessage) == "" {
|
||||
auth.StatusMessage = "disabled via management API"
|
||||
}
|
||||
return
|
||||
}
|
||||
auth.Status = coreauth.StatusActive
|
||||
auth.StatusMessage = ""
|
||||
}
|
||||
|
||||
func (h *Handler) removeAuth(ctx context.Context, id string) {
|
||||
if h == nil || h.authManager == nil {
|
||||
return
|
||||
}
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
return
|
||||
}
|
||||
if _, ok := h.authManager.GetByID(id); ok {
|
||||
h.authManager.Remove(ctx, id)
|
||||
return
|
||||
}
|
||||
authID := h.authIDForPath(id)
|
||||
if authID == "" {
|
||||
return
|
||||
}
|
||||
h.authManager.Remove(ctx, authID)
|
||||
}
|
||||
|
||||
func (h *Handler) removeAuthsForPath(ctx context.Context, path string, fallbackID string) {
|
||||
if h == nil || h.authManager == nil {
|
||||
return
|
||||
}
|
||||
removed := false
|
||||
for _, auth := range h.authManager.List() {
|
||||
if auth == nil {
|
||||
continue
|
||||
}
|
||||
if sameAuthFilePath(authAttribute(auth, "path"), path) || sameAuthFilePath(authAttribute(auth, coreauth.AttributeVirtualSource), path) {
|
||||
h.removeAuth(ctx, auth.ID)
|
||||
removed = true
|
||||
}
|
||||
}
|
||||
if removed {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(fallbackID) != "" {
|
||||
h.removeAuth(ctx, fallbackID)
|
||||
return
|
||||
}
|
||||
h.removeAuth(ctx, path)
|
||||
}
|
||||
|
||||
func sameAuthFilePath(left, right string) bool {
|
||||
left = cleanAuthFilePath(left)
|
||||
right = cleanAuthFilePath(right)
|
||||
if left == "" || right == "" {
|
||||
return false
|
||||
}
|
||||
if runtime.GOOS == "windows" {
|
||||
return strings.EqualFold(left, right)
|
||||
}
|
||||
return left == right
|
||||
}
|
||||
|
||||
func cleanAuthFilePath(path string) string {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return ""
|
||||
}
|
||||
if abs, errAbs := filepath.Abs(path); errAbs == nil && strings.TrimSpace(abs) != "" {
|
||||
path = abs
|
||||
}
|
||||
return filepath.Clean(path)
|
||||
}
|
||||
|
||||
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 {
|
||||
store = sdkAuth.GetTokenStore()
|
||||
h.tokenStore = store
|
||||
}
|
||||
if h.cfg != nil {
|
||||
if dirSetter, ok := store.(interface{ SetBaseDir(string) }); ok {
|
||||
dirSetter.SetBaseDir(h.cfg.AuthDir)
|
||||
}
|
||||
}
|
||||
return store
|
||||
}
|
||||
|
||||
func (h *Handler) mergeExistingAuthFileMetadata(record *coreauth.Auth) {
|
||||
if h == nil || record == nil {
|
||||
return
|
||||
}
|
||||
var existingMap map[string]any
|
||||
|
||||
if h.cfg != nil && strings.TrimSpace(h.cfg.AuthDir) != "" {
|
||||
targetFile := record.FileName
|
||||
if targetFile == "" {
|
||||
targetFile = record.ID
|
||||
}
|
||||
if targetFile != "" {
|
||||
fullPath := filepath.Join(h.cfg.AuthDir, targetFile)
|
||||
if raw, errRead := os.ReadFile(fullPath); errRead == nil && len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &existingMap)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if existingMap == nil && h.authManager != nil {
|
||||
if existing, ok := h.authManager.GetByID(record.ID); ok && existing != nil && existing.Metadata != nil {
|
||||
existingMap = existing.Metadata
|
||||
} else {
|
||||
for _, auth := range h.authManager.List() {
|
||||
if auth != nil && auth.FileName == record.FileName && auth.Metadata != nil {
|
||||
existingMap = auth.Metadata
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(existingMap) > 0 {
|
||||
coreauth.MergeExistingAuthMetadata(record, existingMap)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) saveTokenRecord(ctx context.Context, record *coreauth.Auth) (string, error) {
|
||||
if record == nil {
|
||||
return "", fmt.Errorf("token record is nil")
|
||||
}
|
||||
h.mergeExistingAuthFileMetadata(record)
|
||||
store := h.tokenStoreWithBaseDir()
|
||||
if store == nil {
|
||||
return "", fmt.Errorf("token store unavailable")
|
||||
}
|
||||
if h.postAuthHook != nil {
|
||||
if err := h.postAuthHook(ctx, record); err != nil {
|
||||
return "", fmt.Errorf("post-auth hook failed: %w", err)
|
||||
}
|
||||
}
|
||||
savedPath, errSave := store.Save(ctx, record)
|
||||
if errSave != nil {
|
||||
return savedPath, errSave
|
||||
}
|
||||
if h.postAuthPersistHook != nil {
|
||||
if errHook := h.postAuthPersistHook(ctx, record); errHook != nil {
|
||||
return savedPath, fmt.Errorf("post-auth persist hook failed: %w", errHook)
|
||||
}
|
||||
}
|
||||
return savedPath, nil
|
||||
}
|
||||
|
|
@ -0,0 +1,260 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestListAuthFilesFiltersByNameAndAuthIndex(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "shared-codex.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex"}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
registerAuthForLookupTest(t, manager, &coreauth.Auth{
|
||||
ID: "auth-a",
|
||||
Index: "idx-a",
|
||||
FileName: fileName,
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"path": filePath,
|
||||
},
|
||||
})
|
||||
registerAuthForLookupTest(t, manager, &coreauth.Auth{
|
||||
ID: "auth-b",
|
||||
Index: "idx-b",
|
||||
FileName: fileName,
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"path": filePath,
|
||||
},
|
||||
})
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v0/management/auth-files?name=shared-codex.json&auth_index=idx-b", nil)
|
||||
ctx.Request = req
|
||||
|
||||
h.ListAuthFiles(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
var payload struct {
|
||||
Files []map[string]any `json:"files"`
|
||||
}
|
||||
if errDecode := json.Unmarshal(rec.Body.Bytes(), &payload); errDecode != nil {
|
||||
t.Fatalf("decode response: %v", errDecode)
|
||||
}
|
||||
if len(payload.Files) != 1 {
|
||||
t.Fatalf("files len = %d, want 1 payload=%s", len(payload.Files), rec.Body.String())
|
||||
}
|
||||
if got := payload.Files[0]["id"]; got != "auth-b" {
|
||||
t.Fatalf("id = %#v, want auth-b", got)
|
||||
}
|
||||
if got := payload.Files[0]["auth_index"]; got != "idx-b" {
|
||||
t.Fatalf("auth_index = %#v, want idx-b", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAuthFilesFromDiskFiltersByNameAndRejectsAuthIndex(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
for _, file := range []struct {
|
||||
name string
|
||||
body string
|
||||
}{
|
||||
{name: "alpha.json", body: `{"type":"codex","email":"alpha@example.com"}`},
|
||||
{name: "beta.json", body: `{"type":"codex","email":"beta@example.com"}`},
|
||||
} {
|
||||
if errWrite := os.WriteFile(filepath.Join(authDir, file.name), []byte(file.body), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write auth file %s: %v", file.name, errWrite)
|
||||
}
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files?name=beta.json", nil)
|
||||
|
||||
h.ListAuthFiles(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
var payload struct {
|
||||
Files []map[string]any `json:"files"`
|
||||
}
|
||||
if errDecode := json.Unmarshal(rec.Body.Bytes(), &payload); errDecode != nil {
|
||||
t.Fatalf("decode response: %v", errDecode)
|
||||
}
|
||||
if len(payload.Files) != 1 || payload.Files[0]["name"] != "beta.json" {
|
||||
t.Fatalf("files = %#v, want only beta.json", payload.Files)
|
||||
}
|
||||
|
||||
rec = httptest.NewRecorder()
|
||||
ctx, _ = gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files?name=beta.json&auth_index=idx-b", nil)
|
||||
|
||||
h.ListAuthFiles(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
payload.Files = nil
|
||||
if errDecode := json.Unmarshal(rec.Body.Bytes(), &payload); errDecode != nil {
|
||||
t.Fatalf("decode auth_index response: %v", errDecode)
|
||||
}
|
||||
if len(payload.Files) != 0 {
|
||||
t.Fatalf("files = %#v, want no disk fallback matches for auth_index", payload.Files)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchAuthFileStatusVerifiesAuthIndex(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
registerAuthForLookupTest(t, manager, &coreauth.Auth{
|
||||
ID: "auth-a",
|
||||
Index: "idx-a",
|
||||
FileName: "shared-codex.json",
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
})
|
||||
registerAuthForLookupTest(t, manager, &coreauth.Auth{
|
||||
ID: "auth-b",
|
||||
Index: "idx-b",
|
||||
FileName: "shared-codex.json",
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
})
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"shared-codex.json","auth_index":"idx-b","disabled":true}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
|
||||
h.PatchAuthFileStatus(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
authA, okA := manager.GetByID("auth-a")
|
||||
authB, okB := manager.GetByID("auth-b")
|
||||
if !okA || !okB {
|
||||
t.Fatalf("expected both auth records to exist")
|
||||
}
|
||||
if authA.Disabled || authA.Status == coreauth.StatusDisabled {
|
||||
t.Fatalf("auth-a was modified: %+v", authA)
|
||||
}
|
||||
if !authB.Disabled || authB.Status != coreauth.StatusDisabled {
|
||||
t.Fatalf("auth-b was not disabled: %+v", authB)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchAuthFileStatusRejectsMismatchedAuthIndex(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
registerAuthForLookupTest(t, manager, &coreauth.Auth{
|
||||
ID: "auth-a",
|
||||
Index: "idx-a",
|
||||
FileName: "shared-codex.json",
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
})
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"shared-codex.json","auth_index":"idx-missing","disabled":true}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
|
||||
h.PatchAuthFileStatus(ctx)
|
||||
|
||||
if rec.Code != http.StatusNotFound {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusNotFound, rec.Body.String())
|
||||
}
|
||||
authA, ok := manager.GetByID("auth-a")
|
||||
if !ok {
|
||||
t.Fatalf("expected auth-a to exist")
|
||||
}
|
||||
if authA.Disabled || authA.Status == coreauth.StatusDisabled {
|
||||
t.Fatalf("auth-a was modified: %+v", authA)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthFileLookupAndEntryBuildConcurrentEnsureIndex(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "concurrent-codex.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex"}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
auth := &coreauth.Auth{
|
||||
ID: "auth-concurrent",
|
||||
Index: "idx-concurrent",
|
||||
FileName: fileName,
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"path": filePath,
|
||||
},
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 32; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < 100; j++ {
|
||||
if !matchesAuthFileLookup(auth, fileName, "idx-concurrent") {
|
||||
t.Errorf("auth lookup did not match")
|
||||
}
|
||||
entry := h.buildAuthFileEntry(auth)
|
||||
if entry == nil {
|
||||
t.Errorf("entry is nil")
|
||||
continue
|
||||
}
|
||||
if got := entry["auth_index"]; got != "idx-concurrent" {
|
||||
t.Errorf("auth_index = %#v, want idx-concurrent", got)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func registerAuthForLookupTest(t *testing.T, manager *coreauth.Manager, auth *coreauth.Auth) {
|
||||
t.Helper()
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth %q: %v", auth.ID, errRegister)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,220 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
anthropicCallbackPort = 54545
|
||||
codexCallbackPort = 1455
|
||||
)
|
||||
|
||||
type callbackForwarder struct {
|
||||
provider string
|
||||
server *http.Server
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
func isWebUIRequest(c *gin.Context) bool {
|
||||
raw := strings.TrimSpace(c.Query("is_webui"))
|
||||
if raw == "" {
|
||||
return false
|
||||
}
|
||||
switch strings.ToLower(raw) {
|
||||
case "1", "true", "yes", "on":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func startCallbackForwarder(port int, provider, targetBase string) (*callbackForwarder, error) {
|
||||
callbackForwardersMu.Lock()
|
||||
prev := callbackForwarders[port]
|
||||
if prev != nil {
|
||||
delete(callbackForwarders, port)
|
||||
}
|
||||
callbackForwardersMu.Unlock()
|
||||
|
||||
if prev != nil {
|
||||
stopForwarderInstance(port, prev)
|
||||
}
|
||||
|
||||
addr := fmt.Sprintf("0.0.0.0:%d", port)
|
||||
ln, err := net.Listen("tcp", addr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to listen on %s: %w", addr, err)
|
||||
}
|
||||
|
||||
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
target := targetBase
|
||||
if raw := r.URL.RawQuery; raw != "" {
|
||||
if strings.Contains(target, "?") {
|
||||
target = target + "&" + raw
|
||||
} else {
|
||||
target = target + "?" + raw
|
||||
}
|
||||
}
|
||||
w.Header().Set("Cache-Control", "no-store")
|
||||
http.Redirect(w, r, target, http.StatusFound)
|
||||
})
|
||||
|
||||
srv := &http.Server{
|
||||
Handler: handler,
|
||||
ReadHeaderTimeout: 5 * time.Second,
|
||||
WriteTimeout: 5 * time.Second,
|
||||
}
|
||||
done := make(chan struct{})
|
||||
|
||||
go func() {
|
||||
if errServe := srv.Serve(ln); errServe != nil && !errors.Is(errServe, http.ErrServerClosed) {
|
||||
log.WithError(errServe).Warnf("callback forwarder for %s stopped unexpectedly", provider)
|
||||
}
|
||||
close(done)
|
||||
}()
|
||||
|
||||
forwarder := &callbackForwarder{
|
||||
provider: provider,
|
||||
server: srv,
|
||||
done: done,
|
||||
}
|
||||
|
||||
callbackForwardersMu.Lock()
|
||||
callbackForwarders[port] = forwarder
|
||||
callbackForwardersMu.Unlock()
|
||||
|
||||
log.Infof("callback forwarder for %s listening on %s", provider, addr)
|
||||
|
||||
return forwarder, nil
|
||||
}
|
||||
|
||||
func stopCallbackForwarderInstance(port int, forwarder *callbackForwarder) {
|
||||
if forwarder == nil {
|
||||
return
|
||||
}
|
||||
callbackForwardersMu.Lock()
|
||||
if current := callbackForwarders[port]; current == forwarder {
|
||||
delete(callbackForwarders, port)
|
||||
}
|
||||
callbackForwardersMu.Unlock()
|
||||
|
||||
stopForwarderInstance(port, forwarder)
|
||||
}
|
||||
|
||||
func stopForwarderInstance(port int, forwarder *callbackForwarder) {
|
||||
if forwarder == nil || forwarder.server == nil {
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := forwarder.server.Shutdown(ctx); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
log.WithError(err).Warnf("failed to shut down callback forwarder on port %d", port)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-forwarder.done:
|
||||
case <-time.After(2 * time.Second):
|
||||
}
|
||||
|
||||
log.Infof("callback forwarder on port %d stopped", port)
|
||||
}
|
||||
|
||||
func (h *Handler) managementCallbackURL(path string) (string, error) {
|
||||
if h == nil || h.cfg == nil || h.cfg.Port <= 0 {
|
||||
return "", fmt.Errorf("server port is not configured")
|
||||
}
|
||||
if !strings.HasPrefix(path, "/") {
|
||||
path = "/" + path
|
||||
}
|
||||
scheme := "http"
|
||||
if h.cfg.TLS.Enable {
|
||||
scheme = "https"
|
||||
}
|
||||
return fmt.Sprintf("%s://127.0.0.1:%d%s", scheme, h.cfg.Port, path), nil
|
||||
}
|
||||
|
||||
func pluginAuthProviderFromPath(path string) (string, bool) {
|
||||
path = strings.TrimSpace(path)
|
||||
const prefix = "/v0/management/"
|
||||
const suffix = "-auth-url"
|
||||
if !strings.HasPrefix(path, prefix) || !strings.HasSuffix(path, suffix) {
|
||||
return "", false
|
||||
}
|
||||
provider := strings.TrimSuffix(strings.TrimPrefix(path, prefix), suffix)
|
||||
provider = strings.ToLower(strings.TrimSpace(provider))
|
||||
if provider == "" {
|
||||
return "", false
|
||||
}
|
||||
for _, r := range provider {
|
||||
switch {
|
||||
case r >= 'a' && r <= 'z':
|
||||
case r >= '0' && r <= '9':
|
||||
case r == '-':
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
return provider, true
|
||||
}
|
||||
|
||||
func (h *Handler) ServePluginAuthURL(c *gin.Context) bool {
|
||||
if h == nil || c == nil || c.Request == nil || c.Request.URL == nil {
|
||||
return false
|
||||
}
|
||||
h.mu.Lock()
|
||||
host := h.pluginHost
|
||||
h.mu.Unlock()
|
||||
if host == nil {
|
||||
return false
|
||||
}
|
||||
provider, ok := pluginAuthProviderFromPath(c.Request.URL.Path)
|
||||
if !ok || !host.HasAuthProvider(provider) {
|
||||
return false
|
||||
}
|
||||
|
||||
ctx := PopulateAuthContext(context.Background(), c)
|
||||
baseURL, errBaseURL := h.managementCallbackURL("/v0/management/oauth-callback")
|
||||
if errBaseURL != nil {
|
||||
log.WithError(errBaseURL).Error("failed to compute plugin auth callback URL")
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"})
|
||||
return true
|
||||
}
|
||||
resp, handled, errStart := host.StartLogin(ctx, provider, baseURL)
|
||||
if !handled {
|
||||
return false
|
||||
}
|
||||
if errStart != nil {
|
||||
log.WithError(errStart).Error("failed to start plugin auth login")
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"})
|
||||
return true
|
||||
}
|
||||
state := strings.TrimSpace(resp.State)
|
||||
if state == "" {
|
||||
log.WithField("provider", provider).Error("plugin auth provider returned empty state")
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "invalid oauth state"})
|
||||
return true
|
||||
}
|
||||
if errState := ValidateOAuthState(state); errState != nil {
|
||||
log.WithError(errState).WithField("provider", provider).Error("plugin auth provider returned invalid state")
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "invalid oauth state"})
|
||||
return true
|
||||
}
|
||||
if errRegister := RegisterPluginOAuthSession(state, provider, resp.Metadata); errRegister != nil {
|
||||
log.WithError(errRegister).WithField("provider", provider).Error("failed to register plugin oauth session")
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "failed to generate authorization url"})
|
||||
return true
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok", "url": resp.URL, "state": state})
|
||||
return true
|
||||
}
|
||||
|
|
@ -0,0 +1,600 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
fileauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestPatchAuthFileFields_MergeHeadersAndDeleteEmptyValues(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
store := &memoryAuthStore{}
|
||||
manager := coreauth.NewManager(store, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: "test.json",
|
||||
FileName: "test.json",
|
||||
Provider: "claude",
|
||||
Attributes: map[string]string{
|
||||
"path": "/tmp/test.json",
|
||||
"header:X-Old": "old",
|
||||
"header:X-Remove": "gone",
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "claude",
|
||||
"headers": map[string]any{
|
||||
"X-Old": "old",
|
||||
"X-Remove": "gone",
|
||||
},
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
|
||||
body := `{"name":"test.json","prefix":"p1","proxy_url":"http://proxy.local","headers":{"X-Old":"new","X-New":"v","X-Remove":" ","X-Nope":""}}`
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/fields", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
h.PatchAuthFileFields(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
updated, ok := manager.GetByID("test.json")
|
||||
if !ok || updated == nil {
|
||||
t.Fatalf("expected auth record to exist after patch")
|
||||
}
|
||||
|
||||
if updated.Prefix != "p1" {
|
||||
t.Fatalf("prefix = %q, want %q", updated.Prefix, "p1")
|
||||
}
|
||||
if updated.ProxyURL != "http://proxy.local" {
|
||||
t.Fatalf("proxy_url = %q, want %q", updated.ProxyURL, "http://proxy.local")
|
||||
}
|
||||
|
||||
if updated.Metadata == nil {
|
||||
t.Fatalf("expected metadata to be non-nil")
|
||||
}
|
||||
if got, _ := updated.Metadata["prefix"].(string); got != "p1" {
|
||||
t.Fatalf("metadata.prefix = %q, want %q", got, "p1")
|
||||
}
|
||||
if got, _ := updated.Metadata["proxy_url"].(string); got != "http://proxy.local" {
|
||||
t.Fatalf("metadata.proxy_url = %q, want %q", got, "http://proxy.local")
|
||||
}
|
||||
|
||||
headersMeta, ok := updated.Metadata["headers"].(map[string]any)
|
||||
if !ok {
|
||||
raw, _ := json.Marshal(updated.Metadata["headers"])
|
||||
t.Fatalf("metadata.headers = %T (%s), want map[string]any", updated.Metadata["headers"], string(raw))
|
||||
}
|
||||
if got := headersMeta["X-Old"]; got != "new" {
|
||||
t.Fatalf("metadata.headers.X-Old = %#v, want %q", got, "new")
|
||||
}
|
||||
if got := headersMeta["X-New"]; got != "v" {
|
||||
t.Fatalf("metadata.headers.X-New = %#v, want %q", got, "v")
|
||||
}
|
||||
if _, ok := headersMeta["X-Remove"]; ok {
|
||||
t.Fatalf("expected metadata.headers.X-Remove to be deleted")
|
||||
}
|
||||
if _, ok := headersMeta["X-Nope"]; ok {
|
||||
t.Fatalf("expected metadata.headers.X-Nope to be absent")
|
||||
}
|
||||
|
||||
if got := updated.Attributes["header:X-Old"]; got != "new" {
|
||||
t.Fatalf("attrs header:X-Old = %q, want %q", got, "new")
|
||||
}
|
||||
if got := updated.Attributes["header:X-New"]; got != "v" {
|
||||
t.Fatalf("attrs header:X-New = %q, want %q", got, "v")
|
||||
}
|
||||
if _, ok := updated.Attributes["header:X-Remove"]; ok {
|
||||
t.Fatalf("expected attrs header:X-Remove to be deleted")
|
||||
}
|
||||
if _, ok := updated.Attributes["header:X-Nope"]; ok {
|
||||
t.Fatalf("expected attrs header:X-Nope to be absent")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchAuthFileFields_HeadersEmptyMapIsNoop(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
store := &memoryAuthStore{}
|
||||
manager := coreauth.NewManager(store, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: "noop.json",
|
||||
FileName: "noop.json",
|
||||
Provider: "claude",
|
||||
Attributes: map[string]string{
|
||||
"path": "/tmp/noop.json",
|
||||
"header:X-Kee": "1",
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "claude",
|
||||
"headers": map[string]any{
|
||||
"X-Kee": "1",
|
||||
},
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
|
||||
body := `{"name":"noop.json","note":"hello","headers":{}}`
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/fields", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
h.PatchAuthFileFields(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
updated, ok := manager.GetByID("noop.json")
|
||||
if !ok || updated == nil {
|
||||
t.Fatalf("expected auth record to exist after patch")
|
||||
}
|
||||
if got := updated.Attributes["header:X-Kee"]; got != "1" {
|
||||
t.Fatalf("attrs header:X-Kee = %q, want %q", got, "1")
|
||||
}
|
||||
headersMeta, ok := updated.Metadata["headers"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("expected metadata.headers to remain a map, got %T", updated.Metadata["headers"])
|
||||
}
|
||||
if got := headersMeta["X-Kee"]; got != "1" {
|
||||
t.Fatalf("metadata.headers.X-Kee = %#v, want %q", got, "1")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchAuthFileFields_WebsocketsFalseIsUpdate(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
store := &memoryAuthStore{}
|
||||
manager := coreauth.NewManager(store, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: "codex.json",
|
||||
FileName: "codex.json",
|
||||
Provider: "codex",
|
||||
Attributes: map[string]string{
|
||||
"path": "/tmp/codex.json",
|
||||
"websockets": "true",
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "codex",
|
||||
"websockets": true,
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
|
||||
body := `{"name":"codex.json","websockets":false}`
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/fields", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
h.PatchAuthFileFields(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
updated, ok := manager.GetByID("codex.json")
|
||||
if !ok || updated == nil {
|
||||
t.Fatalf("expected auth record to exist after patch")
|
||||
}
|
||||
if got := updated.Attributes["websockets"]; got != "false" {
|
||||
t.Fatalf("attrs websockets = %q, want %q", got, "false")
|
||||
}
|
||||
if got, ok := updated.Metadata["websockets"].(bool); !ok || got {
|
||||
t.Fatalf("metadata.websockets = %#v, want false", updated.Metadata["websockets"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchAuthFileFields_ArbitraryFieldsPersistToFile(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "generic.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
store := fileauth.NewFileTokenStore()
|
||||
store.SetBaseDir(authDir)
|
||||
manager := coreauth.NewManager(store, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
FileName: fileName,
|
||||
Provider: "codex",
|
||||
Attributes: map[string]string{
|
||||
"path": filePath,
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "codex",
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
|
||||
body := `{"name":"generic.json","abc":true,"nested.cde":true,"fgh":{"ijk":true}}`
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/fields", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
h.PatchAuthFileFields(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
raw, errRead := os.ReadFile(filePath)
|
||||
if errRead != nil {
|
||||
t.Fatalf("failed to read updated auth file: %v", errRead)
|
||||
}
|
||||
var data map[string]any
|
||||
if errUnmarshal := json.Unmarshal(raw, &data); errUnmarshal != nil {
|
||||
t.Fatalf("failed to unmarshal updated auth file: %v", errUnmarshal)
|
||||
}
|
||||
if got := data["abc"]; got != true {
|
||||
t.Fatalf("abc = %#v, want true", got)
|
||||
}
|
||||
nested, ok := data["nested"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("nested = %#v, want object", data["nested"])
|
||||
}
|
||||
if got := nested["cde"]; got != true {
|
||||
t.Fatalf("nested.cde = %#v, want true", got)
|
||||
}
|
||||
fgh, ok := data["fgh"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("fgh = %#v, want object", data["fgh"])
|
||||
}
|
||||
if got := fgh["ijk"]; got != true {
|
||||
t.Fatalf("fgh.ijk = %#v, want true", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchAuthFileFields_WeightPersistsAndSyncsRuntime(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "weighted.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
store := fileauth.NewFileTokenStore()
|
||||
store.SetBaseDir(authDir)
|
||||
manager := coreauth.NewManager(store, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
FileName: fileName,
|
||||
Provider: "codex",
|
||||
Attributes: map[string]string{"path": filePath},
|
||||
Metadata: map[string]any{"type": "codex"},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("Register() error = %v", errRegister)
|
||||
}
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
|
||||
patch := func(weight string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
body := `{"name":"weighted.json","weight":` + weight + `}`
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/fields", strings.NewReader(body))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
h.PatchAuthFileFields(ctx)
|
||||
return rec
|
||||
}
|
||||
|
||||
if rec := patch("7"); rec.Code != http.StatusOK {
|
||||
t.Fatalf("update status = %d, want 200; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
updated, ok := manager.GetByID(fileName)
|
||||
if !ok || updated.Attributes[coreauth.AttributeWeight] != "7" {
|
||||
t.Fatalf("runtime weight = %#v, want 7", updated)
|
||||
}
|
||||
raw, errRead := os.ReadFile(filePath)
|
||||
if errRead != nil {
|
||||
t.Fatalf("ReadFile() error = %v", errRead)
|
||||
}
|
||||
var persisted map[string]any
|
||||
if errUnmarshal := json.Unmarshal(raw, &persisted); errUnmarshal != nil {
|
||||
t.Fatalf("Unmarshal() error = %v", errUnmarshal)
|
||||
}
|
||||
if persisted["weight"] != float64(7) {
|
||||
t.Fatalf("persisted weight = %#v, want 7", persisted["weight"])
|
||||
}
|
||||
|
||||
if rec := patch("null"); rec.Code != http.StatusOK {
|
||||
t.Fatalf("reset status = %d, want 200; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
updated, _ = manager.GetByID(fileName)
|
||||
if _, exists := updated.Attributes[coreauth.AttributeWeight]; exists {
|
||||
t.Fatal("runtime weight remains after reset")
|
||||
}
|
||||
raw, errRead = os.ReadFile(filePath)
|
||||
if errRead != nil {
|
||||
t.Fatalf("ReadFile() after reset error = %v", errRead)
|
||||
}
|
||||
persisted = nil
|
||||
if errUnmarshal := json.Unmarshal(raw, &persisted); errUnmarshal != nil {
|
||||
t.Fatalf("Unmarshal() after reset error = %v", errUnmarshal)
|
||||
}
|
||||
if _, exists := persisted["weight"]; exists {
|
||||
t.Fatal("persisted weight remains after reset")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchAuthFileFields_RejectsInvalidWeights(t *testing.T) {
|
||||
store := &memoryAuthStore{}
|
||||
manager := coreauth.NewManager(store, nil, nil)
|
||||
record := &coreauth.Auth{ID: "auth.json", FileName: "auth.json", Provider: "codex", Metadata: map[string]any{"type": "codex"}}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("Register() error = %v", errRegister)
|
||||
}
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{}, manager)
|
||||
|
||||
for _, weight := range []string{"1.5", "1000001", "9223372036854775808", `"7"`} {
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
body := `{"name":"auth.json","weight":` + weight + `}`
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/fields", strings.NewReader(body))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
h.PatchAuthFileFields(ctx)
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("weight %s status = %d, want 400; body=%s", weight, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchAuthFileFields_RequestRetryRoundTrip(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "request-retry.json"
|
||||
store := fileauth.NewFileTokenStore()
|
||||
store.SetBaseDir(authDir)
|
||||
manager := coreauth.NewManager(store, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
FileName: fileName,
|
||||
Provider: "codex",
|
||||
Attributes: map[string]string{
|
||||
"path": filepath.Join(authDir, fileName),
|
||||
},
|
||||
Metadata: map[string]any{"type": "codex"},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("Register() error = %v", errRegister)
|
||||
}
|
||||
|
||||
handler := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
engine := gin.New()
|
||||
engine.GET("/auth-files", handler.ListAuthFiles)
|
||||
engine.PATCH("/auth-files/fields", handler.PatchAuthFileFields)
|
||||
|
||||
patch := func(body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
response := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPatch, "/auth-files/fields", strings.NewReader(body))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
engine.ServeHTTP(response, request)
|
||||
return response
|
||||
}
|
||||
getRequestRetry := func() *int {
|
||||
t.Helper()
|
||||
response := httptest.NewRecorder()
|
||||
engine.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/auth-files", nil))
|
||||
if response.Code != http.StatusOK {
|
||||
t.Fatalf("GET status = %d body=%s", response.Code, response.Body.String())
|
||||
}
|
||||
var payload struct {
|
||||
Files []struct {
|
||||
Name string `json:"name"`
|
||||
RequestRetry *int `json:"request_retry"`
|
||||
} `json:"files"`
|
||||
}
|
||||
if errDecode := json.Unmarshal(response.Body.Bytes(), &payload); errDecode != nil {
|
||||
t.Fatalf("decode GET response: %v", errDecode)
|
||||
}
|
||||
if len(payload.Files) != 1 || payload.Files[0].Name != fileName {
|
||||
t.Fatalf("GET files = %#v", payload.Files)
|
||||
}
|
||||
return payload.Files[0].RequestRetry
|
||||
}
|
||||
|
||||
if response := patch(`{"name":"request-retry.json","request-retry":2}`); response.Code != http.StatusOK {
|
||||
t.Fatalf("PATCH status = %d body=%s", response.Code, response.Body.String())
|
||||
}
|
||||
updated, ok := manager.GetByID(fileName)
|
||||
if !ok {
|
||||
t.Fatal("updated auth is missing")
|
||||
}
|
||||
if retry, okRetry := updated.RequestRetryOverride(); !okRetry || retry != 2 {
|
||||
t.Fatalf("RequestRetryOverride() = (%d, %t), want (2, true)", retry, okRetry)
|
||||
}
|
||||
if _, exists := updated.Metadata["request-retry"]; exists {
|
||||
t.Fatalf("legacy request-retry metadata remains: %#v", updated.Metadata)
|
||||
}
|
||||
persistedData, errRead := os.ReadFile(filepath.Join(authDir, fileName))
|
||||
if errRead != nil {
|
||||
t.Fatalf("read persisted auth: %v", errRead)
|
||||
}
|
||||
var persisted map[string]any
|
||||
if errUnmarshal := json.Unmarshal(persistedData, &persisted); errUnmarshal != nil {
|
||||
t.Fatalf("decode persisted auth: %v", errUnmarshal)
|
||||
}
|
||||
if persisted["request_retry"] != float64(2) {
|
||||
t.Fatalf("persisted request_retry = %#v, want 2", persisted["request_retry"])
|
||||
}
|
||||
if _, exists := persisted["request-retry"]; exists {
|
||||
t.Fatalf("persisted legacy request-retry remains: %#v", persisted)
|
||||
}
|
||||
if retry := getRequestRetry(); retry == nil || *retry != 2 {
|
||||
t.Fatalf("GET request_retry = %#v, want 2", retry)
|
||||
}
|
||||
|
||||
if response := patch(`{"name":"request-retry.json","request_retry":0}`); response.Code != http.StatusOK {
|
||||
t.Fatalf("PATCH underscore status = %d body=%s", response.Code, response.Body.String())
|
||||
}
|
||||
if retry := getRequestRetry(); retry == nil || *retry != 0 {
|
||||
t.Fatalf("GET request_retry = %#v, want explicit 0", retry)
|
||||
}
|
||||
|
||||
if response := patch(`{"name":"request-retry.json","request-retry":-1}`); response.Code != http.StatusOK {
|
||||
t.Fatalf("PATCH negative status = %d body=%s", response.Code, response.Body.String())
|
||||
}
|
||||
if retry := getRequestRetry(); retry != nil {
|
||||
t.Fatalf("GET request_retry after negative clear = %#v, want omitted", retry)
|
||||
}
|
||||
|
||||
if response := patch(`{"name":"request-retry.json","request_retry":2}`); response.Code != http.StatusOK {
|
||||
t.Fatalf("PATCH reset status = %d body=%s", response.Code, response.Body.String())
|
||||
}
|
||||
if response := patch(`{"name":"request-retry.json","request-retry":2,"request_retry":3}`); response.Code != http.StatusOK {
|
||||
t.Fatalf("PATCH canonical precedence status = %d body=%s", response.Code, response.Body.String())
|
||||
}
|
||||
if retry := getRequestRetry(); retry == nil || *retry != 3 {
|
||||
t.Fatalf("GET request_retry after alias conflict = %#v, want canonical 3", retry)
|
||||
}
|
||||
if response := patch(`{"name":"request-retry.json","request_retry":2}`); response.Code != http.StatusOK {
|
||||
t.Fatalf("PATCH second reset status = %d body=%s", response.Code, response.Body.String())
|
||||
}
|
||||
for _, body := range []string{
|
||||
`{"name":"request-retry.json","request-retry":"2"}`,
|
||||
`{"name":"request-retry.json","request-retry":1.5}`,
|
||||
`{"name":"request-retry.json","request_retry.child":2}`,
|
||||
`{"name":"request-retry.json","request_retry .child":2}`,
|
||||
`{"name":"request-retry.json","request-retry .child":2}`,
|
||||
} {
|
||||
if response := patch(body); response.Code != http.StatusBadRequest {
|
||||
t.Fatalf("PATCH %s status = %d, want 400 body=%s", body, response.Code, response.Body.String())
|
||||
}
|
||||
if retry := getRequestRetry(); retry == nil || *retry != 2 {
|
||||
t.Fatalf("invalid PATCH changed request_retry to %#v", retry)
|
||||
}
|
||||
}
|
||||
|
||||
if response := patch(`{"name":"request-retry.json","request_retry":null}`); response.Code != http.StatusOK {
|
||||
t.Fatalf("PATCH null status = %d body=%s", response.Code, response.Body.String())
|
||||
}
|
||||
if retry := getRequestRetry(); retry != nil {
|
||||
t.Fatalf("GET request_retry after null clear = %#v, want omitted", retry)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthFileRequestRetryFromJSON(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
raw string
|
||||
want int
|
||||
ok bool
|
||||
}{
|
||||
{name: "canonical", raw: `{"request_retry":2}`, want: 2, ok: true},
|
||||
{name: "legacy", raw: `{"request-retry":2}`, want: 2, ok: true},
|
||||
{name: "negative inherits", raw: `{"request_retry":-1}`},
|
||||
{name: "string integer compatibility", raw: `{"request_retry":"2"}`, want: 2, ok: true},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
got, ok := authFileRequestRetryFromJSON([]byte(test.raw))
|
||||
if got != test.want || ok != test.ok {
|
||||
t.Fatalf("authFileRequestRetryFromJSON(%s) = (%d, %t), want (%d, %t)", test.raw, got, ok, test.want, test.ok)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeAuthFilePatchFieldsCanonicalizesLegacyRoots(t *testing.T) {
|
||||
fields := map[string]json.RawMessage{
|
||||
"request-retry": json.RawMessage(`2`),
|
||||
" disable-cooling ": json.RawMessage(`true`),
|
||||
"fingerprint-profile.value": json.RawMessage(`"x"`),
|
||||
"provider-specific": json.RawMessage(`"preserved"`),
|
||||
}
|
||||
|
||||
normalized, errNormalize := normalizeAuthFilePatchFields(fields)
|
||||
if errNormalize != nil {
|
||||
t.Fatalf("normalizeAuthFilePatchFields() error = %v", errNormalize)
|
||||
}
|
||||
for _, key := range []string{"request_retry", "disable_cooling", "fingerprint_profile.value", "provider-specific"} {
|
||||
if _, exists := normalized[key]; !exists {
|
||||
t.Fatalf("normalized fields missing %q: %#v", key, normalized)
|
||||
}
|
||||
}
|
||||
|
||||
canonicalWins, errCanonicalWins := normalizeAuthFilePatchFields(map[string]json.RawMessage{
|
||||
"request-retry": json.RawMessage(`2`),
|
||||
"request_retry": json.RawMessage(`3`),
|
||||
})
|
||||
if errCanonicalWins != nil {
|
||||
t.Fatalf("normalizeAuthFilePatchFields() canonical precedence error = %v", errCanonicalWins)
|
||||
}
|
||||
if got := string(canonicalWins["request_retry"]); got != "3" {
|
||||
t.Fatalf("normalized request_retry = %s, want canonical value 3", got)
|
||||
}
|
||||
|
||||
_, errNestedDuplicate := normalizeAuthFilePatchFields(map[string]json.RawMessage{
|
||||
"disable_cooling.value": json.RawMessage(`true`),
|
||||
"disable_cooling . value": json.RawMessage(`false`),
|
||||
})
|
||||
if errNestedDuplicate == nil {
|
||||
t.Fatal("normalizeAuthFilePatchFields() accepted equivalent nested paths")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetSourceAuthFileDisabledNormalizesLegacyMetadata(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "legacy.json")
|
||||
if errWrite := os.WriteFile(path, []byte(`{"type":"codex","request-retry":2,"disable-cooling":true}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("write legacy auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
if errDisable := setSourceAuthFileDisabled(path, true); errDisable != nil {
|
||||
t.Fatalf("setSourceAuthFileDisabled() error = %v", errDisable)
|
||||
}
|
||||
persistedData, errRead := os.ReadFile(path)
|
||||
if errRead != nil {
|
||||
t.Fatalf("read persisted auth file: %v", errRead)
|
||||
}
|
||||
var persisted map[string]any
|
||||
if errUnmarshal := json.Unmarshal(persistedData, &persisted); errUnmarshal != nil {
|
||||
t.Fatalf("decode persisted auth file: %v", errUnmarshal)
|
||||
}
|
||||
if got := persisted["request_retry"]; got != float64(2) {
|
||||
t.Fatalf("persisted request_retry = %#v, want 2", got)
|
||||
}
|
||||
if got := persisted["disable_cooling"]; got != true {
|
||||
t.Fatalf("persisted disable_cooling = %#v, want true", got)
|
||||
}
|
||||
if got := persisted["disabled"]; got != true {
|
||||
t.Fatalf("persisted disabled = %#v, want true", got)
|
||||
}
|
||||
for _, legacy := range []string{"request-retry", "disable-cooling"} {
|
||||
if _, exists := persisted[legacy]; exists {
|
||||
t.Fatalf("persisted metadata retained %q: %#v", legacy, persisted)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,259 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
)
|
||||
|
||||
func TestPluginLoginPollAuthsExpandsMultipleAuths(t *testing.T) {
|
||||
host := pluginhost.New()
|
||||
resp := pluginapi.AuthLoginPollResponse{
|
||||
Status: pluginapi.AuthLoginStatusSuccess,
|
||||
Auths: []pluginapi.AuthData{
|
||||
{
|
||||
Provider: "gemini-cli",
|
||||
ID: "geminicli.json",
|
||||
FileName: "geminicli.json",
|
||||
StorageJSON: []byte(`{"type":"gemini-cli"}`),
|
||||
},
|
||||
{
|
||||
Provider: "gemini-cli",
|
||||
ID: "geminicli-project-a.json",
|
||||
FileName: "geminicli-project-a.json",
|
||||
StorageJSON: []byte(`{"type":"gemini-cli","project_id":"project-a"}`),
|
||||
Metadata: map[string]any{"project_id": "project-a"},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
records := pluginLoginPollAuths(host, resp)
|
||||
if len(records) != 2 {
|
||||
t.Fatalf("pluginLoginPollAuths() len = %d, want two records", len(records))
|
||||
}
|
||||
if records[0].ID != "geminicli.json" || records[1].ID != "geminicli-project-a.json" {
|
||||
t.Fatalf("records = %#v, want both plugin auths", records)
|
||||
}
|
||||
if gotProject := records[1].Metadata["project_id"]; gotProject != "project-a" {
|
||||
t.Fatalf("project_id = %#v, want project-a", gotProject)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSavePluginLoginRecordsRollsBackSavedAuthsOnFailure(t *testing.T) {
|
||||
store := &pluginLoginRollbackStore{failAt: 2}
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, nil)
|
||||
h.tokenStore = store
|
||||
|
||||
records := []*coreauth.Auth{
|
||||
{
|
||||
ID: "geminicli.json",
|
||||
FileName: "geminicli.json",
|
||||
Provider: "gemini-cli",
|
||||
Metadata: map[string]any{"type": "gemini-cli"},
|
||||
},
|
||||
{
|
||||
ID: "geminicli-project-a.json",
|
||||
FileName: "geminicli-project-a.json",
|
||||
Provider: "gemini-cli",
|
||||
Metadata: map[string]any{"type": "gemini-cli", "project_id": "project-a"},
|
||||
},
|
||||
}
|
||||
|
||||
errSave := h.savePluginLoginRecords(context.Background(), records)
|
||||
if errSave == nil {
|
||||
t.Fatal("savePluginLoginRecords() error = nil, want rollback-triggering error")
|
||||
}
|
||||
if len(store.saved) != 2 {
|
||||
t.Fatalf("saved len = %d, want two attempted saves", len(store.saved))
|
||||
}
|
||||
if !store.deleted["geminicli.json"] || !store.deleted["geminicli-project-a.json"] {
|
||||
t.Fatalf("deleted = %#v, want both saved auths rolled back", store.deleted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchPluginVirtualAuthStatusReturnsConflictForVirtualChild(t *testing.T) {
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
auth := pluginVirtualAuthForTest(t.TempDir(), "source.json", "auth-1")
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register virtual auth: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"auth-1","disabled":true}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
|
||||
h.PatchAuthFileStatus(ctx)
|
||||
|
||||
if rec.Code != http.StatusConflict {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusConflict, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchPluginVirtualSourceStatusDisablesAllExpandedAuths(t *testing.T) {
|
||||
authDir := t.TempDir()
|
||||
fileName := "source.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
if errWrite := os.WriteFile(filePath, []byte(`{"type":"gemini-cli","disabled":false}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("write source auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
for _, id := range []string{"source.json", "virtual-project-a"} {
|
||||
auth := pluginVirtualAuthForTest(authDir, fileName, id)
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register virtual auth %s: %v", id, errRegister)
|
||||
}
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/status", strings.NewReader(`{"name":"source.json","disabled":true}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
|
||||
h.PatchAuthFileStatus(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
raw, errRead := os.ReadFile(filePath)
|
||||
if errRead != nil {
|
||||
t.Fatalf("read source auth file: %v", errRead)
|
||||
}
|
||||
if !strings.Contains(string(raw), `"disabled":true`) {
|
||||
t.Fatalf("source auth file = %s, want disabled:true", string(raw))
|
||||
}
|
||||
for _, id := range []string{"source.json", "virtual-project-a"} {
|
||||
auth, ok := manager.GetByID(id)
|
||||
if !ok || auth == nil {
|
||||
t.Fatalf("expected auth %s to remain registered", id)
|
||||
}
|
||||
if !auth.Disabled || auth.Status != coreauth.StatusDisabled {
|
||||
t.Fatalf("auth %s disabled/status = %v/%s, want disabled", id, auth.Disabled, auth.Status)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchPluginVirtualAuthFieldsReturnsConflict(t *testing.T) {
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
auth := pluginVirtualAuthForTest(t.TempDir(), "source.json", "auth-1")
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register virtual auth: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/fields", strings.NewReader(`{"name":"auth-1","note":"hello"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
|
||||
h.PatchAuthFileFields(ctx)
|
||||
|
||||
if rec.Code != http.StatusConflict {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusConflict, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeletePluginVirtualSourceRemovesExpandedRuntimeAuths(t *testing.T) {
|
||||
authDir := t.TempDir()
|
||||
fileName := "source.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
if errWrite := os.WriteFile(filePath, []byte(`{"type":"gemini-cli"}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("write source auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
for _, id := range []string{"auth-1", "auth-2"} {
|
||||
auth := pluginVirtualAuthForTest(authDir, fileName, id)
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register virtual auth %s: %v", id, errRegister)
|
||||
}
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
h.tokenStore = &memoryAuthStore{}
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodDelete, "/v0/management/auth-files?name="+url.QueryEscape(fileName), nil)
|
||||
ctx.Request = req
|
||||
|
||||
h.DeleteAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if _, errStat := os.Stat(filePath); !os.IsNotExist(errStat) {
|
||||
t.Fatalf("expected source auth file to be removed, stat err: %v", errStat)
|
||||
}
|
||||
for _, id := range []string{"auth-1", "auth-2"} {
|
||||
if _, ok := manager.GetByID(id); ok {
|
||||
t.Fatalf("expected virtual auth %s to be removed", id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func pluginVirtualAuthForTest(authDir, fileName, id string) *coreauth.Auth {
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
auth := &coreauth.Auth{
|
||||
ID: id,
|
||||
FileName: fileName,
|
||||
Provider: "gemini-cli",
|
||||
Attributes: map[string]string{
|
||||
"path": filePath,
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "gemini-cli",
|
||||
},
|
||||
}
|
||||
coreauth.MarkPluginVirtualAuth(auth, filePath, 0)
|
||||
return auth
|
||||
}
|
||||
|
||||
type pluginLoginRollbackStore struct {
|
||||
failAt int
|
||||
saved []string
|
||||
deleted map[string]bool
|
||||
}
|
||||
|
||||
func (s *pluginLoginRollbackStore) List(context.Context) ([]*coreauth.Auth, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (s *pluginLoginRollbackStore) Save(_ context.Context, auth *coreauth.Auth) (string, error) {
|
||||
path := strings.TrimSpace(auth.FileName)
|
||||
if path == "" {
|
||||
path = strings.TrimSpace(auth.ID)
|
||||
}
|
||||
s.saved = append(s.saved, path)
|
||||
if len(s.saved) == s.failAt {
|
||||
return path, errors.New("save failed after write")
|
||||
}
|
||||
return path, nil
|
||||
}
|
||||
|
||||
func (s *pluginLoginRollbackStore) Delete(_ context.Context, id string) error {
|
||||
if s.deleted == nil {
|
||||
s.deleted = make(map[string]bool)
|
||||
}
|
||||
s.deleted[id] = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *pluginLoginRollbackStore) SetBaseDir(string) {}
|
||||
|
|
@ -0,0 +1,155 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestListAuthFiles_IncludesProjectIDFromManager(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "antigravity-user@example.com-project-a.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
if errWrite := os.WriteFile(filePath, []byte(`{"type":"antigravity","email":"user@example.com","project_id":"project-a"}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
FileName: fileName,
|
||||
Provider: "antigravity",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"path": filePath,
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "antigravity",
|
||||
"email": "user@example.com",
|
||||
"project_id": "project-a",
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
h.tokenStore = &memoryAuthStore{}
|
||||
|
||||
entry := firstAuthFileEntry(t, h)
|
||||
if got := entry["project_id"]; got != "project-a" {
|
||||
t.Fatalf("expected project_id %q, got %#v", "project-a", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAuthFilesFromDisk_IncludesProjectID(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
filePath := filepath.Join(authDir, "antigravity-user@example.com-project-a.json")
|
||||
if errWrite := os.WriteFile(filePath, []byte(`{"type":"antigravity","email":"user@example.com","project_id":"project-a"}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
|
||||
entry := firstAuthFileEntry(t, h)
|
||||
if got := entry["project_id"]; got != "project-a" {
|
||||
t.Fatalf("expected project_id %q, got %#v", "project-a", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAuthFiles_IncludesWebsocketsFromManager(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "codex-user@example.com-pro.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex","email":"user@example.com"}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
FileName: fileName,
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"path": filePath,
|
||||
"websockets": "true",
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "codex",
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
h.tokenStore = &memoryAuthStore{}
|
||||
|
||||
entry := firstAuthFileEntry(t, h)
|
||||
if got := entry["websockets"]; got != true {
|
||||
t.Fatalf("expected websockets true, got %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAuthFilesFromDisk_IncludesWebsockets(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
authDir := t.TempDir()
|
||||
filePath := filepath.Join(authDir, "codex-user@example.com-pro.json")
|
||||
if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex","email":"user@example.com","websockets":false}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
|
||||
entry := firstAuthFileEntry(t, h)
|
||||
if got := entry["websockets"]; got != false {
|
||||
t.Fatalf("expected websockets false, got %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func firstAuthFileEntry(t *testing.T, h *Handler) map[string]any {
|
||||
t.Helper()
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ginCtx, _ := gin.CreateTestContext(rec)
|
||||
ginCtx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files", nil)
|
||||
|
||||
h.ListAuthFiles(ginCtx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected list status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if errUnmarshal := json.Unmarshal(rec.Body.Bytes(), &payload); errUnmarshal != nil {
|
||||
t.Fatalf("failed to decode list payload: %v", errUnmarshal)
|
||||
}
|
||||
filesRaw, ok := payload["files"].([]any)
|
||||
if !ok {
|
||||
t.Fatalf("expected files array, payload: %#v", payload)
|
||||
}
|
||||
if len(filesRaw) != 1 {
|
||||
t.Fatalf("expected 1 auth entry, got %d", len(filesRaw))
|
||||
}
|
||||
fileEntry, ok := filesRaw[0].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("expected file entry object, got %#v", filesRaw[0])
|
||||
}
|
||||
return fileEntry
|
||||
}
|
||||
|
|
@ -0,0 +1,888 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/auth/antigravity"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/auth/claude"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/auth/codex"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/auth/kimi"
|
||||
xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/misc"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type codexOAuthService interface {
|
||||
GenerateAuthURL(state string, pkceCodes *codex.PKCECodes) (string, error)
|
||||
ExchangeCodeForTokens(ctx context.Context, code string, pkceCodes *codex.PKCECodes) (*codex.CodexAuthBundle, error)
|
||||
CreateTokenStorage(bundle *codex.CodexAuthBundle) *codex.CodexTokenStorage
|
||||
}
|
||||
|
||||
func (h *Handler) RequestAnthropicToken(c *gin.Context) {
|
||||
ctx := context.Background()
|
||||
ctx = PopulateAuthContext(ctx, c)
|
||||
|
||||
fmt.Println("Initializing Claude authentication...")
|
||||
|
||||
// Generate PKCE codes
|
||||
pkceCodes, err := claude.GeneratePKCECodes()
|
||||
if err != nil {
|
||||
log.Errorf("Failed to generate PKCE codes: %v", err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate PKCE codes"})
|
||||
return
|
||||
}
|
||||
|
||||
// Generate random state parameter
|
||||
state, err := misc.GenerateRandomState()
|
||||
if err != nil {
|
||||
log.Errorf("Failed to generate state parameter: %v", err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"})
|
||||
return
|
||||
}
|
||||
|
||||
// Initialize Claude auth service
|
||||
anthropicAuth := claude.NewClaudeAuth(h.cfg)
|
||||
|
||||
// Generate authorization URL (then override redirect_uri to reuse server port)
|
||||
authURL, state, err := anthropicAuth.GenerateAuthURL(state, pkceCodes)
|
||||
if err != nil {
|
||||
log.Errorf("Failed to generate authorization URL: %v", err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"})
|
||||
return
|
||||
}
|
||||
|
||||
RegisterOAuthSession(state, "anthropic")
|
||||
|
||||
isWebUI := isWebUIRequest(c)
|
||||
var forwarder *callbackForwarder
|
||||
if isWebUI {
|
||||
targetURL, errTarget := h.managementCallbackURL("/anthropic/callback")
|
||||
if errTarget != nil {
|
||||
log.WithError(errTarget).Error("failed to compute anthropic callback target")
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"})
|
||||
return
|
||||
}
|
||||
var errStart error
|
||||
if forwarder, errStart = startCallbackForwarder(anthropicCallbackPort, "anthropic", targetURL); errStart != nil {
|
||||
log.WithError(errStart).Error("failed to start anthropic callback forwarder")
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
go func() {
|
||||
if isWebUI {
|
||||
defer stopCallbackForwarderInstance(anthropicCallbackPort, forwarder)
|
||||
}
|
||||
|
||||
// Helper: wait for callback file
|
||||
waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-anthropic-%s.oauth", state))
|
||||
waitForFile := func(path string, timeout time.Duration) (map[string]string, error) {
|
||||
deadline := time.Now().Add(timeout)
|
||||
for {
|
||||
if !IsOAuthSessionPending(state, "anthropic") {
|
||||
return nil, errOAuthSessionNotPending
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
SetOAuthSessionError(state, "Timeout waiting for OAuth callback")
|
||||
return nil, fmt.Errorf("timeout waiting for OAuth callback")
|
||||
}
|
||||
data, errRead := os.ReadFile(path)
|
||||
if errRead == nil {
|
||||
var m map[string]string
|
||||
_ = json.Unmarshal(data, &m)
|
||||
_ = os.Remove(path)
|
||||
return m, nil
|
||||
}
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Println("Waiting for authentication callback...")
|
||||
// Wait up to 5 minutes
|
||||
resultMap, errWait := waitForFile(waitFile, 5*time.Minute)
|
||||
if errWait != nil {
|
||||
if errors.Is(errWait, errOAuthSessionNotPending) {
|
||||
return
|
||||
}
|
||||
authErr := claude.NewAuthenticationError(claude.ErrCallbackTimeout, errWait)
|
||||
log.Error(claude.GetUserFriendlyMessage(authErr))
|
||||
return
|
||||
}
|
||||
if errStr := resultMap["error"]; errStr != "" {
|
||||
oauthErr := claude.NewOAuthError(errStr, "", http.StatusBadRequest)
|
||||
log.Error(claude.GetUserFriendlyMessage(oauthErr))
|
||||
SetOAuthSessionError(state, "Bad request")
|
||||
return
|
||||
}
|
||||
if resultMap["state"] != state {
|
||||
authErr := claude.NewAuthenticationError(claude.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, resultMap["state"]))
|
||||
log.Error(claude.GetUserFriendlyMessage(authErr))
|
||||
SetOAuthSessionError(state, "State code error")
|
||||
return
|
||||
}
|
||||
|
||||
// Parse code (Claude may append state after '#')
|
||||
rawCode := resultMap["code"]
|
||||
code := strings.Split(rawCode, "#")[0]
|
||||
|
||||
// Exchange code for tokens using internal auth service
|
||||
bundle, errExchange := anthropicAuth.ExchangeCodeForTokens(ctx, code, state, pkceCodes)
|
||||
if errExchange != nil {
|
||||
authErr := claude.NewAuthenticationError(claude.ErrCodeExchangeFailed, errExchange)
|
||||
log.Errorf("Failed to exchange authorization code for tokens: %v", authErr)
|
||||
SetOAuthSessionError(state, "Failed to exchange authorization code for tokens")
|
||||
return
|
||||
}
|
||||
|
||||
// Create token storage
|
||||
tokenStorage := anthropicAuth.CreateTokenStorage(bundle)
|
||||
metadata := map[string]any{"email": tokenStorage.Email}
|
||||
if tokenStorage.AccountUUID != "" {
|
||||
metadata["account_uuid"] = tokenStorage.AccountUUID
|
||||
}
|
||||
if tokenStorage.OrganizationUUID != "" {
|
||||
metadata["organization_uuid"] = tokenStorage.OrganizationUUID
|
||||
}
|
||||
if tokenStorage.OrganizationName != "" {
|
||||
metadata["organization_name"] = tokenStorage.OrganizationName
|
||||
}
|
||||
if len(tokenStorage.DeviceIDs) > 0 {
|
||||
metadata[claude.ClaudeDeviceIDsMetadataKey] = append([]string(nil), tokenStorage.DeviceIDs...)
|
||||
}
|
||||
record := &coreauth.Auth{
|
||||
ID: fmt.Sprintf("claude-%s.json", tokenStorage.Email),
|
||||
Provider: "claude",
|
||||
FileName: fmt.Sprintf("claude-%s.json", tokenStorage.Email),
|
||||
Storage: tokenStorage,
|
||||
Metadata: metadata,
|
||||
}
|
||||
if errGuard := guardOAuthSessionPendingForSave(state, "anthropic"); errGuard != nil {
|
||||
return
|
||||
}
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if errSave != nil {
|
||||
log.Errorf("Failed to save authentication tokens: %v", errSave)
|
||||
SetOAuthSessionError(state, "Failed to save authentication tokens")
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
|
||||
if bundle.APIKey != "" {
|
||||
fmt.Println("API key obtained and saved")
|
||||
}
|
||||
fmt.Println("You can now use Claude services through this CLI")
|
||||
CompleteOAuthSession(state)
|
||||
}()
|
||||
|
||||
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
|
||||
}
|
||||
|
||||
func (h *Handler) RequestCodexToken(c *gin.Context) {
|
||||
ctx := context.Background()
|
||||
ctx = PopulateAuthContext(ctx, c)
|
||||
|
||||
fmt.Println("Initializing Codex authentication...")
|
||||
|
||||
// Generate PKCE codes
|
||||
pkceCodes, err := codex.GeneratePKCECodes()
|
||||
if err != nil {
|
||||
log.Errorf("Failed to generate PKCE codes: %v", err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate PKCE codes"})
|
||||
return
|
||||
}
|
||||
|
||||
// Generate random state parameter
|
||||
state, err := misc.GenerateRandomState()
|
||||
if err != nil {
|
||||
log.Errorf("Failed to generate state parameter: %v", err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"})
|
||||
return
|
||||
}
|
||||
|
||||
// Initialize Codex auth service
|
||||
openaiAuth := newCodexOAuthService(h.cfg)
|
||||
|
||||
// Generate authorization URL
|
||||
authURL, err := openaiAuth.GenerateAuthURL(state, pkceCodes)
|
||||
if err != nil {
|
||||
log.Errorf("Failed to generate authorization URL: %v", err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"})
|
||||
return
|
||||
}
|
||||
|
||||
RegisterOAuthSession(state, "codex")
|
||||
|
||||
isWebUI := isWebUIRequest(c)
|
||||
var forwarder *callbackForwarder
|
||||
if isWebUI {
|
||||
targetURL, errTarget := h.managementCallbackURL("/codex/callback")
|
||||
if errTarget != nil {
|
||||
log.WithError(errTarget).Error("failed to compute codex callback target")
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"})
|
||||
return
|
||||
}
|
||||
var errStart error
|
||||
if forwarder, errStart = startCallbackForwarder(codexCallbackPort, "codex", targetURL); errStart != nil {
|
||||
log.WithError(errStart).Error("failed to start codex callback forwarder")
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
go func() {
|
||||
if isWebUI {
|
||||
defer stopCallbackForwarderInstance(codexCallbackPort, forwarder)
|
||||
}
|
||||
|
||||
// Wait for callback file
|
||||
waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-codex-%s.oauth", state))
|
||||
deadline := time.Now().Add(5 * time.Minute)
|
||||
var code string
|
||||
for {
|
||||
if !IsOAuthSessionPending(state, "codex") {
|
||||
return
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
authErr := codex.NewAuthenticationError(codex.ErrCallbackTimeout, fmt.Errorf("timeout waiting for OAuth callback"))
|
||||
log.Error(codex.GetUserFriendlyMessage(authErr))
|
||||
SetOAuthSessionError(state, "Timeout waiting for OAuth callback")
|
||||
return
|
||||
}
|
||||
if data, errR := os.ReadFile(waitFile); errR == nil {
|
||||
var m map[string]string
|
||||
_ = json.Unmarshal(data, &m)
|
||||
_ = os.Remove(waitFile)
|
||||
if errStr := m["error"]; errStr != "" {
|
||||
oauthErr := codex.NewOAuthError(errStr, "", http.StatusBadRequest)
|
||||
log.Error(codex.GetUserFriendlyMessage(oauthErr))
|
||||
SetOAuthSessionError(state, "Bad Request")
|
||||
return
|
||||
}
|
||||
if m["state"] != state {
|
||||
authErr := codex.NewAuthenticationError(codex.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, m["state"]))
|
||||
SetOAuthSessionError(state, "State code error")
|
||||
log.Error(codex.GetUserFriendlyMessage(authErr))
|
||||
return
|
||||
}
|
||||
code = m["code"]
|
||||
break
|
||||
}
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
}
|
||||
|
||||
log.Debug("Authorization code received, exchanging for tokens...")
|
||||
// Exchange code for tokens using internal auth service
|
||||
bundle, errExchange := openaiAuth.ExchangeCodeForTokens(ctx, code, pkceCodes)
|
||||
if errExchange != nil {
|
||||
authErr := codex.NewAuthenticationError(codex.ErrCodeExchangeFailed, errExchange)
|
||||
SetOAuthSessionError(state, oauthSessionErrorWithCause("Failed to exchange authorization code for tokens", errExchange))
|
||||
log.Errorf("Failed to exchange authorization code for tokens: %v", authErr)
|
||||
return
|
||||
}
|
||||
|
||||
// Extract additional info for filename generation
|
||||
claims, _ := codex.ParseJWTToken(bundle.TokenData.IDToken)
|
||||
planType := ""
|
||||
hashAccountID := ""
|
||||
if claims != nil {
|
||||
planType = strings.TrimSpace(claims.CodexAuthInfo.ChatgptPlanType)
|
||||
if accountID := claims.GetAccountID(); accountID != "" {
|
||||
digest := sha256.Sum256([]byte(accountID))
|
||||
hashAccountID = hex.EncodeToString(digest[:])[:8]
|
||||
}
|
||||
}
|
||||
|
||||
// Create token storage and persist
|
||||
tokenStorage := openaiAuth.CreateTokenStorage(bundle)
|
||||
fileName := codex.CredentialFileName(tokenStorage.Email, planType, hashAccountID, true)
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
Provider: "codex",
|
||||
FileName: fileName,
|
||||
Storage: tokenStorage,
|
||||
Metadata: map[string]any{
|
||||
"email": tokenStorage.Email,
|
||||
"account_id": tokenStorage.AccountID,
|
||||
},
|
||||
}
|
||||
if errGuard := guardOAuthSessionPendingForSave(state, "codex"); errGuard != nil {
|
||||
return
|
||||
}
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if errSave != nil {
|
||||
SetOAuthSessionError(state, "Failed to save authentication tokens")
|
||||
log.Errorf("Failed to save authentication tokens: %v", errSave)
|
||||
return
|
||||
}
|
||||
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
|
||||
if bundle.APIKey != "" {
|
||||
fmt.Println("API key obtained and saved")
|
||||
}
|
||||
fmt.Println("You can now use Codex services through this CLI")
|
||||
CompleteOAuthSession(state)
|
||||
}()
|
||||
|
||||
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
|
||||
}
|
||||
|
||||
func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
ctx := context.Background()
|
||||
ctx = PopulateAuthContext(ctx, c)
|
||||
|
||||
fmt.Println("Initializing Antigravity authentication...")
|
||||
|
||||
authSvc := antigravity.NewAntigravityAuth(h.cfg, nil)
|
||||
|
||||
state, errState := misc.GenerateRandomState()
|
||||
if errState != nil {
|
||||
log.Errorf("Failed to generate state parameter: %v", errState)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"})
|
||||
return
|
||||
}
|
||||
|
||||
redirectURI := fmt.Sprintf("http://localhost:%d/oauth-callback", antigravity.CallbackPort)
|
||||
authURL := authSvc.BuildAuthURL(state, redirectURI)
|
||||
|
||||
RegisterOAuthSession(state, "antigravity")
|
||||
|
||||
isWebUI := isWebUIRequest(c)
|
||||
var forwarder *callbackForwarder
|
||||
if isWebUI {
|
||||
targetURL, errTarget := h.managementCallbackURL("/antigravity/callback")
|
||||
if errTarget != nil {
|
||||
log.WithError(errTarget).Error("failed to compute antigravity callback target")
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"})
|
||||
return
|
||||
}
|
||||
var errStart error
|
||||
if forwarder, errStart = startCallbackForwarder(antigravity.CallbackPort, "antigravity", targetURL); errStart != nil {
|
||||
log.WithError(errStart).Error("failed to start antigravity callback forwarder")
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
go func() {
|
||||
if isWebUI {
|
||||
defer stopCallbackForwarderInstance(antigravity.CallbackPort, forwarder)
|
||||
}
|
||||
|
||||
waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-antigravity-%s.oauth", state))
|
||||
deadline := time.Now().Add(5 * time.Minute)
|
||||
var authCode string
|
||||
for {
|
||||
if !IsOAuthSessionPending(state, "antigravity") {
|
||||
return
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
log.Error("oauth flow timed out")
|
||||
SetOAuthSessionError(state, "OAuth flow timed out")
|
||||
return
|
||||
}
|
||||
if data, errReadFile := os.ReadFile(waitFile); errReadFile == nil {
|
||||
var payload map[string]string
|
||||
_ = json.Unmarshal(data, &payload)
|
||||
_ = os.Remove(waitFile)
|
||||
if errStr := strings.TrimSpace(payload["error"]); errStr != "" {
|
||||
log.Errorf("Authentication failed: %s", errStr)
|
||||
SetOAuthSessionError(state, "Authentication failed")
|
||||
return
|
||||
}
|
||||
if payloadState := strings.TrimSpace(payload["state"]); payloadState != "" && payloadState != state {
|
||||
log.Errorf("Authentication failed: state mismatch")
|
||||
SetOAuthSessionError(state, "Authentication failed: state mismatch")
|
||||
return
|
||||
}
|
||||
authCode = strings.TrimSpace(payload["code"])
|
||||
if authCode == "" {
|
||||
log.Error("Authentication failed: code not found")
|
||||
SetOAuthSessionError(state, "Authentication failed: code not found")
|
||||
return
|
||||
}
|
||||
break
|
||||
}
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
}
|
||||
|
||||
tokenResp, errToken := authSvc.ExchangeCodeForTokens(ctx, authCode, redirectURI)
|
||||
if errToken != nil {
|
||||
log.Errorf("Failed to exchange token: %v", errToken)
|
||||
SetOAuthSessionError(state, "Failed to exchange token")
|
||||
return
|
||||
}
|
||||
|
||||
accessToken := strings.TrimSpace(tokenResp.AccessToken)
|
||||
if accessToken == "" {
|
||||
log.Error("antigravity: token exchange returned empty access token")
|
||||
SetOAuthSessionError(state, "Failed to exchange token")
|
||||
return
|
||||
}
|
||||
|
||||
email, errInfo := authSvc.FetchUserInfo(ctx, accessToken)
|
||||
if errInfo != nil {
|
||||
log.Errorf("Failed to fetch user info: %v", errInfo)
|
||||
SetOAuthSessionError(state, "Failed to fetch user info")
|
||||
return
|
||||
}
|
||||
email = strings.TrimSpace(email)
|
||||
if email == "" {
|
||||
log.Error("antigravity: user info returned empty email")
|
||||
SetOAuthSessionError(state, "Failed to fetch user info")
|
||||
return
|
||||
}
|
||||
|
||||
projectID := ""
|
||||
if accessToken != "" {
|
||||
fetchedProjectID, errProject := authSvc.FetchProjectID(ctx, accessToken)
|
||||
if errProject != nil {
|
||||
log.Warnf("antigravity: failed to fetch project ID: %v", errProject)
|
||||
} else {
|
||||
projectID = fetchedProjectID
|
||||
log.Infof("antigravity: obtained project ID %s", util.HideAPIKey(projectID))
|
||||
}
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
metadata := map[string]any{
|
||||
"type": "antigravity",
|
||||
"access_token": tokenResp.AccessToken,
|
||||
"refresh_token": tokenResp.RefreshToken,
|
||||
"expires_in": tokenResp.ExpiresIn,
|
||||
"timestamp": now.UnixMilli(),
|
||||
"expired": now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339),
|
||||
}
|
||||
if email != "" {
|
||||
metadata["email"] = email
|
||||
}
|
||||
if projectID != "" {
|
||||
metadata["project_id"] = projectID
|
||||
}
|
||||
|
||||
fileName := antigravity.CredentialFileName(email)
|
||||
label := strings.TrimSpace(email)
|
||||
if label == "" {
|
||||
label = "antigravity"
|
||||
}
|
||||
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
Provider: "antigravity",
|
||||
FileName: fileName,
|
||||
Label: label,
|
||||
Metadata: metadata,
|
||||
}
|
||||
if errGuard := guardOAuthSessionPendingForSave(state, "antigravity"); errGuard != nil {
|
||||
return
|
||||
}
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if errSave != nil {
|
||||
log.Errorf("Failed to save token to file: %v", errSave)
|
||||
SetOAuthSessionError(state, "Failed to save token to file")
|
||||
return
|
||||
}
|
||||
|
||||
CompleteOAuthSession(state)
|
||||
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
|
||||
if projectID != "" {
|
||||
fmt.Printf("Using GCP project: %s\n", util.HideAPIKey(projectID))
|
||||
}
|
||||
fmt.Println("You can now use Antigravity services through this CLI")
|
||||
}()
|
||||
|
||||
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
|
||||
}
|
||||
|
||||
func (h *Handler) RequestXAIToken(c *gin.Context) {
|
||||
ctx := context.Background()
|
||||
ctx = PopulateAuthContext(ctx, c)
|
||||
|
||||
fmt.Println("Initializing xAI authentication...")
|
||||
|
||||
state := fmt.Sprintf("xai-%d", time.Now().UnixNano())
|
||||
authSvc := xaiauth.NewXAIAuth(h.cfg)
|
||||
|
||||
deviceFlow, errStartDeviceFlow := authSvc.StartDeviceFlow(ctx)
|
||||
if errStartDeviceFlow != nil {
|
||||
log.Errorf("Failed to start xAI device flow: %v", errStartDeviceFlow)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start device authorization flow"})
|
||||
return
|
||||
}
|
||||
authURL := strings.TrimSpace(deviceFlow.VerificationURIComplete)
|
||||
if authURL == "" {
|
||||
authURL = strings.TrimSpace(deviceFlow.VerificationURI)
|
||||
}
|
||||
|
||||
RegisterOAuthSession(state, "xai")
|
||||
|
||||
go func() {
|
||||
pollCtx, cancelPoll := context.WithCancel(ctx)
|
||||
defer cancelPoll()
|
||||
go watchOAuthSessionCancel(pollCtx, cancelPoll, state, "xai")
|
||||
|
||||
fmt.Println("Waiting for xAI authentication...")
|
||||
bundle, errWaitForAuthorization := authSvc.WaitForAuthorization(pollCtx, deviceFlow)
|
||||
if errWaitForAuthorization != nil {
|
||||
if !IsOAuthSessionPending(state, "xai") {
|
||||
return
|
||||
}
|
||||
log.Errorf("xAI authentication failed: %v", errWaitForAuthorization)
|
||||
SetOAuthSessionError(state, oauthSessionErrorWithCause("Authentication failed", errWaitForAuthorization))
|
||||
return
|
||||
}
|
||||
if !IsOAuthSessionPending(state, "xai") {
|
||||
return
|
||||
}
|
||||
|
||||
tokenStorage := authSvc.CreateTokenStorage(bundle)
|
||||
if tokenStorage == nil || strings.TrimSpace(tokenStorage.AccessToken) == "" {
|
||||
log.Error("xAI token exchange returned empty access token")
|
||||
SetOAuthSessionError(state, "Failed to exchange token")
|
||||
return
|
||||
}
|
||||
|
||||
fileName := xaiauth.CredentialFileName(tokenStorage.Email, tokenStorage.Subject)
|
||||
label := strings.TrimSpace(tokenStorage.Email)
|
||||
if label == "" {
|
||||
label = "xAI"
|
||||
}
|
||||
|
||||
metadata := map[string]any{
|
||||
"type": "xai",
|
||||
"access_token": tokenStorage.AccessToken,
|
||||
"refresh_token": tokenStorage.RefreshToken,
|
||||
"id_token": tokenStorage.IDToken,
|
||||
"token_type": tokenStorage.TokenType,
|
||||
"expires_in": tokenStorage.ExpiresIn,
|
||||
"expired": tokenStorage.Expire,
|
||||
"last_refresh": tokenStorage.LastRefresh,
|
||||
"base_url": tokenStorage.BaseURL,
|
||||
"token_endpoint": tokenStorage.TokenEndpoint,
|
||||
"auth_kind": "oauth",
|
||||
}
|
||||
if tokenStorage.Email != "" {
|
||||
metadata["email"] = tokenStorage.Email
|
||||
}
|
||||
if tokenStorage.Subject != "" {
|
||||
metadata["sub"] = tokenStorage.Subject
|
||||
}
|
||||
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
Provider: "xai",
|
||||
FileName: fileName,
|
||||
Label: label,
|
||||
Storage: tokenStorage,
|
||||
Metadata: metadata,
|
||||
Attributes: map[string]string{
|
||||
"auth_kind": "oauth",
|
||||
"base_url": tokenStorage.BaseURL,
|
||||
},
|
||||
}
|
||||
if errGuard := guardOAuthSessionPendingForSave(state, "xai"); errGuard != nil {
|
||||
return
|
||||
}
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if errSave != nil {
|
||||
log.Errorf("Failed to save xAI token to file: %v", errSave)
|
||||
SetOAuthSessionError(state, "Failed to save token to file")
|
||||
return
|
||||
}
|
||||
|
||||
CompleteOAuthSession(state)
|
||||
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
|
||||
fmt.Println("You can now use xAI services through this CLI")
|
||||
}()
|
||||
|
||||
response := gin.H{"status": "ok", "url": authURL, "state": state, "flow": "device"}
|
||||
if userCode := strings.TrimSpace(deviceFlow.UserCode); userCode != "" {
|
||||
response["user_code"] = userCode
|
||||
}
|
||||
if deviceFlow.ExpiresIn > 0 {
|
||||
response["expires_in"] = deviceFlow.ExpiresIn
|
||||
} else {
|
||||
response["expires_in"] = int(xaiauth.MaxPollDuration / time.Second)
|
||||
}
|
||||
c.JSON(200, response)
|
||||
}
|
||||
|
||||
func (h *Handler) RequestKimiToken(c *gin.Context) {
|
||||
ctx := context.Background()
|
||||
ctx = PopulateAuthContext(ctx, c)
|
||||
|
||||
fmt.Println("Initializing Kimi authentication...")
|
||||
|
||||
state := fmt.Sprintf("kmi-%d", time.Now().UnixNano())
|
||||
// Initialize Kimi auth service
|
||||
kimiAuth := kimi.NewKimiAuth(h.cfg)
|
||||
|
||||
// Generate authorization URL
|
||||
deviceFlow, errStartDeviceFlow := kimiAuth.StartDeviceFlow(ctx)
|
||||
if errStartDeviceFlow != nil {
|
||||
log.Errorf("Failed to generate authorization URL: %v", errStartDeviceFlow)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"})
|
||||
return
|
||||
}
|
||||
authURL := deviceFlow.VerificationURIComplete
|
||||
if authURL == "" {
|
||||
authURL = deviceFlow.VerificationURI
|
||||
}
|
||||
|
||||
RegisterOAuthSession(state, "kimi")
|
||||
|
||||
go func() {
|
||||
pollCtx, cancelPoll := context.WithCancel(ctx)
|
||||
defer cancelPoll()
|
||||
go watchOAuthSessionCancel(pollCtx, cancelPoll, state, "kimi")
|
||||
|
||||
fmt.Println("Waiting for authentication...")
|
||||
authBundle, errWaitForAuthorization := kimiAuth.WaitForAuthorization(pollCtx, deviceFlow)
|
||||
if errWaitForAuthorization != nil {
|
||||
if !IsOAuthSessionPending(state, "kimi") {
|
||||
return
|
||||
}
|
||||
SetOAuthSessionError(state, oauthSessionErrorWithCause("Authentication failed", errWaitForAuthorization))
|
||||
fmt.Printf("Authentication failed: %v\n", errWaitForAuthorization)
|
||||
return
|
||||
}
|
||||
if !IsOAuthSessionPending(state, "kimi") {
|
||||
return
|
||||
}
|
||||
|
||||
// Create token storage
|
||||
tokenStorage := kimiAuth.CreateTokenStorage(authBundle)
|
||||
|
||||
metadata := map[string]any{
|
||||
"type": "kimi",
|
||||
"access_token": authBundle.TokenData.AccessToken,
|
||||
"refresh_token": authBundle.TokenData.RefreshToken,
|
||||
"token_type": authBundle.TokenData.TokenType,
|
||||
"scope": authBundle.TokenData.Scope,
|
||||
"timestamp": time.Now().UnixMilli(),
|
||||
}
|
||||
if authBundle.TokenData.ExpiresAt > 0 {
|
||||
expired := time.Unix(authBundle.TokenData.ExpiresAt, 0).UTC().Format(time.RFC3339)
|
||||
metadata["expired"] = expired
|
||||
}
|
||||
if strings.TrimSpace(authBundle.DeviceID) != "" {
|
||||
metadata["device_id"] = strings.TrimSpace(authBundle.DeviceID)
|
||||
}
|
||||
|
||||
fileName := fmt.Sprintf("kimi-%d.json", time.Now().UnixMilli())
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
Provider: "kimi",
|
||||
FileName: fileName,
|
||||
Label: "Kimi User",
|
||||
Storage: tokenStorage,
|
||||
Metadata: metadata,
|
||||
}
|
||||
if errGuard := guardOAuthSessionPendingForSave(state, "kimi"); errGuard != nil {
|
||||
return
|
||||
}
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if errSave != nil {
|
||||
log.Errorf("Failed to save authentication tokens: %v", errSave)
|
||||
SetOAuthSessionError(state, "Failed to save authentication tokens")
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
|
||||
fmt.Println("You can now use Kimi services through this CLI")
|
||||
CompleteOAuthSession(state)
|
||||
}()
|
||||
|
||||
response := gin.H{"status": "ok", "url": authURL, "state": state, "flow": "device"}
|
||||
if userCode := strings.TrimSpace(deviceFlow.UserCode); userCode != "" {
|
||||
response["user_code"] = userCode
|
||||
}
|
||||
if deviceFlow.ExpiresIn > 0 {
|
||||
response["expires_in"] = deviceFlow.ExpiresIn
|
||||
}
|
||||
c.JSON(200, response)
|
||||
}
|
||||
|
||||
// watchOAuthSessionCancel cancels pollCtx once the OAuth session is no longer pending.
|
||||
func watchOAuthSessionCancel(pollCtx context.Context, cancel context.CancelFunc, state, provider string) {
|
||||
if cancel == nil {
|
||||
return
|
||||
}
|
||||
ticker := time.NewTicker(2 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-pollCtx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if !IsOAuthSessionPending(state, provider) {
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// CancelAuthSession cancels a pending OAuth session identified by state.
|
||||
// Protected by management auth. Safe for both callback and device-code flows:
|
||||
// waiters check IsOAuthSessionPending and exit without saving credentials.
|
||||
func (h *Handler) CancelAuthSession(c *gin.Context) {
|
||||
state := strings.TrimSpace(c.Query("state"))
|
||||
if state == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "missing state"})
|
||||
return
|
||||
}
|
||||
if err := ValidateOAuthState(state); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "invalid state"})
|
||||
return
|
||||
}
|
||||
cancelled := CancelOAuthSession(state)
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok", "cancelled": cancelled})
|
||||
}
|
||||
|
||||
func (h *Handler) GetAuthStatus(c *gin.Context) {
|
||||
state := strings.TrimSpace(c.Query("state"))
|
||||
if state == "" {
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||||
return
|
||||
}
|
||||
if err := ValidateOAuthState(state); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "invalid state"})
|
||||
return
|
||||
}
|
||||
|
||||
provider, status, isPlugin, metadata, completed, ok := GetOAuthSessionDetails(state)
|
||||
if !ok {
|
||||
c.JSON(http.StatusOK, gin.H{"status": "error", "error": "unknown or expired state"})
|
||||
return
|
||||
}
|
||||
if completed {
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||||
return
|
||||
}
|
||||
if status != "" {
|
||||
c.JSON(http.StatusOK, gin.H{"status": "error", "error": status})
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
host := h.pluginHost
|
||||
h.mu.Unlock()
|
||||
if isPlugin && host != nil && host.HasAuthProvider(provider) {
|
||||
ctx := PopulateAuthContext(context.Background(), c)
|
||||
resp, handled, errPoll := host.PollLogin(ctx, provider, state, metadata)
|
||||
if handled {
|
||||
if errPoll != nil {
|
||||
message := strings.TrimSpace(errPoll.Error())
|
||||
if message == "" {
|
||||
message = "Authentication failed"
|
||||
}
|
||||
SetOAuthSessionError(state, message)
|
||||
c.JSON(http.StatusOK, gin.H{"status": "error", "error": message})
|
||||
return
|
||||
}
|
||||
switch resp.Status {
|
||||
case "", pluginapi.AuthLoginStatusPending:
|
||||
c.JSON(http.StatusOK, gin.H{"status": "wait"})
|
||||
return
|
||||
case pluginapi.AuthLoginStatusError:
|
||||
message := strings.TrimSpace(resp.Message)
|
||||
if message == "" {
|
||||
message = "Authentication failed"
|
||||
}
|
||||
SetOAuthSessionError(state, message)
|
||||
c.JSON(http.StatusOK, gin.H{"status": "error", "error": message})
|
||||
return
|
||||
case pluginapi.AuthLoginStatusSuccess:
|
||||
records := pluginLoginPollAuths(host, resp)
|
||||
if len(records) == 0 {
|
||||
SetOAuthSessionError(state, "Authentication failed")
|
||||
c.JSON(http.StatusOK, gin.H{"status": "error", "error": "Authentication failed"})
|
||||
return
|
||||
}
|
||||
if errSave := h.savePluginLoginRecords(ctx, records); errSave != nil {
|
||||
log.WithError(errSave).WithField("provider", provider).Error("failed to save plugin auth tokens")
|
||||
SetOAuthSessionError(state, "Failed to save authentication tokens")
|
||||
c.JSON(http.StatusOK, gin.H{"status": "error", "error": "Failed to save authentication tokens"})
|
||||
return
|
||||
}
|
||||
CompleteOAuthSession(state)
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||||
return
|
||||
default:
|
||||
c.JSON(http.StatusOK, gin.H{"status": "wait"})
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"status": "wait"})
|
||||
}
|
||||
|
||||
func pluginLoginPollAuths(host *pluginhost.Host, resp pluginapi.AuthLoginPollResponse) []*coreauth.Auth {
|
||||
if host == nil {
|
||||
return nil
|
||||
}
|
||||
authDatas := resp.Auths
|
||||
if len(authDatas) == 0 {
|
||||
authDatas = []pluginapi.AuthData{resp.Auth}
|
||||
}
|
||||
records := make([]*coreauth.Auth, 0, len(authDatas))
|
||||
for _, authData := range authDatas {
|
||||
record := host.AuthDataToCoreAuth(authData, "", "")
|
||||
if record == nil {
|
||||
return nil
|
||||
}
|
||||
records = append(records, record)
|
||||
}
|
||||
return records
|
||||
}
|
||||
|
||||
func (h *Handler) savePluginLoginRecords(ctx context.Context, records []*coreauth.Auth) error {
|
||||
savedPaths := make([]string, 0, len(records))
|
||||
for _, record := range records {
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if strings.TrimSpace(savedPath) != "" {
|
||||
savedPaths = append(savedPaths, savedPath)
|
||||
}
|
||||
if errSave != nil {
|
||||
h.rollbackSavedTokenRecords(ctx, savedPaths)
|
||||
return errSave
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *Handler) rollbackSavedTokenRecords(ctx context.Context, savedPaths []string) {
|
||||
for i := len(savedPaths) - 1; i >= 0; i-- {
|
||||
path := strings.TrimSpace(savedPaths[i])
|
||||
if path == "" {
|
||||
continue
|
||||
}
|
||||
if errDelete := h.deleteTokenRecord(ctx, path); errDelete != nil {
|
||||
log.WithError(errDelete).WithField("path", path).Warn("failed to roll back plugin auth token")
|
||||
}
|
||||
h.removeAuthsForPath(ctx, path, path)
|
||||
}
|
||||
}
|
||||
|
||||
// PopulateAuthContext extracts request info and adds it to the context
|
||||
func PopulateAuthContext(ctx context.Context, c *gin.Context) context.Context {
|
||||
info := &coreauth.RequestInfo{
|
||||
Query: c.Request.URL.Query(),
|
||||
Headers: c.Request.Header,
|
||||
}
|
||||
return coreauth.WithRequestInfo(ctx, info)
|
||||
}
|
||||
|
|
@ -0,0 +1,93 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestListAuthFiles_IncludesRecentRequestsBuckets(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: "runtime-only-auth-1",
|
||||
Provider: "codex",
|
||||
Attributes: map[string]string{
|
||||
"runtime_only": "true",
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "codex",
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
h.tokenStore = &memoryAuthStore{}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ginCtx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v0/management/auth-files", nil)
|
||||
ginCtx.Request = req
|
||||
|
||||
h.ListAuthFiles(ginCtx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected list status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if errUnmarshal := json.Unmarshal(rec.Body.Bytes(), &payload); errUnmarshal != nil {
|
||||
t.Fatalf("failed to decode list payload: %v", errUnmarshal)
|
||||
}
|
||||
filesRaw, ok := payload["files"].([]any)
|
||||
if !ok {
|
||||
t.Fatalf("expected files array, payload: %#v", payload)
|
||||
}
|
||||
if len(filesRaw) != 1 {
|
||||
t.Fatalf("expected 1 auth entry, got %d", len(filesRaw))
|
||||
}
|
||||
|
||||
fileEntry, ok := filesRaw[0].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("expected file entry object, got %#v", filesRaw[0])
|
||||
}
|
||||
|
||||
if _, ok := fileEntry["success"].(float64); !ok {
|
||||
t.Fatalf("expected success number, got %#v", fileEntry["success"])
|
||||
}
|
||||
if _, ok := fileEntry["failed"].(float64); !ok {
|
||||
t.Fatalf("expected failed number, got %#v", fileEntry["failed"])
|
||||
}
|
||||
|
||||
recentRaw, ok := fileEntry["recent_requests"].([]any)
|
||||
if !ok {
|
||||
t.Fatalf("expected recent_requests array, got %#v", fileEntry["recent_requests"])
|
||||
}
|
||||
if len(recentRaw) != 20 {
|
||||
t.Fatalf("expected 20 recent_requests buckets, got %d", len(recentRaw))
|
||||
}
|
||||
for idx, item := range recentRaw {
|
||||
bucket, ok := item.(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("expected bucket object at %d, got %#v", idx, item)
|
||||
}
|
||||
if _, ok := bucket["time"].(string); !ok {
|
||||
t.Fatalf("expected bucket time string at %d, got %#v", idx, bucket["time"])
|
||||
}
|
||||
if _, ok := bucket["success"].(float64); !ok {
|
||||
t.Fatalf("expected bucket success number at %d, got %#v", idx, bucket["success"])
|
||||
}
|
||||
if _, ok := bucket["failed"].(float64); !ok {
|
||||
t.Fatalf("expected bucket failed number at %d, got %#v", idx, bucket["failed"])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,199 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/auth/codex"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestSaveTokenRecord_PreservesExistingAuthFileSettings(t *testing.T) {
|
||||
authDir := t.TempDir()
|
||||
fileName := "codex-user@example.com.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
|
||||
// User configured fields on existing OAuth account
|
||||
initialContent := map[string]any{
|
||||
"type": "codex",
|
||||
"email": "user@example.com",
|
||||
"access_token": "old-access",
|
||||
"refresh_token": "old-refresh",
|
||||
"prefix": "custom-prefix",
|
||||
"websockets": false,
|
||||
"note": "my important account",
|
||||
"proxy_url": "http://127.0.0.1:8080",
|
||||
"weight": float64(5),
|
||||
"headers": map[string]any{"User-Agent": "Custom"},
|
||||
"models": []any{"o3-mini"},
|
||||
"thinking": map[string]any{"enabled": true},
|
||||
"priority": float64(2),
|
||||
}
|
||||
raw, errMarshal := json.Marshal(initialContent)
|
||||
if errMarshal != nil {
|
||||
t.Fatalf("marshal initial error: %v", errMarshal)
|
||||
}
|
||||
if errWrite := os.WriteFile(filePath, raw, 0o600); errWrite != nil {
|
||||
t.Fatalf("write initial file error: %v", errWrite)
|
||||
}
|
||||
|
||||
cfg := &config.Config{
|
||||
AuthDir: authDir,
|
||||
}
|
||||
h := NewHandler(cfg, "", nil)
|
||||
|
||||
// Re-login arrives with new OAuth tokens
|
||||
tokenStorage := &codex.CodexTokenStorage{
|
||||
Type: "codex",
|
||||
Email: "user@example.com",
|
||||
AccessToken: "new-access-token",
|
||||
RefreshToken: "new-refresh-token",
|
||||
IDToken: "new-id-token",
|
||||
AccountID: "act-123",
|
||||
Expire: "2026-12-31T23:59:59Z",
|
||||
}
|
||||
newRecord := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
Provider: "codex",
|
||||
FileName: fileName,
|
||||
Storage: tokenStorage,
|
||||
Metadata: map[string]any{
|
||||
"email": tokenStorage.Email,
|
||||
"account_id": tokenStorage.AccountID,
|
||||
},
|
||||
}
|
||||
|
||||
savedPath, errSave := h.saveTokenRecord(context.Background(), newRecord)
|
||||
if errSave != nil {
|
||||
t.Fatalf("saveTokenRecord error: %v", errSave)
|
||||
}
|
||||
if savedPath != filePath {
|
||||
t.Fatalf("savedPath = %s, want %s", savedPath, filePath)
|
||||
}
|
||||
|
||||
savedRaw, errRead := os.ReadFile(filePath)
|
||||
if errRead != nil {
|
||||
t.Fatalf("ReadFile error: %v", errRead)
|
||||
}
|
||||
var saved map[string]any
|
||||
if errUnmarshal := json.Unmarshal(savedRaw, &saved); errUnmarshal != nil {
|
||||
t.Fatalf("Unmarshal error: %v", errUnmarshal)
|
||||
}
|
||||
|
||||
// Verify new OAuth token data was updated
|
||||
if saved["access_token"] != "new-access-token" {
|
||||
t.Errorf("access_token = %v, want new-access-token", saved["access_token"])
|
||||
}
|
||||
if saved["refresh_token"] != "new-refresh-token" {
|
||||
t.Errorf("refresh_token = %v, want new-refresh-token", saved["refresh_token"])
|
||||
}
|
||||
|
||||
// Verify user-configured fields were preserved
|
||||
if saved["prefix"] != "custom-prefix" {
|
||||
t.Errorf("prefix = %v, want custom-prefix", saved["prefix"])
|
||||
}
|
||||
if saved["websockets"] != false {
|
||||
t.Errorf("websockets = %v, want false", saved["websockets"])
|
||||
}
|
||||
if saved["note"] != "my important account" {
|
||||
t.Errorf("note = %v, want my important account", saved["note"])
|
||||
}
|
||||
if saved["proxy_url"] != "http://127.0.0.1:8080" {
|
||||
t.Errorf("proxy_url = %v, want http://127.0.0.1:8080", saved["proxy_url"])
|
||||
}
|
||||
if saved["weight"] != float64(5) {
|
||||
t.Errorf("weight = %v, want 5", saved["weight"])
|
||||
}
|
||||
if !reflect.DeepEqual(saved["headers"], map[string]any{"User-Agent": "Custom"}) {
|
||||
t.Errorf("headers = %#v, want map[User-Agent:Custom]", saved["headers"])
|
||||
}
|
||||
if !reflect.DeepEqual(saved["models"], []any{"o3-mini"}) {
|
||||
t.Errorf("models = %#v, want [o3-mini]", saved["models"])
|
||||
}
|
||||
if !reflect.DeepEqual(saved["thinking"], map[string]any{"enabled": true}) {
|
||||
t.Errorf("thinking = %#v, want map[enabled:true]", saved["thinking"])
|
||||
}
|
||||
if saved["priority"] != float64(2) {
|
||||
t.Errorf("priority = %v, want 2", saved["priority"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchAuthFileFields_DeletesPluginFields(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
authDir := t.TempDir()
|
||||
fileName := "plugin-auth.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
|
||||
initialContent := map[string]any{
|
||||
"type": "demo-plugin",
|
||||
"token": "tok-123",
|
||||
"weight": float64(10),
|
||||
"headers": map[string]any{"X-Header": "val"},
|
||||
}
|
||||
raw, errMarshal := json.Marshal(initialContent)
|
||||
if errMarshal != nil {
|
||||
t.Fatalf("marshal error: %v", errMarshal)
|
||||
}
|
||||
if errWrite := os.WriteFile(filePath, raw, 0o600); errWrite != nil {
|
||||
t.Fatalf("write error: %v", errWrite)
|
||||
}
|
||||
|
||||
store := sdkAuth.NewFileTokenStore()
|
||||
store.SetBaseDir(authDir)
|
||||
manager := coreauth.NewManager(store, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
FileName: fileName,
|
||||
Provider: "demo-plugin",
|
||||
Metadata: map[string]any{
|
||||
"type": "demo-plugin",
|
||||
"token": "tok-123",
|
||||
"weight": float64(10),
|
||||
"headers": map[string]any{"X-Header": "val"},
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("Register() error = %v", errRegister)
|
||||
}
|
||||
|
||||
cfg := &config.Config{AuthDir: authDir}
|
||||
h := NewHandlerWithoutConfigFilePath(cfg, manager)
|
||||
|
||||
// Patch weight: null to delete weight
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
body := `{"name":"plugin-auth.json","weight":null}`
|
||||
c.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/fields", strings.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
h.PatchAuthFileFields(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("PatchAuthFileFields status = %d, want 200; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
savedRaw, errRead := os.ReadFile(filePath)
|
||||
if errRead != nil {
|
||||
t.Fatalf("ReadFile error: %v", errRead)
|
||||
}
|
||||
var saved map[string]any
|
||||
if errUnmarshal := json.Unmarshal(savedRaw, &saved); errUnmarshal != nil {
|
||||
t.Fatalf("Unmarshal error: %v", errUnmarshal)
|
||||
}
|
||||
|
||||
if _, exists := saved["weight"]; exists {
|
||||
t.Errorf("weight still exists in file after delete: %#v", saved["weight"])
|
||||
}
|
||||
if saved["token"] != "tok-123" {
|
||||
t.Errorf("token = %v, want tok-123", saved["token"])
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,69 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestUploadAuthFile_PreservesPriorityAttributes(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
authDir := t.TempDir()
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
|
||||
content := `{"type":"codex","email":"midai0530@gmail.com","priority":98}`
|
||||
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
part, err := writer.CreateFormFile("file", "codex-midai0530@gmail.com-plus.json")
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create multipart file: %v", err)
|
||||
}
|
||||
if _, err = part.Write([]byte(content)); err != nil {
|
||||
t.Fatalf("failed to write multipart content: %v", err)
|
||||
}
|
||||
if err = writer.Close(); err != nil {
|
||||
t.Fatalf("failed to close multipart writer: %v", err)
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v0/management/auth-files", &body)
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
ctx.Request = req
|
||||
|
||||
h.UploadAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected upload status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err = json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
|
||||
t.Fatalf("failed to decode response: %v", err)
|
||||
}
|
||||
if status, _ := payload["status"].(string); status != "ok" {
|
||||
t.Fatalf("expected status ok, got %#v", payload["status"])
|
||||
}
|
||||
|
||||
auth, ok := manager.GetByID("codex-midai0530@gmail.com-plus.json")
|
||||
if !ok || auth == nil {
|
||||
t.Fatalf("expected uploaded auth record to exist")
|
||||
}
|
||||
if got := auth.Attributes["priority"]; got != "98" {
|
||||
t.Fatalf("priority attribute = %q, want %q", got, "98")
|
||||
}
|
||||
if got := auth.Metadata["priority"]; got != float64(98) {
|
||||
t.Fatalf("priority metadata = %#v, want 98", got)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,132 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/synthesizer"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
const configAPIKeyDisablePattern = "*"
|
||||
|
||||
func setConfigAPIKeyExcludedAll(models []string, disable bool) []string {
|
||||
if disable {
|
||||
for _, item := range models {
|
||||
if strings.TrimSpace(item) == configAPIKeyDisablePattern {
|
||||
return config.NormalizeExcludedModels(models)
|
||||
}
|
||||
}
|
||||
return config.NormalizeExcludedModels(append(append([]string(nil), models...), configAPIKeyDisablePattern))
|
||||
}
|
||||
filtered := make([]string, 0, len(models))
|
||||
for _, item := range models {
|
||||
if strings.TrimSpace(item) == configAPIKeyDisablePattern {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
return config.NormalizeExcludedModels(filtered)
|
||||
}
|
||||
|
||||
func toggleConfigAPIKeyExcludedAll(cfg *config.Config, auth *coreauth.Auth, disable bool) (bool, error) {
|
||||
if cfg == nil || auth == nil || !coreauth.IsConfigAPIKeyAuth(auth) {
|
||||
return false, nil
|
||||
}
|
||||
authID := strings.TrimSpace(auth.ID)
|
||||
if authID == "" {
|
||||
return false, fmt.Errorf("auth id is empty")
|
||||
}
|
||||
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
|
||||
for i := range cfg.GeminiKey {
|
||||
entry := &cfg.GeminiKey[i]
|
||||
key := strings.TrimSpace(entry.APIKey)
|
||||
base := strings.TrimSpace(entry.BaseURL)
|
||||
proxyURL := strings.TrimSpace(entry.ProxyURL)
|
||||
prefix := strings.TrimSpace(entry.Prefix)
|
||||
if key == "" && base == "" {
|
||||
continue
|
||||
}
|
||||
id, _ := idGen.Next("gemini:apikey", key, base, proxyURL, prefix, config.FormatSortedHeaders(entry.Headers))
|
||||
if id == authID {
|
||||
entry.ExcludedModels = setConfigAPIKeyExcludedAll(entry.ExcludedModels, disable)
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
for i := range cfg.InteractionsKey {
|
||||
entry := &cfg.InteractionsKey[i]
|
||||
key := strings.TrimSpace(entry.APIKey)
|
||||
base := strings.TrimSpace(entry.BaseURL)
|
||||
proxyURL := strings.TrimSpace(entry.ProxyURL)
|
||||
prefix := strings.TrimSpace(entry.Prefix)
|
||||
if key == "" && base == "" {
|
||||
continue
|
||||
}
|
||||
id, _ := idGen.Next("gemini-interactions:apikey", key, base, proxyURL, prefix, config.FormatSortedHeaders(entry.Headers))
|
||||
if id == authID {
|
||||
entry.ExcludedModels = setConfigAPIKeyExcludedAll(entry.ExcludedModels, disable)
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
for i := range cfg.ClaudeKey {
|
||||
entry := &cfg.ClaudeKey[i]
|
||||
key := strings.TrimSpace(entry.APIKey)
|
||||
base := strings.TrimSpace(entry.BaseURL)
|
||||
proxyURL := strings.TrimSpace(entry.ProxyURL)
|
||||
prefix := strings.TrimSpace(entry.Prefix)
|
||||
if key == "" && base == "" {
|
||||
continue
|
||||
}
|
||||
id, _ := idGen.Next("claude:apikey", key, base, proxyURL, prefix, config.FormatSortedHeaders(entry.Headers))
|
||||
if id == authID {
|
||||
entry.ExcludedModels = setConfigAPIKeyExcludedAll(entry.ExcludedModels, disable)
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
for i := range cfg.CodexKey {
|
||||
entry := &cfg.CodexKey[i]
|
||||
key := strings.TrimSpace(entry.APIKey)
|
||||
base := strings.TrimSpace(entry.BaseURL)
|
||||
proxyURL := strings.TrimSpace(entry.ProxyURL)
|
||||
prefix := strings.TrimSpace(entry.Prefix)
|
||||
if key == "" && base == "" {
|
||||
continue
|
||||
}
|
||||
id, _ := idGen.Next("codex:apikey", key, base, proxyURL, prefix, config.FormatSortedHeaders(entry.Headers))
|
||||
if id == authID {
|
||||
entry.ExcludedModels = setConfigAPIKeyExcludedAll(entry.ExcludedModels, disable)
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
for i := range cfg.XAIKey {
|
||||
entry := &cfg.XAIKey[i]
|
||||
key := strings.TrimSpace(entry.APIKey)
|
||||
base := strings.TrimSpace(entry.BaseURL)
|
||||
proxyURL := strings.TrimSpace(entry.ProxyURL)
|
||||
prefix := strings.TrimSpace(entry.Prefix)
|
||||
if key == "" && base == "" {
|
||||
continue
|
||||
}
|
||||
id, _ := idGen.Next("xai:apikey", key, base, proxyURL, prefix, config.FormatSortedHeaders(entry.Headers))
|
||||
if id == authID {
|
||||
entry.ExcludedModels = setConfigAPIKeyExcludedAll(entry.ExcludedModels, disable)
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
for i := range cfg.VertexCompatAPIKey {
|
||||
entry := &cfg.VertexCompatAPIKey[i]
|
||||
key := strings.TrimSpace(entry.APIKey)
|
||||
base := strings.TrimSpace(entry.BaseURL)
|
||||
proxy := strings.TrimSpace(entry.ProxyURL)
|
||||
id, _ := idGen.Next("vertex:apikey", key, base, proxy)
|
||||
if id == authID {
|
||||
entry.ExcludedModels = setConfigAPIKeyExcludedAll(entry.ExcludedModels, disable)
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, nil
|
||||
}
|
||||
|
|
@ -0,0 +1,162 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/synthesizer"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestSetConfigAPIKeyExcludedAll(t *testing.T) {
|
||||
gotDisable := setConfigAPIKeyExcludedAll([]string{"gpt-5"}, true)
|
||||
if len(gotDisable) != 2 || gotDisable[0] != "gpt-5" || gotDisable[1] != "*" {
|
||||
t.Fatalf("unexpected disable list: %#v", gotDisable)
|
||||
}
|
||||
gotEnable := setConfigAPIKeyExcludedAll([]string{"gpt-5", "*"}, false)
|
||||
if len(gotEnable) != 1 || gotEnable[0] != "gpt-5" {
|
||||
t.Fatalf("unexpected enable list: %#v", gotEnable)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToggleConfigAPIKeyExcludedAll_XAI(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
XAIKey: []config.XAIKey{{
|
||||
APIKey: "xai-test",
|
||||
BaseURL: "https://api.x.ai/v1",
|
||||
}},
|
||||
}
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
authID, _ := idGen.Next("xai:apikey", "xai-test", "https://api.x.ai/v1", "", "", "")
|
||||
auth := &coreauth.Auth{
|
||||
ID: authID,
|
||||
Provider: "xai",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "xai-test",
|
||||
"base_url": "https://api.x.ai/v1",
|
||||
"source": "config:xai[abc]",
|
||||
},
|
||||
}
|
||||
|
||||
handled, errToggle := toggleConfigAPIKeyExcludedAll(cfg, auth, true)
|
||||
if errToggle != nil || !handled {
|
||||
t.Fatalf("toggle disable: handled=%v err=%v", handled, errToggle)
|
||||
}
|
||||
if len(cfg.XAIKey[0].ExcludedModels) != 1 || cfg.XAIKey[0].ExcludedModels[0] != "*" {
|
||||
t.Fatalf("excluded-models = %#v, want [*]", cfg.XAIKey[0].ExcludedModels)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToggleConfigAPIKeyExcludedAll_Codex(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
CodexKey: []config.CodexKey{{
|
||||
APIKey: "sk-test",
|
||||
BaseURL: "https://example.com/v1",
|
||||
}},
|
||||
}
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
authID, _ := idGen.Next("codex:apikey", "sk-test", "https://example.com/v1", "", "", "")
|
||||
auth := &coreauth.Auth{
|
||||
ID: authID,
|
||||
Provider: "codex",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "sk-test",
|
||||
"base_url": "https://example.com/v1",
|
||||
"source": "config:codex[abc]",
|
||||
},
|
||||
}
|
||||
|
||||
handled, err := toggleConfigAPIKeyExcludedAll(cfg, auth, true)
|
||||
if err != nil || !handled {
|
||||
t.Fatalf("toggle disable: handled=%v err=%v", handled, err)
|
||||
}
|
||||
if len(cfg.CodexKey[0].ExcludedModels) != 1 || cfg.CodexKey[0].ExcludedModels[0] != "*" {
|
||||
t.Fatalf("expected excluded-models [*], got %#v", cfg.CodexKey[0].ExcludedModels)
|
||||
}
|
||||
|
||||
handled, err = toggleConfigAPIKeyExcludedAll(cfg, auth, false)
|
||||
if err != nil || !handled {
|
||||
t.Fatalf("toggle enable: handled=%v err=%v", handled, err)
|
||||
}
|
||||
if len(cfg.CodexKey[0].ExcludedModels) != 0 {
|
||||
t.Fatalf("expected excluded-models cleared, got %#v", cfg.CodexKey[0].ExcludedModels)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToggleConfigAPIKeyExcludedAll_Vertex_NoBaseURL(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
VertexCompatAPIKey: []config.VertexCompatKey{{
|
||||
APIKey: "vertex-key-only",
|
||||
}},
|
||||
}
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
authID, _ := idGen.Next("vertex:apikey", "vertex-key-only", "", "")
|
||||
auth := &coreauth.Auth{
|
||||
ID: authID,
|
||||
Provider: "vertex",
|
||||
Attributes: map[string]string{
|
||||
"auth_kind": "apikey",
|
||||
"api_key": "vertex-key-only",
|
||||
"source": "config:vertex[xyz]",
|
||||
},
|
||||
}
|
||||
|
||||
handled, errToggle := toggleConfigAPIKeyExcludedAll(cfg, auth, true)
|
||||
if errToggle != nil || !handled {
|
||||
t.Fatalf("toggle disable: handled=%v err=%v", handled, errToggle)
|
||||
}
|
||||
if len(cfg.VertexCompatAPIKey[0].ExcludedModels) != 1 || cfg.VertexCompatAPIKey[0].ExcludedModels[0] != "*" {
|
||||
t.Fatalf("excluded-models = %#v, want [*]", cfg.VertexCompatAPIKey[0].ExcludedModels)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToggleConfigAPIKeyExcludedAll_EmptyKeyWithBaseURL(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
ClaudeKey: []config.ClaudeKey{{
|
||||
APIKey: "",
|
||||
BaseURL: "https://custom-claude.example.com",
|
||||
}},
|
||||
GeminiKey: []config.GeminiKey{{
|
||||
APIKey: " ",
|
||||
BaseURL: "https://custom-gemini.example.com",
|
||||
}},
|
||||
}
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
claudeID, _ := idGen.Next("claude:apikey", "", "https://custom-claude.example.com", "", "", "")
|
||||
geminiID, _ := idGen.Next("gemini:apikey", "", "https://custom-gemini.example.com", "", "", "")
|
||||
|
||||
claudeAuth := &coreauth.Auth{
|
||||
ID: claudeID,
|
||||
Provider: "claude",
|
||||
Attributes: map[string]string{
|
||||
"auth_kind": "apikey",
|
||||
"base_url": "https://custom-claude.example.com",
|
||||
"source": "config:claude[abc]",
|
||||
},
|
||||
}
|
||||
geminiAuth := &coreauth.Auth{
|
||||
ID: geminiID,
|
||||
Provider: "gemini",
|
||||
Attributes: map[string]string{
|
||||
"auth_kind": "apikey",
|
||||
"base_url": "https://custom-gemini.example.com",
|
||||
"source": "config:gemini[def]",
|
||||
},
|
||||
}
|
||||
|
||||
handled, err := toggleConfigAPIKeyExcludedAll(cfg, claudeAuth, true)
|
||||
if err != nil || !handled {
|
||||
t.Fatalf("toggle claude: handled=%v err=%v", handled, err)
|
||||
}
|
||||
if len(cfg.ClaudeKey[0].ExcludedModels) != 1 || cfg.ClaudeKey[0].ExcludedModels[0] != "*" {
|
||||
t.Fatalf("claude excluded-models = %#v, want [*]", cfg.ClaudeKey[0].ExcludedModels)
|
||||
}
|
||||
|
||||
handled, err = toggleConfigAPIKeyExcludedAll(cfg, geminiAuth, true)
|
||||
if err != nil || !handled {
|
||||
t.Fatalf("toggle gemini: handled=%v err=%v", handled, err)
|
||||
}
|
||||
if len(cfg.GeminiKey[0].ExcludedModels) != 1 || cfg.GeminiKey[0].ExcludedModels[0] != "*" {
|
||||
t.Fatalf("gemini excluded-models = %#v, want [*]", cfg.GeminiKey[0].ExcludedModels)
|
||||
}
|
||||
}
|
||||
334
backend/internal/api/handlers/management/config_auth_index.go
Normal file
334
backend/internal/api/handlers/management/config_auth_index.go
Normal file
|
|
@ -0,0 +1,334 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/synthesizer"
|
||||
)
|
||||
|
||||
type geminiKeyWithAuthIndex struct {
|
||||
config.GeminiKey
|
||||
AuthIndex string `json:"auth-index,omitempty"`
|
||||
}
|
||||
|
||||
type claudeKeyWithAuthIndex struct {
|
||||
config.ClaudeKey
|
||||
AuthIndex string `json:"auth-index,omitempty"`
|
||||
}
|
||||
|
||||
type codexKeyWithAuthIndex struct {
|
||||
config.CodexKey
|
||||
AuthIndex string `json:"auth-index,omitempty"`
|
||||
}
|
||||
|
||||
type xaiKeyWithAuthIndex struct {
|
||||
config.XAIKey
|
||||
AuthIndex string `json:"auth-index,omitempty"`
|
||||
}
|
||||
|
||||
type vertexCompatKeyWithAuthIndex struct {
|
||||
config.VertexCompatKey
|
||||
AuthIndex string `json:"auth-index,omitempty"`
|
||||
}
|
||||
|
||||
type openAICompatibilityAPIKeyWithAuthIndex struct {
|
||||
config.OpenAICompatibilityAPIKey
|
||||
AuthIndex string `json:"auth-index,omitempty"`
|
||||
}
|
||||
|
||||
type openAICompatibilityWithAuthIndex struct {
|
||||
Name string `json:"name"`
|
||||
Priority int `json:"priority,omitempty"`
|
||||
Disabled bool `json:"disabled"`
|
||||
Prefix string `json:"prefix,omitempty"`
|
||||
BaseURL string `json:"base-url"`
|
||||
APIKeyEntries []openAICompatibilityAPIKeyWithAuthIndex `json:"api-key-entries,omitempty"`
|
||||
Models []config.OpenAICompatibilityModel `json:"models,omitempty"`
|
||||
Headers map[string]string `json:"headers,omitempty"`
|
||||
SupportPromptCacheKey bool `json:"support-prompt-cache-key,omitempty"`
|
||||
DisableCooling *bool `json:"disable-cooling,omitempty"`
|
||||
RequestRetry *int `json:"request-retry,omitempty"`
|
||||
RequestScopedErrors []config.RequestScopedErrorRule `json:"request-scoped-errors,omitempty"`
|
||||
AuthIndex string `json:"auth-index,omitempty"`
|
||||
}
|
||||
|
||||
func (h *Handler) liveAuthIndexByID() map[string]string {
|
||||
out := map[string]string{}
|
||||
if h == nil {
|
||||
return out
|
||||
}
|
||||
h.mu.Lock()
|
||||
manager := h.authManager
|
||||
h.mu.Unlock()
|
||||
if manager == nil {
|
||||
return out
|
||||
}
|
||||
// authManager.List() returns clones, so EnsureIndex only affects these copies.
|
||||
for _, auth := range manager.List() {
|
||||
if auth == nil {
|
||||
continue
|
||||
}
|
||||
id := strings.TrimSpace(auth.ID)
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
idx := strings.TrimSpace(auth.Index)
|
||||
if idx == "" {
|
||||
idx = auth.EnsureIndex()
|
||||
}
|
||||
if idx == "" {
|
||||
continue
|
||||
}
|
||||
out[id] = idx
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *Handler) geminiKeysWithAuthIndex() []geminiKeyWithAuthIndex {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
liveIndexByID := h.liveAuthIndexByID()
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
out := make([]geminiKeyWithAuthIndex, len(h.cfg.GeminiKey))
|
||||
for i := range h.cfg.GeminiKey {
|
||||
entry := h.cfg.GeminiKey[i]
|
||||
authIndex := ""
|
||||
key := strings.TrimSpace(entry.APIKey)
|
||||
base := strings.TrimSpace(entry.BaseURL)
|
||||
proxyURL := strings.TrimSpace(entry.ProxyURL)
|
||||
prefix := strings.TrimSpace(entry.Prefix)
|
||||
if key != "" || base != "" {
|
||||
id, _ := idGen.Next("gemini:apikey", key, base, proxyURL, prefix, config.FormatSortedHeaders(entry.Headers))
|
||||
authIndex = liveIndexByID[id]
|
||||
}
|
||||
out[i] = geminiKeyWithAuthIndex{
|
||||
GeminiKey: entry,
|
||||
AuthIndex: authIndex,
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *Handler) interactionsKeysWithAuthIndex() []geminiKeyWithAuthIndex {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
liveIndexByID := h.liveAuthIndexByID()
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
out := make([]geminiKeyWithAuthIndex, len(h.cfg.InteractionsKey))
|
||||
for i := range h.cfg.InteractionsKey {
|
||||
entry := h.cfg.InteractionsKey[i]
|
||||
authIndex := ""
|
||||
key := strings.TrimSpace(entry.APIKey)
|
||||
base := strings.TrimSpace(entry.BaseURL)
|
||||
proxyURL := strings.TrimSpace(entry.ProxyURL)
|
||||
prefix := strings.TrimSpace(entry.Prefix)
|
||||
if key != "" || base != "" {
|
||||
id, _ := idGen.Next("gemini-interactions:apikey", key, base, proxyURL, prefix, config.FormatSortedHeaders(entry.Headers))
|
||||
authIndex = liveIndexByID[id]
|
||||
}
|
||||
out[i] = geminiKeyWithAuthIndex{
|
||||
GeminiKey: entry,
|
||||
AuthIndex: authIndex,
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *Handler) claudeKeysWithAuthIndex() []claudeKeyWithAuthIndex {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
liveIndexByID := h.liveAuthIndexByID()
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
out := make([]claudeKeyWithAuthIndex, len(h.cfg.ClaudeKey))
|
||||
for i := range h.cfg.ClaudeKey {
|
||||
entry := h.cfg.ClaudeKey[i]
|
||||
authIndex := ""
|
||||
key := strings.TrimSpace(entry.APIKey)
|
||||
base := strings.TrimSpace(entry.BaseURL)
|
||||
proxyURL := strings.TrimSpace(entry.ProxyURL)
|
||||
prefix := strings.TrimSpace(entry.Prefix)
|
||||
if key != "" || base != "" {
|
||||
id, _ := idGen.Next("claude:apikey", key, base, proxyURL, prefix, config.FormatSortedHeaders(entry.Headers))
|
||||
authIndex = liveIndexByID[id]
|
||||
}
|
||||
out[i] = claudeKeyWithAuthIndex{
|
||||
ClaudeKey: entry,
|
||||
AuthIndex: authIndex,
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *Handler) codexKeysWithAuthIndex() []codexKeyWithAuthIndex {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
liveIndexByID := h.liveAuthIndexByID()
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
out := make([]codexKeyWithAuthIndex, len(h.cfg.CodexKey))
|
||||
for i := range h.cfg.CodexKey {
|
||||
entry := h.cfg.CodexKey[i]
|
||||
authIndex := ""
|
||||
key := strings.TrimSpace(entry.APIKey)
|
||||
base := strings.TrimSpace(entry.BaseURL)
|
||||
proxyURL := strings.TrimSpace(entry.ProxyURL)
|
||||
prefix := strings.TrimSpace(entry.Prefix)
|
||||
if key != "" || base != "" {
|
||||
id, _ := idGen.Next("codex:apikey", key, base, proxyURL, prefix, config.FormatSortedHeaders(entry.Headers))
|
||||
authIndex = liveIndexByID[id]
|
||||
}
|
||||
out[i] = codexKeyWithAuthIndex{
|
||||
CodexKey: entry,
|
||||
AuthIndex: authIndex,
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *Handler) xaiKeysWithAuthIndex() []xaiKeyWithAuthIndex {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
liveIndexByID := h.liveAuthIndexByID()
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
out := make([]xaiKeyWithAuthIndex, len(h.cfg.XAIKey))
|
||||
for i := range h.cfg.XAIKey {
|
||||
entry := h.cfg.XAIKey[i]
|
||||
authIndex := ""
|
||||
key := strings.TrimSpace(entry.APIKey)
|
||||
base := strings.TrimSpace(entry.BaseURL)
|
||||
proxyURL := strings.TrimSpace(entry.ProxyURL)
|
||||
prefix := strings.TrimSpace(entry.Prefix)
|
||||
if key != "" || base != "" {
|
||||
id, _ := idGen.Next("xai:apikey", key, base, proxyURL, prefix, config.FormatSortedHeaders(entry.Headers))
|
||||
authIndex = liveIndexByID[id]
|
||||
}
|
||||
out[i] = xaiKeyWithAuthIndex{
|
||||
XAIKey: entry,
|
||||
AuthIndex: authIndex,
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *Handler) vertexCompatKeysWithAuthIndex() []vertexCompatKeyWithAuthIndex {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
liveIndexByID := h.liveAuthIndexByID()
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
out := make([]vertexCompatKeyWithAuthIndex, len(h.cfg.VertexCompatAPIKey))
|
||||
for i := range h.cfg.VertexCompatAPIKey {
|
||||
entry := h.cfg.VertexCompatAPIKey[i]
|
||||
id, _ := idGen.Next("vertex:apikey", entry.APIKey, entry.BaseURL, entry.ProxyURL)
|
||||
authIndex := liveIndexByID[id]
|
||||
out[i] = vertexCompatKeyWithAuthIndex{
|
||||
VertexCompatKey: entry,
|
||||
AuthIndex: authIndex,
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *Handler) openAICompatibilityWithAuthIndex() []openAICompatibilityWithAuthIndex {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
liveIndexByID := h.liveAuthIndexByID()
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
normalized := normalizedOpenAICompatibilityEntries(h.cfg.OpenAICompatibility)
|
||||
out := make([]openAICompatibilityWithAuthIndex, len(normalized))
|
||||
idGen := synthesizer.NewStableIDGenerator()
|
||||
for i := range normalized {
|
||||
entry := normalized[i]
|
||||
providerName := strings.ToLower(strings.TrimSpace(entry.Name))
|
||||
if providerName == "" {
|
||||
providerName = "openai-compatibility"
|
||||
}
|
||||
idKind := fmt.Sprintf("openai-compatibility:%s", providerName)
|
||||
|
||||
response := openAICompatibilityWithAuthIndex{
|
||||
Name: entry.Name,
|
||||
Priority: entry.Priority,
|
||||
Disabled: entry.Disabled,
|
||||
Prefix: entry.Prefix,
|
||||
BaseURL: entry.BaseURL,
|
||||
Models: entry.Models,
|
||||
Headers: entry.Headers,
|
||||
SupportPromptCacheKey: entry.SupportPromptCacheKey,
|
||||
DisableCooling: entry.DisableCooling,
|
||||
RequestRetry: entry.RequestRetry,
|
||||
RequestScopedErrors: entry.RequestScopedErrors,
|
||||
AuthIndex: "",
|
||||
}
|
||||
if len(entry.APIKeyEntries) == 0 {
|
||||
id, _ := idGen.Next(idKind, entry.BaseURL)
|
||||
response.AuthIndex = liveIndexByID[id]
|
||||
} else {
|
||||
response.APIKeyEntries = make([]openAICompatibilityAPIKeyWithAuthIndex, len(entry.APIKeyEntries))
|
||||
for j := range entry.APIKeyEntries {
|
||||
apiKeyEntry := entry.APIKeyEntries[j]
|
||||
id, _ := idGen.Next(idKind, apiKeyEntry.APIKey, entry.BaseURL, apiKeyEntry.ProxyURL)
|
||||
response.APIKeyEntries[j] = openAICompatibilityAPIKeyWithAuthIndex{
|
||||
OpenAICompatibilityAPIKey: apiKeyEntry,
|
||||
AuthIndex: liveIndexByID[id],
|
||||
}
|
||||
}
|
||||
}
|
||||
out[i] = response
|
||||
}
|
||||
return out
|
||||
}
|
||||
338
backend/internal/api/handlers/management/config_basic.go
Normal file
338
backend/internal/api/handlers/management/config_basic.go
Normal file
|
|
@ -0,0 +1,338 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
const (
|
||||
latestReleaseURL = "https://api.github.com/repos/router-for-me/CLIProxyAPI/releases/latest"
|
||||
latestReleaseUserAgent = "CLIProxyAPI"
|
||||
)
|
||||
|
||||
func (h *Handler) GetConfig(c *gin.Context) {
|
||||
if h == nil || h.cfg == nil {
|
||||
c.JSON(200, gin.H{})
|
||||
return
|
||||
}
|
||||
c.JSON(200, new(*h.cfg))
|
||||
}
|
||||
|
||||
type releaseInfo struct {
|
||||
TagName string `json:"tag_name"`
|
||||
Name string `json:"name"`
|
||||
}
|
||||
|
||||
// GetLatestVersion returns the latest release version from GitHub without downloading assets.
|
||||
func (h *Handler) GetLatestVersion(c *gin.Context) {
|
||||
client := &http.Client{Timeout: 10 * time.Second}
|
||||
proxyURL := ""
|
||||
if h != nil && h.cfg != nil {
|
||||
proxyURL = strings.TrimSpace(h.cfg.ProxyURL)
|
||||
}
|
||||
if proxyURL != "" {
|
||||
sdkCfg := &sdkconfig.SDKConfig{ProxyURL: proxyURL}
|
||||
util.SetProxy(sdkCfg, client)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(c.Request.Context(), http.MethodGet, latestReleaseURL, nil)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "request_create_failed", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
req.Header.Set("Accept", "application/vnd.github+json")
|
||||
req.Header.Set("User-Agent", latestReleaseUserAgent)
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "request_failed", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if errClose := resp.Body.Close(); errClose != nil {
|
||||
log.WithError(errClose).Debug("failed to close latest version response body")
|
||||
}
|
||||
}()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "unexpected_status", "message": fmt.Sprintf("status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))})
|
||||
return
|
||||
}
|
||||
|
||||
var info releaseInfo
|
||||
if errDecode := json.NewDecoder(resp.Body).Decode(&info); errDecode != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "decode_failed", "message": errDecode.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
version := strings.TrimSpace(info.TagName)
|
||||
if version == "" {
|
||||
version = strings.TrimSpace(info.Name)
|
||||
}
|
||||
if version == "" {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "invalid_response", "message": "missing release version"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"latest-version": version})
|
||||
}
|
||||
|
||||
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 _, errWrite := f.Write(data); errWrite != nil {
|
||||
_ = f.Close()
|
||||
return errWrite
|
||||
}
|
||||
if errSync := f.Sync(); errSync != nil {
|
||||
_ = f.Close()
|
||||
return errSync
|
||||
}
|
||||
return f.Close()
|
||||
}
|
||||
|
||||
func (h *Handler) PutConfigYAML(c *gin.Context) {
|
||||
body, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_yaml", "message": "cannot read request body"})
|
||||
return
|
||||
}
|
||||
var cfg config.Config
|
||||
if err = yaml.Unmarshal(body, &cfg); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_yaml", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
// Validate config using LoadConfigOptional with optional=false to enforce parsing
|
||||
tmpDir := filepath.Dir(h.configFilePath)
|
||||
tmpFile, err := os.CreateTemp(tmpDir, "config-validate-*.yaml")
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "write_failed", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
tempFile := tmpFile.Name()
|
||||
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 errClose := tmpFile.Close(); errClose != nil {
|
||||
_ = os.Remove(tempFile)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "write_failed", "message": errClose.Error()})
|
||||
return
|
||||
}
|
||||
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()})
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if WriteConfig(h.configFilePath, body) != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "write_failed", "message": "failed to write config"})
|
||||
return
|
||||
}
|
||||
// Reload into handler to keep memory in sync
|
||||
newCfg, err := config.LoadConfig(h.configFilePath)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "reload_failed", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
h.cfg = newCfg
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "changed": []string{"config"}})
|
||||
}
|
||||
|
||||
// GetConfigYAML returns the raw config.yaml file bytes without re-encoding.
|
||||
// It preserves comments and original formatting/styles.
|
||||
func (h *Handler) GetConfigYAML(c *gin.Context) {
|
||||
data, err := os.ReadFile(h.configFilePath)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "not_found", "message": "config file not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "read_failed", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
c.Header("Content-Type", "application/yaml; charset=utf-8")
|
||||
c.Header("Cache-Control", "no-store")
|
||||
c.Header("X-Content-Type-Options", "nosniff")
|
||||
// Write raw bytes as-is
|
||||
_, _ = c.Writer.Write(data)
|
||||
}
|
||||
|
||||
// Debug
|
||||
func (h *Handler) GetDebug(c *gin.Context) { c.JSON(200, gin.H{"debug": h.cfg.Debug}) }
|
||||
func (h *Handler) PutDebug(c *gin.Context) { h.updateBoolField(c, func(v bool) { h.cfg.Debug = v }) }
|
||||
|
||||
// UsageStatisticsEnabled
|
||||
func (h *Handler) GetUsageStatisticsEnabled(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"usage-statistics-enabled": h.cfg.UsageStatisticsEnabled})
|
||||
}
|
||||
func (h *Handler) PutUsageStatisticsEnabled(c *gin.Context) {
|
||||
h.updateBoolField(c, func(v bool) { h.cfg.UsageStatisticsEnabled = v })
|
||||
}
|
||||
|
||||
// UsageStatisticsEnabled
|
||||
func (h *Handler) GetLoggingToFile(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"logging-to-file": h.cfg.LoggingToFile})
|
||||
}
|
||||
func (h *Handler) PutLoggingToFile(c *gin.Context) {
|
||||
h.updateBoolField(c, func(v bool) { h.cfg.LoggingToFile = v })
|
||||
}
|
||||
|
||||
// LogsMaxTotalSizeMB
|
||||
func (h *Handler) GetLogsMaxTotalSizeMB(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"logs-max-total-size-mb": h.cfg.LogsMaxTotalSizeMB})
|
||||
}
|
||||
func (h *Handler) PutLogsMaxTotalSizeMB(c *gin.Context) {
|
||||
var body struct {
|
||||
Value *int `json:"value"`
|
||||
}
|
||||
if errBindJSON := c.ShouldBindJSON(&body); errBindJSON != nil || body.Value == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
value := *body.Value
|
||||
if value < 0 {
|
||||
value = 0
|
||||
}
|
||||
h.cfg.LogsMaxTotalSizeMB = value
|
||||
h.persist(c)
|
||||
}
|
||||
|
||||
// ErrorLogsMaxFiles
|
||||
func (h *Handler) GetErrorLogsMaxFiles(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"error-logs-max-files": h.cfg.ErrorLogsMaxFiles})
|
||||
}
|
||||
func (h *Handler) PutErrorLogsMaxFiles(c *gin.Context) {
|
||||
var body struct {
|
||||
Value *int `json:"value"`
|
||||
}
|
||||
if errBindJSON := c.ShouldBindJSON(&body); errBindJSON != nil || body.Value == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
value := *body.Value
|
||||
if value < 0 {
|
||||
value = 10
|
||||
}
|
||||
h.cfg.ErrorLogsMaxFiles = value
|
||||
h.persist(c)
|
||||
}
|
||||
|
||||
// Request log
|
||||
func (h *Handler) GetRequestLog(c *gin.Context) { c.JSON(200, gin.H{"request-log": h.cfg.RequestLog}) }
|
||||
func (h *Handler) PutRequestLog(c *gin.Context) {
|
||||
h.updateBoolField(c, func(v bool) { h.cfg.RequestLog = v })
|
||||
}
|
||||
|
||||
// Websocket auth
|
||||
func (h *Handler) GetWebsocketAuth(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"ws-auth": h.cfg.WebsocketAuth})
|
||||
}
|
||||
func (h *Handler) PutWebsocketAuth(c *gin.Context) {
|
||||
h.updateBoolField(c, func(v bool) { h.cfg.WebsocketAuth = v })
|
||||
}
|
||||
|
||||
// Request retry
|
||||
func (h *Handler) GetRequestRetry(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"request-retry": h.cfg.RequestRetry})
|
||||
}
|
||||
func (h *Handler) PutRequestRetry(c *gin.Context) {
|
||||
h.updateIntField(c, func(v int) { h.cfg.RequestRetry = v })
|
||||
}
|
||||
|
||||
// Max retry credentials
|
||||
func (h *Handler) GetMaxRetryCredentials(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"max-retry-credentials": h.cfg.MaxRetryCredentials})
|
||||
}
|
||||
func (h *Handler) PutMaxRetryCredentials(c *gin.Context) {
|
||||
h.updateIntField(c, func(v int) { h.cfg.MaxRetryCredentials = v })
|
||||
}
|
||||
|
||||
// Max retry interval
|
||||
func (h *Handler) GetMaxRetryInterval(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"max-retry-interval": h.cfg.MaxRetryInterval})
|
||||
}
|
||||
func (h *Handler) PutMaxRetryInterval(c *gin.Context) {
|
||||
h.updateIntField(c, func(v int) { h.cfg.MaxRetryInterval = v })
|
||||
}
|
||||
|
||||
// ForceModelPrefix
|
||||
func (h *Handler) GetForceModelPrefix(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"force-model-prefix": h.cfg.ForceModelPrefix})
|
||||
}
|
||||
func (h *Handler) PutForceModelPrefix(c *gin.Context) {
|
||||
h.updateBoolField(c, func(v bool) { h.cfg.ForceModelPrefix = v })
|
||||
}
|
||||
|
||||
func normalizeRoutingStrategy(strategy string) (string, bool) {
|
||||
normalized := strings.ToLower(strings.TrimSpace(strategy))
|
||||
switch normalized {
|
||||
case "", "round-robin", "roundrobin", "rr":
|
||||
return "round-robin", true
|
||||
case "weighted-round-robin", "weightedroundrobin", "wrr":
|
||||
return "weighted-round-robin", true
|
||||
case "fill-first", "fillfirst", "ff":
|
||||
return "fill-first", true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
// RoutingStrategy
|
||||
func (h *Handler) GetRoutingStrategy(c *gin.Context) {
|
||||
strategy, ok := normalizeRoutingStrategy(h.cfg.Routing.Strategy)
|
||||
if !ok {
|
||||
c.JSON(200, gin.H{"strategy": strings.TrimSpace(h.cfg.Routing.Strategy)})
|
||||
return
|
||||
}
|
||||
c.JSON(200, gin.H{"strategy": strategy})
|
||||
}
|
||||
func (h *Handler) PutRoutingStrategy(c *gin.Context) {
|
||||
var body struct {
|
||||
Value *string `json:"value"`
|
||||
}
|
||||
if errBindJSON := c.ShouldBindJSON(&body); errBindJSON != nil || body.Value == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
normalized, ok := normalizeRoutingStrategy(*body.Value)
|
||||
if !ok {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid strategy"})
|
||||
return
|
||||
}
|
||||
h.cfg.Routing.Strategy = normalized
|
||||
h.persist(c)
|
||||
}
|
||||
|
||||
// Proxy URL
|
||||
func (h *Handler) GetProxyURL(c *gin.Context) { c.JSON(200, gin.H{"proxy-url": h.cfg.ProxyURL}) }
|
||||
func (h *Handler) PutProxyURL(c *gin.Context) {
|
||||
h.updateStringField(c, func(v string) { h.cfg.ProxyURL = v })
|
||||
}
|
||||
func (h *Handler) DeleteProxyURL(c *gin.Context) {
|
||||
h.cfg.ProxyURL = ""
|
||||
h.persist(c)
|
||||
}
|
||||
|
|
@ -0,0 +1,12 @@
|
|||
package management
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestNormalizeRoutingStrategyWeightedRoundRobin(t *testing.T) {
|
||||
for _, input := range []string{"weighted-round-robin", "weightedroundrobin", "wrr"} {
|
||||
got, ok := normalizeRoutingStrategy(input)
|
||||
if !ok || got != "weighted-round-robin" {
|
||||
t.Fatalf("normalizeRoutingStrategy(%q) = %q, %v; want weighted-round-robin, true", input, got, ok)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,116 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestPatchClaudeKeyFingerprintProfile(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
ClaudeKey: []config.ClaudeKey{
|
||||
{APIKey: "test-claude-key"},
|
||||
},
|
||||
}
|
||||
h := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)}
|
||||
|
||||
// Patch fingerprint-profile to claude-code-cli
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/claude-api-key",
|
||||
strings.NewReader(`{"index":0,"value":{"fingerprint-profile":"claude-code-cli"}}`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
h.PatchClaudeKey(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if got := cfg.ClaudeKey[0].FingerprintProfile; got != "claude-code-cli" {
|
||||
t.Fatalf("FingerprintProfile = %q, want %q", got, "claude-code-cli")
|
||||
}
|
||||
|
||||
// Patch fingerprint-profile back to empty
|
||||
rec = httptest.NewRecorder()
|
||||
ctx, _ = gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/claude-api-key",
|
||||
strings.NewReader(`{"index":0,"value":{"fingerprint-profile":""}}`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
h.PatchClaudeKey(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if got := cfg.ClaudeKey[0].FingerprintProfile; got != "" {
|
||||
t.Fatalf("FingerprintProfile = %q, want empty", got)
|
||||
}
|
||||
|
||||
// A legacy alias is stored in canonical form so the config file and the request
|
||||
// path agree on one spelling.
|
||||
rec = httptest.NewRecorder()
|
||||
ctx, _ = gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/claude-api-key",
|
||||
strings.NewReader(`{"index":0,"value":{"fingerprint-profile":" OAuth-CLI "}}`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
h.PatchClaudeKey(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if got := cfg.ClaudeKey[0].FingerprintProfile; got != "claude-code-cli" {
|
||||
t.Fatalf("FingerprintProfile = %q, want canonical %q", got, "claude-code-cli")
|
||||
}
|
||||
}
|
||||
|
||||
// A typo must fail the write instead of reaching the request path, where it can
|
||||
// only be reported as a warning behind every later request.
|
||||
func TestPatchClaudeKeyRejectsUnknownFingerprintProfile(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
ClaudeKey: []config.ClaudeKey{
|
||||
{APIKey: "test-claude-key", FingerprintProfile: "claude-code-cli"},
|
||||
},
|
||||
}
|
||||
h := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/claude-api-key",
|
||||
strings.NewReader(`{"index":0,"value":{"fingerprint-profile":"claude-code"}}`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
h.PatchClaudeKey(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want 400; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "fingerprint-profile") {
|
||||
t.Fatalf("error body = %s, want it to name the field", rec.Body.String())
|
||||
}
|
||||
if got := cfg.ClaudeKey[0].FingerprintProfile; got != "claude-code-cli" {
|
||||
t.Fatalf("FingerprintProfile = %q, want the rejected patch to leave it unchanged", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPutClaudeKeysRejectsUnknownFingerprintProfile(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
h := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPut, "/v0/management/claude-api-key",
|
||||
strings.NewReader(`[{"api-key":"k1"},{"api-key":"k2","fingerprint-profile":"claude-cli"}]`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
h.PutClaudeKeys(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want 400; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "claude-api-key[1].fingerprint-profile") {
|
||||
t.Fatalf("error body = %s, want the offending index", rec.Body.String())
|
||||
}
|
||||
if len(cfg.ClaudeKey) != 0 {
|
||||
t.Fatalf("ClaudeKey = %+v, want the rejected write to change nothing", cfg.ClaudeKey)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,34 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestPatchCodexKeyUpdatesAlphaSearch(t *testing.T) {
|
||||
h := &Handler{
|
||||
cfg: &config.Config{CodexKey: []config.CodexKey{{
|
||||
APIKey: "codex-key",
|
||||
BaseURL: "https://codex.example.com",
|
||||
}}},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/codex-api-key", strings.NewReader(`{"index":0,"value":{"alpha-search":true}}`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
h.PatchCodexKey(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if !h.cfg.CodexKey[0].AlphaSearch {
|
||||
t.Fatal("alpha-search = false, want true")
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,129 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestPatchDisableCoolingOverrideForEveryFamily(t *testing.T) {
|
||||
initial := true
|
||||
tests := []struct {
|
||||
name string
|
||||
setup func(*config.Config)
|
||||
patch func(*Handler, *gin.Context)
|
||||
get func(*config.Config) *bool
|
||||
}{
|
||||
{
|
||||
name: "gemini",
|
||||
setup: func(cfg *config.Config) {
|
||||
cfg.GeminiKey = []config.GeminiKey{{APIKey: "key", DisableCooling: &initial}}
|
||||
},
|
||||
patch: (*Handler).PatchGeminiKey,
|
||||
get: func(cfg *config.Config) *bool { return cfg.GeminiKey[0].DisableCooling },
|
||||
},
|
||||
{
|
||||
name: "interactions",
|
||||
setup: func(cfg *config.Config) {
|
||||
cfg.InteractionsKey = []config.GeminiKey{{APIKey: "key", DisableCooling: &initial}}
|
||||
},
|
||||
patch: (*Handler).PatchInteractionsKey,
|
||||
get: func(cfg *config.Config) *bool { return cfg.InteractionsKey[0].DisableCooling },
|
||||
},
|
||||
{
|
||||
name: "claude",
|
||||
setup: func(cfg *config.Config) {
|
||||
cfg.ClaudeKey = []config.ClaudeKey{{APIKey: "key", DisableCooling: &initial}}
|
||||
},
|
||||
patch: (*Handler).PatchClaudeKey,
|
||||
get: func(cfg *config.Config) *bool { return cfg.ClaudeKey[0].DisableCooling },
|
||||
},
|
||||
{
|
||||
name: "openai compatibility",
|
||||
setup: func(cfg *config.Config) {
|
||||
cfg.OpenAICompatibility = []config.OpenAICompatibility{{
|
||||
Name: "compat",
|
||||
BaseURL: "https://compat.example.com",
|
||||
APIKeyEntries: []config.OpenAICompatibilityAPIKey{{APIKey: "key"}},
|
||||
DisableCooling: &initial,
|
||||
}}
|
||||
},
|
||||
patch: (*Handler).PatchOpenAICompat,
|
||||
get: func(cfg *config.Config) *bool { return cfg.OpenAICompatibility[0].DisableCooling },
|
||||
},
|
||||
{
|
||||
name: "vertex",
|
||||
setup: func(cfg *config.Config) {
|
||||
cfg.VertexCompatAPIKey = []config.VertexCompatKey{{
|
||||
APIKey: "key",
|
||||
BaseURL: "https://vertex.example.com",
|
||||
DisableCooling: &initial,
|
||||
}}
|
||||
},
|
||||
patch: (*Handler).PatchVertexCompatKey,
|
||||
get: func(cfg *config.Config) *bool { return cfg.VertexCompatAPIKey[0].DisableCooling },
|
||||
},
|
||||
{
|
||||
name: "codex",
|
||||
setup: func(cfg *config.Config) {
|
||||
cfg.CodexKey = []config.CodexKey{{
|
||||
APIKey: "key",
|
||||
BaseURL: "https://codex.example.com",
|
||||
DisableCooling: &initial,
|
||||
}}
|
||||
},
|
||||
patch: (*Handler).PatchCodexKey,
|
||||
get: func(cfg *config.Config) *bool { return cfg.CodexKey[0].DisableCooling },
|
||||
},
|
||||
{
|
||||
name: "xai",
|
||||
setup: func(cfg *config.Config) {
|
||||
cfg.XAIKey = []config.XAIKey{{
|
||||
APIKey: "key",
|
||||
BaseURL: "https://api.x.ai/v1",
|
||||
DisableCooling: &initial,
|
||||
}}
|
||||
},
|
||||
patch: (*Handler).PatchXAIKey,
|
||||
get: func(cfg *config.Config) *bool { return cfg.XAIKey[0].DisableCooling },
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
tc.setup(cfg)
|
||||
h := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)}
|
||||
|
||||
patch := func(value string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
body := fmt.Sprintf(`{"index":0,"value":{"disable-cooling":%s}}`, value)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/key", strings.NewReader(body))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
tc.patch(h, ctx)
|
||||
return rec
|
||||
}
|
||||
|
||||
if rec := patch("false"); rec.Code != http.StatusOK {
|
||||
t.Fatalf("false patch status = %d, want 200; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if override := tc.get(cfg); override == nil || *override {
|
||||
t.Fatalf("disable-cooling = %v, want explicit false", override)
|
||||
}
|
||||
|
||||
if rec := patch("null"); rec.Code != http.StatusOK {
|
||||
t.Fatalf("null patch status = %d, want 200; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if override := tc.get(cfg); override != nil {
|
||||
t.Fatalf("disable-cooling = %v, want inherited value", override)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
1949
backend/internal/api/handlers/management/config_lists.go
Normal file
1949
backend/internal/api/handlers/management/config_lists.go
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -0,0 +1,299 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func writeTestConfigFile(t *testing.T) string {
|
||||
t.Helper()
|
||||
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "config.yaml")
|
||||
if errWrite := os.WriteFile(path, []byte("{}\n"), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write test config: %v", errWrite)
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
func TestDeleteGeminiKey_RequiresBaseURLWhenAPIKeyDuplicated(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
GeminiKey: []config.GeminiKey{
|
||||
{APIKey: "shared-key", BaseURL: "https://a.example.com"},
|
||||
{APIKey: "shared-key", BaseURL: "https://b.example.com"},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodDelete, "/v0/management/gemini-api-key?api-key=shared-key", nil)
|
||||
|
||||
h.DeleteGeminiKey(c)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
if got := len(h.cfg.GeminiKey); got != 2 {
|
||||
t.Fatalf("gemini keys len = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteGeminiKey_DeletesOnlyMatchingBaseURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
GeminiKey: []config.GeminiKey{
|
||||
{APIKey: "shared-key", BaseURL: "https://a.example.com"},
|
||||
{APIKey: "shared-key", BaseURL: "https://b.example.com"},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodDelete, "/v0/management/gemini-api-key?api-key=shared-key&base-url=https://a.example.com", nil)
|
||||
|
||||
h.DeleteGeminiKey(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if got := len(h.cfg.GeminiKey); got != 1 {
|
||||
t.Fatalf("gemini keys len = %d, want 1", got)
|
||||
}
|
||||
if got := h.cfg.GeminiKey[0].BaseURL; got != "https://b.example.com" {
|
||||
t.Fatalf("remaining base-url = %q, want %q", got, "https://b.example.com")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteGeminiStyleKeyRejectsAmbiguousRoutingIdentity(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
interactions bool
|
||||
}{
|
||||
{name: "Gemini"},
|
||||
{name: "Interactions", interactions: true},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
entries := []config.GeminiKey{
|
||||
{APIKey: "shared-key", BaseURL: "https://shared.example.com", Prefix: "team-a"},
|
||||
{APIKey: "shared-key", BaseURL: "https://shared.example.com", Prefix: "team-b"},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
path := "/v0/management/gemini-api-key?api-key=shared-key&base-url=https://shared.example.com"
|
||||
if tc.interactions {
|
||||
cfg.InteractionsKey = entries
|
||||
path = "/v0/management/interactions-api-key?api-key=shared-key&base-url=https://shared.example.com"
|
||||
} else {
|
||||
cfg.GeminiKey = entries
|
||||
}
|
||||
handler := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)}
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodDelete, path, nil)
|
||||
|
||||
if tc.interactions {
|
||||
handler.DeleteInteractionsKey(ctx)
|
||||
} else {
|
||||
handler.DeleteGeminiKey(ctx)
|
||||
}
|
||||
|
||||
if recorder.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusBadRequest, recorder.Body.String())
|
||||
}
|
||||
remaining := cfg.GeminiKey
|
||||
if tc.interactions {
|
||||
remaining = cfg.InteractionsKey
|
||||
}
|
||||
if len(remaining) != 2 {
|
||||
t.Fatalf("remaining credential count = %d, want 2", len(remaining))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchGeminiStyleKeyRoutingIdentity(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
interactions bool
|
||||
firstBase string
|
||||
wantStatus int
|
||||
}{
|
||||
{name: "Gemini unique base URL", firstBase: "https://first.example.com", wantStatus: http.StatusOK},
|
||||
{name: "Gemini ambiguous base URL", firstBase: "https://shared.example.com", wantStatus: http.StatusBadRequest},
|
||||
{name: "Interactions unique base URL", interactions: true, firstBase: "https://first.example.com", wantStatus: http.StatusOK},
|
||||
{name: "Interactions ambiguous base URL", interactions: true, firstBase: "https://shared.example.com", wantStatus: http.StatusBadRequest},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
entries := []config.GeminiKey{
|
||||
{APIKey: "shared-key", BaseURL: tc.firstBase, Prefix: "team-a"},
|
||||
{APIKey: "shared-key", BaseURL: "https://shared.example.com", Prefix: "team-b"},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
path := "/v0/management/gemini-api-key?base-url=https://shared.example.com"
|
||||
if tc.interactions {
|
||||
cfg.InteractionsKey = entries
|
||||
path = "/v0/management/interactions-api-key?base-url=https://shared.example.com"
|
||||
} else {
|
||||
cfg.GeminiKey = entries
|
||||
}
|
||||
handler := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)}
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, path, strings.NewReader(`{"match":"shared-key","value":{"prefix":"updated"}}`))
|
||||
|
||||
if tc.interactions {
|
||||
handler.PatchInteractionsKey(ctx)
|
||||
} else {
|
||||
handler.PatchGeminiKey(ctx)
|
||||
}
|
||||
|
||||
if recorder.Code != tc.wantStatus {
|
||||
t.Fatalf("status = %d, want %d; body=%s", recorder.Code, tc.wantStatus, recorder.Body.String())
|
||||
}
|
||||
remaining := cfg.GeminiKey
|
||||
if tc.interactions {
|
||||
remaining = cfg.InteractionsKey
|
||||
}
|
||||
if tc.wantStatus == http.StatusOK {
|
||||
if remaining[0].Prefix != "team-a" || remaining[1].Prefix != "updated" {
|
||||
t.Fatalf("prefixes = %q, %q; want team-a, updated", remaining[0].Prefix, remaining[1].Prefix)
|
||||
}
|
||||
} else if remaining[0].Prefix != "team-a" || remaining[1].Prefix != "team-b" {
|
||||
t.Fatalf("ambiguous patch changed prefixes to %q, %q", remaining[0].Prefix, remaining[1].Prefix)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteClaudeKey_DeletesEmptyBaseURLWhenExplicitlyProvided(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
ClaudeKey: []config.ClaudeKey{
|
||||
{APIKey: "shared-key", BaseURL: ""},
|
||||
{APIKey: "shared-key", BaseURL: "https://claude.example.com"},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodDelete, "/v0/management/claude-api-key?api-key=shared-key&base-url=", nil)
|
||||
|
||||
h.DeleteClaudeKey(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if got := len(h.cfg.ClaudeKey); got != 1 {
|
||||
t.Fatalf("claude keys len = %d, want 1", got)
|
||||
}
|
||||
if got := h.cfg.ClaudeKey[0].BaseURL; got != "https://claude.example.com" {
|
||||
t.Fatalf("remaining base-url = %q, want %q", got, "https://claude.example.com")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteVertexCompatKey_DeletesOnlyMatchingBaseURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
VertexCompatAPIKey: []config.VertexCompatKey{
|
||||
{APIKey: "shared-key", BaseURL: "https://a.example.com"},
|
||||
{APIKey: "shared-key", BaseURL: "https://b.example.com"},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodDelete, "/v0/management/vertex-api-key?api-key=shared-key&base-url=https://b.example.com", nil)
|
||||
|
||||
h.DeleteVertexCompatKey(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if got := len(h.cfg.VertexCompatAPIKey); got != 1 {
|
||||
t.Fatalf("vertex keys len = %d, want 1", got)
|
||||
}
|
||||
if got := h.cfg.VertexCompatAPIKey[0].BaseURL; got != "https://a.example.com" {
|
||||
t.Fatalf("remaining base-url = %q, want %q", got, "https://a.example.com")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteXAIKey_RequiresBaseURLWhenAPIKeyDuplicated(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
XAIKey: []config.XAIKey{
|
||||
{APIKey: "shared-key", BaseURL: "https://a.example.com"},
|
||||
{APIKey: "shared-key", BaseURL: "https://b.example.com"},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodDelete, "/v0/management/xai-api-key?api-key=shared-key", nil)
|
||||
|
||||
h.DeleteXAIKey(c)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
if got := len(h.cfg.XAIKey); got != 2 {
|
||||
t.Fatalf("xAI keys len = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteCodexKey_RequiresBaseURLWhenAPIKeyDuplicated(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
CodexKey: []config.CodexKey{
|
||||
{APIKey: "shared-key", BaseURL: "https://a.example.com"},
|
||||
{APIKey: "shared-key", BaseURL: "https://b.example.com"},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodDelete, "/v0/management/codex-api-key?api-key=shared-key", nil)
|
||||
|
||||
h.DeleteCodexKey(c)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
if got := len(h.cfg.CodexKey); got != 2 {
|
||||
t.Fatalf("codex keys len = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,67 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestGetOpenAICompatIncludesDisableCooling(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
requestRetry := 0
|
||||
disableCooling := true
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{
|
||||
OpenAICompatibility: []config.OpenAICompatibility{
|
||||
{
|
||||
Name: "Mimo CN",
|
||||
BaseURL: "https://token-plan-cn.xiaomimimo.com/v1",
|
||||
APIKeyEntries: []config.OpenAICompatibilityAPIKey{
|
||||
{APIKey: "test-key"},
|
||||
},
|
||||
Models: []config.OpenAICompatibilityModel{
|
||||
{Name: "mimo-v2.5", Alias: ""},
|
||||
},
|
||||
SupportPromptCacheKey: true,
|
||||
DisableCooling: &disableCooling,
|
||||
RequestRetry: &requestRetry,
|
||||
},
|
||||
},
|
||||
}, nil)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/openai-compatibility", nil)
|
||||
h.GetOpenAICompat(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var body struct {
|
||||
OpenAICompatibility []struct {
|
||||
SupportPromptCacheKey *bool `json:"support-prompt-cache-key"`
|
||||
DisableCooling *bool `json:"disable-cooling"`
|
||||
RequestRetry *int `json:"request-retry"`
|
||||
} `json:"openai-compatibility"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
|
||||
t.Fatalf("failed to decode response: %v", err)
|
||||
}
|
||||
if len(body.OpenAICompatibility) != 1 {
|
||||
t.Fatalf("expected 1 openai-compatibility entry, got %d", len(body.OpenAICompatibility))
|
||||
}
|
||||
if body.OpenAICompatibility[0].SupportPromptCacheKey == nil || !*body.OpenAICompatibility[0].SupportPromptCacheKey {
|
||||
t.Fatalf("expected support-prompt-cache-key to be present and true, got %#v", body.OpenAICompatibility[0].SupportPromptCacheKey)
|
||||
}
|
||||
if body.OpenAICompatibility[0].DisableCooling == nil || !*body.OpenAICompatibility[0].DisableCooling {
|
||||
t.Fatalf("expected disable-cooling to be present and true, got %#v", body.OpenAICompatibility[0].DisableCooling)
|
||||
}
|
||||
if body.OpenAICompatibility[0].RequestRetry == nil || *body.OpenAICompatibility[0].RequestRetry != 0 {
|
||||
t.Fatalf("expected request-retry to be present and 0, got %#v", body.OpenAICompatibility[0].RequestRetry)
|
||||
}
|
||||
}
|
||||
104
backend/internal/api/handlers/management/config_weight_test.go
Normal file
104
backend/internal/api/handlers/management/config_weight_test.go
Normal file
|
|
@ -0,0 +1,104 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestPatchAPIKeyWeightForEveryFamily(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
setup func(*config.Config)
|
||||
patch func(*Handler, *gin.Context)
|
||||
get func(*config.Config) *int
|
||||
}{
|
||||
{name: "gemini", setup: func(cfg *config.Config) { cfg.GeminiKey = []config.GeminiKey{{APIKey: "key"}} }, patch: (*Handler).PatchGeminiKey, get: func(cfg *config.Config) *int { return cfg.GeminiKey[0].Weight }},
|
||||
{name: "interactions", setup: func(cfg *config.Config) { cfg.InteractionsKey = []config.GeminiKey{{APIKey: "key"}} }, patch: (*Handler).PatchInteractionsKey, get: func(cfg *config.Config) *int { return cfg.InteractionsKey[0].Weight }},
|
||||
{name: "claude", setup: func(cfg *config.Config) { cfg.ClaudeKey = []config.ClaudeKey{{APIKey: "key"}} }, patch: (*Handler).PatchClaudeKey, get: func(cfg *config.Config) *int { return cfg.ClaudeKey[0].Weight }},
|
||||
{name: "vertex", setup: func(cfg *config.Config) {
|
||||
cfg.VertexCompatAPIKey = []config.VertexCompatKey{{APIKey: "key", BaseURL: "https://example.com"}}
|
||||
}, patch: (*Handler).PatchVertexCompatKey, get: func(cfg *config.Config) *int { return cfg.VertexCompatAPIKey[0].Weight }},
|
||||
{name: "codex", setup: func(cfg *config.Config) {
|
||||
cfg.CodexKey = []config.CodexKey{{APIKey: "key", BaseURL: "https://example.com"}}
|
||||
}, patch: (*Handler).PatchCodexKey, get: func(cfg *config.Config) *int { return cfg.CodexKey[0].Weight }},
|
||||
{name: "xai", setup: func(cfg *config.Config) {
|
||||
cfg.XAIKey = []config.XAIKey{{APIKey: "key", BaseURL: "https://example.com"}}
|
||||
}, patch: (*Handler).PatchXAIKey, get: func(cfg *config.Config) *int { return cfg.XAIKey[0].Weight }},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
test.setup(cfg)
|
||||
h := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/key", strings.NewReader(`{"index":0,"value":{"weight":7}}`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
test.patch(h, ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if weight := test.get(cfg); weight == nil || *weight != 7 {
|
||||
t.Fatalf("weight = %v, want 7", weight)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchAPIKeyWeightResetAndStrictValidation(t *testing.T) {
|
||||
initial := 5
|
||||
cfg := &config.Config{GeminiKey: []config.GeminiKey{{APIKey: "key", Weight: &initial}}}
|
||||
h := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)}
|
||||
|
||||
patch := func(raw string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
body := fmt.Sprintf(`{"index":0,"value":{"weight":%s}}`, raw)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/gemini-api-key", strings.NewReader(body))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
h.PatchGeminiKey(ctx)
|
||||
return rec
|
||||
}
|
||||
|
||||
for _, invalid := range []string{"1.5", "1000001", "9223372036854775808", `"7"`} {
|
||||
rec := patch(invalid)
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("weight %s status = %d, want 400; body=%s", invalid, rec.Code, rec.Body.String())
|
||||
}
|
||||
if cfg.GeminiKey[0].Weight == nil || *cfg.GeminiKey[0].Weight != initial {
|
||||
t.Fatalf("invalid weight %s changed config", invalid)
|
||||
}
|
||||
}
|
||||
|
||||
if rec := patch("null"); rec.Code != http.StatusOK {
|
||||
t.Fatalf("reset status = %d, want 200; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if cfg.GeminiKey[0].Weight != nil {
|
||||
t.Fatalf("reset weight = %v, want nil default", cfg.GeminiKey[0].Weight)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPutAPIKeyWeightRejectsAboveMaximum(t *testing.T) {
|
||||
h := &Handler{cfg: &config.Config{}, configFilePath: writeTestConfigFile(t)}
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPut, "/v0/management/gemini-api-key", strings.NewReader(`[{"api-key":"key","weight":1000001}]`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
h.PutGeminiKeys(ctx)
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want 400; body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if len(h.cfg.GeminiKey) != 0 {
|
||||
t.Fatal("invalid PUT changed config")
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,57 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestPatchXAIKeyUpdatesExecutionFields(t *testing.T) {
|
||||
disableCooling := false
|
||||
h := &Handler{
|
||||
cfg: &config.Config{XAIKey: []config.XAIKey{{
|
||||
APIKey: "xai-key",
|
||||
Priority: 1,
|
||||
BaseURL: "https://api.x.ai/v1",
|
||||
Websockets: true,
|
||||
DisableCooling: &disableCooling,
|
||||
}}},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/xai-api-key", strings.NewReader(`{
|
||||
"index": 0,
|
||||
"value": {
|
||||
"priority": 7,
|
||||
"websockets": false,
|
||||
"disable-cooling": true,
|
||||
"request-retry": 0
|
||||
}
|
||||
}`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h.PatchXAIKey(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
entry := h.cfg.XAIKey[0]
|
||||
if entry.Priority != 7 {
|
||||
t.Fatalf("priority = %d, want 7", entry.Priority)
|
||||
}
|
||||
if entry.Websockets {
|
||||
t.Fatal("websockets = true, want false")
|
||||
}
|
||||
if entry.DisableCooling == nil || !*entry.DisableCooling {
|
||||
t.Fatalf("disable-cooling = %v, want true", entry.DisableCooling)
|
||||
}
|
||||
if entry.RequestRetry == nil || *entry.RequestRetry != 0 {
|
||||
t.Fatalf("request-retry = %v, want 0", entry.RequestRetry)
|
||||
}
|
||||
}
|
||||
459
backend/internal/api/handlers/management/handler.go
Normal file
459
backend/internal/api/handlers/management/handler.go
Normal file
|
|
@ -0,0 +1,459 @@
|
|||
// Package management provides the management API handlers and middleware
|
||||
// for configuring the server and managing auth files.
|
||||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/subtle"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/buildinfo"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginstore"
|
||||
sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
)
|
||||
|
||||
type attemptInfo struct {
|
||||
count int
|
||||
blockedUntil time.Time
|
||||
lastActivity time.Time // track last activity for cleanup
|
||||
}
|
||||
|
||||
// attemptCleanupInterval controls how often stale IP entries are purged
|
||||
const attemptCleanupInterval = 1 * time.Hour
|
||||
|
||||
// attemptMaxIdleTime controls how long an IP can be idle before cleanup
|
||||
const attemptMaxIdleTime = 2 * time.Hour
|
||||
|
||||
// Handler aggregates config reference, persistence path and helpers.
|
||||
type Handler struct {
|
||||
cfg *config.Config
|
||||
configFilePath string
|
||||
mu sync.Mutex
|
||||
reloadMu sync.Mutex
|
||||
reloadGeneration uint64
|
||||
appliedReloadGeneration uint64
|
||||
attemptsMu sync.Mutex
|
||||
failedAttempts map[string]*attemptInfo // keyed by client IP
|
||||
authManager *coreauth.Manager
|
||||
tokenStore coreauth.Store
|
||||
localPassword string
|
||||
allowRemoteOverride bool
|
||||
envSecret string
|
||||
logDir string
|
||||
postAuthHook coreauth.PostAuthHook
|
||||
postAuthPersistHook coreauth.PostAuthHook
|
||||
pluginHost *pluginhost.Host
|
||||
configReloadHook func(context.Context, *config.Config)
|
||||
pluginStoreRegistryURL string
|
||||
pluginStoreHTTPClient pluginstore.HTTPDoer
|
||||
pluginReleaseCacheMu sync.Mutex
|
||||
pluginReleaseCache map[string]pluginReleaseCacheEntry
|
||||
}
|
||||
|
||||
type configReloadSnapshot struct {
|
||||
cfg *config.Config
|
||||
generation uint64
|
||||
}
|
||||
|
||||
// NewHandler creates a new management handler instance.
|
||||
func NewHandler(cfg *config.Config, configFilePath string, manager *coreauth.Manager) *Handler {
|
||||
envSecret, _ := os.LookupEnv("MANAGEMENT_PASSWORD")
|
||||
envSecret = strings.TrimSpace(envSecret)
|
||||
|
||||
h := &Handler{
|
||||
cfg: cfg,
|
||||
configFilePath: configFilePath,
|
||||
failedAttempts: make(map[string]*attemptInfo),
|
||||
authManager: manager,
|
||||
tokenStore: sdkAuth.GetTokenStore(),
|
||||
allowRemoteOverride: envSecret != "",
|
||||
envSecret: envSecret,
|
||||
}
|
||||
h.startAttemptCleanup()
|
||||
return h
|
||||
}
|
||||
|
||||
// startAttemptCleanup launches a background goroutine that periodically
|
||||
// removes stale IP entries from failedAttempts to prevent memory leaks.
|
||||
func (h *Handler) startAttemptCleanup() {
|
||||
go func() {
|
||||
ticker := time.NewTicker(attemptCleanupInterval)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
h.purgeStaleAttempts()
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// purgeStaleAttempts removes IP entries that have been idle beyond attemptMaxIdleTime
|
||||
// and whose ban (if any) has expired.
|
||||
func (h *Handler) purgeStaleAttempts() {
|
||||
now := time.Now()
|
||||
h.attemptsMu.Lock()
|
||||
defer h.attemptsMu.Unlock()
|
||||
for ip, ai := range h.failedAttempts {
|
||||
// Skip if still banned
|
||||
if !ai.blockedUntil.IsZero() && now.Before(ai.blockedUntil) {
|
||||
continue
|
||||
}
|
||||
// Remove if idle too long
|
||||
if now.Sub(ai.lastActivity) > attemptMaxIdleTime {
|
||||
delete(h.failedAttempts, ip)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// NewHandler creates a new management handler instance.
|
||||
func NewHandlerWithoutConfigFilePath(cfg *config.Config, manager *coreauth.Manager) *Handler {
|
||||
return NewHandler(cfg, "", manager)
|
||||
}
|
||||
|
||||
// SetConfig updates the in-memory config reference when the server hot-reloads.
|
||||
func (h *Handler) SetConfig(cfg *config.Config) {
|
||||
if h == nil {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
h.cfg = cfg
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetAuthManager updates the auth manager reference used by management endpoints.
|
||||
func (h *Handler) SetAuthManager(manager *coreauth.Manager) {
|
||||
if h == nil {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
h.authManager = manager
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetPluginHost updates the plugin host used by plugin-backed management endpoints.
|
||||
func (h *Handler) SetPluginHost(host *pluginhost.Host) {
|
||||
if h == nil {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
h.pluginHost = host
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetConfigReloadHook updates the callback used after management saves config changes.
|
||||
func (h *Handler) SetConfigReloadHook(hook func(context.Context, *config.Config)) {
|
||||
if h == nil {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
h.configReloadHook = hook
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
// reloadSnapshotConfigLocked clones the runtime config and assigns a reload generation.
|
||||
// Callers must hold h.mu.
|
||||
func (h *Handler) reloadSnapshotConfigLocked() configReloadSnapshot {
|
||||
if h == nil || h.cfg == nil {
|
||||
return configReloadSnapshot{}
|
||||
}
|
||||
h.reloadGeneration++
|
||||
return configReloadSnapshot{
|
||||
cfg: h.cfg.CloneForRuntime(),
|
||||
generation: h.reloadGeneration,
|
||||
}
|
||||
}
|
||||
|
||||
// saveConfigAndSnapshotLocked saves h.cfg and returns a full runtime config snapshot.
|
||||
// Callers must hold h.mu.
|
||||
func (h *Handler) saveConfigAndSnapshotLocked(c *gin.Context) (configReloadSnapshot, bool) {
|
||||
if errSave := config.SaveConfigPreserveComments(h.configFilePath, h.cfg); errSave != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to save config: %v", errSave)})
|
||||
return configReloadSnapshot{}, false
|
||||
}
|
||||
return h.reloadSnapshotConfigLocked(), true
|
||||
}
|
||||
|
||||
// reloadConfigAfterManagementSave reloads from an independent config snapshot.
|
||||
// Callers must pass a full Config clone captured immediately after a successful save.
|
||||
func (h *Handler) reloadConfigAfterManagementSave(ctx context.Context, snapshot configReloadSnapshot) {
|
||||
if h == nil || snapshot.cfg == nil || snapshot.generation == 0 {
|
||||
return
|
||||
}
|
||||
h.reloadMu.Lock()
|
||||
defer h.reloadMu.Unlock()
|
||||
|
||||
h.mu.Lock()
|
||||
if snapshot.generation < h.appliedReloadGeneration {
|
||||
h.mu.Unlock()
|
||||
return
|
||||
}
|
||||
hook := h.configReloadHook
|
||||
host := h.pluginHost
|
||||
h.mu.Unlock()
|
||||
if hook != nil {
|
||||
hook(ctx, snapshot.cfg)
|
||||
} else if host != nil {
|
||||
host.ApplyConfig(ctx, snapshot.cfg)
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
if snapshot.generation > h.appliedReloadGeneration {
|
||||
h.appliedReloadGeneration = snapshot.generation
|
||||
}
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
// reloadConfigAfterManagementSaveAsync reloads from an independent config snapshot.
|
||||
// Callers must pass a full Config clone captured immediately after a successful save.
|
||||
func (h *Handler) reloadConfigAfterManagementSaveAsync(ctx context.Context, snapshot configReloadSnapshot) {
|
||||
if h == nil || snapshot.cfg == nil || snapshot.generation == 0 {
|
||||
return
|
||||
}
|
||||
reloadCtx := context.Background()
|
||||
if ctx != nil {
|
||||
reloadCtx = context.WithoutCancel(ctx)
|
||||
}
|
||||
go func() {
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
log.WithField("panic", recovered).Error("management: async config reload panicked")
|
||||
}
|
||||
}()
|
||||
h.reloadConfigAfterManagementSave(reloadCtx, snapshot)
|
||||
}()
|
||||
}
|
||||
|
||||
// SetLocalPassword configures the runtime-local password accepted for localhost requests.
|
||||
func (h *Handler) SetLocalPassword(password string) { h.localPassword = password }
|
||||
|
||||
// SetLogDirectory updates the directory where main.log should be looked up.
|
||||
func (h *Handler) SetLogDirectory(dir string) {
|
||||
if dir == "" {
|
||||
return
|
||||
}
|
||||
if !filepath.IsAbs(dir) {
|
||||
if abs, err := filepath.Abs(dir); err == nil {
|
||||
dir = abs
|
||||
}
|
||||
}
|
||||
h.logDir = dir
|
||||
}
|
||||
|
||||
// SetPostAuthHook registers a hook to be called after auth record creation but before persistence.
|
||||
func (h *Handler) SetPostAuthHook(hook coreauth.PostAuthHook) {
|
||||
h.postAuthHook = hook
|
||||
}
|
||||
|
||||
// SetPostAuthPersistHook registers a hook to be called after auth persistence.
|
||||
func (h *Handler) SetPostAuthPersistHook(hook coreauth.PostAuthHook) {
|
||||
h.postAuthPersistHook = hook
|
||||
}
|
||||
|
||||
// Middleware enforces access control for management endpoints.
|
||||
// All requests (local and remote) require a valid management key.
|
||||
// Additionally, remote access requires allow-remote-management=true.
|
||||
func (h *Handler) Middleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.Header("X-CPA-VERSION", buildinfo.Version)
|
||||
c.Header("X-CPA-COMMIT", buildinfo.Commit)
|
||||
c.Header("X-CPA-BUILD-DATE", buildinfo.BuildDate)
|
||||
c.Header("X-CPA-SUPPORT-PLUGIN", pluginhost.SupportPluginHeaderValue())
|
||||
|
||||
clientIP := c.ClientIP()
|
||||
localClient := clientIP == "127.0.0.1" || clientIP == "::1"
|
||||
|
||||
// Accept either Authorization: Bearer <key> or X-Management-Key
|
||||
var provided string
|
||||
if ah := c.GetHeader("Authorization"); ah != "" {
|
||||
parts := strings.SplitN(ah, " ", 2)
|
||||
if len(parts) == 2 && strings.ToLower(parts[0]) == "bearer" {
|
||||
provided = parts[1]
|
||||
} else {
|
||||
provided = ah
|
||||
}
|
||||
}
|
||||
if provided == "" {
|
||||
provided = c.GetHeader("X-Management-Key")
|
||||
}
|
||||
|
||||
allowed, statusCode, errMsg := h.AuthenticateManagementKey(clientIP, localClient, provided)
|
||||
if !allowed {
|
||||
c.AbortWithStatusJSON(statusCode, gin.H{"error": errMsg})
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// AuthenticateManagementKey verifies the provided management key for the given client.
|
||||
// It mirrors the behaviour of Middleware() so non-HTTP callers can reuse the same logic.
|
||||
func (h *Handler) AuthenticateManagementKey(clientIP string, localClient bool, provided string) (bool, int, string) {
|
||||
const maxFailures = 5
|
||||
const banDuration = 30 * time.Minute
|
||||
|
||||
if h == nil {
|
||||
return false, http.StatusForbidden, "remote management disabled"
|
||||
}
|
||||
|
||||
cfg := h.cfg
|
||||
var (
|
||||
allowRemote bool
|
||||
secretHash string
|
||||
)
|
||||
if cfg != nil {
|
||||
allowRemote = cfg.RemoteManagement.AllowRemote
|
||||
secretHash = cfg.RemoteManagement.SecretKey
|
||||
}
|
||||
if h.allowRemoteOverride {
|
||||
allowRemote = true
|
||||
}
|
||||
envSecret := h.envSecret
|
||||
|
||||
now := time.Now()
|
||||
h.attemptsMu.Lock()
|
||||
ai := h.failedAttempts[clientIP]
|
||||
if ai != nil && !ai.blockedUntil.IsZero() {
|
||||
if now.Before(ai.blockedUntil) {
|
||||
remaining := ai.blockedUntil.Sub(now).Round(time.Second)
|
||||
h.attemptsMu.Unlock()
|
||||
return false, http.StatusForbidden, fmt.Sprintf("IP banned due to too many failed attempts. Try again in %s", remaining)
|
||||
}
|
||||
// Ban expired, reset state
|
||||
ai.blockedUntil = time.Time{}
|
||||
ai.count = 0
|
||||
}
|
||||
h.attemptsMu.Unlock()
|
||||
|
||||
if !localClient && !allowRemote {
|
||||
return false, http.StatusForbidden, "remote management disabled"
|
||||
}
|
||||
|
||||
fail := func() {
|
||||
h.attemptsMu.Lock()
|
||||
aip := h.failedAttempts[clientIP]
|
||||
if aip == nil {
|
||||
aip = &attemptInfo{}
|
||||
h.failedAttempts[clientIP] = aip
|
||||
}
|
||||
aip.count++
|
||||
aip.lastActivity = time.Now()
|
||||
if aip.count >= maxFailures {
|
||||
aip.blockedUntil = time.Now().Add(banDuration)
|
||||
aip.count = 0
|
||||
}
|
||||
h.attemptsMu.Unlock()
|
||||
}
|
||||
|
||||
reset := func() {
|
||||
h.attemptsMu.Lock()
|
||||
if ai := h.failedAttempts[clientIP]; ai != nil {
|
||||
ai.count = 0
|
||||
ai.blockedUntil = time.Time{}
|
||||
}
|
||||
h.attemptsMu.Unlock()
|
||||
}
|
||||
|
||||
if secretHash == "" && envSecret == "" {
|
||||
return false, http.StatusForbidden, "remote management key not set"
|
||||
}
|
||||
|
||||
if provided == "" {
|
||||
fail()
|
||||
return false, http.StatusUnauthorized, "missing management key"
|
||||
}
|
||||
|
||||
if localClient {
|
||||
if lp := h.localPassword; lp != "" {
|
||||
if subtle.ConstantTimeCompare([]byte(provided), []byte(lp)) == 1 {
|
||||
reset()
|
||||
return true, 0, ""
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if envSecret != "" && subtle.ConstantTimeCompare([]byte(provided), []byte(envSecret)) == 1 {
|
||||
reset()
|
||||
return true, 0, ""
|
||||
}
|
||||
|
||||
if secretHash == "" || bcrypt.CompareHashAndPassword([]byte(secretHash), []byte(provided)) != nil {
|
||||
fail()
|
||||
return false, http.StatusUnauthorized, "invalid management key"
|
||||
}
|
||||
|
||||
reset()
|
||||
|
||||
return true, 0, ""
|
||||
}
|
||||
|
||||
// persist saves the current in-memory config to disk.
|
||||
func (h *Handler) persist(c *gin.Context) bool {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
return h.persistLocked(c)
|
||||
}
|
||||
|
||||
// persistLocked saves the current in-memory config to disk.
|
||||
// It expects the caller to hold h.mu.
|
||||
func (h *Handler) persistLocked(c *gin.Context) bool {
|
||||
// Preserve comments when writing
|
||||
if err := config.SaveConfigPreserveComments(h.configFilePath, h.cfg); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to save config: %v", err)})
|
||||
return false
|
||||
}
|
||||
snapshot := h.reloadSnapshotConfigLocked()
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||||
var reqCtx context.Context
|
||||
if c != nil && c.Request != nil {
|
||||
reqCtx = c.Request.Context()
|
||||
}
|
||||
h.reloadConfigAfterManagementSaveAsync(reqCtx, snapshot)
|
||||
return true
|
||||
}
|
||||
|
||||
// Helper methods for simple types
|
||||
func (h *Handler) updateBoolField(c *gin.Context, set func(bool)) {
|
||||
var body struct {
|
||||
Value *bool `json:"value"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil || body.Value == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
set(*body.Value)
|
||||
h.persist(c)
|
||||
}
|
||||
|
||||
func (h *Handler) updateIntField(c *gin.Context, set func(int)) {
|
||||
var body struct {
|
||||
Value *int `json:"value"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil || body.Value == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
set(*body.Value)
|
||||
h.persist(c)
|
||||
}
|
||||
|
||||
func (h *Handler) updateStringField(c *gin.Context, set func(string)) {
|
||||
var body struct {
|
||||
Value *string `json:"value"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil || body.Value == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
set(*body.Value)
|
||||
h.persist(c)
|
||||
}
|
||||
88
backend/internal/api/handlers/management/handler_test.go
Normal file
88
backend/internal/api/handlers/management/handler_test.go
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"
|
||||
)
|
||||
|
||||
func TestAuthenticateManagementKey_LocalhostIPBan_BlocksCorrectKeyDuringBan(t *testing.T) {
|
||||
h := &Handler{
|
||||
cfg: &config.Config{},
|
||||
failedAttempts: make(map[string]*attemptInfo),
|
||||
envSecret: "test-secret",
|
||||
}
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
allowed, statusCode, errMsg := h.AuthenticateManagementKey("127.0.0.1", true, "wrong-secret")
|
||||
if allowed {
|
||||
t.Fatalf("expected auth to be denied at attempt %d", i+1)
|
||||
}
|
||||
if statusCode != http.StatusUnauthorized || errMsg != "invalid management key" {
|
||||
t.Fatalf("unexpected auth failure at attempt %d: status=%d msg=%q", i+1, statusCode, errMsg)
|
||||
}
|
||||
}
|
||||
|
||||
allowed, statusCode, errMsg := h.AuthenticateManagementKey("127.0.0.1", true, "test-secret")
|
||||
if allowed {
|
||||
t.Fatalf("expected correct key to be denied while banned")
|
||||
}
|
||||
if statusCode != http.StatusForbidden {
|
||||
t.Fatalf("expected forbidden status while banned, got %d", statusCode)
|
||||
}
|
||||
if !strings.HasPrefix(errMsg, "IP banned due to too many failed attempts. Try again in") {
|
||||
t.Fatalf("unexpected banned message: %q", errMsg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMiddlewareSetsSupportPluginHeader(t *testing.T) {
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{},
|
||||
failedAttempts: make(map[string]*attemptInfo),
|
||||
envSecret: "test-secret",
|
||||
}
|
||||
middleware := h.Middleware()
|
||||
|
||||
t.Run("invalid key", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v0/management/config", nil)
|
||||
c.Request.RemoteAddr = "127.0.0.1:12345"
|
||||
c.Request.Header.Set("X-Management-Key", "wrong-secret")
|
||||
|
||||
middleware(c)
|
||||
|
||||
if rec.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("status = %d, want %d", rec.Code, http.StatusUnauthorized)
|
||||
}
|
||||
if got := rec.Header().Get("X-CPA-SUPPORT-PLUGIN"); got != pluginhost.SupportPluginHeaderValue() {
|
||||
t.Fatalf("X-CPA-SUPPORT-PLUGIN = %q, want %q", got, pluginhost.SupportPluginHeaderValue())
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("valid key", func(t *testing.T) {
|
||||
engine := gin.New()
|
||||
engine.GET("/v0/management/config", middleware, func(c *gin.Context) {
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/v0/management/config", nil)
|
||||
req.RemoteAddr = "127.0.0.1:12345"
|
||||
req.Header.Set("X-Management-Key", "test-secret")
|
||||
engine.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK)
|
||||
}
|
||||
if got := rec.Header().Get("X-CPA-SUPPORT-PLUGIN"); got != pluginhost.SupportPluginHeaderValue() {
|
||||
t.Fatalf("X-CPA-SUPPORT-PLUGIN = %q, want %q", got, pluginhost.SupportPluginHeaderValue())
|
||||
}
|
||||
})
|
||||
}
|
||||
1310
backend/internal/api/handlers/management/logs.go
Normal file
1310
backend/internal/api/handlers/management/logs.go
Normal file
File diff suppressed because it is too large
Load diff
736
backend/internal/api/handlers/management/logs_test.go
Normal file
736
backend/internal/api/handlers/management/logs_test.go
Normal file
|
|
@ -0,0 +1,736 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestDecodeLogCursorRejectsUnsafeFiles(t *testing.T) {
|
||||
unsafeNames := []string{
|
||||
"",
|
||||
".",
|
||||
"..",
|
||||
"../secret",
|
||||
"nested/main.log",
|
||||
`nested\main.log`,
|
||||
"error.log",
|
||||
}
|
||||
|
||||
for _, name := range unsafeNames {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
raw := mustEncodeRawCursor(t, logCursor{
|
||||
Version: logCursorVersion,
|
||||
File: name,
|
||||
Fingerprint: "fingerprint",
|
||||
})
|
||||
if _, err := decodeLogCursor(raw); err == nil {
|
||||
t.Fatalf("decodeLogCursor(%q) succeeded, want error", name)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
for _, name := range []string{defaultLogFileName, defaultLogFileName + ".1", "main-2026-06-15T10-00-00.log"} {
|
||||
t.Run("allowed_"+name, func(t *testing.T) {
|
||||
raw := mustEncodeRawCursor(t, logCursor{
|
||||
Version: logCursorVersion,
|
||||
File: name,
|
||||
Fingerprint: "fingerprint",
|
||||
})
|
||||
if _, err := decodeLogCursor(raw); err != nil {
|
||||
t.Fatalf("decodeLogCursor(%q) error = %v", name, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogCursorRoundTripOmitsAbsolutePath(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, defaultLogFileName)
|
||||
if err := os.WriteFile(path, []byte("line one\nline two\n"), 0o644); err != nil {
|
||||
t.Fatalf("write log file: %v", err)
|
||||
}
|
||||
|
||||
boundary, errBoundary := completeLogBoundary(path)
|
||||
if errBoundary != nil {
|
||||
t.Fatalf("completeLogBoundary() error = %v", errBoundary)
|
||||
}
|
||||
raw, errCursor := newLogCursor(path, boundary, 123)
|
||||
if errCursor != nil {
|
||||
t.Fatalf("newLogCursor() error = %v", errCursor)
|
||||
}
|
||||
decoded, errDecode := decodeLogCursor(raw)
|
||||
if errDecode != nil {
|
||||
t.Fatalf("decodeLogCursor() error = %v", errDecode)
|
||||
}
|
||||
if decoded.File != defaultLogFileName {
|
||||
t.Fatalf("cursor file = %q, want %q", decoded.File, defaultLogFileName)
|
||||
}
|
||||
if decoded.Offset != boundary {
|
||||
t.Fatalf("cursor offset = %d, want %d", decoded.Offset, boundary)
|
||||
}
|
||||
if decoded.LatestTimestamp != 123 {
|
||||
t.Fatalf("cursor latest timestamp = %d, want 123", decoded.LatestTimestamp)
|
||||
}
|
||||
if strings.Contains(raw, dir) {
|
||||
t.Fatalf("encoded cursor contains log directory %q: %q", dir, raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadCompleteLogLinesSkipsTrailingPartial(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, defaultLogFileName)
|
||||
initial := "first\nsecond\r\npartial"
|
||||
if err := os.WriteFile(path, []byte(initial), 0o644); err != nil {
|
||||
t.Fatalf("write log file: %v", err)
|
||||
}
|
||||
|
||||
read, errRead := readCompleteLogLines(path, 0, -1, 0)
|
||||
if errRead != nil {
|
||||
t.Fatalf("readCompleteLogLines() error = %v", errRead)
|
||||
}
|
||||
wantLines := []string{"first", "second"}
|
||||
if !reflect.DeepEqual(read.lines, wantLines) {
|
||||
t.Fatalf("lines = %#v, want %#v", read.lines, wantLines)
|
||||
}
|
||||
wantOffset := int64(len("first\nsecond\r\n"))
|
||||
if read.endOffset != wantOffset {
|
||||
t.Fatalf("endOffset = %d, want %d", read.endOffset, wantOffset)
|
||||
}
|
||||
|
||||
file, errOpen := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0)
|
||||
if errOpen != nil {
|
||||
t.Fatalf("open log file: %v", errOpen)
|
||||
}
|
||||
if _, errWrite := file.WriteString("\n"); errWrite != nil {
|
||||
_ = file.Close()
|
||||
t.Fatalf("append newline: %v", errWrite)
|
||||
}
|
||||
if errClose := file.Close(); errClose != nil {
|
||||
t.Fatalf("close log file: %v", errClose)
|
||||
}
|
||||
|
||||
next, errNext := readCompleteLogLines(path, read.endOffset, -1, 0)
|
||||
if errNext != nil {
|
||||
t.Fatalf("readCompleteLogLines() after append error = %v", errNext)
|
||||
}
|
||||
if !reflect.DeepEqual(next.lines, []string{"partial"}) {
|
||||
t.Fatalf("next lines = %#v, want partial", next.lines)
|
||||
}
|
||||
if next.endOffset != int64(len(initial)+1) {
|
||||
t.Fatalf("next endOffset = %d, want %d", next.endOffset, len(initial)+1)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsTailLimitReturnsRecentLinesWithCursor(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
lines := []string{
|
||||
"[2026-06-15 10:00:00] first",
|
||||
"[2026-06-15 10:00:01] second",
|
||||
"[2026-06-15 10:00:02] third",
|
||||
"[2026-06-15 10:00:03] fourth",
|
||||
}
|
||||
writeMainLog(t, dir, strings.Join(lines, "\n")+"\n")
|
||||
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=2")
|
||||
wantLines := []string{lines[2], lines[3]}
|
||||
if !reflect.DeepEqual(resp.Lines, wantLines) {
|
||||
t.Fatalf("lines = %#v, want %#v", resp.Lines, wantLines)
|
||||
}
|
||||
if resp.LineCount != len(wantLines) {
|
||||
t.Fatalf("line-count = %d, want returned line count %d", resp.LineCount, len(wantLines))
|
||||
}
|
||||
if resp.NextCursor == "" {
|
||||
t.Fatal("next-cursor is empty")
|
||||
}
|
||||
wantLatest := time.Date(2026, 6, 15, 10, 0, 3, 0, time.Local).Unix()
|
||||
if resp.LatestTimestamp != wantLatest {
|
||||
t.Fatalf("latest-timestamp = %d, want %d", resp.LatestTimestamp, wantLatest)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsTailLimitDoesNotScanOlderFilesForLineCount(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
rotatedPath := filepath.Join(dir, defaultLogFileName+".1")
|
||||
if err := os.WriteFile(rotatedPath, []byte(strings.Repeat("x", logScannerMaxBuffer+1)+"\n"), 0o644); err != nil {
|
||||
t.Fatalf("write rotated log: %v", err)
|
||||
}
|
||||
writeMainLog(t, dir, "[2026-06-15 10:00:00] current\n")
|
||||
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=1")
|
||||
wantLines := []string{"[2026-06-15 10:00:00] current"}
|
||||
if !reflect.DeepEqual(resp.Lines, wantLines) {
|
||||
t.Fatalf("lines = %#v, want %#v", resp.Lines, wantLines)
|
||||
}
|
||||
if resp.LineCount != len(wantLines) {
|
||||
t.Fatalf("line-count = %d, want returned line count %d", resp.LineCount, len(wantLines))
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsNoLimitKeepsFullScanBehavior(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
writeMainLog(t, dir, "complete\npartial")
|
||||
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs")
|
||||
wantLines := []string{"complete", "partial"}
|
||||
if !reflect.DeepEqual(resp.Lines, wantLines) {
|
||||
t.Fatalf("lines = %#v, want %#v", resp.Lines, wantLines)
|
||||
}
|
||||
if resp.LineCount != 2 {
|
||||
t.Fatalf("line-count = %d, want full scan count 2", resp.LineCount)
|
||||
}
|
||||
if resp.NextCursor == "" {
|
||||
t.Fatal("next-cursor is empty")
|
||||
}
|
||||
cursor, errCursor := decodeLogCursor(resp.NextCursor)
|
||||
if errCursor != nil {
|
||||
t.Fatalf("decode next-cursor: %v", errCursor)
|
||||
}
|
||||
if cursor.Offset != int64(len("complete\n")) {
|
||||
t.Fatalf("cursor offset = %d, want complete-line boundary", cursor.Offset)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsAfterKeepsTimestampScanAndReturnsCursor(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
lines := []string{
|
||||
"[2026-06-15 10:00:00] first",
|
||||
"[2026-06-15 10:00:01] second",
|
||||
"[2026-06-15 10:00:02] third",
|
||||
}
|
||||
writeMainLog(t, dir, strings.Join(lines, "\n")+"\n")
|
||||
|
||||
cutoff := time.Date(2026, 6, 15, 10, 0, 0, 0, time.Local).Unix()
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?after="+strconv.FormatInt(cutoff, 10))
|
||||
wantLines := []string{lines[1], lines[2]}
|
||||
if !reflect.DeepEqual(resp.Lines, wantLines) {
|
||||
t.Fatalf("lines = %#v, want %#v", resp.Lines, wantLines)
|
||||
}
|
||||
if resp.LineCount != 3 {
|
||||
t.Fatalf("line-count = %d, want full scan count 3", resp.LineCount)
|
||||
}
|
||||
if resp.NextCursor == "" {
|
||||
t.Fatal("next-cursor is empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsCursorReturnsOnlyNewCompleteLines(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
lines := []string{
|
||||
"[2026-06-15 10:00:00] first",
|
||||
"[2026-06-15 10:00:01] second",
|
||||
"[2026-06-15 10:00:02] third",
|
||||
}
|
||||
writeMainLog(t, dir, strings.Join(lines, "\n")+"\n")
|
||||
initial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=2")
|
||||
if initial.NextCursor == "" {
|
||||
t.Fatal("initial next-cursor is empty")
|
||||
}
|
||||
|
||||
appendMainLog(t, dir, "[2026-06-15 10:00:03] fourth\n")
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=10")
|
||||
wantLines := []string{"[2026-06-15 10:00:03] fourth"}
|
||||
if !reflect.DeepEqual(resp.Lines, wantLines) {
|
||||
t.Fatalf("lines = %#v, want %#v", resp.Lines, wantLines)
|
||||
}
|
||||
if resp.LineCount != 1 {
|
||||
t.Fatalf("line-count = %d, want 1", resp.LineCount)
|
||||
}
|
||||
if resp.CursorReset {
|
||||
t.Fatal("cursor-reset = true, want false")
|
||||
}
|
||||
wantLatest := time.Date(2026, 6, 15, 10, 0, 3, 0, time.Local).Unix()
|
||||
if resp.LatestTimestamp != wantLatest {
|
||||
t.Fatalf("latest-timestamp = %d, want %d", resp.LatestTimestamp, wantLatest)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsCursorRejectsOversizedLine(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
writeMainLog(t, dir, "[2026-06-15 10:00:00] first\n")
|
||||
initial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=1")
|
||||
if initial.NextCursor == "" {
|
||||
t.Fatal("initial next-cursor is empty")
|
||||
}
|
||||
|
||||
appendMainLog(t, dir, strings.Repeat("x", logScannerMaxBuffer+1)+"\n")
|
||||
status, body := performGetLogsRaw(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=1")
|
||||
if status != http.StatusInternalServerError {
|
||||
t.Fatalf("status = %d, want %d", status, http.StatusInternalServerError)
|
||||
}
|
||||
if !strings.Contains(body, "log line exceeds") {
|
||||
t.Fatalf("body = %s, want oversized line error", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsCursorNoNewLinesKeepsCursorStable(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
line := "[2026-06-15 10:00:00] first"
|
||||
writeMainLog(t, dir, line+"\n")
|
||||
initial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=1")
|
||||
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=10")
|
||||
if len(resp.Lines) != 0 {
|
||||
t.Fatalf("lines = %#v, want empty", resp.Lines)
|
||||
}
|
||||
if resp.LineCount != 0 {
|
||||
t.Fatalf("line-count = %d, want 0", resp.LineCount)
|
||||
}
|
||||
if resp.NextCursor != initial.NextCursor {
|
||||
t.Fatalf("next-cursor changed with no complete lines")
|
||||
}
|
||||
if resp.LatestTimestamp != initial.LatestTimestamp {
|
||||
t.Fatalf("latest-timestamp = %d, want %d", resp.LatestTimestamp, initial.LatestTimestamp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsCursorDoesNotAdvancePastTrailingPartial(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
line := "[2026-06-15 10:00:00] first"
|
||||
writeMainLog(t, dir, line+"\n")
|
||||
initial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=1")
|
||||
|
||||
appendMainLog(t, dir, "partial")
|
||||
partial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=10")
|
||||
if len(partial.Lines) != 0 {
|
||||
t.Fatalf("partial lines = %#v, want empty", partial.Lines)
|
||||
}
|
||||
if partial.NextCursor != initial.NextCursor {
|
||||
t.Fatalf("cursor advanced past partial line")
|
||||
}
|
||||
|
||||
appendMainLog(t, dir, "\n")
|
||||
complete := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=10")
|
||||
if !reflect.DeepEqual(complete.Lines, []string{"partial"}) {
|
||||
t.Fatalf("complete lines = %#v, want partial", complete.Lines)
|
||||
}
|
||||
if complete.LatestTimestamp != initial.LatestTimestamp {
|
||||
t.Fatalf("latest-timestamp = %d, want %d", complete.LatestTimestamp, initial.LatestTimestamp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsCursorResetAfterTruncateTailsLimit(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
lines := []string{
|
||||
"[2026-06-15 10:00:00] first",
|
||||
"[2026-06-15 10:00:01] second",
|
||||
"[2026-06-15 10:00:02] third",
|
||||
}
|
||||
writeMainLog(t, dir, strings.Join(lines, "\n")+"\n")
|
||||
initial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=3")
|
||||
|
||||
resetLine := "[2026-06-15 10:00:03] reset"
|
||||
writeMainLog(t, dir, resetLine+"\n")
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=1")
|
||||
if !resp.CursorReset {
|
||||
t.Fatal("cursor-reset = false, want true")
|
||||
}
|
||||
if !reflect.DeepEqual(resp.Lines, []string{resetLine}) {
|
||||
t.Fatalf("lines = %#v, want reset tail", resp.Lines)
|
||||
}
|
||||
if resp.LineCount != 1 {
|
||||
t.Fatalf("line-count = %d, want 1", resp.LineCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsCursorReadsAcrossRotation(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
line1 := "[2026-06-15 10:00:00] first"
|
||||
line2 := "[2026-06-15 10:00:01] second"
|
||||
line3 := "[2026-06-15 10:00:02] third"
|
||||
writeMainLog(t, dir, line1+"\n")
|
||||
initial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=1")
|
||||
|
||||
appendMainLog(t, dir, line2+"\n")
|
||||
if err := os.Rename(filepath.Join(dir, defaultLogFileName), filepath.Join(dir, defaultLogFileName+".1")); err != nil {
|
||||
t.Fatalf("rotate main log: %v", err)
|
||||
}
|
||||
writeMainLog(t, dir, line3+"\n")
|
||||
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=10")
|
||||
wantLines := []string{line2, line3}
|
||||
if !reflect.DeepEqual(resp.Lines, wantLines) {
|
||||
t.Fatalf("lines = %#v, want %#v", resp.Lines, wantLines)
|
||||
}
|
||||
if resp.CursorReset {
|
||||
t.Fatal("cursor-reset = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsCursorReadsRotatedFileWhenNewMainIsSmaller(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
line1 := "[2026-06-15 10:00:00] first line with enough bytes"
|
||||
line2 := "[2026-06-15 10:00:01] second"
|
||||
line3 := "new"
|
||||
writeMainLog(t, dir, line1+"\n")
|
||||
initial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=1")
|
||||
|
||||
appendMainLog(t, dir, line2+"\n")
|
||||
if err := os.Rename(filepath.Join(dir, defaultLogFileName), filepath.Join(dir, defaultLogFileName+".1")); err != nil {
|
||||
t.Fatalf("rotate main log: %v", err)
|
||||
}
|
||||
writeMainLog(t, dir, line3+"\n")
|
||||
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=1")
|
||||
if !reflect.DeepEqual(resp.Lines, []string{line2}) {
|
||||
t.Fatalf("lines = %#v, want rotated unread line", resp.Lines)
|
||||
}
|
||||
if resp.CursorReset {
|
||||
t.Fatal("cursor-reset = true, want false")
|
||||
}
|
||||
|
||||
next := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(resp.NextCursor)+"&limit=1")
|
||||
if !reflect.DeepEqual(next.Lines, []string{line3}) {
|
||||
t.Fatalf("next lines = %#v, want new main line", next.Lines)
|
||||
}
|
||||
if next.CursorReset {
|
||||
t.Fatal("next cursor-reset = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsZeroOffsetCursorWithPartialLineReadsAcrossRotation(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
writeMainLog(t, dir, "partial")
|
||||
initial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=1")
|
||||
if initial.NextCursor == "" {
|
||||
t.Fatal("initial next-cursor is empty")
|
||||
}
|
||||
cursor, errCursor := decodeLogCursor(initial.NextCursor)
|
||||
if errCursor != nil {
|
||||
t.Fatalf("decode initial cursor: %v", errCursor)
|
||||
}
|
||||
if cursor.Offset != 0 || cursor.Size == 0 {
|
||||
t.Fatalf("cursor offset/size = %d/%d, want zero offset with partial size", cursor.Offset, cursor.Size)
|
||||
}
|
||||
|
||||
appendMainLog(t, dir, " complete\n")
|
||||
if err := os.Rename(filepath.Join(dir, defaultLogFileName), filepath.Join(dir, defaultLogFileName+".1")); err != nil {
|
||||
t.Fatalf("rotate main log: %v", err)
|
||||
}
|
||||
writeMainLog(t, dir, "new\n")
|
||||
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=10")
|
||||
wantLines := []string{"partial complete", "new"}
|
||||
if !reflect.DeepEqual(resp.Lines, wantLines) {
|
||||
t.Fatalf("lines = %#v, want %#v", resp.Lines, wantLines)
|
||||
}
|
||||
if resp.CursorReset {
|
||||
t.Fatal("cursor-reset = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsZeroOffsetCursorWithEmptyFileReadsAcrossRotation(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
writeMainLog(t, dir, "")
|
||||
initial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=1")
|
||||
if initial.NextCursor == "" {
|
||||
t.Fatal("initial next-cursor is empty")
|
||||
}
|
||||
cursor, errCursor := decodeLogCursor(initial.NextCursor)
|
||||
if errCursor != nil {
|
||||
t.Fatalf("decode initial cursor: %v", errCursor)
|
||||
}
|
||||
if cursor.Offset != 0 || cursor.Size != 0 {
|
||||
t.Fatalf("cursor offset/size = %d/%d, want empty zero offset", cursor.Offset, cursor.Size)
|
||||
}
|
||||
|
||||
appendMainLog(t, dir, "first\n")
|
||||
mainPath := filepath.Join(dir, defaultLogFileName)
|
||||
nextModTime := time.Unix(0, cursorModTimeUnixNano(cursor)+int64(time.Second))
|
||||
if err := os.Chtimes(mainPath, nextModTime, nextModTime); err != nil {
|
||||
t.Fatalf("update main log mtime: %v", err)
|
||||
}
|
||||
if err := os.Rename(mainPath, filepath.Join(dir, defaultLogFileName+".1")); err != nil {
|
||||
t.Fatalf("rotate main log: %v", err)
|
||||
}
|
||||
writeMainLog(t, dir, "second\n")
|
||||
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=1")
|
||||
if !reflect.DeepEqual(resp.Lines, []string{"first"}) {
|
||||
t.Fatalf("lines = %#v, want first rotated line", resp.Lines)
|
||||
}
|
||||
if resp.CursorReset {
|
||||
t.Fatal("cursor-reset = true, want false")
|
||||
}
|
||||
|
||||
next := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(resp.NextCursor)+"&limit=1")
|
||||
if !reflect.DeepEqual(next.Lines, []string{"second"}) {
|
||||
t.Fatalf("next lines = %#v, want second main line", next.Lines)
|
||||
}
|
||||
if next.CursorReset {
|
||||
t.Fatal("next cursor-reset = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsZeroOffsetCursorWithEmptyFileReadsAcrossTwoRotations(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
writeMainLog(t, dir, "")
|
||||
initial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=1")
|
||||
if initial.NextCursor == "" {
|
||||
t.Fatal("initial next-cursor is empty")
|
||||
}
|
||||
cursor, errCursor := decodeLogCursor(initial.NextCursor)
|
||||
if errCursor != nil {
|
||||
t.Fatalf("decode initial cursor: %v", errCursor)
|
||||
}
|
||||
if cursor.Offset != 0 || cursor.Size != 0 {
|
||||
t.Fatalf("cursor offset/size = %d/%d, want empty zero offset", cursor.Offset, cursor.Size)
|
||||
}
|
||||
|
||||
mainPath := filepath.Join(dir, defaultLogFileName)
|
||||
firstRotatedPath := filepath.Join(dir, defaultLogFileName+".1")
|
||||
secondRotatedPath := filepath.Join(dir, defaultLogFileName+".2")
|
||||
firstModTime := time.Unix(0, cursorModTimeUnixNano(cursor)+int64(time.Second))
|
||||
secondModTime := time.Unix(0, cursorModTimeUnixNano(cursor)+2*int64(time.Second))
|
||||
|
||||
appendMainLog(t, dir, "first\n")
|
||||
if err := os.Chtimes(mainPath, firstModTime, firstModTime); err != nil {
|
||||
t.Fatalf("update first main log mtime: %v", err)
|
||||
}
|
||||
if err := os.Rename(mainPath, firstRotatedPath); err != nil {
|
||||
t.Fatalf("rotate first main log: %v", err)
|
||||
}
|
||||
writeMainLog(t, dir, "second\n")
|
||||
if err := os.Chtimes(mainPath, secondModTime, secondModTime); err != nil {
|
||||
t.Fatalf("update second main log mtime: %v", err)
|
||||
}
|
||||
if err := os.Rename(firstRotatedPath, secondRotatedPath); err != nil {
|
||||
t.Fatalf("advance first rotated log: %v", err)
|
||||
}
|
||||
if err := os.Rename(mainPath, firstRotatedPath); err != nil {
|
||||
t.Fatalf("rotate second main log: %v", err)
|
||||
}
|
||||
writeMainLog(t, dir, "third\n")
|
||||
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=1")
|
||||
if !reflect.DeepEqual(resp.Lines, []string{"first"}) {
|
||||
t.Fatalf("lines = %#v, want oldest rotated line", resp.Lines)
|
||||
}
|
||||
if resp.CursorReset {
|
||||
t.Fatal("cursor-reset = true, want false")
|
||||
}
|
||||
|
||||
next := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(resp.NextCursor)+"&limit=1")
|
||||
if !reflect.DeepEqual(next.Lines, []string{"second"}) {
|
||||
t.Fatalf("next lines = %#v, want newer rotated line", next.Lines)
|
||||
}
|
||||
if next.CursorReset {
|
||||
t.Fatal("next cursor-reset = true, want false")
|
||||
}
|
||||
|
||||
latest := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(next.NextCursor)+"&limit=1")
|
||||
if !reflect.DeepEqual(latest.Lines, []string{"third"}) {
|
||||
t.Fatalf("latest lines = %#v, want main line", latest.Lines)
|
||||
}
|
||||
if latest.CursorReset {
|
||||
t.Fatal("latest cursor-reset = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsZeroOffsetCursorWithEmptyFileResetsWhenRotationModTimeAmbiguous(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
mainPath := filepath.Join(dir, defaultLogFileName)
|
||||
fixedModTime := time.Date(2026, 6, 15, 10, 0, 0, 0, time.Local)
|
||||
writeMainLog(t, dir, "")
|
||||
if err := os.Chtimes(mainPath, fixedModTime, fixedModTime); err != nil {
|
||||
t.Fatalf("set initial main mtime: %v", err)
|
||||
}
|
||||
initial := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?limit=1")
|
||||
if initial.NextCursor == "" {
|
||||
t.Fatal("initial next-cursor is empty")
|
||||
}
|
||||
cursor, errCursor := decodeLogCursor(initial.NextCursor)
|
||||
if errCursor != nil {
|
||||
t.Fatalf("decode initial cursor: %v", errCursor)
|
||||
}
|
||||
if cursor.Offset != 0 || cursor.Size != 0 {
|
||||
t.Fatalf("cursor offset/size = %d/%d, want empty zero offset", cursor.Offset, cursor.Size)
|
||||
}
|
||||
|
||||
first := "[2026-06-15 10:00:01] first"
|
||||
second := "[2026-06-15 10:00:02] second"
|
||||
appendMainLog(t, dir, first+"\n")
|
||||
if err := os.Chtimes(mainPath, fixedModTime, fixedModTime); err != nil {
|
||||
t.Fatalf("set rotated mtime: %v", err)
|
||||
}
|
||||
if err := os.Rename(mainPath, filepath.Join(dir, defaultLogFileName+".1")); err != nil {
|
||||
t.Fatalf("rotate main log: %v", err)
|
||||
}
|
||||
writeMainLog(t, dir, second+"\n")
|
||||
if err := os.Chtimes(mainPath, fixedModTime, fixedModTime); err != nil {
|
||||
t.Fatalf("set new main mtime: %v", err)
|
||||
}
|
||||
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(initial.NextCursor)+"&limit=2")
|
||||
wantLines := []string{first, second}
|
||||
if !reflect.DeepEqual(resp.Lines, wantLines) {
|
||||
t.Fatalf("lines = %#v, want %#v", resp.Lines, wantLines)
|
||||
}
|
||||
if !resp.CursorReset {
|
||||
t.Fatal("cursor-reset = false, want true for ambiguous empty cursor rotation")
|
||||
}
|
||||
if resp.LineCount != len(wantLines) {
|
||||
t.Fatalf("line-count = %d, want returned line count %d", resp.LineCount, len(wantLines))
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsInvalidCursorResetsToTail(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
lines := []string{
|
||||
"[2026-06-15 10:00:00] first",
|
||||
"[2026-06-15 10:00:01] second",
|
||||
}
|
||||
writeMainLog(t, dir, strings.Join(lines, "\n")+"\n")
|
||||
|
||||
cases := []string{
|
||||
"not-base64",
|
||||
mustEncodeRawCursor(t, logCursor{
|
||||
Version: logCursorVersion,
|
||||
File: "../secret",
|
||||
Fingerprint: "fingerprint",
|
||||
}),
|
||||
}
|
||||
for _, raw := range cases {
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(raw)+"&limit=1")
|
||||
if !resp.CursorReset {
|
||||
t.Fatalf("cursor-reset = false for cursor %q", raw)
|
||||
}
|
||||
if !reflect.DeepEqual(resp.Lines, []string{lines[1]}) {
|
||||
t.Fatalf("lines = %#v, want latest line", resp.Lines)
|
||||
}
|
||||
if resp.LineCount != 1 {
|
||||
t.Fatalf("line-count = %d, want 1", resp.LineCount)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsMissingRotatedCursorFileResetsToTail(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
current := "[2026-06-15 10:00:01] current"
|
||||
writeMainLog(t, dir, current+"\n")
|
||||
rotatedPath := filepath.Join(dir, defaultLogFileName+".1")
|
||||
if err := os.WriteFile(rotatedPath, []byte("[2026-06-15 10:00:00] old\n"), 0o644); err != nil {
|
||||
t.Fatalf("write rotated log: %v", err)
|
||||
}
|
||||
cursor, errCursor := newLogCursor(rotatedPath, int64(len("[2026-06-15 10:00:00] old\n")), 0)
|
||||
if errCursor != nil {
|
||||
t.Fatalf("newLogCursor() error = %v", errCursor)
|
||||
}
|
||||
if errRemove := os.Remove(rotatedPath); errRemove != nil {
|
||||
t.Fatalf("remove rotated log: %v", errRemove)
|
||||
}
|
||||
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape(cursor)+"&limit=1")
|
||||
if !resp.CursorReset {
|
||||
t.Fatal("cursor-reset = false, want true")
|
||||
}
|
||||
if !reflect.DeepEqual(resp.Lines, []string{current}) {
|
||||
t.Fatalf("lines = %#v, want current tail", resp.Lines)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsMissingLogDirKeepsOKEmptyResponse(t *testing.T) {
|
||||
dir := filepath.Join(t.TempDir(), "missing")
|
||||
resp := performGetLogs(t, newLogsTestHandler(dir, true), "/v0/management/logs?cursor="+url.QueryEscape("not-base64")+"&limit=1")
|
||||
if len(resp.Lines) != 0 {
|
||||
t.Fatalf("lines = %#v, want empty", resp.Lines)
|
||||
}
|
||||
if resp.LineCount != 0 {
|
||||
t.Fatalf("line-count = %d, want 0", resp.LineCount)
|
||||
}
|
||||
if !resp.CursorReset {
|
||||
t.Fatal("cursor-reset = false, want true for cursor against missing log dir")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetLogsLoggingDisabledKeepsBadRequest(t *testing.T) {
|
||||
status, body := performGetLogsRaw(t, newLogsTestHandler(t.TempDir(), false), "/v0/management/logs?cursor=not-base64&limit=1")
|
||||
if status != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d", status, http.StatusBadRequest)
|
||||
}
|
||||
if !strings.Contains(body, "logging to file disabled") {
|
||||
t.Fatalf("body = %s, want logging disabled error", body)
|
||||
}
|
||||
}
|
||||
|
||||
func mustEncodeRawCursor(t *testing.T, cursor logCursor) string {
|
||||
t.Helper()
|
||||
raw, err := json.Marshal(cursor)
|
||||
if err != nil {
|
||||
t.Fatalf("json.Marshal cursor: %v", err)
|
||||
}
|
||||
return base64.RawURLEncoding.EncodeToString(raw)
|
||||
}
|
||||
|
||||
type logsAPIResponse struct {
|
||||
Lines []string `json:"lines"`
|
||||
LineCount int `json:"line-count"`
|
||||
LatestTimestamp int64 `json:"latest-timestamp"`
|
||||
NextCursor string `json:"next-cursor"`
|
||||
CursorReset bool `json:"cursor-reset"`
|
||||
}
|
||||
|
||||
func newLogsTestHandler(dir string, loggingToFile bool) *Handler {
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{LoggingToFile: loggingToFile}, nil)
|
||||
h.SetLogDirectory(dir)
|
||||
return h
|
||||
}
|
||||
|
||||
func performGetLogs(t *testing.T, h *Handler, target string) logsAPIResponse {
|
||||
t.Helper()
|
||||
status, body := performGetLogsRaw(t, h, target)
|
||||
if status != http.StatusOK {
|
||||
t.Fatalf("GetLogs status = %d, body = %s", status, body)
|
||||
}
|
||||
var resp logsAPIResponse
|
||||
if err := json.Unmarshal([]byte(body), &resp); err != nil {
|
||||
t.Fatalf("decode response: %v", err)
|
||||
}
|
||||
if resp.Lines == nil {
|
||||
resp.Lines = []string{}
|
||||
}
|
||||
return resp
|
||||
}
|
||||
|
||||
func performGetLogsRaw(t *testing.T, h *Handler, target string) (int, string) {
|
||||
t.Helper()
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, target, nil)
|
||||
h.GetLogs(c)
|
||||
return rec.Code, rec.Body.String()
|
||||
}
|
||||
|
||||
func writeMainLog(t *testing.T, dir, content string) {
|
||||
t.Helper()
|
||||
if err := os.WriteFile(filepath.Join(dir, defaultLogFileName), []byte(content), 0o644); err != nil {
|
||||
t.Fatalf("write main log: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func appendMainLog(t *testing.T, dir, content string) {
|
||||
t.Helper()
|
||||
file, errOpen := os.OpenFile(filepath.Join(dir, defaultLogFileName), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
|
||||
if errOpen != nil {
|
||||
t.Fatalf("open main log: %v", errOpen)
|
||||
}
|
||||
if _, errWrite := file.WriteString(content); errWrite != nil {
|
||||
_ = file.Close()
|
||||
t.Fatalf("append main log: %v", errWrite)
|
||||
}
|
||||
if errClose := file.Close(); errClose != nil {
|
||||
t.Fatalf("close main log: %v", errClose)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,33 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
)
|
||||
|
||||
// GetStaticModelDefinitions returns static model metadata for a given channel.
|
||||
// Channel is provided via path param (:channel) or query param (?channel=...).
|
||||
func (h *Handler) GetStaticModelDefinitions(c *gin.Context) {
|
||||
channel := strings.TrimSpace(c.Param("channel"))
|
||||
if channel == "" {
|
||||
channel = strings.TrimSpace(c.Query("channel"))
|
||||
}
|
||||
if channel == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "channel is required"})
|
||||
return
|
||||
}
|
||||
|
||||
models := registry.GetStaticModelDefinitionsByChannel(channel)
|
||||
if models == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "unknown channel", "channel": channel})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"channel": strings.ToLower(strings.TrimSpace(channel)),
|
||||
"models": models,
|
||||
})
|
||||
}
|
||||
148
backend/internal/api/handlers/management/oauth_callback.go
Normal file
148
backend/internal/api/handlers/management/oauth_callback.go
Normal file
|
|
@ -0,0 +1,148 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type oauthCallbackRequest struct {
|
||||
Provider string `json:"provider"`
|
||||
RedirectURL string `json:"redirect_url"`
|
||||
Code string `json:"code"`
|
||||
State string `json:"state"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
|
||||
func (h *Handler) PostOAuthCallback(c *gin.Context) {
|
||||
if h == nil || h.cfg == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"status": "error", "error": "handler not initialized"})
|
||||
return
|
||||
}
|
||||
|
||||
var req oauthCallbackRequest
|
||||
if errBindJSON := c.ShouldBindJSON(&req); errBindJSON != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "invalid body"})
|
||||
return
|
||||
}
|
||||
h.handleOAuthCallback(c, req)
|
||||
}
|
||||
|
||||
func (h *Handler) GetOAuthCallback(c *gin.Context) {
|
||||
req := oauthCallbackRequest{
|
||||
Provider: strings.TrimSpace(c.Query("provider")),
|
||||
Code: strings.TrimSpace(c.Query("code")),
|
||||
State: strings.TrimSpace(c.Query("state")),
|
||||
Error: firstNonEmpty(c.Query("error"), c.Query("error_description")),
|
||||
}
|
||||
h.handleOAuthCallback(c, req)
|
||||
}
|
||||
|
||||
func (h *Handler) handleOAuthCallback(c *gin.Context, req oauthCallbackRequest) {
|
||||
if h == nil || h.cfg == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"status": "error", "error": "handler not initialized"})
|
||||
return
|
||||
}
|
||||
|
||||
state := strings.TrimSpace(req.State)
|
||||
code := strings.TrimSpace(req.Code)
|
||||
errMsg := strings.TrimSpace(req.Error)
|
||||
|
||||
if rawRedirect := strings.TrimSpace(req.RedirectURL); rawRedirect != "" {
|
||||
u, errParse := url.Parse(rawRedirect)
|
||||
if errParse != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "invalid redirect_url"})
|
||||
return
|
||||
}
|
||||
q := u.Query()
|
||||
if state == "" {
|
||||
state = strings.TrimSpace(q.Get("state"))
|
||||
}
|
||||
if code == "" {
|
||||
code = strings.TrimSpace(q.Get("code"))
|
||||
}
|
||||
if errMsg == "" {
|
||||
errMsg = strings.TrimSpace(q.Get("error"))
|
||||
if errMsg == "" {
|
||||
errMsg = strings.TrimSpace(q.Get("error_description"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if state == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "state is required"})
|
||||
return
|
||||
}
|
||||
if err := ValidateOAuthState(state); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "invalid state"})
|
||||
return
|
||||
}
|
||||
if code == "" && errMsg == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "code or error is required"})
|
||||
return
|
||||
}
|
||||
|
||||
sessionProvider, sessionStatus, isPlugin, _, completed, ok := GetOAuthSessionDetails(state)
|
||||
if !ok {
|
||||
c.JSON(http.StatusNotFound, gin.H{"status": "error", "error": "unknown or expired state"})
|
||||
return
|
||||
}
|
||||
if completed {
|
||||
c.JSON(http.StatusConflict, gin.H{"status": "error", "error": "oauth flow is already completed"})
|
||||
return
|
||||
}
|
||||
provider := strings.TrimSpace(req.Provider)
|
||||
if provider == "" {
|
||||
provider = sessionProvider
|
||||
}
|
||||
var canonicalProvider string
|
||||
var errNormalize error
|
||||
if isPlugin {
|
||||
canonicalProvider, errNormalize = NormalizePluginOAuthCallbackProvider(provider)
|
||||
} else {
|
||||
canonicalProvider, errNormalize = NormalizeOAuthCallbackProvider(provider)
|
||||
}
|
||||
if errNormalize != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "unsupported provider"})
|
||||
return
|
||||
}
|
||||
if sessionStatus != "" {
|
||||
c.JSON(http.StatusConflict, gin.H{"status": "error", "error": sessionStatus})
|
||||
return
|
||||
}
|
||||
if !strings.EqualFold(sessionProvider, canonicalProvider) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "provider does not match state"})
|
||||
return
|
||||
}
|
||||
|
||||
if _, errWrite := WriteOAuthCallbackFileForPendingSession(h.cfg.AuthDir, canonicalProvider, state, code, errMsg); errWrite != nil {
|
||||
if errors.Is(errWrite, errOAuthSessionNotPending) {
|
||||
_, status, okSession := GetOAuthSession(state)
|
||||
if okSession && status != "" {
|
||||
c.JSON(http.StatusConflict, gin.H{"status": "error", "error": status})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusConflict, gin.H{"status": "error", "error": "oauth flow is not pending"})
|
||||
return
|
||||
}
|
||||
log.WithError(errWrite).Error("failed to persist oauth callback")
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"status": "error", "error": "failed to persist oauth callback"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||||
}
|
||||
|
||||
func firstNonEmpty(values ...string) string {
|
||||
for _, value := range values {
|
||||
trimmed := strings.TrimSpace(value)
|
||||
if trimmed != "" {
|
||||
return trimmed
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
148
backend/internal/api/handlers/management/oauth_callback_test.go
Normal file
148
backend/internal/api/handlers/management/oauth_callback_test.go
Normal file
|
|
@ -0,0 +1,148 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestPostOAuthCallbackCreatesMissingAuthDir(t *testing.T) {
|
||||
|
||||
authDir := filepath.Join(t.TempDir(), "missing-auth")
|
||||
state := "test-antigravity-state"
|
||||
RegisterOAuthSession(state, "antigravity")
|
||||
defer CompleteOAuthSession(state)
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
router := gin.New()
|
||||
router.POST("/v0/management/oauth-callback", h.PostOAuthCallback)
|
||||
|
||||
body := `{"provider":"antigravity","redirect_url":"http://localhost:59788/oauth-callback?state=test-antigravity-state&code=test-code"}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/v0/management/oauth-callback", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
callbackPath := filepath.Join(authDir, ".oauth-antigravity-"+state+".oauth")
|
||||
data, errRead := os.ReadFile(callbackPath)
|
||||
if errRead != nil {
|
||||
t.Fatalf("expected callback file to be written: %v", errRead)
|
||||
}
|
||||
|
||||
var payload oauthCallbackFilePayload
|
||||
if errUnmarshal := json.Unmarshal(data, &payload); errUnmarshal != nil {
|
||||
t.Fatalf("failed to decode callback payload: %v", errUnmarshal)
|
||||
}
|
||||
if payload.State != state || payload.Code != "test-code" || payload.Error != "" {
|
||||
t.Fatalf("unexpected callback payload: %+v", payload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetOAuthCallbackWritesPluginProviderCallback(t *testing.T) {
|
||||
authDir := filepath.Join(t.TempDir(), "missing-auth")
|
||||
state := "test-geminicli-state"
|
||||
if errRegister := RegisterPluginOAuthSession(state, "gemini-cli", nil); errRegister != nil {
|
||||
t.Fatalf("register plugin oauth session: %v", errRegister)
|
||||
}
|
||||
defer CompleteOAuthSession(state)
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
router := gin.New()
|
||||
router.GET("/v0/management/oauth-callback", h.GetOAuthCallback)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v0/management/oauth-callback?state="+state+"&code=test-code", nil)
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
callbackPath := filepath.Join(authDir, ".oauth-gemini-cli-"+state+".oauth")
|
||||
data, errRead := os.ReadFile(callbackPath)
|
||||
if errRead != nil {
|
||||
t.Fatalf("expected callback file to be written: %v", errRead)
|
||||
}
|
||||
|
||||
var payload oauthCallbackFilePayload
|
||||
if errUnmarshal := json.Unmarshal(data, &payload); errUnmarshal != nil {
|
||||
t.Fatalf("failed to decode callback payload: %v", errUnmarshal)
|
||||
}
|
||||
if payload.State != state || payload.Code != "test-code" || payload.Error != "" {
|
||||
t.Fatalf("unexpected callback payload: %+v", payload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetOAuthCallbackDoesNotAliasPluginProvider(t *testing.T) {
|
||||
authDir := filepath.Join(t.TempDir(), "missing-auth")
|
||||
state := "test-openai-plugin-state"
|
||||
if errRegister := RegisterPluginOAuthSession(state, "openai", nil); errRegister != nil {
|
||||
t.Fatalf("register plugin oauth session: %v", errRegister)
|
||||
}
|
||||
defer CompleteOAuthSession(state)
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
router := gin.New()
|
||||
router.GET("/v0/management/oauth-callback", h.GetOAuthCallback)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v0/management/oauth-callback?state="+state+"&code=test-code", nil)
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
callbackPath := filepath.Join(authDir, ".oauth-openai-"+state+".oauth")
|
||||
if _, errRead := os.ReadFile(callbackPath); errRead != nil {
|
||||
t.Fatalf("expected plugin callback provider to stay openai: %v", errRead)
|
||||
}
|
||||
if _, errRead := os.ReadFile(filepath.Join(authDir, ".oauth-codex-"+state+".oauth")); errRead == nil {
|
||||
t.Fatal("unexpected codex callback file for openai plugin provider")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteOAuthCallbackFileForPendingSessionCreatesMissingAuthDirForCallbackProviders(t *testing.T) {
|
||||
// xAI uses device-code flow and no longer writes callback files.
|
||||
providers := []string{"anthropic", "codex", "gemini", "antigravity"}
|
||||
for _, provider := range providers {
|
||||
t.Run(provider, func(t *testing.T) {
|
||||
authDir := filepath.Join(t.TempDir(), "missing-auth")
|
||||
state := provider + "-state"
|
||||
RegisterOAuthSession(state, provider)
|
||||
defer CompleteOAuthSession(state)
|
||||
|
||||
path, errWrite := WriteOAuthCallbackFileForPendingSession(authDir, provider, state, "code-"+provider, "")
|
||||
if errWrite != nil {
|
||||
t.Fatalf("expected callback file write to succeed: %v", errWrite)
|
||||
}
|
||||
|
||||
data, errRead := os.ReadFile(path)
|
||||
if errRead != nil {
|
||||
t.Fatalf("expected callback file to be written: %v", errRead)
|
||||
}
|
||||
|
||||
var payload oauthCallbackFilePayload
|
||||
if errUnmarshal := json.Unmarshal(data, &payload); errUnmarshal != nil {
|
||||
t.Fatalf("failed to decode callback payload: %v", errUnmarshal)
|
||||
}
|
||||
if payload.State != state || payload.Code != "code-"+provider || payload.Error != "" {
|
||||
t.Fatalf("unexpected callback payload: %+v", payload)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,111 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/auth/codex"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
type fakeCodexOAuthService struct{}
|
||||
|
||||
func (f *fakeCodexOAuthService) GenerateAuthURL(state string, pkceCodes *codex.PKCECodes) (string, error) {
|
||||
return "https://auth.example.test/oauth?state=" + state, nil
|
||||
}
|
||||
|
||||
func (f *fakeCodexOAuthService) ExchangeCodeForTokens(ctx context.Context, code string, pkceCodes *codex.PKCECodes) (*codex.CodexAuthBundle, error) {
|
||||
now := time.Now()
|
||||
return &codex.CodexAuthBundle{
|
||||
TokenData: codex.CodexTokenData{
|
||||
IDToken: "invalid-test-id-token",
|
||||
AccessToken: "access-" + code,
|
||||
RefreshToken: "refresh-" + code,
|
||||
Email: "codex-" + code + "@example.test",
|
||||
Expire: now.Add(time.Hour).Format(time.RFC3339),
|
||||
},
|
||||
LastRefresh: now.Format(time.RFC3339),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (f *fakeCodexOAuthService) CreateTokenStorage(bundle *codex.CodexAuthBundle) *codex.CodexTokenStorage {
|
||||
return &codex.CodexTokenStorage{
|
||||
IDToken: bundle.TokenData.IDToken,
|
||||
AccessToken: bundle.TokenData.AccessToken,
|
||||
RefreshToken: bundle.TokenData.RefreshToken,
|
||||
AccountID: bundle.TokenData.AccountID,
|
||||
LastRefresh: bundle.LastRefresh,
|
||||
Email: bundle.TokenData.Email,
|
||||
Expire: bundle.TokenData.Expire,
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestCodexTokenCompletionKeepsConcurrentSessionPending(t *testing.T) {
|
||||
originalNewCodexOAuthService := newCodexOAuthService
|
||||
newCodexOAuthService = func(cfg *config.Config) codexOAuthService {
|
||||
return &fakeCodexOAuthService{}
|
||||
}
|
||||
defer func() {
|
||||
newCodexOAuthService = originalNewCodexOAuthService
|
||||
}()
|
||||
|
||||
authDir := filepath.Join(t.TempDir(), "auths")
|
||||
handler := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
router := gin.New()
|
||||
router.GET("/codex-auth-url", handler.RequestCodexToken)
|
||||
|
||||
firstState := requestCodexTokenState(t, router)
|
||||
secondState := requestCodexTokenState(t, router)
|
||||
defer CompleteOAuthSession(firstState)
|
||||
defer CompleteOAuthSession(secondState)
|
||||
|
||||
if _, errWrite := WriteOAuthCallbackFileForPendingSession(authDir, "codex", firstState, "first-code", ""); errWrite != nil {
|
||||
t.Fatalf("write first callback file: %v", errWrite)
|
||||
}
|
||||
|
||||
waitForOAuthSessionDone(t, firstState)
|
||||
if !IsOAuthSessionPending(secondState, "codex") {
|
||||
t.Fatalf("expected concurrent codex session %s to remain pending after %s completed", secondState, firstState)
|
||||
}
|
||||
}
|
||||
|
||||
func requestCodexTokenState(t *testing.T, router http.Handler) string {
|
||||
t.Helper()
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/codex-auth-url", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var payload struct {
|
||||
State string `json:"state"`
|
||||
}
|
||||
if errDecode := json.Unmarshal(w.Body.Bytes(), &payload); errDecode != nil {
|
||||
t.Fatalf("decode codex auth URL response: %v", errDecode)
|
||||
}
|
||||
if payload.State == "" {
|
||||
t.Fatalf("expected codex auth URL response to include state")
|
||||
}
|
||||
return payload.State
|
||||
}
|
||||
|
||||
func waitForOAuthSessionDone(t *testing.T, state string) {
|
||||
t.Helper()
|
||||
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !IsOAuthSessionPending(state, "codex") {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("timed out waiting for codex session %s to complete", state)
|
||||
}
|
||||
463
backend/internal/api/handlers/management/oauth_sessions.go
Normal file
463
backend/internal/api/handlers/management/oauth_sessions.go
Normal file
|
|
@ -0,0 +1,463 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
// oauthSessionTTL must cover device-code flows (xAI ~30m, Kimi ~15m).
|
||||
oauthSessionTTL = 30 * time.Minute
|
||||
oauthCompletedSessionTTL = time.Minute
|
||||
maxOAuthStateLength = 128
|
||||
)
|
||||
|
||||
const (
|
||||
oauthSessionSourceBuiltin = "builtin"
|
||||
oauthSessionSourcePlugin = "plugin"
|
||||
)
|
||||
|
||||
var (
|
||||
errInvalidOAuthState = errors.New("invalid oauth state")
|
||||
errUnsupportedOAuthFlow = errors.New("unsupported oauth provider")
|
||||
errOAuthSessionNotPending = errors.New("oauth session is not pending")
|
||||
errOAuthSessionExists = errors.New("oauth session already exists")
|
||||
)
|
||||
|
||||
type oauthSession struct {
|
||||
Provider string
|
||||
Status string
|
||||
Source string
|
||||
Metadata map[string]any
|
||||
Completed bool
|
||||
CreatedAt time.Time
|
||||
ExpiresAt time.Time
|
||||
}
|
||||
|
||||
type oauthSessionStore struct {
|
||||
mu sync.RWMutex
|
||||
ttl time.Duration
|
||||
completedTTL time.Duration
|
||||
sessions map[string]oauthSession
|
||||
}
|
||||
|
||||
func newOAuthSessionStore(ttl time.Duration) *oauthSessionStore {
|
||||
if ttl <= 0 {
|
||||
ttl = oauthSessionTTL
|
||||
}
|
||||
completedTTL := oauthCompletedSessionTTL
|
||||
if ttl < completedTTL {
|
||||
completedTTL = ttl
|
||||
}
|
||||
return &oauthSessionStore{
|
||||
ttl: ttl,
|
||||
completedTTL: completedTTL,
|
||||
sessions: make(map[string]oauthSession),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *oauthSessionStore) purgeExpiredLocked(now time.Time) {
|
||||
for state, session := range s.sessions {
|
||||
if !session.ExpiresAt.IsZero() && now.After(session.ExpiresAt) {
|
||||
delete(s.sessions, state)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *oauthSessionStore) Register(state, provider string) {
|
||||
state = strings.TrimSpace(state)
|
||||
provider = strings.ToLower(strings.TrimSpace(provider))
|
||||
if state == "" || provider == "" {
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.purgeExpiredLocked(now)
|
||||
s.sessions[state] = oauthSession{
|
||||
Provider: provider,
|
||||
Status: "",
|
||||
Source: oauthSessionSourceBuiltin,
|
||||
CreatedAt: now,
|
||||
ExpiresAt: now.Add(s.ttl),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *oauthSessionStore) RegisterPlugin(state, provider string, metadata map[string]any) error {
|
||||
state = strings.TrimSpace(state)
|
||||
provider = strings.ToLower(strings.TrimSpace(provider))
|
||||
if state == "" || provider == "" {
|
||||
return fmt.Errorf("%w: empty state or provider", errInvalidOAuthState)
|
||||
}
|
||||
if errState := ValidateOAuthState(state); errState != nil {
|
||||
return errState
|
||||
}
|
||||
now := time.Now()
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.purgeExpiredLocked(now)
|
||||
if _, ok := s.sessions[state]; ok {
|
||||
return errOAuthSessionExists
|
||||
}
|
||||
s.sessions[state] = oauthSession{
|
||||
Provider: provider,
|
||||
Status: "",
|
||||
Source: oauthSessionSourcePlugin,
|
||||
Metadata: cloneOAuthSessionMetadata(metadata),
|
||||
CreatedAt: now,
|
||||
ExpiresAt: now.Add(s.ttl),
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *oauthSessionStore) SetError(state, message string) {
|
||||
state = strings.TrimSpace(state)
|
||||
message = strings.TrimSpace(message)
|
||||
if state == "" {
|
||||
return
|
||||
}
|
||||
if message == "" {
|
||||
message = "Authentication failed"
|
||||
}
|
||||
now := time.Now()
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.purgeExpiredLocked(now)
|
||||
session, ok := s.sessions[state]
|
||||
if !ok || session.Completed {
|
||||
return
|
||||
}
|
||||
session.Status = message
|
||||
session.ExpiresAt = now.Add(s.ttl)
|
||||
s.sessions[state] = session
|
||||
}
|
||||
|
||||
func (s *oauthSessionStore) Complete(state string) {
|
||||
state = strings.TrimSpace(state)
|
||||
if state == "" {
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.purgeExpiredLocked(now)
|
||||
session, ok := s.sessions[state]
|
||||
if !ok || session.Completed {
|
||||
return
|
||||
}
|
||||
session.Status = ""
|
||||
session.Metadata = nil
|
||||
session.Completed = true
|
||||
session.ExpiresAt = now.Add(s.completedTTL)
|
||||
s.sessions[state] = session
|
||||
}
|
||||
|
||||
func (s *oauthSessionStore) CompleteProvider(provider string, source string) int {
|
||||
provider = strings.ToLower(strings.TrimSpace(provider))
|
||||
if provider == "" {
|
||||
return 0
|
||||
}
|
||||
source = strings.TrimSpace(source)
|
||||
now := time.Now()
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.purgeExpiredLocked(now)
|
||||
removed := 0
|
||||
for state, session := range s.sessions {
|
||||
if !session.Completed && strings.EqualFold(session.Provider, provider) && (source == "" || session.Source == source) {
|
||||
session.Status = ""
|
||||
session.Metadata = nil
|
||||
session.Completed = true
|
||||
session.ExpiresAt = now.Add(s.completedTTL)
|
||||
s.sessions[state] = session
|
||||
removed++
|
||||
}
|
||||
}
|
||||
return removed
|
||||
}
|
||||
|
||||
func (s *oauthSessionStore) Get(state string) (oauthSession, bool) {
|
||||
state = strings.TrimSpace(state)
|
||||
now := time.Now()
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.purgeExpiredLocked(now)
|
||||
session, ok := s.sessions[state]
|
||||
session.Metadata = cloneOAuthSessionMetadata(session.Metadata)
|
||||
return session, ok
|
||||
}
|
||||
|
||||
func (s *oauthSessionStore) IsPending(state, provider string) bool {
|
||||
state = strings.TrimSpace(state)
|
||||
provider = strings.ToLower(strings.TrimSpace(provider))
|
||||
now := time.Now()
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.purgeExpiredLocked(now)
|
||||
session, ok := s.sessions[state]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if session.Completed || session.Status != "" {
|
||||
return false
|
||||
}
|
||||
if provider == "" {
|
||||
return true
|
||||
}
|
||||
return strings.EqualFold(session.Provider, provider)
|
||||
}
|
||||
|
||||
// Cancel removes a pending OAuth session so background waiters exit without saving credentials.
|
||||
// Returns true when a pending session was cancelled.
|
||||
func (s *oauthSessionStore) Cancel(state string) bool {
|
||||
state = strings.TrimSpace(state)
|
||||
if state == "" {
|
||||
return false
|
||||
}
|
||||
now := time.Now()
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.purgeExpiredLocked(now)
|
||||
session, ok := s.sessions[state]
|
||||
if !ok || session.Completed || session.Status != "" {
|
||||
return false
|
||||
}
|
||||
delete(s.sessions, state)
|
||||
return true
|
||||
}
|
||||
|
||||
func cloneOAuthSessionMetadata(in map[string]any) map[string]any {
|
||||
if len(in) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]any, len(in))
|
||||
for key, value := range in {
|
||||
out[key] = value
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
var oauthSessions = newOAuthSessionStore(oauthSessionTTL)
|
||||
|
||||
func RegisterOAuthSession(state, provider string) { oauthSessions.Register(state, provider) }
|
||||
|
||||
func RegisterPluginOAuthSession(state, provider string, metadata map[string]any) error {
|
||||
return oauthSessions.RegisterPlugin(state, provider, metadata)
|
||||
}
|
||||
|
||||
func SetOAuthSessionError(state, message string) { oauthSessions.SetError(state, message) }
|
||||
|
||||
func CompleteOAuthSession(state string) { oauthSessions.Complete(state) }
|
||||
|
||||
func CompleteOAuthSessionsByProvider(provider string) int {
|
||||
return oauthSessions.CompleteProvider(provider, oauthSessionSourceBuiltin)
|
||||
}
|
||||
|
||||
func CompletePluginOAuthSessionsByProvider(provider string) int {
|
||||
return oauthSessions.CompleteProvider(provider, oauthSessionSourcePlugin)
|
||||
}
|
||||
|
||||
func GetOAuthSession(state string) (provider string, status string, ok bool) {
|
||||
session, ok := oauthSessions.Get(state)
|
||||
if !ok || session.Completed {
|
||||
return "", "", false
|
||||
}
|
||||
return session.Provider, session.Status, true
|
||||
}
|
||||
|
||||
func GetOAuthSessionDetails(state string) (provider string, status string, isPlugin bool, metadata map[string]any, completed bool, ok bool) {
|
||||
session, ok := oauthSessions.Get(state)
|
||||
if !ok {
|
||||
return "", "", false, nil, false, false
|
||||
}
|
||||
return session.Provider, session.Status, session.Source == oauthSessionSourcePlugin, cloneOAuthSessionMetadata(session.Metadata), session.Completed, true
|
||||
}
|
||||
|
||||
func IsOAuthSessionPending(state, provider string) bool {
|
||||
return oauthSessions.IsPending(state, provider)
|
||||
}
|
||||
|
||||
// guardOAuthSessionPendingForSave returns errOAuthSessionNotPending when the session
|
||||
// is no longer pending (cancelled, completed, errored, or expired).
|
||||
// Call immediately before persisting credentials so a cancel that races with token
|
||||
// exchange or metadata fetch cannot save credentials for a cancelled flow.
|
||||
func guardOAuthSessionPendingForSave(state, provider string) error {
|
||||
if IsOAuthSessionPending(state, provider) {
|
||||
return nil
|
||||
}
|
||||
return errOAuthSessionNotPending
|
||||
}
|
||||
|
||||
// CancelOAuthSession cancels a pending OAuth session by state.
|
||||
// Background callback and device-code waiters observe IsOAuthSessionPending as false and exit without saving credentials.
|
||||
func CancelOAuthSession(state string) bool {
|
||||
return oauthSessions.Cancel(state)
|
||||
}
|
||||
|
||||
func oauthSessionErrorWithCause(message string, cause error) string {
|
||||
message = strings.TrimSpace(message)
|
||||
if message == "" {
|
||||
message = "Authentication failed"
|
||||
}
|
||||
if cause == nil {
|
||||
return message
|
||||
}
|
||||
detail := strings.TrimSpace(cause.Error())
|
||||
if detail == "" {
|
||||
return message
|
||||
}
|
||||
return message + ": " + detail
|
||||
}
|
||||
|
||||
func ValidateOAuthState(state string) error {
|
||||
trimmed := strings.TrimSpace(state)
|
||||
if trimmed == "" {
|
||||
return fmt.Errorf("%w: empty", errInvalidOAuthState)
|
||||
}
|
||||
if len(trimmed) > maxOAuthStateLength {
|
||||
return fmt.Errorf("%w: too long", errInvalidOAuthState)
|
||||
}
|
||||
if strings.Contains(trimmed, "/") || strings.Contains(trimmed, "\\") {
|
||||
return fmt.Errorf("%w: contains path separator", errInvalidOAuthState)
|
||||
}
|
||||
if strings.Contains(trimmed, "..") {
|
||||
return fmt.Errorf("%w: contains '..'", errInvalidOAuthState)
|
||||
}
|
||||
for _, r := range trimmed {
|
||||
switch {
|
||||
case r >= 'a' && r <= 'z':
|
||||
case r >= 'A' && r <= 'Z':
|
||||
case r >= '0' && r <= '9':
|
||||
case r == '-' || r == '_' || r == '.':
|
||||
default:
|
||||
return fmt.Errorf("%w: invalid character", errInvalidOAuthState)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func NormalizeOAuthProvider(provider string) (string, error) {
|
||||
switch strings.ToLower(strings.TrimSpace(provider)) {
|
||||
case "anthropic", "claude":
|
||||
return "anthropic", nil
|
||||
case "codex", "openai":
|
||||
return "codex", nil
|
||||
case "antigravity", "anti-gravity":
|
||||
return "antigravity", nil
|
||||
case "xai", "x-ai", "x.ai", "grok":
|
||||
return "xai", nil
|
||||
default:
|
||||
return "", errUnsupportedOAuthFlow
|
||||
}
|
||||
}
|
||||
|
||||
func NormalizeOAuthCallbackProvider(provider string) (string, error) {
|
||||
if normalized, errNormalize := NormalizeOAuthProvider(provider); errNormalize == nil {
|
||||
return normalized, nil
|
||||
}
|
||||
return NormalizePluginOAuthCallbackProvider(provider)
|
||||
}
|
||||
|
||||
func NormalizePluginOAuthCallbackProvider(provider string) (string, error) {
|
||||
trimmed := strings.ToLower(strings.TrimSpace(provider))
|
||||
if trimmed == "" {
|
||||
return "", errUnsupportedOAuthFlow
|
||||
}
|
||||
for _, r := range trimmed {
|
||||
switch {
|
||||
case r >= 'a' && r <= 'z':
|
||||
case r >= '0' && r <= '9':
|
||||
case r == '-':
|
||||
default:
|
||||
return "", errUnsupportedOAuthFlow
|
||||
}
|
||||
}
|
||||
return trimmed, nil
|
||||
}
|
||||
|
||||
func normalizeOAuthCallbackProviderForPendingSession(provider, state string) (string, error) {
|
||||
session, ok := oauthSessions.Get(state)
|
||||
if ok && session.Source == oauthSessionSourcePlugin {
|
||||
return NormalizePluginOAuthCallbackProvider(provider)
|
||||
}
|
||||
return NormalizeOAuthCallbackProvider(provider)
|
||||
}
|
||||
|
||||
type oauthCallbackFilePayload struct {
|
||||
Code string `json:"code"`
|
||||
State string `json:"state"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
|
||||
func WriteOAuthCallbackFile(authDir, provider, state, code, errorMessage string) (string, error) {
|
||||
canonicalProvider, err := NormalizeOAuthCallbackProvider(provider)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return writeOAuthCallbackFile(authDir, canonicalProvider, state, code, errorMessage)
|
||||
}
|
||||
|
||||
func writeOAuthCallbackFile(authDir, canonicalProvider, state, code, errorMessage string) (string, error) {
|
||||
if strings.TrimSpace(authDir) == "" {
|
||||
return "", fmt.Errorf("auth dir is empty")
|
||||
}
|
||||
canonicalProvider = strings.TrimSpace(canonicalProvider)
|
||||
if canonicalProvider == "" {
|
||||
return "", errUnsupportedOAuthFlow
|
||||
}
|
||||
if err := ValidateOAuthState(state); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
fileName := fmt.Sprintf(".oauth-%s-%s.oauth", canonicalProvider, state)
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
if err := os.MkdirAll(authDir, 0o700); err != nil {
|
||||
return "", fmt.Errorf("create oauth callback dir: %w", err)
|
||||
}
|
||||
payload := oauthCallbackFilePayload{
|
||||
Code: strings.TrimSpace(code),
|
||||
State: strings.TrimSpace(state),
|
||||
Error: strings.TrimSpace(errorMessage),
|
||||
}
|
||||
data, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("marshal oauth callback payload: %w", err)
|
||||
}
|
||||
if err := os.WriteFile(filePath, data, 0o600); err != nil {
|
||||
return "", fmt.Errorf("write oauth callback file: %w", err)
|
||||
}
|
||||
return filePath, nil
|
||||
}
|
||||
|
||||
func WriteOAuthCallbackFileForPendingSession(authDir, provider, state, code, errorMessage string) (string, error) {
|
||||
canonicalProvider, err := normalizeOAuthCallbackProviderForPendingSession(provider, state)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !IsOAuthSessionPending(state, canonicalProvider) {
|
||||
return "", errOAuthSessionNotPending
|
||||
}
|
||||
return writeOAuthCallbackFile(authDir, canonicalProvider, state, code, errorMessage)
|
||||
}
|
||||
344
backend/internal/api/handlers/management/oauth_sessions_test.go
Normal file
344
backend/internal/api/handlers/management/oauth_sessions_test.go
Normal file
|
|
@ -0,0 +1,344 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestOAuthSessionStoreCompleteKeepsShortLivedSession(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
store.Register("completed-state", "codex")
|
||||
|
||||
store.Complete("completed-state")
|
||||
|
||||
if _, ok := store.Get("completed-state"); !ok {
|
||||
t.Fatal("completed OAuth session was deleted instead of retained as a tombstone")
|
||||
}
|
||||
if store.IsPending("completed-state", "codex") {
|
||||
t.Fatal("completed OAuth session remained pending")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOAuthSessionStoreCompleteDoesNotExtendCompletedSession(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
store.Register("completed-state", "codex")
|
||||
store.Complete("completed-state")
|
||||
before, ok := store.Get("completed-state")
|
||||
if !ok {
|
||||
t.Fatal("completed OAuth session tombstone is missing")
|
||||
}
|
||||
|
||||
store.completedTTL = 2 * time.Minute
|
||||
store.Complete("completed-state")
|
||||
after, ok := store.Get("completed-state")
|
||||
if !ok {
|
||||
t.Fatal("completed OAuth session tombstone is missing after repeated completion")
|
||||
}
|
||||
if !after.ExpiresAt.Equal(before.ExpiresAt) {
|
||||
t.Fatalf("repeated completion extended expiry from %s to %s", before.ExpiresAt, after.ExpiresAt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOAuthSessionStoreCompleteProviderSkipsCompletedSessions(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
store.Register("completed-state", "codex")
|
||||
store.Register("pending-state", "codex")
|
||||
store.Complete("completed-state")
|
||||
completedBefore, ok := store.Get("completed-state")
|
||||
if !ok {
|
||||
t.Fatal("completed OAuth session tombstone is missing")
|
||||
}
|
||||
|
||||
store.completedTTL = 2 * time.Minute
|
||||
if got := store.CompleteProvider("codex", oauthSessionSourceBuiltin); got != 1 {
|
||||
t.Fatalf("CompleteProvider() = %d, want 1 newly completed session", got)
|
||||
}
|
||||
completedAfter, ok := store.Get("completed-state")
|
||||
if !ok {
|
||||
t.Fatal("completed OAuth session tombstone is missing after provider completion")
|
||||
}
|
||||
if !completedAfter.ExpiresAt.Equal(completedBefore.ExpiresAt) {
|
||||
t.Fatalf("provider completion extended existing tombstone from %s to %s", completedBefore.ExpiresAt, completedAfter.ExpiresAt)
|
||||
}
|
||||
pendingAfter, ok := store.Get("pending-state")
|
||||
if !ok || !pendingAfter.Completed {
|
||||
t.Fatalf("pending session completed/ok = %t/%t, want true/true", pendingAfter.Completed, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetOAuthSessionHidesCompletedSession(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
replaceOAuthSessionStoreForTest(t, store)
|
||||
store.Register("completed-state", "codex")
|
||||
store.Complete("completed-state")
|
||||
|
||||
provider, status, ok := GetOAuthSession("completed-state")
|
||||
if ok {
|
||||
t.Fatalf("GetOAuthSession() = (%q, %q, true), want completed session hidden", provider, status)
|
||||
}
|
||||
|
||||
_, _, _, _, completed, detailsOK := GetOAuthSessionDetails("completed-state")
|
||||
if !detailsOK || !completed {
|
||||
t.Fatalf("GetOAuthSessionDetails() completed/ok = %t/%t, want true/true", completed, detailsOK)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetAuthStatusRejectsUnknownStateAndAcceptsCompletedState(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
replaceOAuthSessionStoreForTest(t, store)
|
||||
|
||||
handler := &Handler{}
|
||||
router := gin.New()
|
||||
router.GET("/status", handler.GetAuthStatus)
|
||||
|
||||
unknown := performOAuthStatusRequest(t, router, "unknown-state")
|
||||
if unknown.Status != "error" || unknown.Error != "unknown or expired state" {
|
||||
t.Fatalf("unknown state response = %#v, want unknown/expired error", unknown)
|
||||
}
|
||||
|
||||
store.Register("completed-state", "codex")
|
||||
store.Complete("completed-state")
|
||||
completed := performOAuthStatusRequest(t, router, "completed-state")
|
||||
if completed.Status != "ok" || completed.Error != "" {
|
||||
t.Fatalf("completed state response = %#v, want success", completed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOAuthCallbackRejectsCompletedSession(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
replaceOAuthSessionStoreForTest(t, store)
|
||||
store.Register("completed-state", "codex")
|
||||
store.Complete("completed-state")
|
||||
|
||||
handler := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, nil)
|
||||
router := gin.New()
|
||||
router.POST("/oauth-callback", handler.PostOAuthCallback)
|
||||
|
||||
req := httptest.NewRequest(
|
||||
http.MethodPost,
|
||||
"/oauth-callback",
|
||||
strings.NewReader(`{"provider":"codex","state":"completed-state","code":"test-code"}`),
|
||||
)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusConflict {
|
||||
t.Fatalf("completed callback status = %d, want %d; body=%s", w.Code, http.StatusConflict, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
type oauthStatusResponse struct {
|
||||
Status string `json:"status"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
|
||||
func performOAuthStatusRequest(t *testing.T, router http.Handler, state string) oauthStatusResponse {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodGet, "/status?state="+state, nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status request returned %d, want %d; body=%s", w.Code, http.StatusOK, w.Body.String())
|
||||
}
|
||||
var response oauthStatusResponse
|
||||
if errDecode := json.Unmarshal(w.Body.Bytes(), &response); errDecode != nil {
|
||||
t.Fatalf("decode status response: %v", errDecode)
|
||||
}
|
||||
return response
|
||||
}
|
||||
|
||||
func TestOAuthSessionStoreCancelRemovesPendingSession(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
store.Register("pending-state", "xai")
|
||||
|
||||
if !store.Cancel("pending-state") {
|
||||
t.Fatal("Cancel() = false, want true for pending session")
|
||||
}
|
||||
if store.IsPending("pending-state", "xai") {
|
||||
t.Fatal("cancelled session remained pending")
|
||||
}
|
||||
if _, ok := store.Get("pending-state"); ok {
|
||||
t.Fatal("cancelled session still present in store")
|
||||
}
|
||||
if store.Cancel("pending-state") {
|
||||
t.Fatal("second Cancel() = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOAuthSessionStoreCancelIgnoresCompletedAndUnknown(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
store.Register("completed-state", "codex")
|
||||
store.Complete("completed-state")
|
||||
|
||||
if store.Cancel("completed-state") {
|
||||
t.Fatal("Cancel() completed session = true, want false")
|
||||
}
|
||||
if _, ok := store.Get("completed-state"); !ok {
|
||||
t.Fatal("completed tombstone was removed by Cancel")
|
||||
}
|
||||
if store.Cancel("missing-state") {
|
||||
t.Fatal("Cancel() unknown session = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOAuthSessionStoreCancelIgnoresErrorSession(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
store.Register("error-state", "kimi")
|
||||
store.SetError("error-state", "Authentication failed")
|
||||
|
||||
if store.IsPending("error-state", "kimi") {
|
||||
t.Fatal("error session should not be pending")
|
||||
}
|
||||
if store.Cancel("error-state") {
|
||||
t.Fatal("Cancel() error session = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelOAuthSessionAndCallbackRejectAfterCancel(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
replaceOAuthSessionStoreForTest(t, store)
|
||||
store.Register("callback-state", "anthropic")
|
||||
|
||||
if !CancelOAuthSession("callback-state") {
|
||||
t.Fatal("CancelOAuthSession() = false, want true")
|
||||
}
|
||||
if IsOAuthSessionPending("callback-state", "anthropic") {
|
||||
t.Fatal("session still pending after cancel")
|
||||
}
|
||||
|
||||
_, errWrite := WriteOAuthCallbackFileForPendingSession(t.TempDir(), "anthropic", "callback-state", "code", "")
|
||||
if errWrite == nil {
|
||||
t.Fatal("expected callback write to fail after cancel")
|
||||
}
|
||||
if !errors.Is(errWrite, errOAuthSessionNotPending) {
|
||||
t.Fatalf("callback write error = %v, want %v", errWrite, errOAuthSessionNotPending)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGuardOAuthSessionPendingForSave(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
replaceOAuthSessionStoreForTest(t, store)
|
||||
|
||||
providers := []string{"anthropic", "codex", "antigravity", "xai", "kimi"}
|
||||
for _, provider := range providers {
|
||||
state := provider + "-save-guard"
|
||||
store.Register(state, provider)
|
||||
|
||||
if errGuard := guardOAuthSessionPendingForSave(state, provider); errGuard != nil {
|
||||
t.Fatalf("%s pending guard error = %v, want nil", provider, errGuard)
|
||||
}
|
||||
|
||||
if !CancelOAuthSession(state) {
|
||||
t.Fatalf("%s CancelOAuthSession() = false, want true", provider)
|
||||
}
|
||||
if errGuard := guardOAuthSessionPendingForSave(state, provider); !errors.Is(errGuard, errOAuthSessionNotPending) {
|
||||
t.Fatalf("%s after cancel guard error = %v, want %v", provider, errGuard, errOAuthSessionNotPending)
|
||||
}
|
||||
}
|
||||
|
||||
// Completed and errored sessions must also refuse save.
|
||||
store.Register("completed-save", "codex")
|
||||
store.Complete("completed-save")
|
||||
if errGuard := guardOAuthSessionPendingForSave("completed-save", "codex"); !errors.Is(errGuard, errOAuthSessionNotPending) {
|
||||
t.Fatalf("completed guard error = %v, want %v", errGuard, errOAuthSessionNotPending)
|
||||
}
|
||||
|
||||
store.Register("error-save", "anthropic")
|
||||
store.SetError("error-save", "Authentication failed")
|
||||
if errGuard := guardOAuthSessionPendingForSave("error-save", "anthropic"); !errors.Is(errGuard, errOAuthSessionNotPending) {
|
||||
t.Fatalf("error guard error = %v, want %v", errGuard, errOAuthSessionNotPending)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelAuthSessionHandler(t *testing.T) {
|
||||
store := newOAuthSessionStore(time.Minute)
|
||||
replaceOAuthSessionStoreForTest(t, store)
|
||||
store.Register("device-state", "xai")
|
||||
|
||||
handler := &Handler{}
|
||||
router := gin.New()
|
||||
router.DELETE("/oauth-session", handler.CancelAuthSession)
|
||||
|
||||
missing := performOAuthCancelRequest(t, router, "")
|
||||
if missing.status != http.StatusBadRequest {
|
||||
t.Fatalf("missing state status = %d, want %d", missing.status, http.StatusBadRequest)
|
||||
}
|
||||
|
||||
invalid := performOAuthCancelRequest(t, router, "bad/state")
|
||||
if invalid.status != http.StatusBadRequest {
|
||||
t.Fatalf("invalid state status = %d, want %d", invalid.status, http.StatusBadRequest)
|
||||
}
|
||||
|
||||
cancelled := performOAuthCancelRequest(t, router, "device-state")
|
||||
if cancelled.status != http.StatusOK || !cancelled.cancelled || cancelled.bodyStatus != "ok" {
|
||||
t.Fatalf("cancel pending response = %#v, want ok/cancelled", cancelled)
|
||||
}
|
||||
if IsOAuthSessionPending("device-state", "xai") {
|
||||
t.Fatal("device session still pending after cancel API")
|
||||
}
|
||||
|
||||
repeat := performOAuthCancelRequest(t, router, "device-state")
|
||||
if repeat.status != http.StatusOK || repeat.cancelled {
|
||||
t.Fatalf("repeat cancel response = %#v, want ok with cancelled=false", repeat)
|
||||
}
|
||||
|
||||
// Status after cancel should not report success.
|
||||
statusRouter := gin.New()
|
||||
statusRouter.GET("/status", handler.GetAuthStatus)
|
||||
unknown := performOAuthStatusRequest(t, statusRouter, "device-state")
|
||||
if unknown.Status != "error" || unknown.Error != "unknown or expired state" {
|
||||
t.Fatalf("status after cancel = %#v, want unknown/expired error", unknown)
|
||||
}
|
||||
}
|
||||
|
||||
type oauthCancelResponse struct {
|
||||
status int
|
||||
bodyStatus string
|
||||
cancelled bool
|
||||
}
|
||||
|
||||
func performOAuthCancelRequest(t *testing.T, router http.Handler, state string) oauthCancelResponse {
|
||||
t.Helper()
|
||||
path := "/oauth-session"
|
||||
if state != "" {
|
||||
path += "?state=" + state
|
||||
}
|
||||
req := httptest.NewRequest(http.MethodDelete, path, nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
var body struct {
|
||||
Status string `json:"status"`
|
||||
Cancelled bool `json:"cancelled"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
if w.Body.Len() > 0 {
|
||||
if errDecode := json.Unmarshal(w.Body.Bytes(), &body); errDecode != nil {
|
||||
t.Fatalf("decode cancel response: %v body=%s", errDecode, w.Body.String())
|
||||
}
|
||||
}
|
||||
return oauthCancelResponse{
|
||||
status: w.Code,
|
||||
bodyStatus: body.Status,
|
||||
cancelled: body.Cancelled,
|
||||
}
|
||||
}
|
||||
|
||||
func replaceOAuthSessionStoreForTest(t *testing.T, store *oauthSessionStore) {
|
||||
t.Helper()
|
||||
original := oauthSessions
|
||||
oauthSessions = store
|
||||
t.Cleanup(func() {
|
||||
oauthSessions = original
|
||||
})
|
||||
}
|
||||
937
backend/internal/api/handlers/management/plugin_store.go
Normal file
937
backend/internal/api/handlers/management/plugin_store.go
Normal file
|
|
@ -0,0 +1,937 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/htmlsanitize"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginstore"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
const (
|
||||
// pluginReleaseCacheTTL bounds how long a resolved latest release version is
|
||||
// reused before the GitHub API is queried again.
|
||||
pluginReleaseCacheTTL = 10 * time.Minute
|
||||
// pluginReleaseFailureCacheTTL throttles retries after a failed lookup so a
|
||||
// rate-limited or unreachable API is not hammered on every listing.
|
||||
pluginReleaseFailureCacheTTL = 30 * time.Second
|
||||
)
|
||||
|
||||
type pluginReleaseCacheEntry struct {
|
||||
version string
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
type pluginStoreListResponse struct {
|
||||
PluginsEnabled bool `json:"plugins_enabled"`
|
||||
PluginsDir string `json:"plugins_dir"`
|
||||
Sources []pluginStoreSource `json:"sources"`
|
||||
SourceErrors []pluginStoreSourceErr `json:"source_errors,omitempty"`
|
||||
Plugins []pluginStoreListEntry `json:"plugins"`
|
||||
}
|
||||
|
||||
type pluginStoreSource struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
URL string `json:"url"`
|
||||
}
|
||||
|
||||
type pluginStoreSourceErr struct {
|
||||
SourceID string `json:"source_id"`
|
||||
SourceName string `json:"source_name"`
|
||||
SourceURL string `json:"source_url"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
type pluginStoreListEntry struct {
|
||||
StoreID string `json:"store_id"`
|
||||
SourceID string `json:"source_id"`
|
||||
SourceName string `json:"source_name"`
|
||||
SourceURL string `json:"source_url"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Author string `json:"author"`
|
||||
Version string `json:"version"`
|
||||
Repository string `json:"repository"`
|
||||
InstallType string `json:"install_type"`
|
||||
AuthRequired bool `json:"auth_required"`
|
||||
AuthConfigured bool `json:"auth_configured"`
|
||||
Platforms []pluginStorePlatform `json:"platforms,omitempty"`
|
||||
Logo string `json:"logo,omitempty"`
|
||||
Homepage string `json:"homepage,omitempty"`
|
||||
License string `json:"license,omitempty"`
|
||||
Tags []string `json:"tags,omitempty"`
|
||||
Installed bool `json:"installed"`
|
||||
InstalledVersion string `json:"installed_version"`
|
||||
InstalledSourceID string `json:"installed_source_id,omitempty"`
|
||||
InstallSourceStatus string `json:"install_source_status,omitempty"`
|
||||
Path string `json:"path"`
|
||||
Configured bool `json:"configured"`
|
||||
Registered bool `json:"registered"`
|
||||
Enabled bool `json:"enabled"`
|
||||
EffectiveEnabled bool `json:"effective_enabled"`
|
||||
UpdateAvailable bool `json:"update_available"`
|
||||
}
|
||||
|
||||
type pluginStorePlatform struct {
|
||||
GOOS string `json:"goos"`
|
||||
GOARCH string `json:"goarch"`
|
||||
}
|
||||
|
||||
type pluginInstallResponse struct {
|
||||
Status string `json:"status"`
|
||||
SourceID string `json:"source_id"`
|
||||
SourceName string `json:"source_name"`
|
||||
SourceURL string `json:"source_url"`
|
||||
ID string `json:"id"`
|
||||
Version string `json:"version"`
|
||||
InstallType string `json:"install_type"`
|
||||
Path string `json:"path"`
|
||||
PluginsEnabled bool `json:"plugins_enabled"`
|
||||
RestartRequired bool `json:"restart_required"`
|
||||
}
|
||||
|
||||
type pluginInstallRequest struct {
|
||||
Version string `json:"version"`
|
||||
}
|
||||
|
||||
type pluginLocalStatus struct {
|
||||
Installed bool
|
||||
InstalledVersion string
|
||||
StoreManaged bool
|
||||
InstalledSourceID string
|
||||
InstalledSourceURL string
|
||||
Path string
|
||||
Configured bool
|
||||
Registered bool
|
||||
Enabled bool
|
||||
EffectiveEnabled bool
|
||||
}
|
||||
|
||||
type sourcedPlugin struct {
|
||||
source pluginstore.Source
|
||||
plugin pluginstore.Plugin
|
||||
}
|
||||
|
||||
func (h *Handler) ListPluginStore(c *gin.Context) {
|
||||
pluginsEnabled, pluginsDir, proxyURL, sourceConfigs, storeAuth, configs, host := h.pluginStoreSnapshot()
|
||||
resolvedPluginsDir, errResolvePluginsDir := config.ResolvePluginsDir(pluginsDir)
|
||||
if errResolvePluginsDir != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_directory_invalid", "message": errResolvePluginsDir.Error()})
|
||||
return
|
||||
}
|
||||
pluginsDir = resolvedPluginsDir
|
||||
sources, errSources := h.pluginStoreSources(sourceConfigs)
|
||||
if errSources != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_store_source_invalid", "message": errSources.Error()})
|
||||
return
|
||||
}
|
||||
plugins, sourceErrors := h.fetchSourcedPlugins(c.Request.Context(), proxyURL, storeAuth, sources)
|
||||
if len(plugins) == 0 && len(sourceErrors) > 0 {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "plugin_store_registry_failed", "message": sourceErrors[0].Message})
|
||||
return
|
||||
}
|
||||
statuses, errStatus := pluginLocalStatuses(pluginsEnabled, pluginsDir, configs, host)
|
||||
if errStatus != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_discovery_failed", "message": errStatus.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
latestInput := make([]pluginstore.Plugin, 0, len(plugins))
|
||||
for _, item := range plugins {
|
||||
latestInput = append(latestInput, item.plugin)
|
||||
}
|
||||
client := h.newPluginStoreClient(proxyURL, "", storeAuth)
|
||||
latestVersions := h.latestPluginVersions(c.Request.Context(), client, latestInput)
|
||||
pluginSourceCounts := make(map[string]int, len(plugins))
|
||||
for _, item := range plugins {
|
||||
pluginSourceCounts[item.plugin.ID]++
|
||||
}
|
||||
|
||||
entries := make([]pluginStoreListEntry, 0, len(plugins))
|
||||
for index, item := range plugins {
|
||||
plugin := item.plugin
|
||||
status := statuses[plugin.ID]
|
||||
installedSourceID, installSourceStatus, sourceAllowsUpdate := pluginStoreInstallSourceStatus(
|
||||
status,
|
||||
sources,
|
||||
item.source.ID,
|
||||
pluginSourceCounts[plugin.ID],
|
||||
)
|
||||
installedVersion := status.InstalledVersion
|
||||
// Fall back to the registry version when the latest release is unknown.
|
||||
storeVersion := plugin.Version
|
||||
if latestVersions[index] != "" {
|
||||
storeVersion = latestVersions[index]
|
||||
}
|
||||
entries = append(entries, pluginStoreListEntry{
|
||||
StoreID: htmlsanitize.String(item.source.ID + "/" + plugin.ID),
|
||||
SourceID: htmlsanitize.String(item.source.ID),
|
||||
SourceName: htmlsanitize.String(item.source.Name),
|
||||
SourceURL: htmlsanitize.String(item.source.URL),
|
||||
ID: htmlsanitize.String(plugin.ID),
|
||||
Name: htmlsanitize.String(plugin.Name),
|
||||
Description: htmlsanitize.String(plugin.Description),
|
||||
Author: htmlsanitize.String(plugin.Author),
|
||||
Version: htmlsanitize.String(storeVersion),
|
||||
Repository: htmlsanitize.String(plugin.Repository),
|
||||
InstallType: htmlsanitize.String(pluginstore.PluginInstallType(plugin)),
|
||||
AuthRequired: plugin.AuthRequired,
|
||||
AuthConfigured: pluginAuthConfigured(item.source, plugin, storeAuth),
|
||||
Platforms: sanitizePluginStorePlatforms(pluginstore.PluginPlatforms(plugin)),
|
||||
Logo: htmlsanitize.String(plugin.Logo),
|
||||
Homepage: htmlsanitize.String(plugin.Homepage),
|
||||
License: htmlsanitize.String(plugin.License),
|
||||
Tags: htmlsanitize.Strings(plugin.Tags),
|
||||
Installed: status.Installed,
|
||||
InstalledVersion: htmlsanitize.String(installedVersion),
|
||||
InstalledSourceID: htmlsanitize.String(installedSourceID),
|
||||
InstallSourceStatus: htmlsanitize.String(installSourceStatus),
|
||||
Path: htmlsanitize.String(status.Path),
|
||||
Configured: status.Configured,
|
||||
Registered: status.Registered,
|
||||
Enabled: status.Enabled,
|
||||
EffectiveEnabled: status.EffectiveEnabled,
|
||||
UpdateAvailable: sourceAllowsUpdate && pluginstore.UpdateAvailable(installedVersion, storeVersion),
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, pluginStoreListResponse{
|
||||
PluginsEnabled: pluginsEnabled,
|
||||
PluginsDir: htmlsanitize.String(pluginsDir),
|
||||
Sources: sanitizePluginStoreSources(sources),
|
||||
SourceErrors: sanitizePluginStoreSourceErrors(sourceErrors),
|
||||
Plugins: entries,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *Handler) InstallPluginFromStore(c *gin.Context) {
|
||||
h.installPluginFromStore(c, runtime.GOOS, runtime.GOARCH)
|
||||
}
|
||||
|
||||
func (h *Handler) installPluginFromStore(c *gin.Context, goos, goarch string) {
|
||||
id, okID := pluginIDFromRequest(c)
|
||||
if !okID {
|
||||
return
|
||||
}
|
||||
requestedVersion, errVersionRequest := pluginInstallRequestedVersion(c)
|
||||
if errVersionRequest != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_request", "message": errVersionRequest.Error()})
|
||||
return
|
||||
}
|
||||
installCtx := c.Request.Context()
|
||||
pluginsEnabled, pluginsDir, proxyURL, sourceConfigs, storeAuth, configs, host := h.pluginStoreSnapshot()
|
||||
resolvedPluginsDir, errResolvePluginsDir := config.ResolvePluginsDir(pluginsDir)
|
||||
if errResolvePluginsDir != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_directory_invalid", "message": errResolvePluginsDir.Error()})
|
||||
return
|
||||
}
|
||||
pluginsDir = resolvedPluginsDir
|
||||
sources, errSources := h.pluginStoreSources(sourceConfigs)
|
||||
if errSources != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_store_source_invalid", "message": errSources.Error()})
|
||||
return
|
||||
}
|
||||
source, plugin, client, okPlugin := h.findPluginStoreInstallTarget(installCtx, proxyURL, storeAuth, sources, id, c.Query("source"), c)
|
||||
if !okPlugin {
|
||||
return
|
||||
}
|
||||
if !validatePluginStoreInstallSource(c, configs, sources, id, source.ID) {
|
||||
return
|
||||
}
|
||||
pluginIsBusy := func() bool { return pluginBusy(host, id) }
|
||||
installOptions := pluginstore.InstallOptions{
|
||||
PluginsDir: pluginsDir,
|
||||
GOOS: goos,
|
||||
GOARCH: goarch,
|
||||
PluginLoaded: pluginIsBusy,
|
||||
}
|
||||
var manifest pluginstore.Manifest
|
||||
var result pluginstore.InstallResult
|
||||
var errInstall error
|
||||
switch pluginstore.PluginInstallType(plugin) {
|
||||
case pluginstore.InstallTypeDirect:
|
||||
var errManifest error
|
||||
manifest, errManifest = pluginStoreDirectManifest(source, plugin, requestedVersion)
|
||||
if errManifest != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "plugin_manifest_invalid", "message": errManifest.Error()})
|
||||
return
|
||||
}
|
||||
result, errInstall = client.InstallManifest(installCtx, manifest, installOptions)
|
||||
case pluginstore.InstallTypeGitHubRelease:
|
||||
result, errInstall = installPluginStoreGitHubRelease(installCtx, client, plugin, requestedVersion, installOptions)
|
||||
default:
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "plugin_manifest_invalid", "message": fmt.Sprintf("unsupported install type %q", plugin.Install.Type)})
|
||||
return
|
||||
}
|
||||
if errInstall != nil {
|
||||
if errors.Is(errInstall, pluginstore.ErrLoadedPluginLocked) {
|
||||
c.JSON(http.StatusConflict, gin.H{
|
||||
"error": "plugin_update_requires_restart",
|
||||
"message": "loaded plugin cannot be overwritten while the server is running",
|
||||
"restart_required": true,
|
||||
})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "plugin_install_failed", "message": errInstall.Error()})
|
||||
return
|
||||
}
|
||||
if manifest.ID == "" {
|
||||
var errManifest error
|
||||
manifest, errManifest = pluginStoreManifestForInstall(source, plugin, result)
|
||||
if errManifest != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"error": "plugin_manifest_failed",
|
||||
"message": fmt.Sprintf("plugin file installed at %s but creating store manifest failed: %s", result.Path, errManifest.Error()),
|
||||
"path": result.Path,
|
||||
})
|
||||
return
|
||||
}
|
||||
}
|
||||
restartRequired := false
|
||||
|
||||
h.mu.Lock()
|
||||
if h.cfg == nil {
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"error": "config_unavailable",
|
||||
"message": fmt.Sprintf("plugin file installed at %s but config is unavailable to enable it", result.Path),
|
||||
"path": result.Path,
|
||||
})
|
||||
return
|
||||
}
|
||||
if errEnable := h.enablePluginConfigLocked(id, manifest); errEnable != nil {
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"error": "config_update_failed",
|
||||
"message": fmt.Sprintf("plugin file installed at %s but enabling it in config failed: %s", result.Path, errEnable.Error()),
|
||||
"path": result.Path,
|
||||
})
|
||||
return
|
||||
}
|
||||
if errSave := config.SaveConfigPreserveComments(h.configFilePath, h.cfg); errSave != nil {
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"error": "config_save_failed",
|
||||
"message": fmt.Sprintf("plugin file installed at %s but saving config failed: %s", result.Path, errSave.Error()),
|
||||
"path": result.Path,
|
||||
})
|
||||
return
|
||||
}
|
||||
cfgSnapshot := h.reloadSnapshotConfigLocked()
|
||||
h.mu.Unlock()
|
||||
|
||||
h.reloadConfigAfterManagementSaveAsync(c.Request.Context(), cfgSnapshot)
|
||||
log.WithFields(log.Fields{
|
||||
"plugin_id": result.ID,
|
||||
"plugin_name": plugin.Name,
|
||||
"source_id": source.ID,
|
||||
"version": result.Version,
|
||||
"install_type": result.InstallType,
|
||||
"path": result.Path,
|
||||
"overwritten": result.Overwritten,
|
||||
}).Info("pluginstore: plugin installed")
|
||||
|
||||
c.JSON(http.StatusOK, pluginInstallResponse{
|
||||
Status: "installed",
|
||||
SourceID: htmlsanitize.String(source.ID),
|
||||
SourceName: htmlsanitize.String(source.Name),
|
||||
SourceURL: htmlsanitize.String(source.URL),
|
||||
ID: htmlsanitize.String(result.ID),
|
||||
Version: htmlsanitize.String(result.Version),
|
||||
InstallType: htmlsanitize.String(result.InstallType),
|
||||
Path: htmlsanitize.String(result.Path),
|
||||
PluginsEnabled: pluginsEnabled,
|
||||
RestartRequired: restartRequired,
|
||||
})
|
||||
}
|
||||
|
||||
func pluginStoreDirectManifest(source pluginstore.Source, plugin pluginstore.Plugin, requestedVersion string) (pluginstore.Manifest, error) {
|
||||
version := normalizePluginStoreRequestedVersion(requestedVersion)
|
||||
if version == "" {
|
||||
version = normalizePluginStoreRequestedVersion(plugin.Version)
|
||||
}
|
||||
if normalizePluginStoreRequestedVersion(plugin.Version) == version {
|
||||
plugin.Version = version
|
||||
return pluginstore.ManifestFromPlugin(source, plugin)
|
||||
}
|
||||
for _, candidate := range plugin.Versions {
|
||||
if normalizePluginStoreRequestedVersion(candidate.Version) != version {
|
||||
continue
|
||||
}
|
||||
plugin.Version = version
|
||||
plugin.Install = candidate.Install
|
||||
if strings.TrimSpace(plugin.Install.Type) == "" {
|
||||
plugin.Install.Type = pluginstore.InstallTypeDirect
|
||||
}
|
||||
return pluginstore.ManifestFromPlugin(source, plugin)
|
||||
}
|
||||
return pluginstore.Manifest{}, fmt.Errorf("direct plugin version %q not found", version)
|
||||
}
|
||||
|
||||
func installPluginStoreGitHubRelease(ctx context.Context, client pluginstore.Client, plugin pluginstore.Plugin, requestedVersion string, options pluginstore.InstallOptions) (pluginstore.InstallResult, error) {
|
||||
version := normalizePluginStoreRequestedVersion(requestedVersion)
|
||||
if version == "" {
|
||||
return client.Install(ctx, plugin, options)
|
||||
}
|
||||
tags := pluginStoreReleaseTagCandidates(requestedVersion)
|
||||
errs := make([]error, 0, len(tags))
|
||||
for _, tag := range tags {
|
||||
result, errInstall := client.InstallVersion(ctx, plugin, tag, version, options)
|
||||
if errInstall == nil {
|
||||
return result, nil
|
||||
}
|
||||
errs = append(errs, fmt.Errorf("%s: %w", tag, errInstall))
|
||||
}
|
||||
return pluginstore.InstallResult{}, fmt.Errorf("install release by tag: %w", errors.Join(errs...))
|
||||
}
|
||||
|
||||
func pluginStoreManifestForInstall(source pluginstore.Source, plugin pluginstore.Plugin, result pluginstore.InstallResult) (pluginstore.Manifest, error) {
|
||||
installType := strings.TrimSpace(result.InstallType)
|
||||
if installType == "" {
|
||||
installType = pluginstore.PluginInstallType(plugin)
|
||||
}
|
||||
switch installType {
|
||||
case pluginstore.InstallTypeDirect:
|
||||
plugin.Version = strings.TrimSpace(result.Version)
|
||||
plugin.Install = pluginstore.NormalizeInstallPlan(plugin.Install)
|
||||
return pluginstore.ManifestFromPlugin(source, plugin)
|
||||
case pluginstore.InstallTypeGitHubRelease:
|
||||
releaseTag := strings.TrimSpace(result.ReleaseTag)
|
||||
if releaseTag == "" {
|
||||
return pluginstore.Manifest{}, fmt.Errorf("release tag is required")
|
||||
}
|
||||
return pluginstore.ManifestFromRelease(source, plugin, pluginstore.Release{TagName: releaseTag})
|
||||
default:
|
||||
return pluginstore.Manifest{}, fmt.Errorf("unsupported install type %q", result.InstallType)
|
||||
}
|
||||
}
|
||||
|
||||
func pluginInstallRequestedVersion(c *gin.Context) (string, error) {
|
||||
requestedVersion := strings.TrimSpace(c.Query("version"))
|
||||
if c == nil || c.Request == nil || c.Request.Body == nil || c.Request.Body == http.NoBody {
|
||||
return requestedVersion, nil
|
||||
}
|
||||
body, errRead := io.ReadAll(c.Request.Body)
|
||||
if errRead != nil {
|
||||
return "", fmt.Errorf("read install request: %w", errRead)
|
||||
}
|
||||
if strings.TrimSpace(string(body)) == "" {
|
||||
return requestedVersion, nil
|
||||
}
|
||||
var req pluginInstallRequest
|
||||
if errDecode := json.Unmarshal(body, &req); errDecode != nil {
|
||||
return "", fmt.Errorf("decode install request: %w", errDecode)
|
||||
}
|
||||
bodyVersion := strings.TrimSpace(req.Version)
|
||||
if requestedVersion == "" {
|
||||
return bodyVersion, nil
|
||||
}
|
||||
if bodyVersion == "" || normalizePluginStoreRequestedVersion(bodyVersion) == normalizePluginStoreRequestedVersion(requestedVersion) {
|
||||
return requestedVersion, nil
|
||||
}
|
||||
return "", fmt.Errorf("version query %q does not match request body version %q", requestedVersion, bodyVersion)
|
||||
}
|
||||
|
||||
func pluginStoreReleaseTagCandidates(version string) []string {
|
||||
version = strings.TrimSpace(version)
|
||||
if version == "" {
|
||||
return nil
|
||||
}
|
||||
if strings.HasPrefix(strings.ToLower(version), "v") {
|
||||
return []string{version, strings.TrimSpace(version[1:])}
|
||||
}
|
||||
return []string{version, "v" + version}
|
||||
}
|
||||
|
||||
func normalizePluginStoreRequestedVersion(version string) string {
|
||||
version = strings.TrimSpace(version)
|
||||
if strings.HasPrefix(strings.ToLower(version), "v") {
|
||||
return strings.TrimSpace(version[1:])
|
||||
}
|
||||
return version
|
||||
}
|
||||
|
||||
// enablePluginConfigLocked sets plugins.configs.<id>.enabled and store while
|
||||
// preserving the rest of the plugin's raw configuration. Callers must hold h.mu.
|
||||
func (h *Handler) enablePluginConfigLocked(id string, storeManifest pluginstore.Manifest) error {
|
||||
ensurePluginConfigMap(h.cfg)
|
||||
node := pluginConfigNode(h.cfg.Plugins.Configs[id])
|
||||
storeNode, errStoreNode := pluginStoreManifestYAMLNode(storeManifest)
|
||||
if errStoreNode != nil {
|
||||
return errStoreNode
|
||||
}
|
||||
setYAMLMappingValue(node, "enabled", boolYAMLNode(true))
|
||||
setYAMLMappingValue(node, "store", storeNode)
|
||||
updated, errConfig := pluginInstanceConfigFromNode(node)
|
||||
if errConfig != nil {
|
||||
return fmt.Errorf("decode plugin config: %w", errConfig)
|
||||
}
|
||||
h.cfg.Plugins.Configs[id] = updated
|
||||
return nil
|
||||
}
|
||||
|
||||
func pluginStoreManifestYAMLNode(manifest pluginstore.Manifest) (*yaml.Node, error) {
|
||||
var node yaml.Node
|
||||
if errEncode := node.Encode(manifest); errEncode != nil {
|
||||
return nil, fmt.Errorf("encode store manifest: %w", errEncode)
|
||||
}
|
||||
return &node, nil
|
||||
}
|
||||
|
||||
func (h *Handler) pluginStoreSnapshot() (bool, string, string, []string, []pluginstore.AuthConfig, map[string]config.PluginInstanceConfig, *pluginhost.Host) {
|
||||
if h == nil {
|
||||
return false, "plugins", "", nil, nil, map[string]config.PluginInstanceConfig{}, nil
|
||||
}
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.cfg == nil {
|
||||
return false, "plugins", "", nil, nil, map[string]config.PluginInstanceConfig{}, nil
|
||||
}
|
||||
pluginsEnabled := h.cfg.Plugins.Enabled
|
||||
pluginsDir := normalizedPluginsDir(h.cfg.Plugins.Dir)
|
||||
proxyURL := strings.TrimSpace(h.cfg.ProxyURL)
|
||||
sourceConfigs := append([]string(nil), h.cfg.Plugins.StoreSources...)
|
||||
storeAuth := append([]pluginstore.AuthConfig(nil), h.cfg.Plugins.StoreAuth...)
|
||||
configs := make(map[string]config.PluginInstanceConfig, len(h.cfg.Plugins.Configs))
|
||||
for id, item := range h.cfg.Plugins.Configs {
|
||||
configs[id] = item
|
||||
}
|
||||
return pluginsEnabled, pluginsDir, proxyURL, sourceConfigs, storeAuth, configs, h.pluginHost
|
||||
}
|
||||
|
||||
func (h *Handler) pluginStoreSources(sourceConfigs []string) ([]pluginstore.Source, error) {
|
||||
if h != nil && strings.TrimSpace(h.pluginStoreRegistryURL) != "" {
|
||||
source := pluginstore.DefaultSource()
|
||||
source.URL = strings.TrimSpace(h.pluginStoreRegistryURL)
|
||||
return []pluginstore.Source{source}, nil
|
||||
}
|
||||
return pluginstore.NormalizeSources(sourceConfigs)
|
||||
}
|
||||
|
||||
func (h *Handler) newPluginStoreClient(proxyURL string, registryURL string, storeAuth []pluginstore.AuthConfig) pluginstore.Client {
|
||||
registryURL = strings.TrimSpace(registryURL)
|
||||
var httpClient pluginstore.HTTPDoer
|
||||
if h != nil {
|
||||
httpClient = h.pluginStoreHTTPClient
|
||||
}
|
||||
if registryURL == "" {
|
||||
registryURL = pluginstore.DefaultRegistryURL
|
||||
}
|
||||
if httpClient != nil {
|
||||
return pluginstore.Client{HTTPClient: httpClient, RegistryURL: registryURL, Auth: storeAuth}
|
||||
}
|
||||
client := &http.Client{}
|
||||
if strings.TrimSpace(proxyURL) != "" {
|
||||
util.SetProxy(&sdkconfig.SDKConfig{ProxyURL: strings.TrimSpace(proxyURL)}, client)
|
||||
}
|
||||
return pluginstore.Client{HTTPClient: client, RegistryURL: registryURL, Auth: storeAuth}
|
||||
}
|
||||
|
||||
func (h *Handler) fetchSourcedPlugins(ctx context.Context, proxyURL string, storeAuth []pluginstore.AuthConfig, sources []pluginstore.Source) ([]sourcedPlugin, []pluginStoreSourceErr) {
|
||||
plugins := make([]sourcedPlugin, 0)
|
||||
sourceErrors := make([]pluginStoreSourceErr, 0)
|
||||
for _, source := range sources {
|
||||
client := h.newPluginStoreClient(proxyURL, source.URL, storeAuth)
|
||||
registry, errRegistry := client.FetchRegistry(ctx)
|
||||
if errRegistry != nil {
|
||||
sourceErrors = append(sourceErrors, pluginStoreSourceErr{
|
||||
SourceID: source.ID,
|
||||
SourceName: source.Name,
|
||||
SourceURL: source.URL,
|
||||
Message: errRegistry.Error(),
|
||||
})
|
||||
continue
|
||||
}
|
||||
for _, plugin := range registry.Plugins {
|
||||
plugins = append(plugins, sourcedPlugin{source: source, plugin: plugin})
|
||||
}
|
||||
}
|
||||
return plugins, sourceErrors
|
||||
}
|
||||
|
||||
func (h *Handler) findPluginStoreInstallTarget(ctx context.Context, proxyURL string, storeAuth []pluginstore.AuthConfig, sources []pluginstore.Source, id string, requestedSourceID string, c *gin.Context) (pluginstore.Source, pluginstore.Plugin, pluginstore.Client, bool) {
|
||||
requestedSourceID = strings.TrimSpace(requestedSourceID)
|
||||
if requestedSourceID != "" {
|
||||
for _, source := range sources {
|
||||
if source.ID != requestedSourceID {
|
||||
continue
|
||||
}
|
||||
client := h.newPluginStoreClient(proxyURL, source.URL, storeAuth)
|
||||
registry, errRegistry := client.FetchRegistry(ctx)
|
||||
if errRegistry != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "plugin_store_registry_failed", "message": errRegistry.Error()})
|
||||
return pluginstore.Source{}, pluginstore.Plugin{}, pluginstore.Client{}, false
|
||||
}
|
||||
plugin, okPlugin := registry.PluginByID(id)
|
||||
if !okPlugin {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "plugin_not_found", "message": "plugin not found in registry source"})
|
||||
return pluginstore.Source{}, pluginstore.Plugin{}, pluginstore.Client{}, false
|
||||
}
|
||||
return source, plugin, client, true
|
||||
}
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "plugin_store_source_not_found", "message": "plugin store source not found"})
|
||||
return pluginstore.Source{}, pluginstore.Plugin{}, pluginstore.Client{}, false
|
||||
}
|
||||
|
||||
plugins, sourceErrors := h.fetchSourcedPlugins(ctx, proxyURL, storeAuth, sources)
|
||||
matches := make([]sourcedPlugin, 0)
|
||||
for _, item := range plugins {
|
||||
if item.plugin.ID == id {
|
||||
matches = append(matches, item)
|
||||
}
|
||||
}
|
||||
if len(matches) == 0 {
|
||||
if len(plugins) == 0 && len(sourceErrors) > 0 {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "plugin_store_registry_failed", "message": sourceErrors[0].Message})
|
||||
return pluginstore.Source{}, pluginstore.Plugin{}, pluginstore.Client{}, false
|
||||
}
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "plugin_not_found", "message": "plugin not found in registry"})
|
||||
return pluginstore.Source{}, pluginstore.Plugin{}, pluginstore.Client{}, false
|
||||
}
|
||||
if len(matches) > 1 {
|
||||
c.JSON(http.StatusConflict, gin.H{
|
||||
"error": "plugin_store_source_required",
|
||||
"message": "multiple plugin store sources contain this plugin id; specify source",
|
||||
"sources": sanitizePluginStoreSources(sourcedPluginSources(matches)),
|
||||
})
|
||||
return pluginstore.Source{}, pluginstore.Plugin{}, pluginstore.Client{}, false
|
||||
}
|
||||
match := matches[0]
|
||||
return match.source, match.plugin, h.newPluginStoreClient(proxyURL, match.source.URL, storeAuth), true
|
||||
}
|
||||
|
||||
func sourcedPluginSources(plugins []sourcedPlugin) []pluginstore.Source {
|
||||
sources := make([]pluginstore.Source, 0, len(plugins))
|
||||
for _, item := range plugins {
|
||||
sources = append(sources, item.source)
|
||||
}
|
||||
return sources
|
||||
}
|
||||
|
||||
func sanitizePluginStoreSources(sources []pluginstore.Source) []pluginStoreSource {
|
||||
out := make([]pluginStoreSource, 0, len(sources))
|
||||
for _, source := range sources {
|
||||
out = append(out, pluginStoreSource{
|
||||
ID: htmlsanitize.String(source.ID),
|
||||
Name: htmlsanitize.String(source.Name),
|
||||
URL: htmlsanitize.String(source.URL),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func sanitizePluginStoreSourceErrors(sourceErrors []pluginStoreSourceErr) []pluginStoreSourceErr {
|
||||
if len(sourceErrors) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]pluginStoreSourceErr, 0, len(sourceErrors))
|
||||
for _, sourceError := range sourceErrors {
|
||||
out = append(out, pluginStoreSourceErr{
|
||||
SourceID: htmlsanitize.String(sourceError.SourceID),
|
||||
SourceName: htmlsanitize.String(sourceError.SourceName),
|
||||
SourceURL: htmlsanitize.String(sourceError.SourceURL),
|
||||
Message: htmlsanitize.String(sourceError.Message),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func sanitizePluginStorePlatforms(platforms []pluginstore.Platform) []pluginStorePlatform {
|
||||
if len(platforms) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]pluginStorePlatform, 0, len(platforms))
|
||||
for _, platform := range platforms {
|
||||
out = append(out, pluginStorePlatform{
|
||||
GOOS: htmlsanitize.String(platform.GOOS),
|
||||
GOARCH: htmlsanitize.String(platform.GOARCH),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func pluginAuthConfigured(source pluginstore.Source, plugin pluginstore.Plugin, storeAuth []pluginstore.AuthConfig) bool {
|
||||
return pluginstore.PluginAuthConfigured(source, plugin, storeAuth)
|
||||
}
|
||||
|
||||
// latestPluginVersions resolves the latest release version of each registry
|
||||
// plugin concurrently, returning results positionally aligned with plugins.
|
||||
// Unresolved entries are left empty so callers can fall back gracefully.
|
||||
func (h *Handler) latestPluginVersions(ctx context.Context, client pluginstore.Client, plugins []pluginstore.Plugin) []string {
|
||||
versions := make([]string, len(plugins))
|
||||
var wg sync.WaitGroup
|
||||
for index := range plugins {
|
||||
wg.Add(1)
|
||||
go func(index int) {
|
||||
defer wg.Done()
|
||||
versions[index] = h.latestPluginVersion(ctx, client, plugins[index])
|
||||
}(index)
|
||||
}
|
||||
wg.Wait()
|
||||
return versions
|
||||
}
|
||||
|
||||
// latestPluginVersion returns the plugin's latest release version, caching
|
||||
// lookups per repository so repeated listings do not exhaust the GitHub API
|
||||
// rate limit. Failed lookups are cached for a shorter interval and reported
|
||||
// as an empty version.
|
||||
func (h *Handler) latestPluginVersion(ctx context.Context, client pluginstore.Client, plugin pluginstore.Plugin) string {
|
||||
if pluginstore.PluginInstallType(plugin) != pluginstore.InstallTypeGitHubRelease {
|
||||
return ""
|
||||
}
|
||||
repository := strings.TrimSpace(plugin.Repository)
|
||||
if repository == "" {
|
||||
return ""
|
||||
}
|
||||
now := time.Now()
|
||||
h.pluginReleaseCacheMu.Lock()
|
||||
entry, found := h.pluginReleaseCache[repository]
|
||||
h.pluginReleaseCacheMu.Unlock()
|
||||
if found && now.Before(entry.expiresAt) {
|
||||
return entry.version
|
||||
}
|
||||
|
||||
version := ""
|
||||
ttl := pluginReleaseFailureCacheTTL
|
||||
release, errRelease := client.FetchLatestRelease(ctx, plugin)
|
||||
if errRelease != nil {
|
||||
log.WithError(errRelease).WithField("plugin_id", plugin.ID).Warn("pluginstore: failed to fetch latest release")
|
||||
} else if latestVersion, errVersion := pluginstore.ReleaseVersion(release); errVersion != nil {
|
||||
log.WithError(errVersion).WithField("plugin_id", plugin.ID).Warn("pluginstore: invalid latest release tag")
|
||||
} else {
|
||||
version = latestVersion
|
||||
ttl = pluginReleaseCacheTTL
|
||||
}
|
||||
|
||||
h.pluginReleaseCacheMu.Lock()
|
||||
if h.pluginReleaseCache == nil {
|
||||
h.pluginReleaseCache = make(map[string]pluginReleaseCacheEntry)
|
||||
}
|
||||
h.pluginReleaseCache[repository] = pluginReleaseCacheEntry{version: version, expiresAt: now.Add(ttl)}
|
||||
h.pluginReleaseCacheMu.Unlock()
|
||||
return version
|
||||
}
|
||||
|
||||
func pluginLocalStatuses(pluginsEnabled bool, pluginsDir string, configs map[string]config.PluginInstanceConfig, host *pluginhost.Host) (map[string]pluginLocalStatus, error) {
|
||||
statuses := map[string]pluginLocalStatus{}
|
||||
files, errDiscover := pluginhost.DiscoverPluginFiles(pluginsDir, pluginStoreDesiredVersions(configs))
|
||||
if errDiscover != nil {
|
||||
return nil, errDiscover
|
||||
}
|
||||
for _, file := range files {
|
||||
status := statuses[file.ID]
|
||||
status.Installed = true
|
||||
status.Path = file.Path
|
||||
if strings.TrimSpace(file.Version) != "" {
|
||||
status.InstalledVersion = strings.TrimSpace(file.Version)
|
||||
}
|
||||
status.Enabled = true
|
||||
statuses[file.ID] = status
|
||||
}
|
||||
for id, item := range configs {
|
||||
status := statuses[id]
|
||||
status.Configured = true
|
||||
status.Enabled = pluginInstanceEnabled(item)
|
||||
status.InstalledSourceID, status.InstalledSourceURL, status.StoreManaged = pluginStoreConfiguredSource(item)
|
||||
statuses[id] = status
|
||||
}
|
||||
if host != nil {
|
||||
for _, info := range host.RegisteredPlugins() {
|
||||
status := statuses[info.ID]
|
||||
status.Installed = true
|
||||
status.Registered = true
|
||||
status.InstalledVersion = strings.TrimSpace(info.Metadata.Version)
|
||||
if _, configured := configs[info.ID]; !configured && !status.Enabled {
|
||||
status.Enabled = false
|
||||
}
|
||||
statuses[info.ID] = status
|
||||
}
|
||||
}
|
||||
for id, status := range statuses {
|
||||
status.EffectiveEnabled = pluginsEnabled && status.Enabled && status.Registered
|
||||
statuses[id] = status
|
||||
}
|
||||
return statuses, nil
|
||||
}
|
||||
|
||||
func pluginStoreConfiguredSource(item config.PluginInstanceConfig) (sourceID string, sourceURL string, managed bool) {
|
||||
storeNode := pluginStoreConfigNode(item)
|
||||
if storeNode == nil {
|
||||
return "", "", false
|
||||
}
|
||||
var manifest pluginstore.Manifest
|
||||
if errDecode := storeNode.Decode(&manifest); errDecode != nil {
|
||||
return "", "", true
|
||||
}
|
||||
return strings.TrimSpace(manifest.SourceID), strings.TrimSpace(manifest.SourceURL), true
|
||||
}
|
||||
|
||||
func pluginStoreResolveInstalledSource(status pluginLocalStatus, sources []pluginstore.Source) (string, bool) {
|
||||
sourceID := strings.TrimSpace(status.InstalledSourceID)
|
||||
sourceURL := strings.TrimSpace(status.InstalledSourceURL)
|
||||
if sourceID != "" {
|
||||
for _, source := range sources {
|
||||
if strings.TrimSpace(source.ID) != sourceID {
|
||||
continue
|
||||
}
|
||||
if sourceURL != "" && strings.TrimSpace(source.URL) != sourceURL {
|
||||
return "", false
|
||||
}
|
||||
return sourceID, true
|
||||
}
|
||||
return sourceID, true
|
||||
}
|
||||
if sourceURL == "" {
|
||||
return "", false
|
||||
}
|
||||
for _, source := range sources {
|
||||
if strings.TrimSpace(source.URL) == sourceURL {
|
||||
return strings.TrimSpace(source.ID), true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func pluginStoreInstallSourceStatus(status pluginLocalStatus, sources []pluginstore.Source, entrySourceID string, sourceCount int) (installedSourceID string, sourceStatus string, allowUpdate bool) {
|
||||
if !status.Installed && !status.Configured && !status.Registered {
|
||||
return "", "", true
|
||||
}
|
||||
if sourceID, known := pluginStoreResolveInstalledSource(status, sources); known {
|
||||
if sourceID == strings.TrimSpace(entrySourceID) {
|
||||
return sourceID, "matched", true
|
||||
}
|
||||
return sourceID, "different", false
|
||||
}
|
||||
if status.StoreManaged || sourceCount > 1 {
|
||||
return "", "unknown", false
|
||||
}
|
||||
return "", "assumed", true
|
||||
}
|
||||
|
||||
func validatePluginStoreInstallSource(c *gin.Context, configs map[string]config.PluginInstanceConfig, sources []pluginstore.Source, id string, requestedSourceID string) bool {
|
||||
item, configured := configs[id]
|
||||
if !configured {
|
||||
return true
|
||||
}
|
||||
installedSourceID, installedSourceURL, managed := pluginStoreConfiguredSource(item)
|
||||
if !managed {
|
||||
return true
|
||||
}
|
||||
status := pluginLocalStatus{
|
||||
StoreManaged: true,
|
||||
InstalledSourceID: installedSourceID,
|
||||
InstalledSourceURL: installedSourceURL,
|
||||
}
|
||||
resolvedSourceID, known := pluginStoreResolveInstalledSource(status, sources)
|
||||
if !known {
|
||||
c.JSON(http.StatusConflict, gin.H{
|
||||
"error": "plugin_store_installed_source_unknown",
|
||||
"message": "installed plugin source cannot be verified; uninstall it before reinstalling from the store",
|
||||
"requested_source_id": strings.TrimSpace(requestedSourceID),
|
||||
})
|
||||
return false
|
||||
}
|
||||
if resolvedSourceID != strings.TrimSpace(requestedSourceID) {
|
||||
c.JSON(http.StatusConflict, gin.H{
|
||||
"error": "plugin_store_source_conflict",
|
||||
"message": "installed plugin belongs to a different store source; uninstall it before switching sources",
|
||||
"installed_source_id": resolvedSourceID,
|
||||
"requested_source_id": strings.TrimSpace(requestedSourceID),
|
||||
})
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func pluginStoreDesiredVersions(configs map[string]config.PluginInstanceConfig) map[string]string {
|
||||
if len(configs) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]string, len(configs))
|
||||
for id, item := range configs {
|
||||
id = strings.TrimSpace(id)
|
||||
version := pluginStoreDesiredVersion(item)
|
||||
if id == "" || version == "" {
|
||||
continue
|
||||
}
|
||||
out[id] = version
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func pluginStoreDesiredVersion(item config.PluginInstanceConfig) string {
|
||||
storeNode := pluginStoreConfigNode(item)
|
||||
if storeNode == nil {
|
||||
return ""
|
||||
}
|
||||
if version := pluginStoreNormalizeDesiredVersion(pluginStoreYAMLScalar(yamlMappingValue(storeNode, "version"))); version != "" {
|
||||
return version
|
||||
}
|
||||
return pluginStoreNormalizeDesiredVersion(pluginStoreYAMLScalar(yamlMappingValue(storeNode, "release-tag")))
|
||||
}
|
||||
|
||||
func pluginStoreConfigNode(item config.PluginInstanceConfig) *yaml.Node {
|
||||
if item.Raw.Kind != yaml.MappingNode {
|
||||
return nil
|
||||
}
|
||||
return yamlMappingValue(&item.Raw, "store")
|
||||
}
|
||||
|
||||
func yamlMappingValue(node *yaml.Node, key string) *yaml.Node {
|
||||
if node == nil || node.Kind != yaml.MappingNode {
|
||||
return nil
|
||||
}
|
||||
for i := 0; i+1 < len(node.Content); i += 2 {
|
||||
keyNode := node.Content[i]
|
||||
if keyNode == nil || keyNode.Value != key {
|
||||
continue
|
||||
}
|
||||
return node.Content[i+1]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func pluginStoreYAMLScalar(node *yaml.Node) string {
|
||||
if node == nil || node.Kind != yaml.ScalarNode {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(node.Value)
|
||||
}
|
||||
|
||||
func pluginStoreNormalizeDesiredVersion(version string) string {
|
||||
version = strings.TrimSpace(version)
|
||||
if len(version) > 1 && (version[0] == 'v' || version[0] == 'V') {
|
||||
version = version[1:]
|
||||
}
|
||||
if version == "" || version[0] < '0' || version[0] > '9' {
|
||||
return ""
|
||||
}
|
||||
return version
|
||||
}
|
||||
|
||||
func pluginBusy(host *pluginhost.Host, id string) bool {
|
||||
if host == nil {
|
||||
return false
|
||||
}
|
||||
return host.PluginBusy(id)
|
||||
}
|
||||
1436
backend/internal/api/handlers/management/plugin_store_test.go
Normal file
1436
backend/internal/api/handlers/management/plugin_store_test.go
Normal file
File diff suppressed because it is too large
Load diff
713
backend/internal/api/handlers/management/plugins.go
Normal file
713
backend/internal/api/handlers/management/plugins.go
Normal file
|
|
@ -0,0 +1,713 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/htmlsanitize"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
type pluginListResponse struct {
|
||||
PluginsEnabled bool `json:"plugins_enabled"`
|
||||
PluginsDir string `json:"plugins_dir"`
|
||||
Plugins []pluginListEntry `json:"plugins"`
|
||||
}
|
||||
|
||||
type pluginListEntry struct {
|
||||
ID string `json:"id"`
|
||||
Path string `json:"path"`
|
||||
Configured bool `json:"configured"`
|
||||
Registered bool `json:"registered"`
|
||||
Enabled bool `json:"enabled"`
|
||||
EffectiveEnabled bool `json:"effective_enabled"`
|
||||
SupportsOAuth bool `json:"supports_oauth"`
|
||||
OAuthProvider string `json:"oauth_provider"`
|
||||
Logo string `json:"logo"`
|
||||
ConfigFields []pluginConfigFieldInfo `json:"config_fields"`
|
||||
Menus []pluginMenuInfo `json:"menus"`
|
||||
Metadata *pluginMetadataInfo `json:"metadata"`
|
||||
}
|
||||
|
||||
type pluginMetadataInfo struct {
|
||||
Name string `json:"name"`
|
||||
Version string `json:"version"`
|
||||
Author string `json:"author"`
|
||||
GitHubRepository string `json:"github_repository"`
|
||||
Logo string `json:"logo"`
|
||||
ConfigFields []pluginConfigFieldInfo `json:"config_fields"`
|
||||
}
|
||||
|
||||
type pluginConfigFieldInfo struct {
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
EnumValues []string `json:"enum_values"`
|
||||
Description string `json:"description"`
|
||||
}
|
||||
|
||||
type pluginMenuInfo struct {
|
||||
Path string `json:"path"`
|
||||
Menu string `json:"menu"`
|
||||
Description string `json:"description"`
|
||||
}
|
||||
|
||||
// ListPlugins returns discovered, configured, and registered plugin entries.
|
||||
func (h *Handler) ListPlugins(c *gin.Context) {
|
||||
if h == nil || h.cfg == nil {
|
||||
c.JSON(http.StatusOK, pluginListResponse{
|
||||
PluginsDir: "plugins",
|
||||
Plugins: []pluginListEntry{},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
pluginsEnabled := h.cfg.Plugins.Enabled
|
||||
pluginsDir := normalizedPluginsDir(h.cfg.Plugins.Dir)
|
||||
configs := make(map[string]config.PluginInstanceConfig, len(h.cfg.Plugins.Configs))
|
||||
for id, item := range h.cfg.Plugins.Configs {
|
||||
configs[id] = item
|
||||
}
|
||||
host := h.pluginHost
|
||||
h.mu.Unlock()
|
||||
|
||||
resolvedPluginsDir, errResolvePluginsDir := config.ResolvePluginsDir(pluginsDir)
|
||||
if errResolvePluginsDir != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_directory_invalid", "message": errResolvePluginsDir.Error()})
|
||||
return
|
||||
}
|
||||
pluginsDir = resolvedPluginsDir
|
||||
entries := make(map[string]pluginListEntry)
|
||||
files, errDiscover := pluginhost.DiscoverPluginFiles(pluginsDir, pluginStoreDesiredVersions(configs))
|
||||
if errDiscover != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_discovery_failed", "message": errDiscover.Error()})
|
||||
return
|
||||
}
|
||||
for _, file := range files {
|
||||
entries[file.ID] = pluginListEntry{
|
||||
ID: htmlsanitize.String(file.ID),
|
||||
Path: htmlsanitize.String(file.Path),
|
||||
Enabled: false,
|
||||
ConfigFields: []pluginConfigFieldInfo{},
|
||||
Menus: []pluginMenuInfo{},
|
||||
}
|
||||
}
|
||||
for id, item := range configs {
|
||||
entry := entries[id]
|
||||
entry.ID = htmlsanitize.String(id)
|
||||
entry.Configured = true
|
||||
entry.Enabled = pluginInstanceEnabled(item)
|
||||
if entry.ConfigFields == nil {
|
||||
entry.ConfigFields = []pluginConfigFieldInfo{}
|
||||
}
|
||||
if entry.Menus == nil {
|
||||
entry.Menus = []pluginMenuInfo{}
|
||||
}
|
||||
entries[id] = entry
|
||||
}
|
||||
if host != nil {
|
||||
for _, info := range host.RegisteredPlugins() {
|
||||
entry := entries[info.ID]
|
||||
entry.ID = htmlsanitize.String(info.ID)
|
||||
entry.Registered = true
|
||||
entry.SupportsOAuth = info.SupportsOAuth
|
||||
entry.OAuthProvider = htmlsanitize.String(info.OAuthProvider)
|
||||
entry.Logo = htmlsanitize.String(info.Metadata.Logo)
|
||||
entry.ConfigFields = pluginConfigFields(info.Metadata.ConfigFields)
|
||||
entry.Menus = pluginMenus(info.Menus)
|
||||
entry.Metadata = pluginMetadata(info.Metadata)
|
||||
entries[info.ID] = entry
|
||||
}
|
||||
}
|
||||
|
||||
ids := make([]string, 0, len(entries))
|
||||
for id := range entries {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
sort.Strings(ids)
|
||||
out := make([]pluginListEntry, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
entry := entries[id]
|
||||
entry.EffectiveEnabled = pluginsEnabled && entry.Enabled && entry.Registered
|
||||
if entry.ConfigFields == nil {
|
||||
entry.ConfigFields = []pluginConfigFieldInfo{}
|
||||
}
|
||||
if entry.Menus == nil {
|
||||
entry.Menus = []pluginMenuInfo{}
|
||||
}
|
||||
out = append(out, entry)
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, pluginListResponse{
|
||||
PluginsEnabled: pluginsEnabled,
|
||||
PluginsDir: htmlsanitize.String(pluginsDir),
|
||||
Plugins: out,
|
||||
})
|
||||
}
|
||||
|
||||
// GetPluginConfig returns the preserved plugins.configs.<id> object as JSON.
|
||||
func (h *Handler) GetPluginConfig(c *gin.Context) {
|
||||
id, okID := pluginIDFromRequest(c)
|
||||
if !okID {
|
||||
return
|
||||
}
|
||||
if h == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "plugin_not_found", "message": "plugin not found"})
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
if h.cfg == nil {
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "plugin_not_found", "message": "plugin not found"})
|
||||
return
|
||||
}
|
||||
item, configured := h.cfg.Plugins.Configs[id]
|
||||
pluginsDir := normalizedPluginsDir(h.cfg.Plugins.Dir)
|
||||
host := h.pluginHost
|
||||
h.mu.Unlock()
|
||||
|
||||
if configured {
|
||||
body, errBody := pluginConfigJSONObject(item)
|
||||
if errBody != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_config_encode_failed", "message": errBody.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, body)
|
||||
return
|
||||
}
|
||||
|
||||
if pluginRegistered(host, id) {
|
||||
c.JSON(http.StatusOK, gin.H{})
|
||||
return
|
||||
}
|
||||
resolvedPluginsDir, errResolvePluginsDir := config.ResolvePluginsDir(pluginsDir)
|
||||
if errResolvePluginsDir != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_directory_invalid", "message": errResolvePluginsDir.Error()})
|
||||
return
|
||||
}
|
||||
discovered, errDiscover := pluginDiscovered(resolvedPluginsDir, id)
|
||||
if errDiscover != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_discovery_failed", "message": errDiscover.Error()})
|
||||
return
|
||||
}
|
||||
if discovered {
|
||||
c.JSON(http.StatusOK, gin.H{})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "plugin_not_found", "message": "plugin not found"})
|
||||
}
|
||||
|
||||
// PatchPluginEnabled updates plugins.configs.<id>.enabled without touching plugins.enabled.
|
||||
func (h *Handler) PatchPluginEnabled(c *gin.Context) {
|
||||
id, okID := pluginIDFromRequest(c)
|
||||
if !okID {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
if errBindJSON := c.ShouldBindJSON(&body); errBindJSON != nil || body.Enabled == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_body", "message": "enabled is required"})
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
ensurePluginConfigMap(h.cfg)
|
||||
item := h.cfg.Plugins.Configs[id]
|
||||
node := pluginConfigNode(item)
|
||||
setYAMLMappingValue(node, "enabled", boolYAMLNode(*body.Enabled))
|
||||
updated, errConfig := pluginInstanceConfigFromNode(node)
|
||||
if errConfig != nil {
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_config", "message": errConfig.Error()})
|
||||
return
|
||||
}
|
||||
h.cfg.Plugins.Configs[id] = updated
|
||||
cfgSnapshot, okSnapshot := h.saveConfigAndSnapshotLocked(c)
|
||||
h.mu.Unlock()
|
||||
if !okSnapshot {
|
||||
return
|
||||
}
|
||||
|
||||
h.reloadConfigAfterManagementSaveAsync(c.Request.Context(), cfgSnapshot)
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||||
}
|
||||
|
||||
// PutPluginConfig replaces plugins.configs.<id> with the request object.
|
||||
func (h *Handler) PutPluginConfig(c *gin.Context) {
|
||||
id, okID := pluginIDFromRequest(c)
|
||||
if !okID {
|
||||
return
|
||||
}
|
||||
body, okBody := readPluginConfigObject(c)
|
||||
if !okBody {
|
||||
return
|
||||
}
|
||||
node, errNode := yamlNodeFromJSONObject(body)
|
||||
if errNode != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_body", "message": errNode.Error()})
|
||||
return
|
||||
}
|
||||
updated, errConfig := pluginInstanceConfigFromNode(node)
|
||||
if errConfig != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_config", "message": errConfig.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
ensurePluginConfigMap(h.cfg)
|
||||
h.cfg.Plugins.Configs[id] = updated
|
||||
h.persistLocked(c)
|
||||
}
|
||||
|
||||
// PatchPluginConfig shallow-merges plugins.configs.<id> with the request object.
|
||||
func (h *Handler) PatchPluginConfig(c *gin.Context) {
|
||||
id, okID := pluginIDFromRequest(c)
|
||||
if !okID {
|
||||
return
|
||||
}
|
||||
body, okBody := readPluginConfigObject(c)
|
||||
if !okBody {
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
ensurePluginConfigMap(h.cfg)
|
||||
node := pluginConfigNode(h.cfg.Plugins.Configs[id])
|
||||
keys := make([]string, 0, len(body))
|
||||
for key := range body {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
for _, key := range keys {
|
||||
value := body[key]
|
||||
if value == nil {
|
||||
deleteYAMLMappingKey(node, key)
|
||||
continue
|
||||
}
|
||||
valueNode, errNode := yamlNodeFromJSONValue(value)
|
||||
if errNode != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_body", "message": errNode.Error()})
|
||||
return
|
||||
}
|
||||
setYAMLMappingValue(node, key, valueNode)
|
||||
}
|
||||
updated, errConfig := pluginInstanceConfigFromNode(node)
|
||||
if errConfig != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_config", "message": errConfig.Error()})
|
||||
return
|
||||
}
|
||||
h.cfg.Plugins.Configs[id] = updated
|
||||
h.persistLocked(c)
|
||||
}
|
||||
|
||||
// DeletePlugin removes the selected local plugin file and its saved config.
|
||||
func (h *Handler) DeletePlugin(c *gin.Context) {
|
||||
id, okID := pluginIDFromRequest(c)
|
||||
if !okID {
|
||||
return
|
||||
}
|
||||
if h == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "plugin_not_found", "message": "plugin not found"})
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
if h.cfg == nil {
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "plugin_not_found", "message": "plugin not found"})
|
||||
return
|
||||
}
|
||||
pluginsDir := normalizedPluginsDir(h.cfg.Plugins.Dir)
|
||||
item, configured := h.cfg.Plugins.Configs[id]
|
||||
host := h.pluginHost
|
||||
h.mu.Unlock()
|
||||
|
||||
resolvedPluginsDir, errResolvePluginsDir := config.ResolvePluginsDir(pluginsDir)
|
||||
if errResolvePluginsDir != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_directory_invalid", "message": errResolvePluginsDir.Error()})
|
||||
return
|
||||
}
|
||||
pluginsDir = resolvedPluginsDir
|
||||
var desiredVersions map[string]string
|
||||
if configured {
|
||||
desiredVersions = pluginStoreDesiredVersions(map[string]config.PluginInstanceConfig{id: item})
|
||||
}
|
||||
path, errPath := pluginFilePath(pluginsDir, id, desiredVersions)
|
||||
if errPath != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_discovery_failed", "message": errPath.Error()})
|
||||
return
|
||||
}
|
||||
if path == "" && !configured {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "plugin_not_found", "message": "plugin not found"})
|
||||
return
|
||||
}
|
||||
|
||||
if pluginBusy(host, id) && (host == nil || !host.UnloadPlugin(id)) && pluginBusy(host, id) {
|
||||
c.JSON(http.StatusConflict, gin.H{
|
||||
"error": "plugin_delete_requires_restart",
|
||||
"message": "loaded plugin cannot be deleted while the server is running",
|
||||
"restart_required": true,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
fileDeleted := false
|
||||
if path != "" {
|
||||
if errRemove := os.Remove(path); errRemove != nil {
|
||||
if !errors.Is(errRemove, os.ErrNotExist) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "plugin_delete_failed", "message": errRemove.Error()})
|
||||
return
|
||||
}
|
||||
} else {
|
||||
fileDeleted = true
|
||||
}
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
delete(h.cfg.Plugins.Configs, id)
|
||||
if configured {
|
||||
if errSave := config.SaveConfigPreserveComments(h.configFilePath, h.cfg); errSave != nil {
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"error": "config_save_failed",
|
||||
"message": fmt.Sprintf("plugin deleted but saving config failed: %s", errSave.Error()),
|
||||
"file_deleted": fileDeleted,
|
||||
"path": path,
|
||||
})
|
||||
return
|
||||
}
|
||||
}
|
||||
cfgSnapshot := h.reloadSnapshotConfigLocked()
|
||||
h.mu.Unlock()
|
||||
|
||||
h.reloadConfigAfterManagementSaveAsync(c.Request.Context(), cfgSnapshot)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"status": "deleted",
|
||||
"id": htmlsanitize.String(id),
|
||||
"path": htmlsanitize.String(path),
|
||||
"file_deleted": fileDeleted,
|
||||
"configured_removed": configured,
|
||||
"restart_required": false,
|
||||
})
|
||||
}
|
||||
|
||||
func normalizedPluginsDir(dir string) string {
|
||||
dir = strings.TrimSpace(dir)
|
||||
if dir == "" {
|
||||
return "plugins"
|
||||
}
|
||||
return dir
|
||||
}
|
||||
|
||||
func pluginInstanceEnabled(item config.PluginInstanceConfig) bool {
|
||||
if item.Enabled == nil {
|
||||
return false
|
||||
}
|
||||
return *item.Enabled
|
||||
}
|
||||
|
||||
func pluginRegistered(host *pluginhost.Host, id string) bool {
|
||||
if host == nil {
|
||||
return false
|
||||
}
|
||||
for _, info := range host.RegisteredPlugins() {
|
||||
if info.ID == id {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func pluginDiscovered(pluginsDir string, id string) (bool, error) {
|
||||
files, errDiscover := pluginhost.DiscoverPluginFiles(pluginsDir)
|
||||
if errDiscover != nil {
|
||||
return false, errDiscover
|
||||
}
|
||||
for _, file := range files {
|
||||
if file.ID == id {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func pluginFilePath(pluginsDir string, id string, desiredVersions ...map[string]string) (string, error) {
|
||||
files, errDiscover := pluginhost.DiscoverPluginFiles(pluginsDir, desiredVersions...)
|
||||
if errDiscover != nil {
|
||||
return "", errDiscover
|
||||
}
|
||||
for _, file := range files {
|
||||
if file.ID == id {
|
||||
return file.Path, nil
|
||||
}
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
|
||||
func pluginConfigFields(fields []pluginapi.ConfigField) []pluginConfigFieldInfo {
|
||||
out := make([]pluginConfigFieldInfo, 0, len(fields))
|
||||
for _, field := range fields {
|
||||
out = append(out, pluginConfigFieldInfo{
|
||||
Name: htmlsanitize.String(field.Name),
|
||||
Type: htmlsanitize.String(string(field.Type)),
|
||||
EnumValues: htmlsanitize.Strings(field.EnumValues),
|
||||
Description: htmlsanitize.String(field.Description),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func pluginMenus(menus []pluginhost.RegisteredPluginMenu) []pluginMenuInfo {
|
||||
out := make([]pluginMenuInfo, 0, len(menus))
|
||||
for _, menu := range menus {
|
||||
out = append(out, pluginMenuInfo{
|
||||
Path: htmlsanitize.String(menu.Path),
|
||||
Menu: htmlsanitize.String(menu.Menu),
|
||||
Description: htmlsanitize.String(menu.Description),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func pluginMetadata(meta pluginapi.Metadata) *pluginMetadataInfo {
|
||||
return &pluginMetadataInfo{
|
||||
Name: htmlsanitize.String(meta.Name),
|
||||
Version: htmlsanitize.String(meta.Version),
|
||||
Author: htmlsanitize.String(meta.Author),
|
||||
GitHubRepository: htmlsanitize.String(meta.GitHubRepository),
|
||||
Logo: htmlsanitize.String(meta.Logo),
|
||||
ConfigFields: pluginConfigFields(meta.ConfigFields),
|
||||
}
|
||||
}
|
||||
|
||||
func pluginIDFromRequest(c *gin.Context) (string, bool) {
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
if !pluginhost.ValidatePluginID(id) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_plugin_id", "message": "invalid plugin id"})
|
||||
return "", false
|
||||
}
|
||||
return id, true
|
||||
}
|
||||
|
||||
func readPluginConfigObject(c *gin.Context) (map[string]any, bool) {
|
||||
decoder := json.NewDecoder(c.Request.Body)
|
||||
decoder.UseNumber()
|
||||
var body map[string]any
|
||||
if errDecode := decoder.Decode(&body); errDecode != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_body", "message": errDecode.Error()})
|
||||
return nil, false
|
||||
}
|
||||
if body == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_body", "message": "body must be a JSON object"})
|
||||
return nil, false
|
||||
}
|
||||
return body, true
|
||||
}
|
||||
|
||||
func ensurePluginConfigMap(cfg *config.Config) {
|
||||
if cfg == nil {
|
||||
return
|
||||
}
|
||||
cfg.NormalizePluginsConfig()
|
||||
}
|
||||
|
||||
func pluginConfigNode(item config.PluginInstanceConfig) *yaml.Node {
|
||||
if item.Raw.Kind == yaml.MappingNode {
|
||||
return cloneYAMLNode(&item.Raw)
|
||||
}
|
||||
node := emptyYAMLMappingNode()
|
||||
if item.Enabled != nil {
|
||||
setYAMLMappingValue(node, "enabled", boolYAMLNode(*item.Enabled))
|
||||
}
|
||||
if item.Priority != 0 {
|
||||
setYAMLMappingValue(node, "priority", intYAMLNode(item.Priority))
|
||||
}
|
||||
return node
|
||||
}
|
||||
|
||||
func pluginConfigJSONObject(item config.PluginInstanceConfig) (map[string]any, error) {
|
||||
value, errValue := yamlNodeToJSONValue(pluginConfigNode(item))
|
||||
if errValue != nil {
|
||||
return nil, errValue
|
||||
}
|
||||
body, ok := value.(map[string]any)
|
||||
if !ok || body == nil {
|
||||
return map[string]any{}, nil
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func pluginInstanceConfigFromNode(node *yaml.Node) (config.PluginInstanceConfig, error) {
|
||||
if node == nil {
|
||||
node = emptyYAMLMappingNode()
|
||||
}
|
||||
var item config.PluginInstanceConfig
|
||||
if errDecode := node.Decode(&item); errDecode != nil {
|
||||
return config.PluginInstanceConfig{}, errDecode
|
||||
}
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func yamlNodeFromJSONObject(body map[string]any) (*yaml.Node, error) {
|
||||
node := emptyYAMLMappingNode()
|
||||
keys := make([]string, 0, len(body))
|
||||
for key := range body {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
for _, key := range keys {
|
||||
valueNode, errNode := yamlNodeFromJSONValue(body[key])
|
||||
if errNode != nil {
|
||||
return nil, fmt.Errorf("%s: %w", key, errNode)
|
||||
}
|
||||
setYAMLMappingValue(node, key, valueNode)
|
||||
}
|
||||
return node, nil
|
||||
}
|
||||
|
||||
func yamlNodeFromJSONValue(value any) (*yaml.Node, error) {
|
||||
switch typed := value.(type) {
|
||||
case nil:
|
||||
return &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!null", Value: "null"}, nil
|
||||
case string:
|
||||
return &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: typed}, nil
|
||||
case bool:
|
||||
return boolYAMLNode(typed), nil
|
||||
case json.Number:
|
||||
if _, errInt64 := typed.Int64(); errInt64 == nil {
|
||||
return &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!int", Value: typed.String()}, nil
|
||||
}
|
||||
if _, errFloat64 := typed.Float64(); errFloat64 == nil {
|
||||
return &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!float", Value: typed.String()}, nil
|
||||
}
|
||||
return nil, fmt.Errorf("invalid number %q", typed.String())
|
||||
case float64:
|
||||
return &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!float", Value: strconv.FormatFloat(typed, 'f', -1, 64)}, nil
|
||||
case []any:
|
||||
node := &yaml.Node{Kind: yaml.SequenceNode, Tag: "!!seq"}
|
||||
for _, item := range typed {
|
||||
child, errChild := yamlNodeFromJSONValue(item)
|
||||
if errChild != nil {
|
||||
return nil, errChild
|
||||
}
|
||||
node.Content = append(node.Content, child)
|
||||
}
|
||||
return node, nil
|
||||
case map[string]any:
|
||||
return yamlNodeFromJSONObject(typed)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported value type %T", value)
|
||||
}
|
||||
}
|
||||
|
||||
func yamlNodeToJSONValue(node *yaml.Node) (any, error) {
|
||||
if node == nil {
|
||||
return nil, nil
|
||||
}
|
||||
switch node.Kind {
|
||||
case yaml.MappingNode:
|
||||
out := make(map[string]any, len(node.Content)/2)
|
||||
for index := 0; index+1 < len(node.Content); index += 2 {
|
||||
key := node.Content[index]
|
||||
value := node.Content[index+1]
|
||||
if key == nil {
|
||||
continue
|
||||
}
|
||||
child, errChild := yamlNodeToJSONValue(value)
|
||||
if errChild != nil {
|
||||
return nil, fmt.Errorf("%s: %w", key.Value, errChild)
|
||||
}
|
||||
out[key.Value] = child
|
||||
}
|
||||
return out, nil
|
||||
case yaml.SequenceNode:
|
||||
out := make([]any, 0, len(node.Content))
|
||||
for _, childNode := range node.Content {
|
||||
child, errChild := yamlNodeToJSONValue(childNode)
|
||||
if errChild != nil {
|
||||
return nil, errChild
|
||||
}
|
||||
out = append(out, child)
|
||||
}
|
||||
return out, nil
|
||||
case yaml.ScalarNode:
|
||||
if node.Tag == "!!str" || node.Tag == "" {
|
||||
return node.Value, nil
|
||||
}
|
||||
var value any
|
||||
if errDecode := node.Decode(&value); errDecode != nil {
|
||||
return nil, errDecode
|
||||
}
|
||||
return value, nil
|
||||
case yaml.AliasNode:
|
||||
return yamlNodeToJSONValue(node.Alias)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported YAML node kind %d", node.Kind)
|
||||
}
|
||||
}
|
||||
|
||||
func emptyYAMLMappingNode() *yaml.Node {
|
||||
return &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
|
||||
}
|
||||
|
||||
func boolYAMLNode(value bool) *yaml.Node {
|
||||
return &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!bool", Value: strconv.FormatBool(value)}
|
||||
}
|
||||
|
||||
func intYAMLNode(value int) *yaml.Node {
|
||||
return &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!int", Value: strconv.Itoa(value)}
|
||||
}
|
||||
|
||||
func setYAMLMappingValue(mapping *yaml.Node, key string, value *yaml.Node) {
|
||||
if mapping.Kind != yaml.MappingNode {
|
||||
*mapping = *emptyYAMLMappingNode()
|
||||
}
|
||||
for index := 0; index+1 < len(mapping.Content); index += 2 {
|
||||
if mapping.Content[index] != nil && mapping.Content[index].Value == key {
|
||||
mapping.Content[index+1] = value
|
||||
return
|
||||
}
|
||||
}
|
||||
mapping.Content = append(mapping.Content, &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}, value)
|
||||
}
|
||||
|
||||
func deleteYAMLMappingKey(mapping *yaml.Node, key string) {
|
||||
if mapping == nil || mapping.Kind != yaml.MappingNode {
|
||||
return
|
||||
}
|
||||
for index := 0; index+1 < len(mapping.Content); index += 2 {
|
||||
if mapping.Content[index] != nil && mapping.Content[index].Value == key {
|
||||
mapping.Content = append(mapping.Content[:index], mapping.Content[index+2:]...)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func cloneYAMLNode(node *yaml.Node) *yaml.Node {
|
||||
if node == nil {
|
||||
return nil
|
||||
}
|
||||
out := *node
|
||||
if len(node.Content) > 0 {
|
||||
out.Content = make([]*yaml.Node, 0, len(node.Content))
|
||||
for _, child := range node.Content {
|
||||
out.Content = append(out.Content, cloneYAMLNode(child))
|
||||
}
|
||||
}
|
||||
return &out
|
||||
}
|
||||
852
backend/internal/api/handlers/management/plugins_test.go
Normal file
852
backend/internal/api/handlers/management/plugins_test.go
Normal file
|
|
@ -0,0 +1,852 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"html"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
func waitForAsyncReload(t *testing.T, reloads <-chan *config.Config) *config.Config {
|
||||
t.Helper()
|
||||
select {
|
||||
case cfg := <-reloads:
|
||||
return cfg
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("timed out waiting for async config reload")
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func waitForReloadDone(t *testing.T, done <-chan struct{}) {
|
||||
t.Helper()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("timed out waiting for config reload hook to finish")
|
||||
}
|
||||
}
|
||||
|
||||
func captureConfigReload(h *Handler) (<-chan *config.Config, <-chan struct{}) {
|
||||
reloads := make(chan *config.Config, 1)
|
||||
done := make(chan struct{})
|
||||
h.SetConfigReloadHook(func(_ context.Context, cfg *config.Config) {
|
||||
defer close(done)
|
||||
reloads <- cfg
|
||||
})
|
||||
return reloads, done
|
||||
}
|
||||
|
||||
func TestConfigReloadGenerationSkipsOlderSnapshot(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Configs: map[string]config.PluginInstanceConfig{
|
||||
"sample": pluginConfigFromYAML(t, "enabled: true\nmode: old\n"),
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
reloadedModes := make([]string, 0, 1)
|
||||
h.SetConfigReloadHook(func(_ context.Context, cfg *config.Config) {
|
||||
reloadedModes = append(reloadedModes, pluginRawScalarValue(t, cfg.Plugins.Configs["sample"], "mode"))
|
||||
})
|
||||
|
||||
h.mu.Lock()
|
||||
older := h.reloadSnapshotConfigLocked()
|
||||
item := h.cfg.Plugins.Configs["sample"]
|
||||
setPluginRawScalarValue(t, &item.Raw, "mode", "new")
|
||||
h.cfg.Plugins.Configs["sample"] = item
|
||||
newer := h.reloadSnapshotConfigLocked()
|
||||
h.mu.Unlock()
|
||||
|
||||
h.reloadConfigAfterManagementSave(context.Background(), newer)
|
||||
h.reloadConfigAfterManagementSave(context.Background(), older)
|
||||
|
||||
if len(reloadedModes) != 1 || reloadedModes[0] != "new" {
|
||||
t.Fatalf("reloaded modes = %#v, want only new snapshot", reloadedModes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListPluginsIncludesScannedAndConfiguredPlugins(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
pluginsDir := writeManagementPluginFile(t, "scanned")
|
||||
disabled := false
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Enabled: false,
|
||||
Dir: pluginsDir,
|
||||
Configs: map[string]config.PluginInstanceConfig{
|
||||
"configured-only": {Enabled: &disabled},
|
||||
},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v0/management/plugins", nil)
|
||||
|
||||
h.ListPlugins(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
|
||||
var body struct {
|
||||
PluginsEnabled bool `json:"plugins_enabled"`
|
||||
Plugins []struct {
|
||||
ID string `json:"id"`
|
||||
Path string `json:"path"`
|
||||
Configured bool `json:"configured"`
|
||||
Registered bool `json:"registered"`
|
||||
Enabled bool `json:"enabled"`
|
||||
EffectiveEnabled bool `json:"effective_enabled"`
|
||||
SupportsOAuth bool `json:"supports_oauth"`
|
||||
OAuthProvider string `json:"oauth_provider"`
|
||||
Logo string `json:"logo"`
|
||||
ConfigFields []any `json:"config_fields"`
|
||||
Menus []any `json:"menus"`
|
||||
} `json:"plugins"`
|
||||
}
|
||||
if errDecode := json.Unmarshal(rec.Body.Bytes(), &body); errDecode != nil {
|
||||
t.Fatalf("decode response: %v; body=%s", errDecode, rec.Body.String())
|
||||
}
|
||||
if body.PluginsEnabled {
|
||||
t.Fatal("plugins_enabled = true, want false")
|
||||
}
|
||||
entries := map[string]struct {
|
||||
Configured bool
|
||||
Registered bool
|
||||
Enabled bool
|
||||
EffectiveEnabled bool
|
||||
Path string
|
||||
}{}
|
||||
for _, item := range body.Plugins {
|
||||
entries[item.ID] = struct {
|
||||
Configured bool
|
||||
Registered bool
|
||||
Enabled bool
|
||||
EffectiveEnabled bool
|
||||
Path string
|
||||
}{
|
||||
Configured: item.Configured,
|
||||
Registered: item.Registered,
|
||||
Enabled: item.Enabled,
|
||||
EffectiveEnabled: item.EffectiveEnabled,
|
||||
Path: item.Path,
|
||||
}
|
||||
if item.Registered ||
|
||||
item.SupportsOAuth ||
|
||||
item.OAuthProvider != "" ||
|
||||
item.Logo != "" ||
|
||||
len(item.ConfigFields) != 0 ||
|
||||
len(item.Menus) != 0 {
|
||||
t.Fatalf("unregistered plugin entry has runtime fields: %#v", item)
|
||||
}
|
||||
}
|
||||
if got, ok := entries["scanned"]; !ok || got.Configured || got.Enabled || got.EffectiveEnabled || got.Path == "" {
|
||||
t.Fatalf("scanned entry = %#v, exists=%v", got, ok)
|
||||
}
|
||||
if got, ok := entries["configured-only"]; !ok || !got.Configured || got.Enabled || got.EffectiveEnabled || got.Path != "" {
|
||||
t.Fatalf("configured-only entry = %#v, exists=%v", got, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListPluginsUsesConfiguredStoreVersionWhenFilesCoexist(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
pluginsDir := t.TempDir()
|
||||
archDir := filepath.Join(pluginsDir, runtime.GOOS, runtime.GOARCH)
|
||||
if errMkdirAll := os.MkdirAll(archDir, 0o755); errMkdirAll != nil {
|
||||
t.Fatalf("MkdirAll(%s) error = %v", archDir, errMkdirAll)
|
||||
}
|
||||
extension := managementPluginExtension(runtime.GOOS)
|
||||
pinnedPath := filepath.Join(archDir, "sample-provider-v0.1.0"+extension)
|
||||
newerPath := filepath.Join(archDir, "sample-provider-v0.2.0"+extension)
|
||||
for _, path := range []string{pinnedPath, newerPath} {
|
||||
if errWriteFile := os.WriteFile(path, []byte("x"), 0o644); errWriteFile != nil {
|
||||
t.Fatalf("WriteFile(%s) error = %v", path, errWriteFile)
|
||||
}
|
||||
}
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Enabled: true,
|
||||
Dir: pluginsDir,
|
||||
Configs: map[string]config.PluginInstanceConfig{
|
||||
"sample-provider": pluginConfigFromYAML(t, "enabled: true\nstore:\n version: 0.1.0\n"),
|
||||
},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v0/management/plugins", nil)
|
||||
|
||||
h.ListPlugins(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
var body pluginListResponse
|
||||
if errDecode := json.Unmarshal(rec.Body.Bytes(), &body); errDecode != nil {
|
||||
t.Fatalf("decode response: %v; body=%s", errDecode, rec.Body.String())
|
||||
}
|
||||
for _, entry := range body.Plugins {
|
||||
if entry.ID != "sample-provider" {
|
||||
continue
|
||||
}
|
||||
if entry.Path != pinnedPath || !entry.Configured || !entry.Enabled {
|
||||
t.Fatalf("plugin entry = %#v, want pinned path %s", entry, pinnedPath)
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Fatalf("sample-provider entry missing: %#v", body.Plugins)
|
||||
}
|
||||
|
||||
func TestGetPluginConfigReturnsPreservedRawConfig(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Configs: map[string]config.PluginInstanceConfig{
|
||||
"sample": pluginConfigFromYAML(t, `
|
||||
enabled: false
|
||||
priority: 7
|
||||
mode: safe
|
||||
allowed_models:
|
||||
- gemini-2.5-pro
|
||||
- claude-sonnet-4
|
||||
options:
|
||||
retries: 2
|
||||
strict: true
|
||||
`),
|
||||
},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = gin.Params{{Key: "id", Value: "sample"}}
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v0/management/plugins/sample/config", nil)
|
||||
|
||||
h.GetPluginConfig(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
|
||||
var body struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
Priority int `json:"priority"`
|
||||
Mode string `json:"mode"`
|
||||
AllowedModels []string `json:"allowed_models"`
|
||||
Options map[string]any `json:"options"`
|
||||
}
|
||||
if errDecode := json.Unmarshal(rec.Body.Bytes(), &body); errDecode != nil {
|
||||
t.Fatalf("decode response: %v; body=%s", errDecode, rec.Body.String())
|
||||
}
|
||||
if body.Enabled || body.Priority != 7 || body.Mode != "safe" {
|
||||
t.Fatalf("base fields = enabled %v priority %d mode %q, want false 7 safe", body.Enabled, body.Priority, body.Mode)
|
||||
}
|
||||
if len(body.AllowedModels) != 2 || body.AllowedModels[0] != "gemini-2.5-pro" || body.AllowedModels[1] != "claude-sonnet-4" {
|
||||
t.Fatalf("allowed_models = %#v", body.AllowedModels)
|
||||
}
|
||||
if body.Options["retries"] != float64(2) || body.Options["strict"] != true {
|
||||
t.Fatalf("options = %#v", body.Options)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetPluginConfigReturnsEmptyObjectForKnownUnconfiguredPlugin(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
pluginsDir := writeManagementPluginFile(t, "scanned")
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Dir: pluginsDir,
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = gin.Params{{Key: "id", Value: "scanned"}}
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v0/management/plugins/scanned/config", nil)
|
||||
|
||||
h.GetPluginConfig(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
var body map[string]any
|
||||
if errDecode := json.Unmarshal(rec.Body.Bytes(), &body); errDecode != nil {
|
||||
t.Fatalf("decode response: %v; body=%s", errDecode, rec.Body.String())
|
||||
}
|
||||
if len(body) != 0 {
|
||||
t.Fatalf("body = %#v, want empty object", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetPluginConfigReturnsNotFoundForUnknownPlugin(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = gin.Params{{Key: "id", Value: "missing"}}
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v0/management/plugins/missing/config", nil)
|
||||
|
||||
h.GetPluginConfig(c)
|
||||
|
||||
if rec.Code != http.StatusNotFound {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusNotFound, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchPluginEnabledUpdatesOnlyPluginConfig(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Enabled: false,
|
||||
Configs: map[string]config.PluginInstanceConfig{
|
||||
"sample": pluginConfigFromYAML(t, "enabled: false\npriority: 2\nmode: safe\n"),
|
||||
},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
reloads, reloadDone := captureConfigReload(h)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = gin.Params{{Key: "id", Value: "sample"}}
|
||||
c.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/plugins/sample/enabled", strings.NewReader(`{"enabled":true}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h.PatchPluginEnabled(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
cfgSnapshot := waitForAsyncReload(t, reloads)
|
||||
waitForReloadDone(t, reloadDone)
|
||||
if cfgSnapshot == h.cfg {
|
||||
t.Fatalf("reload config = handler config %p, want independent snapshot", h.cfg)
|
||||
}
|
||||
if cfgSnapshot.Plugins.Enabled {
|
||||
t.Fatal("snapshot global Plugins.Enabled changed to true")
|
||||
}
|
||||
snapshotItem := cfgSnapshot.Plugins.Configs["sample"]
|
||||
if snapshotItem.Enabled == nil || !*snapshotItem.Enabled {
|
||||
t.Fatalf("snapshot sample enabled = %#v, want true", snapshotItem.Enabled)
|
||||
}
|
||||
if raw := marshalPluginRaw(t, snapshotItem); !strings.Contains(raw, "mode: safe") {
|
||||
t.Fatalf("snapshot raw config lost custom field:\n%s", raw)
|
||||
}
|
||||
if h.cfg.Plugins.Enabled {
|
||||
t.Fatal("global Plugins.Enabled changed to true")
|
||||
}
|
||||
item := h.cfg.Plugins.Configs["sample"]
|
||||
if item.Enabled == nil || !*item.Enabled {
|
||||
t.Fatalf("sample enabled = %#v, want true", item.Enabled)
|
||||
}
|
||||
raw := marshalPluginRaw(t, item)
|
||||
if !strings.Contains(raw, "mode: safe") {
|
||||
t.Fatalf("raw config lost custom field:\n%s", raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchPluginEnabledReloadSnapshotRawImmutability(t *testing.T) {
|
||||
t.Parallel()
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Configs: map[string]config.PluginInstanceConfig{
|
||||
"sample": pluginConfigFromYAML(t, "enabled: false\nmode: first\n"),
|
||||
},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
reloads := make(chan *config.Config, 1)
|
||||
releaseReload := make(chan struct{})
|
||||
reloadDone := make(chan struct{})
|
||||
h.SetConfigReloadHook(func(_ context.Context, cfg *config.Config) {
|
||||
defer close(reloadDone)
|
||||
reloads <- cfg
|
||||
<-releaseReload
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = gin.Params{{Key: "id", Value: "sample"}}
|
||||
c.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/plugins/sample/enabled", strings.NewReader(`{"enabled":true}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h.PatchPluginEnabled(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
close(releaseReload)
|
||||
waitForReloadDone(t, reloadDone)
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
cfgSnapshot := waitForAsyncReload(t, reloads)
|
||||
|
||||
h.mu.Lock()
|
||||
item := h.cfg.Plugins.Configs["sample"]
|
||||
setPluginRawScalarValue(t, &item.Raw, "mode", "second")
|
||||
h.cfg.Plugins.Configs["sample"] = item
|
||||
h.mu.Unlock()
|
||||
|
||||
if cfgSnapshot == h.cfg {
|
||||
t.Fatalf("reload config = handler config %p, want independent snapshot", h.cfg)
|
||||
}
|
||||
snapshotItem := cfgSnapshot.Plugins.Configs["sample"]
|
||||
if snapshotItem.Enabled == nil || !*snapshotItem.Enabled {
|
||||
t.Fatalf("snapshot sample enabled = %#v, want true", snapshotItem.Enabled)
|
||||
}
|
||||
if got := pluginRawScalarValue(t, snapshotItem, "mode"); got != "first" {
|
||||
t.Fatalf("snapshot raw mode = %q, want first", got)
|
||||
}
|
||||
h.mu.Lock()
|
||||
handlerItem := h.cfg.Plugins.Configs["sample"]
|
||||
h.mu.Unlock()
|
||||
if got := pluginRawScalarValue(t, handlerItem, "mode"); got != "second" {
|
||||
t.Fatalf("handler raw mode = %q, want second", got)
|
||||
}
|
||||
|
||||
close(releaseReload)
|
||||
waitForReloadDone(t, reloadDone)
|
||||
}
|
||||
|
||||
func TestPutPluginConfigReplacesPluginConfig(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Configs: map[string]config.PluginInstanceConfig{
|
||||
"sample": pluginConfigFromYAML(t, "enabled: false\nmode: safe\nold: true\n"),
|
||||
},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = gin.Params{{Key: "id", Value: "sample"}}
|
||||
c.Request = httptest.NewRequest(http.MethodPut, "/v0/management/plugins/sample/config", bytes.NewBufferString(`{"enabled":true,"priority":7,"mode":"fast"}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h.PutPluginConfig(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
item := h.cfg.Plugins.Configs["sample"]
|
||||
if item.Enabled == nil || !*item.Enabled || item.Priority != 7 {
|
||||
t.Fatalf("plugin host fields = enabled %#v priority %d, want true priority 7", item.Enabled, item.Priority)
|
||||
}
|
||||
raw := marshalPluginRaw(t, item)
|
||||
if !strings.Contains(raw, "mode: fast") || strings.Contains(raw, "old:") {
|
||||
t.Fatalf("raw config =\n%s", raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatchPluginConfigMergesAndDeletesFields(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Configs: map[string]config.PluginInstanceConfig{
|
||||
"sample": pluginConfigFromYAML(t, "enabled: false\npriority: 3\nmode: safe\nremove: yes\n"),
|
||||
},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = gin.Params{{Key: "id", Value: "sample"}}
|
||||
c.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/plugins/sample/config", strings.NewReader(`{"mode":"fast","remove":null,"count":3}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h.PatchPluginConfig(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
item := h.cfg.Plugins.Configs["sample"]
|
||||
if item.Enabled == nil || *item.Enabled || item.Priority != 3 {
|
||||
t.Fatalf("plugin host fields = enabled %#v priority %d, want false priority 3", item.Enabled, item.Priority)
|
||||
}
|
||||
raw := marshalPluginRaw(t, item)
|
||||
if !strings.Contains(raw, "mode: fast") || !strings.Contains(raw, "count: 3") || strings.Contains(raw, "remove:") {
|
||||
t.Fatalf("raw config =\n%s", raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeletePluginRejectsUnresolvedPluginsDir(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
t.Setenv("HOME", "")
|
||||
t.Setenv("USERPROFILE", "")
|
||||
t.Chdir(workspace)
|
||||
|
||||
literalPluginsDir := filepath.Join(workspace, "~", ".cli-proxy-api", "plugins")
|
||||
targetDir := filepath.Join(literalPluginsDir, runtime.GOOS, runtime.GOARCH)
|
||||
if errMkdir := os.MkdirAll(targetDir, 0o755); errMkdir != nil {
|
||||
t.Fatalf("MkdirAll(%s) error = %v", targetDir, errMkdir)
|
||||
}
|
||||
target := filepath.Join(targetDir, "sample"+managementPluginExtension(runtime.GOOS))
|
||||
if errWrite := os.WriteFile(target, []byte("library-data"), 0o644); errWrite != nil {
|
||||
t.Fatalf("WriteFile(%s) error = %v", target, errWrite)
|
||||
}
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Dir: "~/.cli-proxy-api/plugins",
|
||||
Configs: map[string]config.PluginInstanceConfig{
|
||||
"sample": pluginConfigFromYAML(t, "enabled: false\n"),
|
||||
},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = gin.Params{{Key: "id", Value: "sample"}}
|
||||
c.Request = httptest.NewRequest(http.MethodDelete, "/v0/management/plugins/sample", nil)
|
||||
|
||||
h.DeletePlugin(c)
|
||||
|
||||
if rec.Code != http.StatusInternalServerError {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusInternalServerError, rec.Body.String())
|
||||
}
|
||||
var body map[string]any
|
||||
if errDecode := json.Unmarshal(rec.Body.Bytes(), &body); errDecode != nil {
|
||||
t.Fatalf("Unmarshal() error = %v; body=%s", errDecode, rec.Body.String())
|
||||
}
|
||||
if body["error"] != "plugin_directory_invalid" {
|
||||
t.Fatalf("error = %#v, want plugin_directory_invalid", body["error"])
|
||||
}
|
||||
if _, errStat := os.Stat(target); errStat != nil {
|
||||
t.Fatalf("literal tilde target stat error = %v, want retained", errStat)
|
||||
}
|
||||
if _, configured := h.cfg.Plugins.Configs["sample"]; !configured {
|
||||
t.Fatal("plugin config removed after directory resolution failure")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeletePluginRemovesDiscoveredFileAndConfig(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
pluginsDir := writeManagementPluginFile(t, "sample")
|
||||
configPath := filepath.Join(t.TempDir(), "config.yaml")
|
||||
if errWrite := os.WriteFile(configPath, []byte("plugins:\n configs:\n sample:\n enabled: true\n mode: safe\n keep:\n enabled: true\n mode: retained\n"), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write test config: %v", errWrite)
|
||||
}
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Dir: pluginsDir,
|
||||
Configs: map[string]config.PluginInstanceConfig{
|
||||
"sample": pluginConfigFromYAML(t, "enabled: true\nmode: safe\n"),
|
||||
"keep": pluginConfigFromYAML(t, "enabled: true\nmode: retained\n"),
|
||||
},
|
||||
},
|
||||
},
|
||||
configFilePath: configPath,
|
||||
}
|
||||
reloads := make(chan *config.Config, 1)
|
||||
releaseReload := make(chan struct{})
|
||||
reloadDone := make(chan struct{})
|
||||
h.SetConfigReloadHook(func(_ context.Context, cfg *config.Config) {
|
||||
defer close(reloadDone)
|
||||
reloads <- cfg
|
||||
<-releaseReload
|
||||
})
|
||||
|
||||
path, errPath := pluginFilePath(pluginsDir, "sample")
|
||||
if errPath != nil {
|
||||
t.Fatalf("pluginFilePath() error = %v", errPath)
|
||||
}
|
||||
if path == "" {
|
||||
t.Fatal("plugin path is empty")
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = gin.Params{{Key: "id", Value: "sample"}}
|
||||
c.Request = httptest.NewRequest(http.MethodDelete, "/v0/management/plugins/sample", nil)
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
h.DeletePlugin(c)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("DeletePlugin blocked waiting for config reload")
|
||||
}
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if _, ok := h.cfg.Plugins.Configs["sample"]; ok {
|
||||
t.Fatal("plugin config still exists after delete")
|
||||
}
|
||||
if _, ok := h.cfg.Plugins.Configs["keep"]; !ok {
|
||||
t.Fatal("retained plugin config was removed")
|
||||
}
|
||||
data, errReadConfig := os.ReadFile(configPath)
|
||||
if errReadConfig != nil {
|
||||
t.Fatalf("failed to read saved config: %v", errReadConfig)
|
||||
}
|
||||
text := string(data)
|
||||
if strings.Contains(text, "sample:") || strings.Contains(text, "mode: safe") {
|
||||
t.Fatalf("saved config still contains removed plugin:\n%s", text)
|
||||
}
|
||||
if !strings.Contains(text, "keep:") || !strings.Contains(text, "mode: retained") {
|
||||
t.Fatalf("saved config lost retained plugin:\n%s", text)
|
||||
}
|
||||
if _, errStat := os.Stat(path); !os.IsNotExist(errStat) {
|
||||
t.Fatalf("plugin file stat error = %v, want not exist", errStat)
|
||||
}
|
||||
cfgSnapshot := waitForAsyncReload(t, reloads)
|
||||
if cfgSnapshot == h.cfg {
|
||||
close(releaseReload)
|
||||
waitForReloadDone(t, reloadDone)
|
||||
t.Fatalf("reload config = handler config %p, want independent snapshot", h.cfg)
|
||||
}
|
||||
if _, ok := cfgSnapshot.Plugins.Configs["sample"]; ok {
|
||||
close(releaseReload)
|
||||
waitForReloadDone(t, reloadDone)
|
||||
t.Fatal("snapshot plugin config still exists after delete")
|
||||
}
|
||||
close(releaseReload)
|
||||
waitForReloadDone(t, reloadDone)
|
||||
}
|
||||
|
||||
func TestDeletePluginUsesConfiguredStoreVersionWhenFilesCoexist(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
pluginsDir := t.TempDir()
|
||||
archDir := filepath.Join(pluginsDir, runtime.GOOS, runtime.GOARCH)
|
||||
if errMkdirAll := os.MkdirAll(archDir, 0o755); errMkdirAll != nil {
|
||||
t.Fatalf("MkdirAll(%s) error = %v", archDir, errMkdirAll)
|
||||
}
|
||||
extension := managementPluginExtension(runtime.GOOS)
|
||||
pinnedPath := filepath.Join(archDir, "sample-provider-v0.1.0"+extension)
|
||||
newerPath := filepath.Join(archDir, "sample-provider-v0.2.0"+extension)
|
||||
for _, path := range []string{pinnedPath, newerPath} {
|
||||
if errWriteFile := os.WriteFile(path, []byte("x"), 0o644); errWriteFile != nil {
|
||||
t.Fatalf("WriteFile(%s) error = %v", path, errWriteFile)
|
||||
}
|
||||
}
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
Plugins: config.PluginsConfig{
|
||||
Dir: pluginsDir,
|
||||
Configs: map[string]config.PluginInstanceConfig{
|
||||
"sample-provider": pluginConfigFromYAML(t, "enabled: true\nstore:\n version: 0.1.0\n"),
|
||||
},
|
||||
},
|
||||
},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = gin.Params{{Key: "id", Value: "sample-provider"}}
|
||||
c.Request = httptest.NewRequest(http.MethodDelete, "/v0/management/plugins/sample-provider", nil)
|
||||
|
||||
h.DeletePlugin(c)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if _, ok := h.cfg.Plugins.Configs["sample-provider"]; ok {
|
||||
t.Fatal("plugin config still exists after delete")
|
||||
}
|
||||
if _, errStat := os.Stat(pinnedPath); !os.IsNotExist(errStat) {
|
||||
t.Fatalf("pinned plugin stat error = %v, want not exist", errStat)
|
||||
}
|
||||
if _, errStat := os.Stat(newerPath); errStat != nil {
|
||||
t.Fatalf("newer plugin stat error = %v, want still exists", errStat)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeletePluginReturnsNotFoundForUnknownPlugin(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{},
|
||||
configFilePath: writeTestConfigFile(t),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = gin.Params{{Key: "id", Value: "missing"}}
|
||||
c.Request = httptest.NewRequest(http.MethodDelete, "/v0/management/plugins/missing", nil)
|
||||
|
||||
h.DeletePlugin(c)
|
||||
|
||||
if rec.Code != http.StatusNotFound {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusNotFound, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPluginDisplayFieldsEscapeHTML(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
fields := pluginConfigFields([]pluginapi.ConfigField{{
|
||||
Name: `<img src=x onerror=alert(1)>`,
|
||||
Type: pluginapi.ConfigFieldTypeEnum,
|
||||
EnumValues: []string{`<fast>`, `safe & sound`},
|
||||
Description: `"quoted" 'single' <b>mode</b>`,
|
||||
}})
|
||||
if len(fields) != 1 {
|
||||
t.Fatalf("fields len = %d, want 1", len(fields))
|
||||
}
|
||||
if fields[0].Name != html.EscapeString(`<img src=x onerror=alert(1)>`) {
|
||||
t.Fatalf("field name = %q, want escaped", fields[0].Name)
|
||||
}
|
||||
if fields[0].EnumValues[0] != html.EscapeString(`<fast>`) || fields[0].EnumValues[1] != html.EscapeString(`safe & sound`) {
|
||||
t.Fatalf("enum values = %#v, want escaped values", fields[0].EnumValues)
|
||||
}
|
||||
if fields[0].Description != html.EscapeString(`"quoted" 'single' <b>mode</b>`) {
|
||||
t.Fatalf("description = %q, want escaped", fields[0].Description)
|
||||
}
|
||||
|
||||
menus := pluginMenus([]pluginhost.RegisteredPluginMenu{{
|
||||
Path: `/v0/resource/plugins/sample/<status>`,
|
||||
Menu: `<b>Status</b>`,
|
||||
Description: `Shows <script>alert(1)</script>.`,
|
||||
}})
|
||||
if len(menus) != 1 {
|
||||
t.Fatalf("menus len = %d, want 1", len(menus))
|
||||
}
|
||||
if menus[0].Path != html.EscapeString(`/v0/resource/plugins/sample/<status>`) ||
|
||||
menus[0].Menu != html.EscapeString(`<b>Status</b>`) ||
|
||||
menus[0].Description != html.EscapeString(`Shows <script>alert(1)</script>.`) {
|
||||
t.Fatalf("menu = %#v, want escaped strings", menus[0])
|
||||
}
|
||||
|
||||
meta := pluginMetadata(pluginapi.Metadata{
|
||||
Name: `<script>alert(1)</script>`,
|
||||
Version: `1.0.0&evil=true`,
|
||||
Author: `"attacker"`,
|
||||
GitHubRepository: `https://example.com/repo?x=<script>`,
|
||||
Logo: `<svg onload=alert(1)>`,
|
||||
})
|
||||
if meta.Name != html.EscapeString(`<script>alert(1)</script>`) ||
|
||||
meta.Version != html.EscapeString(`1.0.0&evil=true`) ||
|
||||
meta.Author != html.EscapeString(`"attacker"`) ||
|
||||
meta.GitHubRepository != html.EscapeString(`https://example.com/repo?x=<script>`) ||
|
||||
meta.Logo != html.EscapeString(`<svg onload=alert(1)>`) {
|
||||
t.Fatalf("metadata = %#v, want escaped strings", meta)
|
||||
}
|
||||
}
|
||||
|
||||
func writeManagementPluginFile(t *testing.T, id string) string {
|
||||
t.Helper()
|
||||
root := t.TempDir()
|
||||
archDir := filepath.Join(root, runtime.GOOS, runtime.GOARCH)
|
||||
if errMkdirAll := os.MkdirAll(archDir, 0o755); errMkdirAll != nil {
|
||||
t.Fatalf("MkdirAll() error = %v", errMkdirAll)
|
||||
}
|
||||
path := filepath.Join(archDir, id+managementPluginExtension(runtime.GOOS))
|
||||
if errWriteFile := os.WriteFile(path, []byte("x"), 0o644); errWriteFile != nil {
|
||||
t.Fatalf("WriteFile(%s) error = %v", path, errWriteFile)
|
||||
}
|
||||
return root
|
||||
}
|
||||
|
||||
func managementPluginExtension(goos string) string {
|
||||
switch goos {
|
||||
case "darwin":
|
||||
return ".dylib"
|
||||
case "windows":
|
||||
return ".dll"
|
||||
default:
|
||||
return ".so"
|
||||
}
|
||||
}
|
||||
|
||||
func pluginConfigFromYAML(t *testing.T, text string) config.PluginInstanceConfig {
|
||||
t.Helper()
|
||||
var item config.PluginInstanceConfig
|
||||
if errUnmarshal := yaml.Unmarshal([]byte(text), &item); errUnmarshal != nil {
|
||||
t.Fatalf("unmarshal plugin config: %v", errUnmarshal)
|
||||
}
|
||||
return item
|
||||
}
|
||||
|
||||
func marshalPluginRaw(t *testing.T, item config.PluginInstanceConfig) string {
|
||||
t.Helper()
|
||||
data, errMarshal := yaml.Marshal(&item.Raw)
|
||||
if errMarshal != nil {
|
||||
t.Fatalf("marshal plugin raw: %v", errMarshal)
|
||||
}
|
||||
return string(data)
|
||||
}
|
||||
|
||||
func pluginRawScalarValue(t *testing.T, item config.PluginInstanceConfig, key string) string {
|
||||
t.Helper()
|
||||
for i := 0; i+1 < len(item.Raw.Content); i += 2 {
|
||||
if item.Raw.Content[i] != nil && item.Raw.Content[i].Value == key && item.Raw.Content[i+1] != nil {
|
||||
return item.Raw.Content[i+1].Value
|
||||
}
|
||||
}
|
||||
t.Fatalf("plugin raw missing scalar key %q", key)
|
||||
return ""
|
||||
}
|
||||
|
||||
func setPluginRawScalarValue(t *testing.T, node *yaml.Node, key, value string) {
|
||||
t.Helper()
|
||||
for i := 0; i+1 < len(node.Content); i += 2 {
|
||||
if node.Content[i] != nil && node.Content[i].Value == key && node.Content[i+1] != nil {
|
||||
node.Content[i+1].Value = value
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatalf("plugin raw missing scalar key %q", key)
|
||||
}
|
||||
69
backend/internal/api/handlers/management/quota.go
Normal file
69
backend/internal/api/handlers/management/quota.go
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// Quota exceeded toggles
|
||||
func (h *Handler) GetSwitchProject(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"switch-project": h.cfg.QuotaExceeded.SwitchProject})
|
||||
}
|
||||
func (h *Handler) PutSwitchProject(c *gin.Context) {
|
||||
h.updateBoolField(c, func(v bool) { h.cfg.QuotaExceeded.SwitchProject = v })
|
||||
}
|
||||
|
||||
func (h *Handler) GetSwitchPreviewModel(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"switch-preview-model": h.cfg.QuotaExceeded.SwitchPreviewModel})
|
||||
}
|
||||
func (h *Handler) PutSwitchPreviewModel(c *gin.Context) {
|
||||
h.updateBoolField(c, func(v bool) { h.cfg.QuotaExceeded.SwitchPreviewModel = v })
|
||||
}
|
||||
|
||||
// ResetQuota clears quota/cooldown routing state for one auth index.
|
||||
func (h *Handler) ResetQuota(c *gin.Context) {
|
||||
if h.authManager == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"})
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
AuthIndex string `json:"auth_index"`
|
||||
}
|
||||
if errBindJSON := c.ShouldBindJSON(&req); errBindJSON != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
|
||||
return
|
||||
}
|
||||
|
||||
authIndex := strings.TrimSpace(req.AuthIndex)
|
||||
if authIndex == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "auth_index is required"})
|
||||
return
|
||||
}
|
||||
|
||||
auth := h.authByIndex(authIndex)
|
||||
if auth == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "auth not found"})
|
||||
return
|
||||
}
|
||||
|
||||
updated, models, errReset := h.authManager.ResetQuota(c.Request.Context(), auth.ID)
|
||||
if errReset != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to reset quota: %v", errReset)})
|
||||
return
|
||||
}
|
||||
if updated == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "auth not found"})
|
||||
return
|
||||
}
|
||||
updated.EnsureIndex()
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"status": "ok",
|
||||
"auth_index": updated.Index,
|
||||
"models": models,
|
||||
})
|
||||
}
|
||||
134
backend/internal/api/handlers/management/quota_test.go
Normal file
134
backend/internal/api/handlers/management/quota_test.go
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestResetQuota_UsesAuthIndex(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
next := time.Now().Add(time.Hour)
|
||||
auth := &coreauth.Auth{
|
||||
ID: "reset-auth-id",
|
||||
FileName: "reset-auth-file.json",
|
||||
Provider: "claude",
|
||||
Status: coreauth.StatusError,
|
||||
StatusMessage: "quota exhausted",
|
||||
Unavailable: true,
|
||||
NextRetryAfter: next,
|
||||
Quota: coreauth.QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: next, BackoffLevel: 2},
|
||||
ModelStates: map[string]*coreauth.ModelState{
|
||||
"claude-reset-model": {
|
||||
Status: coreauth.StatusError,
|
||||
StatusMessage: "quota exhausted",
|
||||
Unavailable: true,
|
||||
NextRetryAfter: next,
|
||||
Quota: coreauth.QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: next, BackoffLevel: 2},
|
||||
},
|
||||
},
|
||||
}
|
||||
authIndex := auth.EnsureIndex()
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v0/management/reset-quota", strings.NewReader(`{"auth_index":"`+authIndex+`"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
h.ResetQuota(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if errUnmarshal := json.Unmarshal(rec.Body.Bytes(), &payload); errUnmarshal != nil {
|
||||
t.Fatalf("failed to decode response: %v", errUnmarshal)
|
||||
}
|
||||
if payload["auth_index"] != authIndex {
|
||||
t.Fatalf("auth_index = %#v, want %q", payload["auth_index"], authIndex)
|
||||
}
|
||||
|
||||
updated, ok := manager.GetByID("reset-auth-id")
|
||||
if !ok || updated == nil {
|
||||
t.Fatalf("expected auth record to exist after reset")
|
||||
}
|
||||
if updated.Status != coreauth.StatusActive || updated.StatusMessage != "" || updated.Unavailable || !updated.NextRetryAfter.IsZero() {
|
||||
t.Fatalf("updated auth state = status %q message %q unavailable %v next %v", updated.Status, updated.StatusMessage, updated.Unavailable, updated.NextRetryAfter)
|
||||
}
|
||||
if updated.Quota.Exceeded || updated.Quota.Reason != "" || !updated.Quota.NextRecoverAt.IsZero() || updated.Quota.BackoffLevel != 0 {
|
||||
t.Fatalf("updated auth quota = %+v, want cleared", updated.Quota)
|
||||
}
|
||||
state := updated.ModelStates["claude-reset-model"]
|
||||
if state == nil {
|
||||
t.Fatalf("expected model state to remain")
|
||||
}
|
||||
if state.Status != coreauth.StatusActive || state.StatusMessage != "" || state.Unavailable || !state.NextRetryAfter.IsZero() {
|
||||
t.Fatalf("updated model state = status %q message %q unavailable %v next %v", state.Status, state.StatusMessage, state.Unavailable, state.NextRetryAfter)
|
||||
}
|
||||
if state.Quota.Exceeded || state.Quota.Reason != "" || !state.Quota.NextRecoverAt.IsZero() || state.Quota.BackoffLevel != 0 {
|
||||
t.Fatalf("updated model quota = %+v, want cleared", state.Quota)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResetQuota_DoesNotAcceptAuthIDOrFileName(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
auth := &coreauth.Auth{
|
||||
ID: "reset-auth-id-only",
|
||||
FileName: "reset-auth-file-only.json",
|
||||
Provider: "claude",
|
||||
Status: coreauth.StatusError,
|
||||
}
|
||||
authIndex := auth.EnsureIndex()
|
||||
if authIndex == auth.ID || authIndex == auth.FileName {
|
||||
t.Fatalf("test auth_index unexpectedly matches id or file name: %q", authIndex)
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
wantCode int
|
||||
}{
|
||||
{name: "auth_id field ignored", body: `{"auth_id":"reset-auth-id-only"}`, wantCode: http.StatusBadRequest},
|
||||
{name: "id field ignored", body: `{"id":"reset-auth-id-only"}`, wantCode: http.StatusBadRequest},
|
||||
{name: "file name is not an index", body: `{"auth_index":"reset-auth-file-only.json"}`, wantCode: http.StatusNotFound},
|
||||
{name: "auth id is not an index", body: `{"auth_index":"reset-auth-id-only"}`, wantCode: http.StatusNotFound},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v0/management/reset-quota", strings.NewReader(tt.body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
h.ResetQuota(ctx)
|
||||
|
||||
if rec.Code != tt.wantCode {
|
||||
t.Fatalf("status = %d, want %d with body %s", rec.Code, tt.wantCode, rec.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
13
backend/internal/api/handlers/management/test_main_test.go
Normal file
13
backend/internal/api/handlers/management/test_main_test.go
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
49
backend/internal/api/handlers/management/test_store_test.go
Normal file
49
backend/internal/api/handlers/management/test_store_test.go
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
type memoryAuthStore struct {
|
||||
mu sync.Mutex
|
||||
items map[string]*coreauth.Auth
|
||||
}
|
||||
|
||||
func (s *memoryAuthStore) List(_ context.Context) ([]*coreauth.Auth, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
out := make([]*coreauth.Auth, 0, len(s.items))
|
||||
for _, item := range s.items {
|
||||
out = append(out, item)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *memoryAuthStore) Save(_ context.Context, auth *coreauth.Auth) (string, error) {
|
||||
if auth == nil {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.items == nil {
|
||||
s.items = make(map[string]*coreauth.Auth)
|
||||
}
|
||||
s.items[auth.ID] = auth
|
||||
return auth.ID, nil
|
||||
}
|
||||
|
||||
func (s *memoryAuthStore) Delete(_ context.Context, id string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
delete(s.items, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *memoryAuthStore) SetBaseDir(string) {}
|
||||
55
backend/internal/api/handlers/management/usage.go
Normal file
55
backend/internal/api/handlers/management/usage.go
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue"
|
||||
)
|
||||
|
||||
type usageQueueRecord []byte
|
||||
|
||||
func (r usageQueueRecord) MarshalJSON() ([]byte, error) {
|
||||
if json.Valid(r) {
|
||||
return append([]byte(nil), r...), nil
|
||||
}
|
||||
return json.Marshal(string(r))
|
||||
}
|
||||
|
||||
// GetUsageQueue pops queued usage records from the usage queue.
|
||||
func (h *Handler) GetUsageQueue(c *gin.Context) {
|
||||
if h == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "handler unavailable"})
|
||||
return
|
||||
}
|
||||
|
||||
count, errCount := parseUsageQueueCount(c.Query("count"))
|
||||
if errCount != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": errCount.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
items := redisqueue.PopOldest(count)
|
||||
records := make([]usageQueueRecord, 0, len(items))
|
||||
for _, item := range items {
|
||||
records = append(records, usageQueueRecord(append([]byte(nil), item...)))
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, records)
|
||||
}
|
||||
|
||||
func parseUsageQueueCount(value string) (int, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return 1, nil
|
||||
}
|
||||
count, errCount := strconv.Atoi(value)
|
||||
if errCount != nil || count <= 0 {
|
||||
return 0, errors.New("count must be a positive integer")
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
96
backend/internal/api/handlers/management/usage_test.go
Normal file
96
backend/internal/api/handlers/management/usage_test.go
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue"
|
||||
)
|
||||
|
||||
func TestGetUsageQueuePopsRequestedRecords(t *testing.T) {
|
||||
withManagementUsageQueue(t, func() {
|
||||
redisqueue.Enqueue([]byte(`{"id":1}`))
|
||||
redisqueue.Enqueue([]byte(`{"id":2}`))
|
||||
redisqueue.Enqueue([]byte(`{"id":3}`))
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ginCtx, _ := gin.CreateTestContext(rec)
|
||||
ginCtx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/usage-queue?count=2", nil)
|
||||
|
||||
h := &Handler{}
|
||||
h.GetUsageQueue(ginCtx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
|
||||
var payload []json.RawMessage
|
||||
if errUnmarshal := json.Unmarshal(rec.Body.Bytes(), &payload); errUnmarshal != nil {
|
||||
t.Fatalf("unmarshal response: %v", errUnmarshal)
|
||||
}
|
||||
if len(payload) != 2 {
|
||||
t.Fatalf("response records = %d, want 2", len(payload))
|
||||
}
|
||||
requireRecordID(t, payload[0], 1)
|
||||
requireRecordID(t, payload[1], 2)
|
||||
|
||||
remaining := redisqueue.PopOldest(10)
|
||||
if len(remaining) != 1 || string(remaining[0]) != `{"id":3}` {
|
||||
t.Fatalf("remaining queue = %q, want third item only", remaining)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetUsageQueueInvalidCountDoesNotPop(t *testing.T) {
|
||||
withManagementUsageQueue(t, func() {
|
||||
redisqueue.Enqueue([]byte(`{"id":1}`))
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ginCtx, _ := gin.CreateTestContext(rec)
|
||||
ginCtx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/usage-queue?count=0", nil)
|
||||
|
||||
h := &Handler{}
|
||||
h.GetUsageQueue(ginCtx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
|
||||
remaining := redisqueue.PopOldest(10)
|
||||
if len(remaining) != 1 || string(remaining[0]) != `{"id":1}` {
|
||||
t.Fatalf("remaining queue = %q, want original item", remaining)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func withManagementUsageQueue(t *testing.T, fn func()) {
|
||||
t.Helper()
|
||||
|
||||
prevQueueEnabled := redisqueue.Enabled()
|
||||
redisqueue.SetEnabled(false)
|
||||
redisqueue.SetEnabled(true)
|
||||
|
||||
defer func() {
|
||||
redisqueue.SetEnabled(false)
|
||||
redisqueue.SetEnabled(prevQueueEnabled)
|
||||
}()
|
||||
|
||||
fn()
|
||||
}
|
||||
|
||||
func requireRecordID(t *testing.T, raw json.RawMessage, want int) {
|
||||
t.Helper()
|
||||
|
||||
var payload struct {
|
||||
ID int `json:"id"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(raw, &payload); errUnmarshal != nil {
|
||||
t.Fatalf("unmarshal record: %v", errUnmarshal)
|
||||
}
|
||||
if payload.ID != want {
|
||||
t.Fatalf("record id = %d, want %d", payload.ID, want)
|
||||
}
|
||||
}
|
||||
156
backend/internal/api/handlers/management/vertex_import.go
Normal file
156
backend/internal/api/handlers/management/vertex_import.go
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/auth/vertex"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
// ImportVertexCredential handles uploading a Vertex service account JSON and saving it as an auth record.
|
||||
func (h *Handler) ImportVertexCredential(c *gin.Context) {
|
||||
if h == nil || h.cfg == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "config unavailable"})
|
||||
return
|
||||
}
|
||||
if h.cfg.AuthDir == "" {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "auth directory not configured"})
|
||||
return
|
||||
}
|
||||
|
||||
fileHeader, err := c.FormFile("file")
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "file required"})
|
||||
return
|
||||
}
|
||||
|
||||
file, err := fileHeader.Open()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("failed to read file: %v", err)})
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
data, err := io.ReadAll(file)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("failed to read file: %v", err)})
|
||||
return
|
||||
}
|
||||
|
||||
var serviceAccount map[string]any
|
||||
if err := json.Unmarshal(data, &serviceAccount); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
normalizedSA, err := vertex.NormalizeServiceAccountMap(serviceAccount)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid service account", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
serviceAccount = normalizedSA
|
||||
|
||||
projectID := strings.TrimSpace(valueAsString(serviceAccount["project_id"]))
|
||||
if projectID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "project_id missing"})
|
||||
return
|
||||
}
|
||||
email := strings.TrimSpace(valueAsString(serviceAccount["client_email"]))
|
||||
|
||||
location := strings.TrimSpace(c.PostForm("location"))
|
||||
if location == "" {
|
||||
location = strings.TrimSpace(c.Query("location"))
|
||||
}
|
||||
if location == "" {
|
||||
location = "us-central1"
|
||||
}
|
||||
|
||||
fileName := fmt.Sprintf("vertex-%s.json", sanitizeVertexFilePart(projectID))
|
||||
label := labelForVertex(projectID, email)
|
||||
storage := &vertex.VertexCredentialStorage{
|
||||
ServiceAccount: serviceAccount,
|
||||
ProjectID: projectID,
|
||||
Email: email,
|
||||
Location: location,
|
||||
Type: "vertex",
|
||||
}
|
||||
metadata := map[string]any{
|
||||
"service_account": serviceAccount,
|
||||
"project_id": projectID,
|
||||
"email": email,
|
||||
"location": location,
|
||||
"type": "vertex",
|
||||
"label": label,
|
||||
}
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
Provider: "vertex",
|
||||
FileName: fileName,
|
||||
Storage: storage,
|
||||
Label: label,
|
||||
Metadata: metadata,
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
if reqCtx := c.Request.Context(); reqCtx != nil {
|
||||
ctx = reqCtx
|
||||
}
|
||||
savedPath, err := h.saveTokenRecord(ctx, record)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "save_failed", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"status": "ok",
|
||||
"auth-file": savedPath,
|
||||
"project_id": projectID,
|
||||
"email": email,
|
||||
"location": location,
|
||||
})
|
||||
}
|
||||
|
||||
func valueAsString(v any) string {
|
||||
if v == nil {
|
||||
return ""
|
||||
}
|
||||
switch t := v.(type) {
|
||||
case string:
|
||||
return t
|
||||
default:
|
||||
return fmt.Sprint(t)
|
||||
}
|
||||
}
|
||||
|
||||
func sanitizeVertexFilePart(s string) string {
|
||||
out := strings.TrimSpace(s)
|
||||
replacers := []string{"/", "_", "\\", "_", ":", "_", " ", "-"}
|
||||
for i := 0; i < len(replacers); i += 2 {
|
||||
out = strings.ReplaceAll(out, replacers[i], replacers[i+1])
|
||||
}
|
||||
if out == "" {
|
||||
return "vertex"
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func labelForVertex(projectID, email string) string {
|
||||
p := strings.TrimSpace(projectID)
|
||||
e := strings.TrimSpace(email)
|
||||
if p != "" && e != "" {
|
||||
return fmt.Sprintf("%s (%s)", p, e)
|
||||
}
|
||||
if p != "" {
|
||||
return p
|
||||
}
|
||||
if e != "" {
|
||||
return e
|
||||
}
|
||||
return "vertex"
|
||||
}
|
||||
465
backend/internal/api/middleware/request_logging.go
Normal file
465
backend/internal/api/middleware/request_logging.go
Normal file
|
|
@ -0,0 +1,465 @@
|
|||
// Package middleware provides HTTP middleware components for the CLI Proxy API server.
|
||||
// This file contains the request logging middleware that captures comprehensive
|
||||
// request and response data when enabled through configuration.
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/klauspost/compress/zstd"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
maxErrorOnlyCapturedRequestBodyBytes int64 = 1 << 20 // 1 MiB
|
||||
maxDeferredErrorRequestBodyBytes int64 = 32 << 20 // 32 MiB
|
||||
)
|
||||
|
||||
// RequestLoggingMiddleware creates a Gin middleware that logs HTTP requests and responses.
|
||||
// It captures detailed information about the request and response, including headers and body,
|
||||
// and uses the provided RequestLogger to record this data. When full request logging is disabled,
|
||||
// large and unknown-size bodies are spooled to disk and retained only for error logs.
|
||||
func RequestLoggingMiddleware(logger logging.RequestLogger) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if logger == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
if shouldSkipMethodForRequestLogging(c.Request) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
path := c.Request.URL.Path
|
||||
if !shouldLogRequest(path) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
loggerEnabled := logger.IsEnabled()
|
||||
captureBody := shouldCaptureRequestBody(loggerEnabled, c.Request)
|
||||
|
||||
// Capture request information
|
||||
requestInfo, err := captureRequestInfo(c, captureBody)
|
||||
if err != nil {
|
||||
// Log error but continue processing
|
||||
// In a real implementation, you might want to use a proper logger here
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
// Create response writer wrapper
|
||||
wrapper := NewResponseWriterWrapper(c.Writer, logger, requestInfo)
|
||||
if !loggerEnabled {
|
||||
wrapper.logOnErrorOnly = true
|
||||
}
|
||||
c.Writer = wrapper
|
||||
attachRequestLogSources(c, logger, loggerEnabled)
|
||||
attachDeferredRequestBodyCapture(c.Request, logger, requestInfo, loggerEnabled, captureBody)
|
||||
|
||||
// Process the request
|
||||
c.Next()
|
||||
|
||||
// Finalize logging after request processing
|
||||
if err = wrapper.Finalize(c); err != nil {
|
||||
// Log error but don't interrupt the response
|
||||
// In a real implementation, you might want to use a proper logger here
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type fileBodySourceFactory interface {
|
||||
NewFileBodySource(prefix string) (*logging.FileBodySource, error)
|
||||
}
|
||||
|
||||
type deferredRequestBodyCapture struct {
|
||||
body io.ReadCloser
|
||||
file *os.File
|
||||
source *logging.FileBodySource
|
||||
contentLength int64
|
||||
bytesRead int64
|
||||
bytesCaptured int64
|
||||
captureErr error
|
||||
finished bool
|
||||
sawEOF bool
|
||||
truncated bool
|
||||
}
|
||||
|
||||
func attachDeferredRequestBodyCapture(req *http.Request, logger logging.RequestLogger, requestInfo *RequestInfo, loggerEnabled, bodyCaptured bool) *deferredRequestBodyCapture {
|
||||
if loggerEnabled || bodyCaptured || req == nil || req.Body == nil || req.Body == http.NoBody || req.ContentLength == 0 || requestInfo == nil {
|
||||
return nil
|
||||
}
|
||||
contentType := strings.ToLower(strings.TrimSpace(req.Header.Get("Content-Type")))
|
||||
if strings.HasPrefix(contentType, "multipart/form-data") {
|
||||
return nil
|
||||
}
|
||||
factory, ok := logger.(fileBodySourceFactory)
|
||||
if !ok || factory == nil {
|
||||
return nil
|
||||
}
|
||||
source, errSource := factory.NewFileBodySource("request-body")
|
||||
if errSource != nil {
|
||||
return nil
|
||||
}
|
||||
file, errPart := source.CreatePart("body")
|
||||
if errPart != nil {
|
||||
_ = source.Cleanup()
|
||||
return nil
|
||||
}
|
||||
capture := &deferredRequestBodyCapture{
|
||||
body: req.Body,
|
||||
file: file,
|
||||
source: source,
|
||||
contentLength: req.ContentLength,
|
||||
}
|
||||
req.Body = capture
|
||||
requestInfo.deferredBodyCapture = capture
|
||||
return capture
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) Read(payload []byte) (int, error) {
|
||||
if c == nil || c.body == nil {
|
||||
return 0, io.EOF
|
||||
}
|
||||
n, errRead := c.body.Read(payload)
|
||||
if errRead == io.EOF {
|
||||
c.sawEOF = true
|
||||
}
|
||||
if n == 0 {
|
||||
return n, errRead
|
||||
}
|
||||
c.bytesRead += int64(n)
|
||||
if c.file == nil || c.captureErr != nil {
|
||||
return n, errRead
|
||||
}
|
||||
|
||||
remaining := maxDeferredErrorRequestBodyBytes - c.bytesCaptured
|
||||
if remaining <= 0 {
|
||||
c.truncated = true
|
||||
return n, errRead
|
||||
}
|
||||
writeLength := int64(n)
|
||||
if writeLength > remaining {
|
||||
writeLength = remaining
|
||||
c.truncated = true
|
||||
}
|
||||
written, errWrite := c.file.Write(payload[:int(writeLength)])
|
||||
c.bytesCaptured += int64(written)
|
||||
if errWrite != nil {
|
||||
c.captureErr = errWrite
|
||||
} else if int64(written) != writeLength {
|
||||
c.captureErr = io.ErrShortWrite
|
||||
}
|
||||
if c.captureErr != nil {
|
||||
if errClose := c.file.Close(); errClose != nil {
|
||||
c.captureErr = fmt.Errorf("%v; close capture file: %w", c.captureErr, errClose)
|
||||
}
|
||||
c.file = nil
|
||||
}
|
||||
return n, errRead
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) Close() error {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
_ = c.Finish()
|
||||
if c.body == nil {
|
||||
return nil
|
||||
}
|
||||
return c.body.Close()
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) Finish() error {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
if c.finished {
|
||||
return c.captureErr
|
||||
}
|
||||
c.finished = true
|
||||
if c.file != nil {
|
||||
if errClose := c.file.Close(); errClose != nil && c.captureErr == nil {
|
||||
c.captureErr = errClose
|
||||
}
|
||||
c.file = nil
|
||||
}
|
||||
return c.captureErr
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) Bytes() ([]byte, string, error) {
|
||||
if c == nil || c.source == nil {
|
||||
return nil, "", nil
|
||||
}
|
||||
if errFinish := c.Finish(); errFinish != nil {
|
||||
return nil, "", errFinish
|
||||
}
|
||||
body, errBytes := c.source.Bytes()
|
||||
if errBytes != nil {
|
||||
return nil, "", errBytes
|
||||
}
|
||||
return body, c.statusMarker(), nil
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) statusMarker() string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
var markers []string
|
||||
if c.truncated {
|
||||
markers = append(markers, fmt.Sprintf("[REQUEST BODY TRUNCATED: captured first %d bytes]", c.bytesCaptured))
|
||||
}
|
||||
complete := c.sawEOF || (c.contentLength >= 0 && c.bytesRead >= c.contentLength)
|
||||
if !complete {
|
||||
if c.contentLength >= 0 {
|
||||
markers = append(markers, fmt.Sprintf("[REQUEST BODY CAPTURE INCOMPLETE: consumed %d of %d bytes]", c.bytesRead, c.contentLength))
|
||||
} else {
|
||||
markers = append(markers, fmt.Sprintf("[REQUEST BODY CAPTURE INCOMPLETE: consumed %d bytes from an unknown-length body]", c.bytesRead))
|
||||
}
|
||||
}
|
||||
return strings.Join(markers, "\n")
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) Cleanup() {
|
||||
if c == nil || c.source == nil {
|
||||
return
|
||||
}
|
||||
if errFinish := c.Finish(); errFinish != nil {
|
||||
log.WithError(errFinish).Warn("failed to finish deferred request body capture")
|
||||
}
|
||||
if errCleanup := c.source.Cleanup(); errCleanup != nil {
|
||||
log.WithError(errCleanup).Warn("failed to clean up deferred request body capture")
|
||||
}
|
||||
c.source = nil
|
||||
}
|
||||
|
||||
func attachRequestLogSources(c *gin.Context, logger logging.RequestLogger, loggerEnabled bool) {
|
||||
if c == nil || !loggerEnabled {
|
||||
return
|
||||
}
|
||||
factory, ok := logger.(fileBodySourceFactory)
|
||||
if !ok || factory == nil {
|
||||
return
|
||||
}
|
||||
if source, errSource := factory.NewFileBodySource("api-request"); errSource == nil {
|
||||
c.Set(logging.APIRequestSourceContextKey, source)
|
||||
}
|
||||
if source, errSource := factory.NewFileBodySource("api-response"); errSource == nil {
|
||||
c.Set(logging.APIResponseSourceContextKey, source)
|
||||
}
|
||||
if !isResponsesWebsocketUpgrade(c.Request) {
|
||||
return
|
||||
}
|
||||
if source, errSource := factory.NewFileBodySource("websocket-timeline"); errSource == nil {
|
||||
c.Set(logging.WebsocketTimelineSourceContextKey, source)
|
||||
}
|
||||
if source, errSource := factory.NewFileBodySource("api-websocket-timeline"); errSource == nil {
|
||||
c.Set(logging.APIWebsocketTimelineSourceContextKey, source)
|
||||
}
|
||||
}
|
||||
|
||||
func shouldSkipMethodForRequestLogging(req *http.Request) bool {
|
||||
if req == nil {
|
||||
return true
|
||||
}
|
||||
if req.Method != http.MethodGet {
|
||||
return false
|
||||
}
|
||||
return !isResponsesWebsocketUpgrade(req)
|
||||
}
|
||||
|
||||
func isResponsesWebsocketUpgrade(req *http.Request) bool {
|
||||
if req == nil || req.URL == nil {
|
||||
return false
|
||||
}
|
||||
if req.URL.Path != "/v1/responses" && req.URL.Path != "/backend-api/codex/responses" {
|
||||
return false
|
||||
}
|
||||
return strings.EqualFold(strings.TrimSpace(req.Header.Get("Upgrade")), "websocket")
|
||||
}
|
||||
|
||||
func shouldCaptureRequestBody(loggerEnabled bool, req *http.Request) bool {
|
||||
if loggerEnabled {
|
||||
return true
|
||||
}
|
||||
if req == nil || req.Body == nil {
|
||||
return false
|
||||
}
|
||||
contentType := strings.ToLower(strings.TrimSpace(req.Header.Get("Content-Type")))
|
||||
if strings.HasPrefix(contentType, "multipart/form-data") {
|
||||
return false
|
||||
}
|
||||
if req.ContentLength <= 0 {
|
||||
return false
|
||||
}
|
||||
return req.ContentLength <= maxErrorOnlyCapturedRequestBodyBytes
|
||||
}
|
||||
|
||||
// captureRequestInfo extracts relevant information from the incoming HTTP request.
|
||||
// It captures the URL, method, headers, and body. The request body is read and then
|
||||
// restored so that it can be processed by subsequent handlers.
|
||||
func captureRequestInfo(c *gin.Context, captureBody bool) (*RequestInfo, error) {
|
||||
// Capture URL with sensitive query parameters masked
|
||||
maskedQuery := util.MaskSensitiveQuery(c.Request.URL.RawQuery)
|
||||
url := c.Request.URL.Path
|
||||
if maskedQuery != "" {
|
||||
url += "?" + maskedQuery
|
||||
}
|
||||
|
||||
// Capture method
|
||||
method := c.Request.Method
|
||||
|
||||
// Capture headers
|
||||
headers := make(map[string][]string)
|
||||
for key, values := range c.Request.Header {
|
||||
headers[key] = values
|
||||
}
|
||||
|
||||
// Capture request body
|
||||
var body []byte
|
||||
if captureBody && c.Request.Body != nil {
|
||||
// Read the body
|
||||
bodyBytes, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Restore the body for the actual request processing
|
||||
c.Request.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
|
||||
body = decodeCapturedRequestBodyForLog(bodyBytes, c.Request.Header.Get("Content-Encoding"))
|
||||
}
|
||||
|
||||
return &RequestInfo{
|
||||
URL: url,
|
||||
Method: method,
|
||||
Headers: headers,
|
||||
Body: body,
|
||||
RequestID: logging.GetGinRequestID(c),
|
||||
Timestamp: time.Now(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func decodeCapturedRequestBodyForLog(raw []byte, encoding string) []byte {
|
||||
if len(raw) == 0 {
|
||||
return raw
|
||||
}
|
||||
|
||||
decoded, errDecode := decodeCapturedRequestBody(raw, encoding)
|
||||
if errDecode != nil {
|
||||
return raw
|
||||
}
|
||||
return decoded
|
||||
}
|
||||
|
||||
func decodeCapturedRequestBodyForLogWithLimit(raw []byte, encoding string, limit int64) []byte {
|
||||
if len(raw) == 0 || limit <= 0 {
|
||||
return raw
|
||||
}
|
||||
encoding = strings.TrimSpace(encoding)
|
||||
if encoding == "" || strings.EqualFold(encoding, "identity") {
|
||||
return raw
|
||||
}
|
||||
|
||||
parts := strings.Split(encoding, ",")
|
||||
body := raw
|
||||
for i := len(parts) - 1; i >= 0; i-- {
|
||||
enc := strings.ToLower(strings.TrimSpace(parts[i]))
|
||||
switch enc {
|
||||
case "", "identity":
|
||||
continue
|
||||
case "zstd":
|
||||
decoded, truncated, errDecode := decodeCapturedZstdRequestBodyWithLimit(body, limit)
|
||||
if errDecode != nil {
|
||||
return raw
|
||||
}
|
||||
body = decoded
|
||||
if truncated {
|
||||
if len(body) > 0 && !bytes.HasSuffix(body, []byte("\n")) {
|
||||
body = append(body, '\n')
|
||||
}
|
||||
return append(body, "[DECOMPRESSED REQUEST BODY TRUNCATED]"...)
|
||||
}
|
||||
default:
|
||||
return raw
|
||||
}
|
||||
}
|
||||
return body
|
||||
}
|
||||
|
||||
func decodeCapturedRequestBody(raw []byte, encoding string) ([]byte, error) {
|
||||
encoding = strings.TrimSpace(encoding)
|
||||
if encoding == "" || strings.EqualFold(encoding, "identity") {
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
parts := strings.Split(encoding, ",")
|
||||
body := raw
|
||||
for i := len(parts) - 1; i >= 0; i-- {
|
||||
enc := strings.ToLower(strings.TrimSpace(parts[i]))
|
||||
switch enc {
|
||||
case "", "identity":
|
||||
continue
|
||||
case "zstd":
|
||||
decoded, errDecode := decodeCapturedZstdRequestBody(body)
|
||||
if errDecode != nil {
|
||||
return nil, errDecode
|
||||
}
|
||||
body = decoded
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported request content encoding: %s", enc)
|
||||
}
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func decodeCapturedZstdRequestBody(raw []byte) ([]byte, error) {
|
||||
decoder, errNewReader := zstd.NewReader(bytes.NewReader(raw))
|
||||
if errNewReader != nil {
|
||||
return nil, fmt.Errorf("failed to create zstd request decoder: %w", errNewReader)
|
||||
}
|
||||
defer decoder.Close()
|
||||
|
||||
decoded, errRead := io.ReadAll(decoder)
|
||||
if errRead != nil {
|
||||
return nil, fmt.Errorf("failed to decode zstd request body: %w", errRead)
|
||||
}
|
||||
return decoded, nil
|
||||
}
|
||||
|
||||
func decodeCapturedZstdRequestBodyWithLimit(raw []byte, limit int64) ([]byte, bool, error) {
|
||||
decoder, errNewReader := zstd.NewReader(bytes.NewReader(raw))
|
||||
if errNewReader != nil {
|
||||
return nil, false, fmt.Errorf("failed to create zstd request decoder: %w", errNewReader)
|
||||
}
|
||||
defer decoder.Close()
|
||||
|
||||
decoded, errRead := io.ReadAll(io.LimitReader(decoder, limit+1))
|
||||
if errRead != nil {
|
||||
return nil, false, fmt.Errorf("failed to decode zstd request body: %w", errRead)
|
||||
}
|
||||
if int64(len(decoded)) > limit {
|
||||
return decoded[:limit], true, nil
|
||||
}
|
||||
return decoded, false, nil
|
||||
}
|
||||
|
||||
// shouldLogRequest determines whether the request should be logged.
|
||||
// It skips management endpoints to avoid leaking secrets but allows
|
||||
// all other routes, including module-provided ones, to honor request-log.
|
||||
func shouldLogRequest(path string) bool {
|
||||
if strings.HasPrefix(path, "/v0/management") || strings.HasPrefix(path, "/management") {
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
507
backend/internal/api/middleware/request_logging_test.go
Normal file
507
backend/internal/api/middleware/request_logging_test.go
Normal file
|
|
@ -0,0 +1,507 @@
|
|||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/klauspost/compress/zstd"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
|
||||
)
|
||||
|
||||
func TestShouldSkipMethodForRequestLogging(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
req *http.Request
|
||||
skip bool
|
||||
}{
|
||||
{
|
||||
name: "nil request",
|
||||
req: nil,
|
||||
skip: true,
|
||||
},
|
||||
{
|
||||
name: "post request should not skip",
|
||||
req: &http.Request{
|
||||
Method: http.MethodPost,
|
||||
URL: &url.URL{Path: "/v1/responses"},
|
||||
},
|
||||
skip: false,
|
||||
},
|
||||
{
|
||||
name: "plain get should skip",
|
||||
req: &http.Request{
|
||||
Method: http.MethodGet,
|
||||
URL: &url.URL{Path: "/v1/models"},
|
||||
Header: http.Header{},
|
||||
},
|
||||
skip: true,
|
||||
},
|
||||
{
|
||||
name: "responses websocket upgrade should not skip",
|
||||
req: &http.Request{
|
||||
Method: http.MethodGet,
|
||||
URL: &url.URL{Path: "/v1/responses"},
|
||||
Header: http.Header{"Upgrade": []string{"websocket"}},
|
||||
},
|
||||
skip: false,
|
||||
},
|
||||
{
|
||||
name: "codex responses websocket upgrade should not skip",
|
||||
req: &http.Request{
|
||||
Method: http.MethodGet,
|
||||
URL: &url.URL{Path: "/backend-api/codex/responses"},
|
||||
Header: http.Header{"Upgrade": []string{"websocket"}},
|
||||
},
|
||||
skip: false,
|
||||
},
|
||||
{
|
||||
name: "responses get without upgrade should skip",
|
||||
req: &http.Request{
|
||||
Method: http.MethodGet,
|
||||
URL: &url.URL{Path: "/v1/responses"},
|
||||
Header: http.Header{},
|
||||
},
|
||||
skip: true,
|
||||
},
|
||||
}
|
||||
|
||||
for i := range tests {
|
||||
got := shouldSkipMethodForRequestLogging(tests[i].req)
|
||||
if got != tests[i].skip {
|
||||
t.Fatalf("%s: got skip=%t, want %t", tests[i].name, got, tests[i].skip)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldCaptureRequestBody(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
loggerEnabled bool
|
||||
req *http.Request
|
||||
want bool
|
||||
}{
|
||||
{
|
||||
name: "logger enabled always captures",
|
||||
loggerEnabled: true,
|
||||
req: &http.Request{
|
||||
Body: io.NopCloser(strings.NewReader("{}")),
|
||||
ContentLength: -1,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
},
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "nil request",
|
||||
loggerEnabled: false,
|
||||
req: nil,
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "small known size json in error-only mode",
|
||||
loggerEnabled: false,
|
||||
req: &http.Request{
|
||||
Body: io.NopCloser(strings.NewReader("{}")),
|
||||
ContentLength: 2,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
},
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "large known size skipped in error-only mode",
|
||||
loggerEnabled: false,
|
||||
req: &http.Request{
|
||||
Body: io.NopCloser(strings.NewReader("x")),
|
||||
ContentLength: maxErrorOnlyCapturedRequestBodyBytes + 1,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
},
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "unknown size skipped in error-only mode",
|
||||
loggerEnabled: false,
|
||||
req: &http.Request{
|
||||
Body: io.NopCloser(strings.NewReader("x")),
|
||||
ContentLength: -1,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
},
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "multipart skipped in error-only mode",
|
||||
loggerEnabled: false,
|
||||
req: &http.Request{
|
||||
Body: io.NopCloser(strings.NewReader("x")),
|
||||
ContentLength: 1,
|
||||
Header: http.Header{"Content-Type": []string{"multipart/form-data; boundary=abc"}},
|
||||
},
|
||||
want: false,
|
||||
},
|
||||
}
|
||||
|
||||
for i := range tests {
|
||||
got := shouldCaptureRequestBody(tests[i].loggerEnabled, tests[i].req)
|
||||
if got != tests[i].want {
|
||||
t.Fatalf("%s: got %t, want %t", tests[i].name, got, tests[i].want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeferredRequestBodyCaptureDoesNotDrainUnreadBody(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
logger := logging.NewFileRequestLogger(false, t.TempDir(), "", 10)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader("remaining-body"))
|
||||
request.ContentLength = -1
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
requestInfo := &RequestInfo{Headers: map[string][]string{"Content-Type": {"application/json"}}}
|
||||
capture := attachDeferredRequestBodyCapture(request, logger, requestInfo, false, false)
|
||||
if capture == nil {
|
||||
t.Fatal("deferred request body capture was not attached")
|
||||
}
|
||||
defer capture.Cleanup()
|
||||
|
||||
firstByte := make([]byte, 1)
|
||||
if _, errRead := request.Body.Read(firstByte); errRead != nil {
|
||||
t.Fatalf("read first request byte: %v", errRead)
|
||||
}
|
||||
captured, marker, errCaptured := capture.Bytes()
|
||||
if errCaptured != nil {
|
||||
t.Fatalf("read captured body: %v", errCaptured)
|
||||
}
|
||||
if string(captured) != "r" {
|
||||
t.Fatalf("captured body = %q, want %q", string(captured), "r")
|
||||
}
|
||||
if !strings.Contains(marker, "REQUEST BODY CAPTURE INCOMPLETE") {
|
||||
t.Fatalf("capture marker = %q, want incomplete marker", marker)
|
||||
}
|
||||
remaining, errRemaining := io.ReadAll(capture.body)
|
||||
if errRemaining != nil {
|
||||
t.Fatalf("read remaining body: %v", errRemaining)
|
||||
}
|
||||
if string(remaining) != "emaining-body" {
|
||||
t.Fatalf("remaining body = %q, want %q", string(remaining), "emaining-body")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestLoggingMiddlewareCapturesLargeErrorRequestAndDeferredAPIRequest(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
logsDir := t.TempDir()
|
||||
logger := logging.NewFileRequestLogger(false, logsDir, "", 10)
|
||||
payload := append([]byte(`{"marker":"large-error-body","padding":"`), bytes.Repeat([]byte("x"), int(maxErrorOnlyCapturedRequestBodyBytes))...)
|
||||
payload = append(payload, []byte(`"}`)...)
|
||||
upstreamBody := []byte(`{"model":"upstream-model","input":"translated"}`)
|
||||
|
||||
router := gin.New()
|
||||
router.Use(RequestLoggingMiddleware(logger))
|
||||
router.POST("/v1/responses", func(c *gin.Context) {
|
||||
body, errRead := io.ReadAll(c.Request.Body)
|
||||
if errRead != nil {
|
||||
c.Status(http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
if !bytes.Equal(body, payload) {
|
||||
c.Status(http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
executorCtx := context.WithValue(context.Background(), "gin", c)
|
||||
helps.RecordAPIRequest(executorCtx, &config.Config{}, helps.UpstreamRequestLog{
|
||||
URL: "https://api.example.com/v1/responses",
|
||||
Method: http.MethodPost,
|
||||
Headers: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: upstreamBody,
|
||||
})
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "upstream rejected request"})
|
||||
})
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(payload))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
|
||||
if response.Code != http.StatusBadRequest {
|
||||
t.Fatalf("response status = %d, want %d", response.Code, http.StatusBadRequest)
|
||||
}
|
||||
entries, errReadDir := os.ReadDir(logsDir)
|
||||
if errReadDir != nil {
|
||||
t.Fatalf("read logs dir: %v", errReadDir)
|
||||
}
|
||||
var logPath string
|
||||
for _, entry := range entries {
|
||||
if strings.HasPrefix(entry.Name(), "error-") && strings.HasSuffix(entry.Name(), ".log") {
|
||||
logPath = logsDir + string(os.PathSeparator) + entry.Name()
|
||||
break
|
||||
}
|
||||
}
|
||||
if logPath == "" {
|
||||
t.Fatal("forced error log was not created")
|
||||
}
|
||||
content, errReadLog := os.ReadFile(logPath)
|
||||
if errReadLog != nil {
|
||||
t.Fatalf("read error log: %v", errReadLog)
|
||||
}
|
||||
if !bytes.Contains(content, payload) {
|
||||
t.Fatal("error log does not contain the complete large request body")
|
||||
}
|
||||
if !bytes.Contains(content, []byte("=== API REQUEST 1 ===")) {
|
||||
t.Fatal("error log does not contain the deferred API request section")
|
||||
}
|
||||
if !bytes.Contains(content, upstreamBody) {
|
||||
t.Fatal("error log does not contain the deferred upstream request body")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAttachRequestLogSourcesUsesLoggerLogsDir(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
logsDir := t.TempDir()
|
||||
logger := logging.NewFileRequestLogger(true, logsDir, "", 0)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/backend-api/codex/responses", nil)
|
||||
c.Request.Header.Set("Upgrade", "websocket")
|
||||
|
||||
attachRequestLogSources(c, logger, true)
|
||||
defer cleanupFileBodySourcesFromContext(c)
|
||||
|
||||
for _, key := range []string{
|
||||
logging.WebsocketTimelineSourceContextKey,
|
||||
logging.APIWebsocketTimelineSourceContextKey,
|
||||
} {
|
||||
value, exists := c.Get(key)
|
||||
if !exists {
|
||||
t.Fatalf("expected %s source to be attached", key)
|
||||
}
|
||||
source, ok := value.(*logging.FileBodySource)
|
||||
if !ok || source == nil {
|
||||
t.Fatalf("%s source type = %T", key, value)
|
||||
}
|
||||
file, errPart := source.CreatePart("probe")
|
||||
if errPart != nil {
|
||||
t.Fatalf("CreatePart(%s): %v", key, errPart)
|
||||
}
|
||||
path := file.Name()
|
||||
if errClose := file.Close(); errClose != nil {
|
||||
t.Fatalf("close part: %v", errClose)
|
||||
}
|
||||
if !strings.HasPrefix(path, logsDir+string(os.PathSeparator)) {
|
||||
t.Fatalf("%s part path %s is not under logs dir %s", key, path, logsDir)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func cleanupFileBodySourcesFromContext(c *gin.Context) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, key := range []string{
|
||||
logging.WebsocketTimelineSourceContextKey,
|
||||
logging.APIWebsocketTimelineSourceContextKey,
|
||||
} {
|
||||
value, exists := c.Get(key)
|
||||
if !exists {
|
||||
continue
|
||||
}
|
||||
if source, ok := value.(*logging.FileBodySource); ok && source != nil {
|
||||
_ = source.Cleanup()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeCapturedRequestBodyForLogWithLimitTruncatesZstdExpansion(t *testing.T) {
|
||||
payload := bytes.Repeat([]byte("x"), 1024)
|
||||
var compressed bytes.Buffer
|
||||
encoder, errNewWriter := zstd.NewWriter(&compressed)
|
||||
if errNewWriter != nil {
|
||||
t.Fatalf("zstd.NewWriter: %v", errNewWriter)
|
||||
}
|
||||
if _, errWrite := encoder.Write(payload); errWrite != nil {
|
||||
t.Fatalf("zstd write: %v", errWrite)
|
||||
}
|
||||
if errClose := encoder.Close(); errClose != nil {
|
||||
t.Fatalf("zstd close: %v", errClose)
|
||||
}
|
||||
|
||||
decoded := decodeCapturedRequestBodyForLogWithLimit(compressed.Bytes(), "zstd", 64)
|
||||
if len(decoded) > 128 {
|
||||
t.Fatalf("limited decoded body length = %d, want bounded output", len(decoded))
|
||||
}
|
||||
if !bytes.Contains(decoded, []byte("DECOMPRESSED REQUEST BODY TRUNCATED")) {
|
||||
t.Fatalf("decoded body = %q, want truncation marker", string(decoded))
|
||||
}
|
||||
}
|
||||
|
||||
func TestCaptureRequestInfoDecodesZstdRequestBodyForLog(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
payload := []byte(`{"model":"test-model","stream":true}`)
|
||||
var compressed bytes.Buffer
|
||||
encoder, errNewWriter := zstd.NewWriter(&compressed)
|
||||
if errNewWriter != nil {
|
||||
t.Fatalf("zstd.NewWriter: %v", errNewWriter)
|
||||
}
|
||||
if _, errWrite := encoder.Write(payload); errWrite != nil {
|
||||
t.Fatalf("zstd write: %v", errWrite)
|
||||
}
|
||||
if errClose := encoder.Close(); errClose != nil {
|
||||
t.Fatalf("zstd close: %v", errClose)
|
||||
}
|
||||
compressedBytes := compressed.Bytes()
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(compressedBytes))
|
||||
req.Header.Set("Content-Encoding", "zstd")
|
||||
c.Request = req
|
||||
|
||||
info, errCapture := captureRequestInfo(c, true)
|
||||
if errCapture != nil {
|
||||
t.Fatalf("captureRequestInfo: %v", errCapture)
|
||||
}
|
||||
if !bytes.Equal(info.Body, payload) {
|
||||
t.Fatalf("logged request body = %q, want %q", string(info.Body), string(payload))
|
||||
}
|
||||
|
||||
restoredBody, errRead := io.ReadAll(c.Request.Body)
|
||||
if errRead != nil {
|
||||
t.Fatalf("read restored request body: %v", errRead)
|
||||
}
|
||||
if !bytes.Equal(restoredBody, compressedBytes) {
|
||||
t.Fatal("request body was not restored with the original compressed bytes")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestLoggingMiddleware_ClientCancellationExclusion(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
t.Run("499 status does not create error log when request-log is false", func(t *testing.T) {
|
||||
logsDir := t.TempDir()
|
||||
logger := logging.NewFileRequestLogger(false, logsDir, "", 10)
|
||||
|
||||
router := gin.New()
|
||||
router.Use(RequestLoggingMiddleware(logger))
|
||||
router.POST("/v1/responses", func(c *gin.Context) {
|
||||
c.AbortWithStatus(clienterror.StatusClientClosedRequest)
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-4"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
|
||||
if resp.Code != clienterror.StatusClientClosedRequest {
|
||||
t.Fatalf("status = %d, want %d", resp.Code, clienterror.StatusClientClosedRequest)
|
||||
}
|
||||
|
||||
entries, errRead := os.ReadDir(logsDir)
|
||||
if errRead != nil {
|
||||
t.Fatalf("read logs dir: %v", errRead)
|
||||
}
|
||||
if len(entries) != 0 {
|
||||
t.Fatalf("expected 0 log files for 499 cancellation in error-only mode, got %d files", len(entries))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("context canceled does not create error log when request-log is false", func(t *testing.T) {
|
||||
logsDir := t.TempDir()
|
||||
logger := logging.NewFileRequestLogger(false, logsDir, "", 10)
|
||||
|
||||
router := gin.New()
|
||||
router.Use(RequestLoggingMiddleware(logger))
|
||||
router.POST("/v1/responses", func(c *gin.Context) {
|
||||
// Simulate client closing connection mid-flight
|
||||
ctx, cancel := context.WithCancel(c.Request.Context())
|
||||
cancel()
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-4"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
|
||||
entries, errRead := os.ReadDir(logsDir)
|
||||
if errRead != nil {
|
||||
t.Fatalf("read logs dir: %v", errRead)
|
||||
}
|
||||
if len(entries) != 0 {
|
||||
t.Fatalf("expected 0 log files for canceled context in error-only mode, got %d files", len(entries))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("400 bad request creates error log when request-log is false", func(t *testing.T) {
|
||||
logsDir := t.TempDir()
|
||||
logger := logging.NewFileRequestLogger(false, logsDir, "", 10)
|
||||
|
||||
router := gin.New()
|
||||
router.Use(RequestLoggingMiddleware(logger))
|
||||
router.POST("/v1/responses", func(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid parameter"})
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"bad":"param"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
|
||||
if resp.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d", resp.Code, http.StatusBadRequest)
|
||||
}
|
||||
|
||||
entries, errRead := os.ReadDir(logsDir)
|
||||
if errRead != nil {
|
||||
t.Fatalf("read logs dir: %v", errRead)
|
||||
}
|
||||
var errorLogCount int
|
||||
for _, entry := range entries {
|
||||
if strings.HasPrefix(entry.Name(), "error-") && strings.HasSuffix(entry.Name(), ".log") {
|
||||
errorLogCount++
|
||||
}
|
||||
}
|
||||
if errorLogCount != 1 {
|
||||
t.Fatalf("expected 1 error log file for 400 Bad Request, got %d", errorLogCount)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("499 status logs standard request when request-log is true", func(t *testing.T) {
|
||||
logsDir := t.TempDir()
|
||||
logger := logging.NewFileRequestLogger(true, logsDir, "", 10)
|
||||
|
||||
router := gin.New()
|
||||
router.Use(RequestLoggingMiddleware(logger))
|
||||
router.POST("/v1/responses", func(c *gin.Context) {
|
||||
c.AbortWithStatus(clienterror.StatusClientClosedRequest)
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-4"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
|
||||
entries, errRead := os.ReadDir(logsDir)
|
||||
if errRead != nil {
|
||||
t.Fatalf("read logs dir: %v", errRead)
|
||||
}
|
||||
var standardLogCount int
|
||||
for _, entry := range entries {
|
||||
if !strings.HasPrefix(entry.Name(), "error-") && strings.HasSuffix(entry.Name(), ".log") {
|
||||
standardLogCount++
|
||||
}
|
||||
}
|
||||
if standardLogCount != 1 {
|
||||
t.Fatalf("expected 1 standard request log file when request-log=true, got %d", standardLogCount)
|
||||
}
|
||||
})
|
||||
}
|
||||
761
backend/internal/api/middleware/response_writer.go
Normal file
761
backend/internal/api/middleware/response_writer.go
Normal file
|
|
@ -0,0 +1,761 @@
|
|||
// Package middleware provides Gin HTTP middleware for the CLI Proxy API server.
|
||||
// It includes a sophisticated response writer wrapper designed to capture and log request and response data,
|
||||
// including support for streaming responses, without impacting latency.
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const requestBodyOverrideContextKey = "REQUEST_BODY_OVERRIDE"
|
||||
const responseBodyOverrideContextKey = "RESPONSE_BODY_OVERRIDE"
|
||||
const websocketTimelineOverrideContextKey = "WEBSOCKET_TIMELINE_OVERRIDE"
|
||||
|
||||
// RequestInfo holds essential details of an incoming HTTP request for logging purposes.
|
||||
type RequestInfo struct {
|
||||
URL string // URL is the request URL.
|
||||
Method string // Method is the HTTP method (e.g., GET or POST).
|
||||
Headers map[string][]string // Headers contains the request headers.
|
||||
Body []byte // Body is the raw request body.
|
||||
RequestID string // RequestID is the unique identifier for the request.
|
||||
Timestamp time.Time // Timestamp is when the request was received.
|
||||
deferredBodyCapture *deferredRequestBodyCapture // deferredBodyCapture spools large error-only request bodies.
|
||||
}
|
||||
|
||||
// ResponseWriterWrapper wraps the standard gin.ResponseWriter to intercept and log response data.
|
||||
// It is designed to handle both standard and streaming responses, ensuring that logging operations do not block the client response.
|
||||
type ResponseWriterWrapper struct {
|
||||
gin.ResponseWriter
|
||||
body *bytes.Buffer // body is a buffer to store the response body for non-streaming responses.
|
||||
isStreaming bool // isStreaming indicates whether the response is a streaming type (e.g., text/event-stream).
|
||||
streamWriter logging.StreamingLogWriter // streamWriter is a writer for handling streaming log entries.
|
||||
chunkChannel chan []byte // chunkChannel is a channel for asynchronously passing response chunks to the logger.
|
||||
streamDone chan struct{} // streamDone signals when the streaming goroutine completes.
|
||||
logger logging.RequestLogger // logger is the instance of the request logger service.
|
||||
requestInfo *RequestInfo // requestInfo holds the details of the original request.
|
||||
statusCode int // statusCode stores the HTTP status code of the response.
|
||||
headers map[string][]string // headers stores the response headers.
|
||||
logOnErrorOnly bool // logOnErrorOnly enables logging only when an error response is detected.
|
||||
firstChunkTimestamp time.Time // firstChunkTimestamp captures TTFB for streaming responses.
|
||||
}
|
||||
|
||||
// NewResponseWriterWrapper creates and initializes a new ResponseWriterWrapper.
|
||||
// It takes the original gin.ResponseWriter, a logger instance, and request information.
|
||||
//
|
||||
// Parameters:
|
||||
// - w: The original gin.ResponseWriter to wrap.
|
||||
// - logger: The logging service to use for recording requests.
|
||||
// - requestInfo: The pre-captured information about the incoming request.
|
||||
//
|
||||
// Returns:
|
||||
// - A pointer to a new ResponseWriterWrapper.
|
||||
func NewResponseWriterWrapper(w gin.ResponseWriter, logger logging.RequestLogger, requestInfo *RequestInfo) *ResponseWriterWrapper {
|
||||
return &ResponseWriterWrapper{
|
||||
ResponseWriter: w,
|
||||
body: &bytes.Buffer{},
|
||||
logger: logger,
|
||||
requestInfo: requestInfo,
|
||||
headers: make(map[string][]string),
|
||||
}
|
||||
}
|
||||
|
||||
// Write wraps the underlying ResponseWriter's Write method to capture response data.
|
||||
// For non-streaming responses, it writes to an internal buffer. For streaming responses,
|
||||
// it sends data chunks to a non-blocking channel for asynchronous logging.
|
||||
// CRITICAL: This method prioritizes writing to the client to ensure zero latency,
|
||||
// handling logging operations subsequently.
|
||||
func (w *ResponseWriterWrapper) Write(data []byte) (int, error) {
|
||||
// Ensure headers are captured before first write
|
||||
// This is critical because Write() may trigger WriteHeader() internally
|
||||
w.ensureHeadersCaptured()
|
||||
|
||||
// CRITICAL: Write to client first (zero latency)
|
||||
n, err := w.ResponseWriter.Write(data)
|
||||
|
||||
// THEN: Handle logging based on response type
|
||||
if w.isStreaming && w.chunkChannel != nil {
|
||||
// Capture TTFB on first chunk (synchronous, before async channel send)
|
||||
if w.firstChunkTimestamp.IsZero() {
|
||||
w.firstChunkTimestamp = time.Now()
|
||||
}
|
||||
// For streaming responses: Send to async logging channel (non-blocking)
|
||||
select {
|
||||
case w.chunkChannel <- append([]byte(nil), data...): // Non-blocking send with copy
|
||||
default: // Channel full, skip logging to avoid blocking
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
if w.shouldBufferResponseBody() {
|
||||
w.body.Write(data)
|
||||
}
|
||||
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) shouldBufferResponseBody() bool {
|
||||
if w.logger != nil && w.logger.IsEnabled() {
|
||||
return true
|
||||
}
|
||||
if !w.logOnErrorOnly {
|
||||
return false
|
||||
}
|
||||
status := w.statusCode
|
||||
if status == 0 {
|
||||
if statusWriter, ok := w.ResponseWriter.(interface{ Status() int }); ok && statusWriter != nil {
|
||||
status = statusWriter.Status()
|
||||
} else {
|
||||
status = http.StatusOK
|
||||
}
|
||||
}
|
||||
return status >= http.StatusBadRequest && status != clienterror.StatusClientClosedRequest
|
||||
}
|
||||
|
||||
// WriteString wraps the underlying ResponseWriter's WriteString method to capture response data.
|
||||
// Some handlers (and fmt/io helpers) write via io.StringWriter; without this override, those writes
|
||||
// bypass Write() and would be missing from request logs.
|
||||
func (w *ResponseWriterWrapper) WriteString(data string) (int, error) {
|
||||
w.ensureHeadersCaptured()
|
||||
|
||||
// CRITICAL: Write to client first (zero latency)
|
||||
n, err := w.ResponseWriter.WriteString(data)
|
||||
|
||||
// THEN: Capture for logging
|
||||
if w.isStreaming && w.chunkChannel != nil {
|
||||
// Capture TTFB on first chunk (synchronous, before async channel send)
|
||||
if w.firstChunkTimestamp.IsZero() {
|
||||
w.firstChunkTimestamp = time.Now()
|
||||
}
|
||||
select {
|
||||
case w.chunkChannel <- []byte(data):
|
||||
default:
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
if w.shouldBufferResponseBody() {
|
||||
w.body.WriteString(data)
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
// WriteHeader wraps the underlying ResponseWriter's WriteHeader method.
|
||||
// It captures the status code, detects if the response is streaming based on the Content-Type header,
|
||||
// and initializes the appropriate logging mechanism (standard or streaming).
|
||||
func (w *ResponseWriterWrapper) WriteHeader(statusCode int) {
|
||||
w.statusCode = statusCode
|
||||
|
||||
// Capture response headers using the new method
|
||||
w.captureCurrentHeaders()
|
||||
|
||||
// Detect streaming based on Content-Type
|
||||
contentType := w.ResponseWriter.Header().Get("Content-Type")
|
||||
w.isStreaming = w.detectStreaming(contentType)
|
||||
|
||||
// If streaming, initialize streaming log writer
|
||||
if w.isStreaming && w.logger.IsEnabled() {
|
||||
streamWriter, err := w.logger.LogStreamingRequest(
|
||||
w.requestInfo.URL,
|
||||
w.requestInfo.Method,
|
||||
w.requestInfo.Headers,
|
||||
w.requestInfo.Body,
|
||||
w.requestInfo.RequestID,
|
||||
)
|
||||
if err == nil {
|
||||
w.streamWriter = streamWriter
|
||||
w.chunkChannel = make(chan []byte, 100) // Buffered channel for async writes
|
||||
doneChan := make(chan struct{})
|
||||
w.streamDone = doneChan
|
||||
|
||||
// Start async chunk processor
|
||||
go w.processStreamingChunks(doneChan)
|
||||
|
||||
// Write status immediately
|
||||
_ = streamWriter.WriteStatus(statusCode, w.headers)
|
||||
}
|
||||
}
|
||||
|
||||
// Call original WriteHeader
|
||||
w.ResponseWriter.WriteHeader(statusCode)
|
||||
}
|
||||
|
||||
// ensureHeadersCaptured is a helper function to make sure response headers are captured.
|
||||
// It is safe to call this method multiple times; it will always refresh the headers
|
||||
// with the latest state from the underlying ResponseWriter.
|
||||
func (w *ResponseWriterWrapper) ensureHeadersCaptured() {
|
||||
// Always capture the current headers to ensure we have the latest state
|
||||
w.captureCurrentHeaders()
|
||||
}
|
||||
|
||||
// captureCurrentHeaders reads all headers from the underlying ResponseWriter and stores them
|
||||
// in the wrapper's headers map. It creates copies of the header values to prevent race conditions.
|
||||
func (w *ResponseWriterWrapper) captureCurrentHeaders() {
|
||||
// Initialize headers map if needed
|
||||
if w.headers == nil {
|
||||
w.headers = make(map[string][]string)
|
||||
}
|
||||
|
||||
// Capture all current headers from the underlying ResponseWriter
|
||||
for key, values := range w.ResponseWriter.Header() {
|
||||
// Make a copy of the values slice to avoid reference issues
|
||||
headerValues := make([]string, len(values))
|
||||
copy(headerValues, values)
|
||||
w.headers[key] = headerValues
|
||||
}
|
||||
}
|
||||
|
||||
// detectStreaming determines if a response should be treated as a streaming response.
|
||||
// It checks for a "text/event-stream" Content-Type or a '"stream": true'
|
||||
// field in the original request body.
|
||||
func (w *ResponseWriterWrapper) detectStreaming(contentType string) bool {
|
||||
// Check Content-Type for Server-Sent Events
|
||||
if strings.Contains(contentType, "text/event-stream") {
|
||||
return true
|
||||
}
|
||||
|
||||
// If a concrete Content-Type is already set (e.g., application/json for error responses),
|
||||
// treat it as non-streaming instead of inferring from the request payload.
|
||||
if strings.TrimSpace(contentType) != "" {
|
||||
return false
|
||||
}
|
||||
|
||||
// Only fall back to request payload hints when Content-Type is not set yet.
|
||||
if w.requestInfo != nil && len(w.requestInfo.Body) > 0 {
|
||||
return bytes.Contains(w.requestInfo.Body, []byte(`"stream": true`)) ||
|
||||
bytes.Contains(w.requestInfo.Body, []byte(`"stream":true`))
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// processStreamingChunks runs in a separate goroutine to process response chunks from the chunkChannel.
|
||||
// It asynchronously writes each chunk to the streaming log writer.
|
||||
func (w *ResponseWriterWrapper) processStreamingChunks(done chan struct{}) {
|
||||
if done == nil {
|
||||
return
|
||||
}
|
||||
|
||||
defer close(done)
|
||||
|
||||
if w.streamWriter == nil || w.chunkChannel == nil {
|
||||
return
|
||||
}
|
||||
|
||||
for chunk := range w.chunkChannel {
|
||||
w.streamWriter.WriteChunkAsync(chunk)
|
||||
}
|
||||
}
|
||||
|
||||
// Finalize completes the logging process for the request and response.
|
||||
// For streaming responses, it closes the chunk channel and the stream writer.
|
||||
// For non-streaming responses, it logs the complete request and response details,
|
||||
// including any API-specific request/response data stored in the Gin context.
|
||||
func (w *ResponseWriterWrapper) Finalize(c *gin.Context) error {
|
||||
if w.requestInfo != nil && w.requestInfo.deferredBodyCapture != nil {
|
||||
defer w.requestInfo.deferredBodyCapture.Cleanup()
|
||||
}
|
||||
if w.logger == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
finalStatusCode := w.statusCode
|
||||
if finalStatusCode == 0 {
|
||||
if statusWriter, ok := w.ResponseWriter.(interface{ Status() int }); ok {
|
||||
finalStatusCode = statusWriter.Status()
|
||||
} else {
|
||||
finalStatusCode = 200
|
||||
}
|
||||
}
|
||||
|
||||
var slicesAPIResponseError []*interfaces.ErrorMessage
|
||||
apiResponseError, isExist := c.Get("API_RESPONSE_ERROR")
|
||||
if isExist {
|
||||
if apiErrors, ok := apiResponseError.([]*interfaces.ErrorMessage); ok {
|
||||
slicesAPIResponseError = apiErrors
|
||||
}
|
||||
}
|
||||
|
||||
hasAPIError := hasActionableError(c, finalStatusCode, slicesAPIResponseError)
|
||||
forceLog := w.logOnErrorOnly && hasAPIError && !w.logger.IsEnabled()
|
||||
websocketTimelineSource := w.extractWebsocketTimelineSource(c)
|
||||
apiRequestSource := w.extractAPIRequestSource(c)
|
||||
apiResponseSource := w.extractAPIResponseSource(c)
|
||||
apiWebsocketTimelineSource := w.extractAPIWebsocketTimelineSource(c)
|
||||
if !w.logger.IsEnabled() && !forceLog {
|
||||
cleanupFileBodySources(websocketTimelineSource, apiRequestSource, apiResponseSource, apiWebsocketTimelineSource)
|
||||
return nil
|
||||
}
|
||||
|
||||
if w.isStreaming && w.streamWriter != nil {
|
||||
if w.chunkChannel != nil {
|
||||
close(w.chunkChannel)
|
||||
w.chunkChannel = nil
|
||||
}
|
||||
|
||||
if w.streamDone != nil {
|
||||
<-w.streamDone
|
||||
w.streamDone = nil
|
||||
}
|
||||
|
||||
w.streamWriter.SetFirstChunkTimestamp(w.firstChunkTimestamp)
|
||||
|
||||
// Write API Request and Response to the streaming log before closing
|
||||
apiRequest := w.extractAPIRequest(c)
|
||||
apiResponse := w.extractAPIResponse(c)
|
||||
if sourceWriter, ok := w.streamWriter.(interface {
|
||||
WriteAPIRequestSource(*logging.FileBodySource) error
|
||||
WriteAPIResponseSource(*logging.FileBodySource) error
|
||||
}); ok {
|
||||
if len(apiRequest) > 0 {
|
||||
_ = w.streamWriter.WriteAPIRequest(apiRequest)
|
||||
}
|
||||
if apiRequestSource != nil && apiRequestSource.HasPayload() {
|
||||
_ = sourceWriter.WriteAPIRequestSource(apiRequestSource)
|
||||
}
|
||||
if len(apiResponse) > 0 {
|
||||
_ = w.streamWriter.WriteAPIResponse(apiResponse)
|
||||
}
|
||||
if apiResponseSource != nil && apiResponseSource.HasPayload() {
|
||||
_ = sourceWriter.WriteAPIResponseSource(apiResponseSource)
|
||||
}
|
||||
} else {
|
||||
var errMerge error
|
||||
apiRequest, errMerge = mergeFileBodySource(apiRequest, apiRequestSource)
|
||||
if errMerge != nil {
|
||||
cleanupFileBodySources(websocketTimelineSource, apiResponseSource, apiWebsocketTimelineSource)
|
||||
return errMerge
|
||||
}
|
||||
apiResponse, errMerge = mergeFileBodySource(apiResponse, apiResponseSource)
|
||||
if errMerge != nil {
|
||||
cleanupFileBodySources(websocketTimelineSource, apiWebsocketTimelineSource)
|
||||
return errMerge
|
||||
}
|
||||
if len(apiRequest) > 0 {
|
||||
_ = w.streamWriter.WriteAPIRequest(apiRequest)
|
||||
}
|
||||
if len(apiResponse) > 0 {
|
||||
_ = w.streamWriter.WriteAPIResponse(apiResponse)
|
||||
}
|
||||
}
|
||||
apiWebsocketTimeline := w.extractAPIWebsocketTimeline(c)
|
||||
var errMerge error
|
||||
apiWebsocketTimeline, errMerge = mergeFileBodySource(apiWebsocketTimeline, apiWebsocketTimelineSource)
|
||||
if errMerge != nil {
|
||||
cleanupFileBodySources(websocketTimelineSource, apiRequestSource, apiResponseSource)
|
||||
return errMerge
|
||||
}
|
||||
if len(apiWebsocketTimeline) > 0 {
|
||||
_ = w.streamWriter.WriteAPIWebsocketTimeline(apiWebsocketTimeline)
|
||||
}
|
||||
if err := w.streamWriter.Close(); err != nil {
|
||||
w.streamWriter = nil
|
||||
cleanupFileBodySources(websocketTimelineSource, apiRequestSource, apiResponseSource)
|
||||
return err
|
||||
}
|
||||
w.streamWriter = nil
|
||||
cleanupFileBodySources(websocketTimelineSource, apiRequestSource, apiResponseSource)
|
||||
return nil
|
||||
}
|
||||
|
||||
apiRequest := w.extractAPIRequest(c)
|
||||
if forceLog && len(apiRequest) == 0 {
|
||||
apiRequest = w.extractDeferredAPIRequest(c)
|
||||
}
|
||||
return w.logRequest(w.extractRequestBody(c), finalStatusCode, w.cloneHeaders(), w.extractResponseBody(c), w.extractWebsocketTimeline(c), websocketTimelineSource, apiRequest, apiRequestSource, w.extractAPIResponse(c), apiResponseSource, w.extractAPIWebsocketTimeline(c), apiWebsocketTimelineSource, w.extractAPIResponseTimestamp(c), slicesAPIResponseError, forceLog)
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) cloneHeaders() map[string][]string {
|
||||
w.ensureHeadersCaptured()
|
||||
|
||||
finalHeaders := make(map[string][]string, len(w.headers))
|
||||
for key, values := range w.headers {
|
||||
headerValues := make([]string, len(values))
|
||||
copy(headerValues, values)
|
||||
finalHeaders[key] = headerValues
|
||||
}
|
||||
|
||||
return finalHeaders
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractAPIRequest(c *gin.Context) []byte {
|
||||
apiRequest, isExist := c.Get("API_REQUEST")
|
||||
if !isExist {
|
||||
return nil
|
||||
}
|
||||
data, ok := apiRequest.([]byte)
|
||||
if !ok || len(data) == 0 {
|
||||
return nil
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractDeferredAPIRequest(c *gin.Context) []byte {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
value, exists := c.Get(logging.DeferredAPIRequestContextKey)
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
requests, ok := value.([]logging.DeferredAPIRequest)
|
||||
if !ok || len(requests) == 0 {
|
||||
return nil
|
||||
}
|
||||
var body bytes.Buffer
|
||||
for _, buildRequest := range requests {
|
||||
if buildRequest == nil {
|
||||
continue
|
||||
}
|
||||
body.Write(buildRequest())
|
||||
}
|
||||
return body.Bytes()
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractAPIResponse(c *gin.Context) []byte {
|
||||
apiResponse, isExist := c.Get("API_RESPONSE")
|
||||
if !isExist {
|
||||
return nil
|
||||
}
|
||||
data, ok := apiResponse.([]byte)
|
||||
if !ok || len(data) == 0 {
|
||||
return nil
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractAPIRequestSource(c *gin.Context) *logging.FileBodySource {
|
||||
return extractFileBodySource(c, logging.APIRequestSourceContextKey)
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractAPIResponseSource(c *gin.Context) *logging.FileBodySource {
|
||||
return extractFileBodySource(c, logging.APIResponseSourceContextKey)
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractAPIWebsocketTimeline(c *gin.Context) []byte {
|
||||
apiTimeline, isExist := c.Get("API_WEBSOCKET_TIMELINE")
|
||||
if !isExist {
|
||||
return nil
|
||||
}
|
||||
data, ok := apiTimeline.([]byte)
|
||||
if !ok || len(data) == 0 {
|
||||
return nil
|
||||
}
|
||||
return bytes.Clone(data)
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractAPIWebsocketTimelineSource(c *gin.Context) *logging.FileBodySource {
|
||||
return extractFileBodySource(c, logging.APIWebsocketTimelineSourceContextKey)
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractAPIResponseTimestamp(c *gin.Context) time.Time {
|
||||
ts, isExist := c.Get("API_RESPONSE_TIMESTAMP")
|
||||
if !isExist {
|
||||
return time.Time{}
|
||||
}
|
||||
if t, ok := ts.(time.Time); ok {
|
||||
return t
|
||||
}
|
||||
return time.Time{}
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractRequestBody(c *gin.Context) []byte {
|
||||
if body := extractBodyOverride(c, requestBodyOverrideContextKey); len(body) > 0 {
|
||||
return body
|
||||
}
|
||||
if w.requestInfo == nil {
|
||||
return nil
|
||||
}
|
||||
if len(w.requestInfo.Body) > 0 {
|
||||
return w.requestInfo.Body
|
||||
}
|
||||
if w.requestInfo.deferredBodyCapture == nil {
|
||||
return nil
|
||||
}
|
||||
body, statusMarker, errRead := w.requestInfo.deferredBodyCapture.Bytes()
|
||||
if errRead != nil {
|
||||
log.WithError(errRead).Warn("failed to read deferred request body capture")
|
||||
return nil
|
||||
}
|
||||
encoding := ""
|
||||
for key, values := range w.requestInfo.Headers {
|
||||
if strings.EqualFold(key, "Content-Encoding") && len(values) > 0 {
|
||||
encoding = values[0]
|
||||
break
|
||||
}
|
||||
}
|
||||
body = decodeCapturedRequestBodyForLogWithLimit(body, encoding, maxDeferredErrorRequestBodyBytes)
|
||||
if statusMarker == "" {
|
||||
return body
|
||||
}
|
||||
if len(body) > 0 && !bytes.HasSuffix(body, []byte("\n")) {
|
||||
body = append(body, '\n')
|
||||
}
|
||||
return append(body, statusMarker...)
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractResponseBody(c *gin.Context) []byte {
|
||||
if body := extractBodyOverride(c, responseBodyOverrideContextKey); len(body) > 0 {
|
||||
return body
|
||||
}
|
||||
if w.body == nil || w.body.Len() == 0 {
|
||||
return nil
|
||||
}
|
||||
return bytes.Clone(w.body.Bytes())
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractWebsocketTimeline(c *gin.Context) []byte {
|
||||
return extractBodyOverride(c, websocketTimelineOverrideContextKey)
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractWebsocketTimelineSource(c *gin.Context) *logging.FileBodySource {
|
||||
return extractFileBodySource(c, logging.WebsocketTimelineSourceContextKey)
|
||||
}
|
||||
|
||||
func extractFileBodySource(c *gin.Context, key string) *logging.FileBodySource {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
value, exists := c.Get(key)
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
source, ok := value.(*logging.FileBodySource)
|
||||
if !ok || source == nil {
|
||||
return nil
|
||||
}
|
||||
return source
|
||||
}
|
||||
|
||||
func extractBodyOverride(c *gin.Context, key string) []byte {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
bodyOverride, isExist := c.Get(key)
|
||||
if !isExist {
|
||||
return nil
|
||||
}
|
||||
switch value := bodyOverride.(type) {
|
||||
case []byte:
|
||||
if len(value) > 0 {
|
||||
return bytes.Clone(value)
|
||||
}
|
||||
case string:
|
||||
if strings.TrimSpace(value) != "" {
|
||||
return []byte(value)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) logRequest(requestBody []byte, statusCode int, headers map[string][]string, body, websocketTimeline []byte, websocketTimelineSource *logging.FileBodySource, apiRequestBody []byte, apiRequestSource *logging.FileBodySource, apiResponseBody []byte, apiResponseSource *logging.FileBodySource, apiWebsocketTimeline []byte, apiWebsocketTimelineSource *logging.FileBodySource, apiResponseTimestamp time.Time, apiResponseErrors []*interfaces.ErrorMessage, forceLog bool) error {
|
||||
if w.requestInfo == nil {
|
||||
cleanupFileBodySources(websocketTimelineSource, apiRequestSource, apiResponseSource, apiWebsocketTimelineSource)
|
||||
return nil
|
||||
}
|
||||
|
||||
if loggerWithAllSources, ok := w.logger.(interface {
|
||||
LogRequestWithOptionsAndAllSources(string, string, map[string][]string, []byte, int, map[string][]string, []byte, []byte, *logging.FileBodySource, []byte, *logging.FileBodySource, []byte, *logging.FileBodySource, []byte, *logging.FileBodySource, []*interfaces.ErrorMessage, bool, string, time.Time, time.Time) error
|
||||
}); ok {
|
||||
return loggerWithAllSources.LogRequestWithOptionsAndAllSources(
|
||||
w.requestInfo.URL,
|
||||
w.requestInfo.Method,
|
||||
w.requestInfo.Headers,
|
||||
requestBody,
|
||||
statusCode,
|
||||
headers,
|
||||
body,
|
||||
websocketTimeline,
|
||||
websocketTimelineSource,
|
||||
apiRequestBody,
|
||||
apiRequestSource,
|
||||
apiResponseBody,
|
||||
apiResponseSource,
|
||||
apiWebsocketTimeline,
|
||||
apiWebsocketTimelineSource,
|
||||
apiResponseErrors,
|
||||
forceLog,
|
||||
w.requestInfo.RequestID,
|
||||
w.requestInfo.Timestamp,
|
||||
apiResponseTimestamp,
|
||||
)
|
||||
}
|
||||
|
||||
if loggerWithSources, ok := w.logger.(interface {
|
||||
LogRequestWithOptionsAndSources(string, string, map[string][]string, []byte, int, map[string][]string, []byte, []byte, *logging.FileBodySource, []byte, []byte, []byte, *logging.FileBodySource, []*interfaces.ErrorMessage, bool, string, time.Time, time.Time) error
|
||||
}); ok {
|
||||
var errMerge error
|
||||
apiRequestBody, errMerge = mergeFileBodySource(apiRequestBody, apiRequestSource)
|
||||
if errMerge != nil {
|
||||
cleanupFileBodySources(websocketTimelineSource, apiResponseSource, apiWebsocketTimelineSource)
|
||||
return errMerge
|
||||
}
|
||||
apiResponseBody, errMerge = mergeFileBodySource(apiResponseBody, apiResponseSource)
|
||||
if errMerge != nil {
|
||||
cleanupFileBodySources(websocketTimelineSource, apiWebsocketTimelineSource)
|
||||
return errMerge
|
||||
}
|
||||
return loggerWithSources.LogRequestWithOptionsAndSources(
|
||||
w.requestInfo.URL,
|
||||
w.requestInfo.Method,
|
||||
w.requestInfo.Headers,
|
||||
requestBody,
|
||||
statusCode,
|
||||
headers,
|
||||
body,
|
||||
websocketTimeline,
|
||||
websocketTimelineSource,
|
||||
apiRequestBody,
|
||||
apiResponseBody,
|
||||
apiWebsocketTimeline,
|
||||
apiWebsocketTimelineSource,
|
||||
apiResponseErrors,
|
||||
forceLog,
|
||||
w.requestInfo.RequestID,
|
||||
w.requestInfo.Timestamp,
|
||||
apiResponseTimestamp,
|
||||
)
|
||||
}
|
||||
|
||||
var errMerge error
|
||||
websocketTimeline, errMerge = mergeFileBodySource(websocketTimeline, websocketTimelineSource)
|
||||
if errMerge != nil {
|
||||
cleanupFileBodySources(apiRequestSource, apiResponseSource, apiWebsocketTimelineSource)
|
||||
return errMerge
|
||||
}
|
||||
apiRequestBody, errMerge = mergeFileBodySource(apiRequestBody, apiRequestSource)
|
||||
if errMerge != nil {
|
||||
cleanupFileBodySources(apiResponseSource, apiWebsocketTimelineSource)
|
||||
return errMerge
|
||||
}
|
||||
apiResponseBody, errMerge = mergeFileBodySource(apiResponseBody, apiResponseSource)
|
||||
if errMerge != nil {
|
||||
cleanupFileBodySources(apiWebsocketTimelineSource)
|
||||
return errMerge
|
||||
}
|
||||
apiWebsocketTimeline, errMerge = mergeFileBodySource(apiWebsocketTimeline, apiWebsocketTimelineSource)
|
||||
if errMerge != nil {
|
||||
return errMerge
|
||||
}
|
||||
|
||||
if loggerWithOptions, ok := w.logger.(interface {
|
||||
LogRequestWithOptions(string, string, map[string][]string, []byte, int, map[string][]string, []byte, []byte, []byte, []byte, []byte, []*interfaces.ErrorMessage, bool, string, time.Time, time.Time) error
|
||||
}); ok {
|
||||
return loggerWithOptions.LogRequestWithOptions(
|
||||
w.requestInfo.URL,
|
||||
w.requestInfo.Method,
|
||||
w.requestInfo.Headers,
|
||||
requestBody,
|
||||
statusCode,
|
||||
headers,
|
||||
body,
|
||||
websocketTimeline,
|
||||
apiRequestBody,
|
||||
apiResponseBody,
|
||||
apiWebsocketTimeline,
|
||||
apiResponseErrors,
|
||||
forceLog,
|
||||
w.requestInfo.RequestID,
|
||||
w.requestInfo.Timestamp,
|
||||
apiResponseTimestamp,
|
||||
)
|
||||
}
|
||||
|
||||
return w.logger.LogRequest(
|
||||
w.requestInfo.URL,
|
||||
w.requestInfo.Method,
|
||||
w.requestInfo.Headers,
|
||||
requestBody,
|
||||
statusCode,
|
||||
headers,
|
||||
body,
|
||||
websocketTimeline,
|
||||
apiRequestBody,
|
||||
apiResponseBody,
|
||||
apiWebsocketTimeline,
|
||||
apiResponseErrors,
|
||||
w.requestInfo.RequestID,
|
||||
w.requestInfo.Timestamp,
|
||||
apiResponseTimestamp,
|
||||
)
|
||||
}
|
||||
|
||||
func mergeFileBodySource(payload []byte, source *logging.FileBodySource) ([]byte, error) {
|
||||
if source == nil {
|
||||
return payload, nil
|
||||
}
|
||||
defer cleanupFileBodySources(source)
|
||||
if !source.HasPayload() {
|
||||
return payload, nil
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
if len(payload) > 0 {
|
||||
buf.Write(payload)
|
||||
if !bytes.HasSuffix(payload, []byte("\n")) {
|
||||
buf.WriteByte('\n')
|
||||
}
|
||||
buf.WriteByte('\n')
|
||||
}
|
||||
if errWrite := source.WriteTo(&buf); errWrite != nil {
|
||||
return nil, errWrite
|
||||
}
|
||||
return buf.Bytes(), nil
|
||||
}
|
||||
|
||||
func cleanupFileBodySources(sources ...*logging.FileBodySource) {
|
||||
for _, source := range sources {
|
||||
if source == nil {
|
||||
continue
|
||||
}
|
||||
if errCleanup := source.Cleanup(); errCleanup != nil {
|
||||
log.WithError(errCleanup).Warn("failed to clean up log part files")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func isClientCancellationErrorMessage(errMsg *interfaces.ErrorMessage) bool {
|
||||
if errMsg == nil {
|
||||
return true
|
||||
}
|
||||
return clienterror.IsClientCancellation(errMsg.StatusCode, errMsg.Error)
|
||||
}
|
||||
|
||||
func hasActionableAPIResponseErrors(apiErrors []*interfaces.ErrorMessage) bool {
|
||||
for _, err := range apiErrors {
|
||||
if !isClientCancellationErrorMessage(err) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isContextCanceled(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil {
|
||||
return false
|
||||
}
|
||||
ctx := c.Request.Context()
|
||||
return ctx != nil && errors.Is(ctx.Err(), context.Canceled)
|
||||
}
|
||||
|
||||
func hasActionableError(c *gin.Context, statusCode int, apiErrors []*interfaces.ErrorMessage) bool {
|
||||
if hasActionableAPIResponseErrors(apiErrors) {
|
||||
return true
|
||||
}
|
||||
if statusCode == clienterror.StatusClientClosedRequest {
|
||||
return false
|
||||
}
|
||||
if isContextCanceled(c) && statusCode < http.StatusBadRequest {
|
||||
return false
|
||||
}
|
||||
return statusCode >= http.StatusBadRequest
|
||||
}
|
||||
380
backend/internal/api/middleware/response_writer_test.go
Normal file
380
backend/internal/api/middleware/response_writer_test.go
Normal file
|
|
@ -0,0 +1,380 @@
|
|||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
)
|
||||
|
||||
func TestExtractRequestBodyPrefersOverride(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
wrapper := &ResponseWriterWrapper{
|
||||
requestInfo: &RequestInfo{Body: []byte("original-body")},
|
||||
}
|
||||
|
||||
body := wrapper.extractRequestBody(c)
|
||||
if string(body) != "original-body" {
|
||||
t.Fatalf("request body = %q, want %q", string(body), "original-body")
|
||||
}
|
||||
|
||||
c.Set(requestBodyOverrideContextKey, []byte("override-body"))
|
||||
body = wrapper.extractRequestBody(c)
|
||||
if string(body) != "override-body" {
|
||||
t.Fatalf("request body = %q, want %q", string(body), "override-body")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractRequestBodySupportsStringOverride(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
wrapper := &ResponseWriterWrapper{body: &bytes.Buffer{}}
|
||||
c.Set(requestBodyOverrideContextKey, "override-as-string")
|
||||
|
||||
body := wrapper.extractRequestBody(c)
|
||||
if string(body) != "override-as-string" {
|
||||
t.Fatalf("request body = %q, want %q", string(body), "override-as-string")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractResponseBodyPrefersOverride(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
wrapper := &ResponseWriterWrapper{body: &bytes.Buffer{}}
|
||||
wrapper.body.WriteString("original-response")
|
||||
|
||||
body := wrapper.extractResponseBody(c)
|
||||
if string(body) != "original-response" {
|
||||
t.Fatalf("response body = %q, want %q", string(body), "original-response")
|
||||
}
|
||||
|
||||
c.Set(responseBodyOverrideContextKey, []byte("override-response"))
|
||||
body = wrapper.extractResponseBody(c)
|
||||
if string(body) != "override-response" {
|
||||
t.Fatalf("response body = %q, want %q", string(body), "override-response")
|
||||
}
|
||||
|
||||
body[0] = 'X'
|
||||
if got := wrapper.extractResponseBody(c); string(got) != "override-response" {
|
||||
t.Fatalf("response override should be cloned, got %q", string(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractResponseBodySupportsStringOverride(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
wrapper := &ResponseWriterWrapper{}
|
||||
c.Set(responseBodyOverrideContextKey, "override-response-as-string")
|
||||
|
||||
body := wrapper.extractResponseBody(c)
|
||||
if string(body) != "override-response-as-string" {
|
||||
t.Fatalf("response body = %q, want %q", string(body), "override-response-as-string")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractBodyOverrideClonesBytes(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
override := []byte("body-override")
|
||||
c.Set(requestBodyOverrideContextKey, override)
|
||||
|
||||
body := extractBodyOverride(c, requestBodyOverrideContextKey)
|
||||
if !bytes.Equal(body, override) {
|
||||
t.Fatalf("body override = %q, want %q", string(body), string(override))
|
||||
}
|
||||
|
||||
body[0] = 'X'
|
||||
if !bytes.Equal(override, []byte("body-override")) {
|
||||
t.Fatalf("override mutated: %q", string(override))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractWebsocketTimelineUsesOverride(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
wrapper := &ResponseWriterWrapper{}
|
||||
if got := wrapper.extractWebsocketTimeline(c); got != nil {
|
||||
t.Fatalf("expected nil websocket timeline, got %q", string(got))
|
||||
}
|
||||
|
||||
c.Set(websocketTimelineOverrideContextKey, []byte("timeline"))
|
||||
body := wrapper.extractWebsocketTimeline(c)
|
||||
if string(body) != "timeline" {
|
||||
t.Fatalf("websocket timeline = %q, want %q", string(body), "timeline")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFinalizeStreamingWritesAPIWebsocketTimeline(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
streamWriter := &testStreamingLogWriter{}
|
||||
wrapper := &ResponseWriterWrapper{
|
||||
ResponseWriter: c.Writer,
|
||||
logger: &testRequestLogger{enabled: true},
|
||||
requestInfo: &RequestInfo{
|
||||
URL: "/v1/responses",
|
||||
Method: "POST",
|
||||
Headers: map[string][]string{"Content-Type": {"application/json"}},
|
||||
RequestID: "req-1",
|
||||
Timestamp: time.Date(2026, time.April, 1, 12, 0, 0, 0, time.UTC),
|
||||
},
|
||||
isStreaming: true,
|
||||
streamWriter: streamWriter,
|
||||
}
|
||||
|
||||
c.Set("API_WEBSOCKET_TIMELINE", []byte("Timestamp: 2026-04-01T12:00:00Z\nEvent: api.websocket.request\n{}"))
|
||||
|
||||
if err := wrapper.Finalize(c); err != nil {
|
||||
t.Fatalf("Finalize error: %v", err)
|
||||
}
|
||||
if string(streamWriter.apiWebsocketTimeline) != "Timestamp: 2026-04-01T12:00:00Z\nEvent: api.websocket.request\n{}" {
|
||||
t.Fatalf("stream writer websocket timeline = %q", string(streamWriter.apiWebsocketTimeline))
|
||||
}
|
||||
if !streamWriter.closed {
|
||||
t.Fatal("expected stream writer to be closed")
|
||||
}
|
||||
}
|
||||
|
||||
type testRequestLogger struct {
|
||||
enabled bool
|
||||
}
|
||||
|
||||
func (l *testRequestLogger) LogRequest(string, string, map[string][]string, []byte, int, map[string][]string, []byte, []byte, []byte, []byte, []byte, []*interfaces.ErrorMessage, string, time.Time, time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *testRequestLogger) LogStreamingRequest(string, string, map[string][]string, []byte, string) (logging.StreamingLogWriter, error) {
|
||||
return &testStreamingLogWriter{}, nil
|
||||
}
|
||||
|
||||
func (l *testRequestLogger) IsEnabled() bool {
|
||||
return l.enabled
|
||||
}
|
||||
|
||||
type testStreamingLogWriter struct {
|
||||
apiWebsocketTimeline []byte
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (w *testStreamingLogWriter) WriteChunkAsync([]byte) {}
|
||||
|
||||
func (w *testStreamingLogWriter) WriteStatus(int, map[string][]string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *testStreamingLogWriter) WriteAPIRequest([]byte) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *testStreamingLogWriter) WriteAPIResponse([]byte) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *testStreamingLogWriter) WriteAPIWebsocketTimeline(apiWebsocketTimeline []byte) error {
|
||||
w.apiWebsocketTimeline = bytes.Clone(apiWebsocketTimeline)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *testStreamingLogWriter) SetFirstChunkTimestamp(time.Time) {}
|
||||
|
||||
func (w *testStreamingLogWriter) Close() error {
|
||||
w.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestHasActionableError(t *testing.T) {
|
||||
canceledCtx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
statusCode int
|
||||
ctx context.Context
|
||||
apiErrors []*interfaces.ErrorMessage
|
||||
want bool
|
||||
}{
|
||||
{
|
||||
name: "200 ok without errors",
|
||||
statusCode: http.StatusOK,
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "499 client closed request",
|
||||
statusCode: clienterror.StatusClientClosedRequest,
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "499 with context canceled api error",
|
||||
statusCode: clienterror.StatusClientClosedRequest,
|
||||
apiErrors: []*interfaces.ErrorMessage{{StatusCode: clienterror.StatusClientClosedRequest, Error: context.Canceled}},
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "200 with canceled context",
|
||||
statusCode: http.StatusOK,
|
||||
ctx: canceledCtx,
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "0 with canceled context",
|
||||
statusCode: 0,
|
||||
ctx: canceledCtx,
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "400 bad request",
|
||||
statusCode: http.StatusBadRequest,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "429 rate limit",
|
||||
statusCode: http.StatusTooManyRequests,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "500 internal server error",
|
||||
statusCode: http.StatusInternalServerError,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "503 with canceled context",
|
||||
statusCode: http.StatusServiceUnavailable,
|
||||
ctx: canceledCtx,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "200 with actionable upstream api error",
|
||||
statusCode: http.StatusOK,
|
||||
apiErrors: []*interfaces.ErrorMessage{{StatusCode: http.StatusBadGateway, Error: errors.New("upstream failed")}},
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "200 with non-actionable cancellation api error",
|
||||
statusCode: http.StatusOK,
|
||||
apiErrors: []*interfaces.ErrorMessage{{StatusCode: 0, Error: fmt.Errorf("read: %w", context.Canceled)}},
|
||||
want: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
if tc.ctx != nil {
|
||||
req = req.WithContext(tc.ctx)
|
||||
}
|
||||
c.Request = req
|
||||
|
||||
got := hasActionableError(c, tc.statusCode, tc.apiErrors)
|
||||
if got != tc.want {
|
||||
t.Fatalf("hasActionableError(status=%d, errors=%v) = %t, want %t", tc.statusCode, tc.apiErrors, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type recordingRequestLogger struct {
|
||||
loggedCalls []int
|
||||
enabled bool
|
||||
}
|
||||
|
||||
func (l *recordingRequestLogger) LogRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiResponseErrors []*interfaces.ErrorMessage, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error {
|
||||
l.loggedCalls = append(l.loggedCalls, statusCode)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *recordingRequestLogger) LogStreamingRequest(string, string, map[string][]string, []byte, string) (logging.StreamingLogWriter, error) {
|
||||
return &testStreamingLogWriter{}, nil
|
||||
}
|
||||
|
||||
func (l *recordingRequestLogger) IsEnabled() bool {
|
||||
return l.enabled
|
||||
}
|
||||
|
||||
func (l *recordingRequestLogger) LogRequestWithOptions(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error {
|
||||
if force || l.enabled {
|
||||
l.loggedCalls = append(l.loggedCalls, statusCode)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestFinalizeExcludes499FromForceLog(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
logger := &recordingRequestLogger{enabled: false}
|
||||
wrapper := &ResponseWriterWrapper{
|
||||
ResponseWriter: c.Writer,
|
||||
logger: logger,
|
||||
logOnErrorOnly: true,
|
||||
statusCode: clienterror.StatusClientClosedRequest,
|
||||
requestInfo: &RequestInfo{
|
||||
URL: "/v1/responses",
|
||||
Method: "POST",
|
||||
RequestID: "req-499",
|
||||
Timestamp: time.Now(),
|
||||
},
|
||||
}
|
||||
|
||||
if err := wrapper.Finalize(c); err != nil {
|
||||
t.Fatalf("Finalize error: %v", err)
|
||||
}
|
||||
if len(logger.loggedCalls) != 0 {
|
||||
t.Fatalf("expected 0 logged calls for 499 cancellation, got %d: %v", len(logger.loggedCalls), logger.loggedCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFinalizeIncludes500InForceLog(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
logger := &recordingRequestLogger{enabled: false}
|
||||
wrapper := &ResponseWriterWrapper{
|
||||
ResponseWriter: c.Writer,
|
||||
logger: logger,
|
||||
logOnErrorOnly: true,
|
||||
statusCode: http.StatusInternalServerError,
|
||||
requestInfo: &RequestInfo{
|
||||
URL: "/v1/responses",
|
||||
Method: "POST",
|
||||
RequestID: "req-500",
|
||||
Timestamp: time.Now(),
|
||||
},
|
||||
}
|
||||
|
||||
if err := wrapper.Finalize(c); err != nil {
|
||||
t.Fatalf("Finalize error: %v", err)
|
||||
}
|
||||
if len(logger.loggedCalls) != 1 || logger.loggedCalls[0] != http.StatusInternalServerError {
|
||||
t.Fatalf("expected 1 logged call for 500 status, got: %v", logger.loggedCalls)
|
||||
}
|
||||
}
|
||||
68
backend/internal/api/mux_listener.go
Normal file
68
backend/internal/api/mux_listener.go
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"net"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type muxListener struct {
|
||||
addr net.Addr
|
||||
connCh chan net.Conn
|
||||
closeCh chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func newMuxListener(addr net.Addr, buffer int) *muxListener {
|
||||
if buffer <= 0 {
|
||||
buffer = 1
|
||||
}
|
||||
return &muxListener{
|
||||
addr: addr,
|
||||
connCh: make(chan net.Conn, buffer),
|
||||
closeCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (l *muxListener) Put(conn net.Conn) error {
|
||||
if conn == nil {
|
||||
return nil
|
||||
}
|
||||
select {
|
||||
case <-l.closeCh:
|
||||
return net.ErrClosed
|
||||
case l.connCh <- conn:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (l *muxListener) Accept() (net.Conn, error) {
|
||||
select {
|
||||
case <-l.closeCh:
|
||||
return nil, net.ErrClosed
|
||||
case conn := <-l.connCh:
|
||||
if conn == nil {
|
||||
return nil, net.ErrClosed
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (l *muxListener) Close() error {
|
||||
if l == nil {
|
||||
return nil
|
||||
}
|
||||
l.once.Do(func() {
|
||||
close(l.closeCh)
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *muxListener) Addr() net.Addr {
|
||||
if l == nil {
|
||||
return &net.TCPAddr{}
|
||||
}
|
||||
if l.addr == nil {
|
||||
return &net.TCPAddr{}
|
||||
}
|
||||
return l.addr
|
||||
}
|
||||
125
backend/internal/api/protocol_multiplexer.go
Normal file
125
backend/internal/api/protocol_multiplexer.go
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
func normalizeHTTPServeError(err error) error {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
if errors.Is(err, net.ErrClosed) {
|
||||
return nil
|
||||
}
|
||||
if errors.Is(err, http.ErrServerClosed) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func normalizeListenerError(err error) error {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
if errors.Is(err, net.ErrClosed) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Server) acceptMuxConnections(listener net.Listener, httpListener *muxListener) error {
|
||||
if s == nil || listener == nil {
|
||||
return net.ErrClosed
|
||||
}
|
||||
|
||||
for {
|
||||
conn, errAccept := listener.Accept()
|
||||
if errAccept != nil {
|
||||
return errAccept
|
||||
}
|
||||
if conn == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// Dispatch each connection to a goroutine so that slow/idle clients
|
||||
// cannot block the accept loop. Previously, TLS handshake and
|
||||
// reader.Peek(1) were performed inline; an idle TCP connection that
|
||||
// never sent bytes would block Peek indefinitely, preventing all
|
||||
// subsequent connections from being accepted (issue #3267).
|
||||
go s.routeMuxConnection(conn, httpListener)
|
||||
}
|
||||
}
|
||||
|
||||
// routeMuxConnection performs per-connection protocol detection and routing.
|
||||
func (s *Server) routeMuxConnection(conn net.Conn, httpListener *muxListener) {
|
||||
// Set a read deadline so that idle connections that never send bytes do not
|
||||
// leak goroutines and file descriptors. The deadline is cleared once the
|
||||
// connection is successfully routed to its handler.
|
||||
const muxSniffDeadline = 10 * time.Second
|
||||
_ = conn.SetReadDeadline(time.Now().Add(muxSniffDeadline))
|
||||
|
||||
tlsConn, ok := conn.(*tls.Conn)
|
||||
if ok {
|
||||
if errHandshake := tlsConn.Handshake(); errHandshake != nil {
|
||||
if errClose := conn.Close(); errClose != nil {
|
||||
log.Errorf("failed to close connection after TLS handshake error: %v", errClose)
|
||||
}
|
||||
return
|
||||
}
|
||||
proto := strings.TrimSpace(tlsConn.ConnectionState().NegotiatedProtocol)
|
||||
if proto == "h2" || proto == "http/1.1" {
|
||||
if httpListener == nil {
|
||||
if errClose := conn.Close(); errClose != nil {
|
||||
log.Errorf("failed to close connection: %v", errClose)
|
||||
}
|
||||
return
|
||||
}
|
||||
if errPut := httpListener.Put(tlsConn); errPut != nil {
|
||||
if errClose := conn.Close(); errClose != nil {
|
||||
log.Errorf("failed to close connection after HTTP routing failure: %v", errClose)
|
||||
}
|
||||
} else {
|
||||
_ = conn.SetReadDeadline(time.Time{})
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
reader := bufio.NewReader(conn)
|
||||
prefix, errPeek := reader.Peek(1)
|
||||
if errPeek != nil {
|
||||
if errClose := conn.Close(); errClose != nil {
|
||||
log.Errorf("failed to close connection after protocol peek failure: %v", errClose)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if isRedisRESPPrefix(prefix[0]) {
|
||||
_ = conn.SetReadDeadline(time.Time{})
|
||||
s.handleRedisConnection(conn, reader)
|
||||
return
|
||||
}
|
||||
|
||||
if httpListener == nil {
|
||||
if errClose := conn.Close(); errClose != nil {
|
||||
log.Errorf("failed to close connection without HTTP listener: %v", errClose)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if errPut := httpListener.Put(&bufferedConn{Conn: conn, reader: reader}); errPut != nil {
|
||||
if errClose := conn.Close(); errClose != nil {
|
||||
log.Errorf("failed to close connection after HTTP routing failure: %v", errClose)
|
||||
}
|
||||
} else {
|
||||
_ = conn.SetReadDeadline(time.Time{})
|
||||
}
|
||||
}
|
||||
65
backend/internal/api/protocol_multiplexer_test.go
Normal file
65
backend/internal/api/protocol_multiplexer_test.go
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestAcceptMuxNotBlockedByIdleConnection(t *testing.T) {
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("failed to listen: %v", err)
|
||||
}
|
||||
defer listener.Close()
|
||||
|
||||
var routed atomic.Int32
|
||||
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
routed.Add(1)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
})
|
||||
srv := httptest.NewUnstartedServer(handler)
|
||||
defer srv.Close()
|
||||
|
||||
muxLn := newMuxListener(listener.Addr(), 1024)
|
||||
server := &Server{managementRoutesEnabled: atomic.Bool{}}
|
||||
server.managementRoutesEnabled.Store(false)
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- server.acceptMuxConnections(listener, muxLn)
|
||||
}()
|
||||
|
||||
srv.Listener = muxLn
|
||||
srv.Start()
|
||||
|
||||
// Open an idle TCP connection that never sends any bytes.
|
||||
idleConn, err := net.DialTimeout("tcp", listener.Addr().String(), 2*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to dial idle connection: %v", err)
|
||||
}
|
||||
defer idleConn.Close()
|
||||
|
||||
// Give the accept loop time to pick up the idle connection.
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
// Send a real HTTP request. Before the fix, the accept loop would be
|
||||
// blocked on Peek(1) for the idle connection, causing this request to
|
||||
// time out.
|
||||
client := &http.Client{Timeout: 3 * time.Second}
|
||||
resp, err := client.Get("http://" + listener.Addr().String() + "/")
|
||||
if err != nil {
|
||||
listener.Close()
|
||||
t.Fatalf("HTTP request failed (accept loop may be blocked by idle connection): %v", err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
|
||||
listener.Close()
|
||||
|
||||
if routed.Load() == 0 {
|
||||
t.Error("expected at least one request to be routed")
|
||||
}
|
||||
}
|
||||
606
backend/internal/api/redis_queue_protocol.go
Normal file
606
backend/internal/api/redis_queue_protocol.go
Normal file
|
|
@ -0,0 +1,606 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
redisUsageChannel = "usage"
|
||||
redisErrorsChannel = "errors"
|
||||
)
|
||||
|
||||
type redisSubscriptionCommand struct {
|
||||
args []string
|
||||
err error
|
||||
}
|
||||
|
||||
func isRedisRESPPrefix(prefix byte) bool {
|
||||
switch prefix {
|
||||
case '*', '$', '+', '-', ':':
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) handleRedisConnection(conn net.Conn, reader *bufio.Reader) {
|
||||
if s == nil || conn == nil {
|
||||
return
|
||||
}
|
||||
if reader == nil {
|
||||
reader = bufio.NewReader(conn)
|
||||
}
|
||||
|
||||
clientIP, localClient := resolveRemoteIP(conn.RemoteAddr())
|
||||
authed := false
|
||||
writer := bufio.NewWriter(conn)
|
||||
defer func() {
|
||||
if errClose := conn.Close(); errClose != nil {
|
||||
log.Errorf("redis connection close error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
|
||||
flush := func() bool {
|
||||
if errFlush := writer.Flush(); errFlush != nil {
|
||||
log.Errorf("redis protocol flush error: %v", errFlush)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
if s.cfg != nil && s.cfg.Home.Enabled {
|
||||
_ = writeRedisError(writer, "ERR redis usage output disabled in home mode")
|
||||
_ = writer.Flush()
|
||||
return
|
||||
}
|
||||
|
||||
for {
|
||||
if !s.managementRoutesEnabled.Load() {
|
||||
return
|
||||
}
|
||||
|
||||
args, errRead := readRESPArray(reader)
|
||||
if errRead != nil {
|
||||
if !errors.Is(errRead, io.EOF) {
|
||||
_ = writeRedisError(writer, "ERR "+errRead.Error())
|
||||
_ = writer.Flush()
|
||||
}
|
||||
return
|
||||
}
|
||||
if len(args) == 0 {
|
||||
_ = writeRedisError(writer, "ERR empty command")
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
cmd := strings.ToUpper(strings.TrimSpace(args[0]))
|
||||
|
||||
if cmd != "AUTH" && !authed {
|
||||
if s.mgmt != nil {
|
||||
_, statusCode, errMsg := s.mgmt.AuthenticateManagementKey(clientIP, localClient, "")
|
||||
if statusCode == http.StatusForbidden && strings.HasPrefix(errMsg, "IP banned due to too many failed attempts") {
|
||||
_ = writeRedisError(writer, "ERR "+errMsg)
|
||||
} else {
|
||||
_ = writeRedisError(writer, "NOAUTH Authentication required.")
|
||||
}
|
||||
} else {
|
||||
_ = writeRedisError(writer, "NOAUTH Authentication required.")
|
||||
}
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
switch cmd {
|
||||
case "AUTH":
|
||||
password, ok := parseAuthPassword(args)
|
||||
if !ok {
|
||||
if s.mgmt != nil {
|
||||
_, statusCode, errMsg := s.mgmt.AuthenticateManagementKey(clientIP, localClient, "")
|
||||
if statusCode == http.StatusForbidden && strings.HasPrefix(errMsg, "IP banned due to too many failed attempts") {
|
||||
_ = writeRedisError(writer, "ERR "+errMsg)
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
}
|
||||
_ = writeRedisError(writer, "ERR wrong number of arguments for 'auth' command")
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if s.mgmt == nil {
|
||||
_ = writeRedisError(writer, "ERR remote management disabled")
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
allowed, _, errMsg := s.mgmt.AuthenticateManagementKey(clientIP, localClient, password)
|
||||
if !allowed {
|
||||
_ = writeRedisError(writer, "ERR "+errMsg)
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
authed = true
|
||||
_ = writeRedisSimpleString(writer, "OK")
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
case "SUBSCRIBE":
|
||||
channel, ok := parseSubscribeChannel(args)
|
||||
if !ok {
|
||||
_ = writeRedisError(writer, "ERR wrong number of arguments for 'subscribe' command")
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
messages, unsubscribe, ok := subscribeRedisChannel(channel)
|
||||
if !ok {
|
||||
_ = writeRedisError(writer, fmt.Sprintf("ERR unsupported channel '%s'", channel))
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if errWrite := writeRedisPubSubSubscribe(writer, channel, 1); errWrite != nil {
|
||||
unsubscribe()
|
||||
log.Errorf("redis protocol subscribe response error: %v", errWrite)
|
||||
return
|
||||
}
|
||||
if !flush() {
|
||||
unsubscribe()
|
||||
return
|
||||
}
|
||||
s.streamRedisSubscription(reader, writer, channel, messages, unsubscribe)
|
||||
return
|
||||
case "LPOP", "RPOP":
|
||||
count, hasCount, ok := parsePopCount(args)
|
||||
if !ok {
|
||||
_ = writeRedisError(writer, "ERR wrong number of arguments for '"+strings.ToLower(cmd)+"' command")
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if count <= 0 {
|
||||
_ = writeRedisError(writer, "ERR value is not an integer or out of range")
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
items, ok := popRedisQueueItems(args[1], count)
|
||||
if !ok {
|
||||
_ = writeRedisError(writer, fmt.Sprintf("ERR unsupported channel '%s'", strings.TrimSpace(args[1])))
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if hasCount {
|
||||
_ = writeRedisArrayOfBulkStrings(writer, items)
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if len(items) == 0 {
|
||||
_ = writeRedisNilBulkString(writer)
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
_ = writeRedisBulkString(writer, items[0])
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
default:
|
||||
_ = writeRedisError(writer, fmt.Sprintf("ERR unknown command '%s'", strings.ToLower(cmd)))
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func subscribeRedisChannel(channel string) (<-chan []byte, func(), bool) {
|
||||
switch strings.ToLower(strings.TrimSpace(channel)) {
|
||||
case redisUsageChannel:
|
||||
messages, unsubscribe := redisqueue.SubscribeUsage()
|
||||
return messages, unsubscribe, true
|
||||
case redisErrorsChannel:
|
||||
messages, unsubscribe := redisqueue.SubscribeErrors()
|
||||
return messages, unsubscribe, true
|
||||
default:
|
||||
return nil, nil, false
|
||||
}
|
||||
}
|
||||
|
||||
func popRedisQueueItems(channel string, count int) ([][]byte, bool) {
|
||||
switch strings.ToLower(strings.TrimSpace(channel)) {
|
||||
case redisUsageChannel:
|
||||
return redisqueue.PopOldest(count), true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) streamRedisSubscription(reader *bufio.Reader, writer *bufio.Writer, channel string, messages <-chan []byte, unsubscribe func()) {
|
||||
if unsubscribe == nil {
|
||||
return
|
||||
}
|
||||
defer unsubscribe()
|
||||
|
||||
done := make(chan struct{})
|
||||
defer close(done)
|
||||
|
||||
commands := make(chan redisSubscriptionCommand, 1)
|
||||
go readRedisSubscriptionCommands(reader, commands, done)
|
||||
|
||||
for {
|
||||
select {
|
||||
case msg, ok := <-messages:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if errWrite := writeRedisPubSubMessage(writer, channel, msg); errWrite != nil {
|
||||
log.Errorf("redis protocol publish message error: %v", errWrite)
|
||||
return
|
||||
}
|
||||
if errFlush := writer.Flush(); errFlush != nil {
|
||||
log.Errorf("redis protocol flush error: %v", errFlush)
|
||||
return
|
||||
}
|
||||
case command, ok := <-commands:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
keepOpen := handleRedisSubscriptionCommand(writer, channel, command)
|
||||
if errFlush := writer.Flush(); errFlush != nil {
|
||||
log.Errorf("redis protocol flush error: %v", errFlush)
|
||||
return
|
||||
}
|
||||
if !keepOpen {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func readRedisSubscriptionCommands(reader *bufio.Reader, commands chan<- redisSubscriptionCommand, done <-chan struct{}) {
|
||||
defer close(commands)
|
||||
|
||||
for {
|
||||
args, errRead := readRESPArray(reader)
|
||||
if errRead != nil {
|
||||
if !errors.Is(errRead, io.EOF) {
|
||||
select {
|
||||
case commands <- redisSubscriptionCommand{err: errRead}:
|
||||
case <-done:
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
select {
|
||||
case commands <- redisSubscriptionCommand{args: args}:
|
||||
case <-done:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func handleRedisSubscriptionCommand(writer *bufio.Writer, channel string, command redisSubscriptionCommand) bool {
|
||||
if command.err != nil {
|
||||
_ = writeRedisError(writer, "ERR "+command.err.Error())
|
||||
return false
|
||||
}
|
||||
if len(command.args) == 0 {
|
||||
_ = writeRedisError(writer, "ERR empty command")
|
||||
return true
|
||||
}
|
||||
|
||||
cmd := strings.ToUpper(strings.TrimSpace(command.args[0]))
|
||||
switch cmd {
|
||||
case "PING":
|
||||
payload := []byte(nil)
|
||||
if len(command.args) > 1 {
|
||||
payload = []byte(command.args[1])
|
||||
}
|
||||
_ = writeRedisPubSubPong(writer, payload)
|
||||
return true
|
||||
case "UNSUBSCRIBE":
|
||||
_ = writeRedisPubSubUnsubscribe(writer, channel, 0)
|
||||
return false
|
||||
case "QUIT":
|
||||
_ = writeRedisSimpleString(writer, "OK")
|
||||
return false
|
||||
default:
|
||||
_ = writeRedisError(writer, fmt.Sprintf("ERR unknown command '%s'", strings.ToLower(cmd)))
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
func resolveRemoteIP(addr net.Addr) (ip string, localClient bool) {
|
||||
if addr == nil {
|
||||
return "", false
|
||||
}
|
||||
|
||||
var host string
|
||||
switch a := addr.(type) {
|
||||
case *net.TCPAddr:
|
||||
if a != nil && a.IP != nil {
|
||||
if ip4 := a.IP.To4(); ip4 != nil {
|
||||
host = ip4.String()
|
||||
} else {
|
||||
host = a.IP.String()
|
||||
}
|
||||
}
|
||||
default:
|
||||
host = addr.String()
|
||||
if h, _, errSplit := net.SplitHostPort(host); errSplit == nil {
|
||||
host = h
|
||||
}
|
||||
host = strings.TrimSpace(host)
|
||||
if raw, _, ok := strings.Cut(host, "%"); ok {
|
||||
host = raw
|
||||
}
|
||||
if parsed := net.ParseIP(host); parsed != nil {
|
||||
if ip4 := parsed.To4(); ip4 != nil {
|
||||
host = ip4.String()
|
||||
} else {
|
||||
host = parsed.String()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
host = strings.TrimSpace(host)
|
||||
localClient = host == "127.0.0.1" || host == "::1"
|
||||
return host, localClient
|
||||
}
|
||||
|
||||
func parseAuthPassword(args []string) (string, bool) {
|
||||
switch len(args) {
|
||||
case 2:
|
||||
return args[1], true
|
||||
case 3:
|
||||
return args[2], true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
func parseSubscribeChannel(args []string) (string, bool) {
|
||||
if len(args) != 2 {
|
||||
return "", false
|
||||
}
|
||||
return strings.TrimSpace(args[1]), true
|
||||
}
|
||||
|
||||
func parsePopCount(args []string) (count int, hasCount bool, ok bool) {
|
||||
if len(args) != 2 && len(args) != 3 {
|
||||
return 0, false, false
|
||||
}
|
||||
if len(args) == 2 {
|
||||
return 1, false, true
|
||||
}
|
||||
parsed, errParse := strconv.Atoi(strings.TrimSpace(args[2]))
|
||||
if errParse != nil {
|
||||
return 0, true, true
|
||||
}
|
||||
return parsed, true, true
|
||||
}
|
||||
|
||||
func readRESPArray(reader *bufio.Reader) ([]string, error) {
|
||||
prefix, errRead := reader.ReadByte()
|
||||
if errRead != nil {
|
||||
return nil, errRead
|
||||
}
|
||||
if prefix != '*' {
|
||||
return nil, fmt.Errorf("protocol error")
|
||||
}
|
||||
line, errLine := readRESPLine(reader)
|
||||
if errLine != nil {
|
||||
return nil, errLine
|
||||
}
|
||||
count, errParse := strconv.Atoi(line)
|
||||
if errParse != nil || count < 0 {
|
||||
return nil, fmt.Errorf("protocol error")
|
||||
}
|
||||
args := make([]string, 0, count)
|
||||
for i := 0; i < count; i++ {
|
||||
value, errString := readRESPString(reader)
|
||||
if errString != nil {
|
||||
return nil, errString
|
||||
}
|
||||
args = append(args, value)
|
||||
}
|
||||
return args, nil
|
||||
}
|
||||
|
||||
func readRESPString(reader *bufio.Reader) (string, error) {
|
||||
prefix, errRead := reader.ReadByte()
|
||||
if errRead != nil {
|
||||
return "", errRead
|
||||
}
|
||||
switch prefix {
|
||||
case '$':
|
||||
return readRESPBulkString(reader)
|
||||
case '+', ':':
|
||||
return readRESPLine(reader)
|
||||
default:
|
||||
return "", fmt.Errorf("protocol error")
|
||||
}
|
||||
}
|
||||
|
||||
func readRESPBulkString(reader *bufio.Reader) (string, error) {
|
||||
line, errLine := readRESPLine(reader)
|
||||
if errLine != nil {
|
||||
return "", errLine
|
||||
}
|
||||
length, errParse := strconv.Atoi(line)
|
||||
if errParse != nil {
|
||||
return "", fmt.Errorf("protocol error")
|
||||
}
|
||||
if length < 0 {
|
||||
return "", nil
|
||||
}
|
||||
buf := make([]byte, length+2)
|
||||
if _, errRead := io.ReadFull(reader, buf); errRead != nil {
|
||||
return "", errRead
|
||||
}
|
||||
if length+2 < 2 || buf[length] != '\r' || buf[length+1] != '\n' {
|
||||
return "", fmt.Errorf("protocol error")
|
||||
}
|
||||
return string(buf[:length]), nil
|
||||
}
|
||||
|
||||
func readRESPLine(reader *bufio.Reader) (string, error) {
|
||||
line, errRead := reader.ReadString('\n')
|
||||
if errRead != nil {
|
||||
return "", errRead
|
||||
}
|
||||
line = strings.TrimSuffix(line, "\n")
|
||||
line = strings.TrimSuffix(line, "\r")
|
||||
return line, nil
|
||||
}
|
||||
|
||||
func writeRedisSimpleString(writer *bufio.Writer, value string) error {
|
||||
if writer == nil {
|
||||
return net.ErrClosed
|
||||
}
|
||||
_, errWrite := writer.WriteString("+" + value + "\r\n")
|
||||
return errWrite
|
||||
}
|
||||
|
||||
func writeRedisError(writer *bufio.Writer, message string) error {
|
||||
if writer == nil {
|
||||
return net.ErrClosed
|
||||
}
|
||||
_, errWrite := writer.WriteString("-" + message + "\r\n")
|
||||
return errWrite
|
||||
}
|
||||
|
||||
func writeRedisNilBulkString(writer *bufio.Writer) error {
|
||||
if writer == nil {
|
||||
return net.ErrClosed
|
||||
}
|
||||
_, errWrite := writer.WriteString("$-1\r\n")
|
||||
return errWrite
|
||||
}
|
||||
|
||||
func writeRedisBulkString(writer *bufio.Writer, payload []byte) error {
|
||||
if writer == nil {
|
||||
return net.ErrClosed
|
||||
}
|
||||
if payload == nil {
|
||||
return writeRedisNilBulkString(writer)
|
||||
}
|
||||
if _, errWrite := writer.WriteString("$" + strconv.Itoa(len(payload)) + "\r\n"); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
if _, errWrite := writer.Write(payload); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
_, errWrite := writer.WriteString("\r\n")
|
||||
return errWrite
|
||||
}
|
||||
|
||||
func writeRedisArrayOfBulkStrings(writer *bufio.Writer, items [][]byte) error {
|
||||
if writer == nil {
|
||||
return net.ErrClosed
|
||||
}
|
||||
if _, errWrite := writer.WriteString("*" + strconv.Itoa(len(items)) + "\r\n"); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
for i := range items {
|
||||
if errWrite := writeRedisBulkString(writer, items[i]); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func writeRedisInteger(writer *bufio.Writer, value int) error {
|
||||
if writer == nil {
|
||||
return net.ErrClosed
|
||||
}
|
||||
_, errWrite := writer.WriteString(":" + strconv.Itoa(value) + "\r\n")
|
||||
return errWrite
|
||||
}
|
||||
|
||||
func writeRedisArrayHeader(writer *bufio.Writer, count int) error {
|
||||
if writer == nil {
|
||||
return net.ErrClosed
|
||||
}
|
||||
_, errWrite := writer.WriteString("*" + strconv.Itoa(count) + "\r\n")
|
||||
return errWrite
|
||||
}
|
||||
|
||||
func writeRedisPubSubSubscribe(writer *bufio.Writer, channel string, count int) error {
|
||||
if errWrite := writeRedisArrayHeader(writer, 3); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
if errWrite := writeRedisBulkString(writer, []byte("subscribe")); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
if errWrite := writeRedisBulkString(writer, []byte(channel)); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
return writeRedisInteger(writer, count)
|
||||
}
|
||||
|
||||
func writeRedisPubSubUnsubscribe(writer *bufio.Writer, channel string, count int) error {
|
||||
if errWrite := writeRedisArrayHeader(writer, 3); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
if errWrite := writeRedisBulkString(writer, []byte("unsubscribe")); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
if errWrite := writeRedisBulkString(writer, []byte(channel)); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
return writeRedisInteger(writer, count)
|
||||
}
|
||||
|
||||
func writeRedisPubSubMessage(writer *bufio.Writer, channel string, payload []byte) error {
|
||||
if errWrite := writeRedisArrayHeader(writer, 3); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
if errWrite := writeRedisBulkString(writer, []byte("message")); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
if errWrite := writeRedisBulkString(writer, []byte(channel)); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
return writeRedisBulkString(writer, payload)
|
||||
}
|
||||
|
||||
func writeRedisPubSubPong(writer *bufio.Writer, payload []byte) error {
|
||||
if errWrite := writeRedisArrayHeader(writer, 2); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
if errWrite := writeRedisBulkString(writer, []byte("pong")); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
return writeRedisBulkString(writer, payload)
|
||||
}
|
||||
516
backend/internal/api/redis_queue_protocol_integration_test.go
Normal file
516
backend/internal/api/redis_queue_protocol_integration_test.go
Normal file
|
|
@ -0,0 +1,516 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue"
|
||||
)
|
||||
|
||||
func startRedisMuxListener(t *testing.T, server *Server) (addr string, stop func()) {
|
||||
t.Helper()
|
||||
|
||||
listener, errListen := net.Listen("tcp", "127.0.0.1:0")
|
||||
if errListen != nil {
|
||||
t.Fatalf("failed to listen: %v", errListen)
|
||||
}
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- server.acceptMuxConnections(listener, nil)
|
||||
}()
|
||||
|
||||
stop = func() {
|
||||
_ = listener.Close()
|
||||
select {
|
||||
case err := <-errCh:
|
||||
if err != nil && !errors.Is(err, net.ErrClosed) {
|
||||
t.Errorf("accept loop returned unexpected error: %v", err)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Errorf("timeout waiting for accept loop to exit")
|
||||
}
|
||||
}
|
||||
|
||||
return listener.Addr().String(), stop
|
||||
}
|
||||
|
||||
func writeTestRESPCommand(conn net.Conn, args ...string) error {
|
||||
if conn == nil {
|
||||
return net.ErrClosed
|
||||
}
|
||||
if len(args) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
var buf bytes.Buffer
|
||||
fmt.Fprintf(&buf, "*%d\r\n", len(args))
|
||||
for _, arg := range args {
|
||||
fmt.Fprintf(&buf, "$%d\r\n%s\r\n", len(arg), arg)
|
||||
}
|
||||
_, err := conn.Write(buf.Bytes())
|
||||
return err
|
||||
}
|
||||
|
||||
func readTestRESPLine(r *bufio.Reader) (string, error) {
|
||||
line, err := r.ReadString('\n')
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !strings.HasSuffix(line, "\r\n") {
|
||||
return "", fmt.Errorf("invalid RESP line terminator: %q", line)
|
||||
}
|
||||
return strings.TrimSuffix(line, "\r\n"), nil
|
||||
}
|
||||
|
||||
func readTestRESPError(r *bufio.Reader) (string, error) {
|
||||
prefix, err := r.ReadByte()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if prefix != '-' {
|
||||
return "", fmt.Errorf("expected error prefix '-', got %q", prefix)
|
||||
}
|
||||
return readTestRESPLine(r)
|
||||
}
|
||||
|
||||
func readTestRESPSimpleString(r *bufio.Reader) (string, error) {
|
||||
prefix, errRead := r.ReadByte()
|
||||
if errRead != nil {
|
||||
return "", errRead
|
||||
}
|
||||
if prefix != '+' {
|
||||
return "", fmt.Errorf("expected simple string prefix '+', got %q", prefix)
|
||||
}
|
||||
return readTestRESPLine(r)
|
||||
}
|
||||
|
||||
func readTestRESPBulkString(r *bufio.Reader) ([]byte, error) {
|
||||
prefix, errRead := r.ReadByte()
|
||||
if errRead != nil {
|
||||
return nil, errRead
|
||||
}
|
||||
if prefix != '$' {
|
||||
return nil, fmt.Errorf("expected bulk string prefix '$', got %q", prefix)
|
||||
}
|
||||
|
||||
line, errLine := readTestRESPLine(r)
|
||||
if errLine != nil {
|
||||
return nil, errLine
|
||||
}
|
||||
length, errParse := strconv.Atoi(line)
|
||||
if errParse != nil {
|
||||
return nil, fmt.Errorf("invalid bulk string length %q: %v", line, errParse)
|
||||
}
|
||||
if length == -1 {
|
||||
return nil, nil
|
||||
}
|
||||
if length < -1 {
|
||||
return nil, fmt.Errorf("invalid bulk string length %d", length)
|
||||
}
|
||||
|
||||
payload := make([]byte, length+2)
|
||||
if _, errRead := io.ReadFull(r, payload); errRead != nil {
|
||||
return nil, errRead
|
||||
}
|
||||
if payload[length] != '\r' || payload[length+1] != '\n' {
|
||||
return nil, fmt.Errorf("invalid bulk string terminator")
|
||||
}
|
||||
return payload[:length], nil
|
||||
}
|
||||
|
||||
func readRESPArrayOfBulkStrings(r *bufio.Reader) ([][]byte, error) {
|
||||
prefix, errRead := r.ReadByte()
|
||||
if errRead != nil {
|
||||
return nil, errRead
|
||||
}
|
||||
if prefix != '*' {
|
||||
return nil, fmt.Errorf("expected array prefix '*', got %q", prefix)
|
||||
}
|
||||
|
||||
line, errLine := readTestRESPLine(r)
|
||||
if errLine != nil {
|
||||
return nil, errLine
|
||||
}
|
||||
count, errParse := strconv.Atoi(line)
|
||||
if errParse != nil {
|
||||
return nil, fmt.Errorf("invalid array length %q: %v", line, errParse)
|
||||
}
|
||||
if count < 0 {
|
||||
return nil, fmt.Errorf("invalid array length %d", count)
|
||||
}
|
||||
|
||||
out := make([][]byte, 0, count)
|
||||
for i := 0; i < count; i++ {
|
||||
item, errItem := readTestRESPBulkString(r)
|
||||
if errItem != nil {
|
||||
return nil, errItem
|
||||
}
|
||||
out = append(out, item)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func readTestRESPPubSubSubscribe(r *bufio.Reader) (string, int, error) {
|
||||
prefix, errRead := r.ReadByte()
|
||||
if errRead != nil {
|
||||
return "", 0, errRead
|
||||
}
|
||||
if prefix != '*' {
|
||||
return "", 0, fmt.Errorf("expected array prefix '*', got %q", prefix)
|
||||
}
|
||||
line, errLine := readTestRESPLine(r)
|
||||
if errLine != nil {
|
||||
return "", 0, errLine
|
||||
}
|
||||
count, errParse := strconv.Atoi(line)
|
||||
if errParse != nil {
|
||||
return "", 0, fmt.Errorf("invalid array length %q: %v", line, errParse)
|
||||
}
|
||||
if count != 3 {
|
||||
return "", 0, fmt.Errorf("subscribe ack length = %d, want 3", count)
|
||||
}
|
||||
kind, errKind := readTestRESPBulkString(r)
|
||||
if errKind != nil {
|
||||
return "", 0, errKind
|
||||
}
|
||||
if string(kind) != "subscribe" {
|
||||
return "", 0, fmt.Errorf("subscribe ack kind = %q", string(kind))
|
||||
}
|
||||
channel, errChannel := readTestRESPBulkString(r)
|
||||
if errChannel != nil {
|
||||
return "", 0, errChannel
|
||||
}
|
||||
prefix, errRead = r.ReadByte()
|
||||
if errRead != nil {
|
||||
return "", 0, errRead
|
||||
}
|
||||
if prefix != ':' {
|
||||
return "", 0, fmt.Errorf("expected integer prefix ':', got %q", prefix)
|
||||
}
|
||||
line, errLine = readTestRESPLine(r)
|
||||
if errLine != nil {
|
||||
return "", 0, errLine
|
||||
}
|
||||
subscriptions, errParse := strconv.Atoi(line)
|
||||
if errParse != nil {
|
||||
return "", 0, fmt.Errorf("invalid subscription count %q: %v", line, errParse)
|
||||
}
|
||||
return string(channel), subscriptions, nil
|
||||
}
|
||||
|
||||
func readTestRESPPubSubMessage(r *bufio.Reader) (string, []byte, error) {
|
||||
items, errItems := readRESPArrayOfBulkStrings(r)
|
||||
if errItems != nil {
|
||||
return "", nil, errItems
|
||||
}
|
||||
if len(items) != 3 {
|
||||
return "", nil, fmt.Errorf("pubsub message length = %d, want 3", len(items))
|
||||
}
|
||||
if string(items[0]) != "message" {
|
||||
return "", nil, fmt.Errorf("pubsub message kind = %q", string(items[0]))
|
||||
}
|
||||
return string(items[1]), items[2], nil
|
||||
}
|
||||
|
||||
func TestRedisProtocol_ManagementDisabled_RejectsConnection(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
redisqueue.SetEnabled(false)
|
||||
|
||||
server := newTestServer(t)
|
||||
if server.managementRoutesEnabled.Load() {
|
||||
t.Fatalf("expected managementRoutesEnabled to be false")
|
||||
}
|
||||
|
||||
addr, stop := startRedisMuxListener(t, server)
|
||||
t.Cleanup(stop)
|
||||
|
||||
conn, errDial := net.DialTimeout("tcp", addr, time.Second)
|
||||
if errDial != nil {
|
||||
t.Fatalf("failed to dial redis listener: %v", errDial)
|
||||
}
|
||||
t.Cleanup(func() { _ = conn.Close() })
|
||||
|
||||
_ = conn.SetDeadline(time.Now().Add(2 * time.Second))
|
||||
if errWrite := writeTestRESPCommand(conn, "PING"); errWrite != nil {
|
||||
t.Fatalf("failed to write RESP command: %v", errWrite)
|
||||
}
|
||||
|
||||
buf := make([]byte, 1)
|
||||
_, errRead := conn.Read(buf)
|
||||
if errRead == nil {
|
||||
t.Fatalf("expected connection to be closed when management is disabled")
|
||||
}
|
||||
if ne, ok := errRead.(net.Error); ok && ne.Timeout() {
|
||||
t.Fatalf("expected connection to be closed when management is disabled, got timeout: %v", errRead)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisProtocol_HomeEnabled_DisablesConnection(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "test-management-password")
|
||||
redisqueue.SetEnabled(false)
|
||||
t.Cleanup(func() { redisqueue.SetEnabled(false) })
|
||||
|
||||
server := newTestServer(t)
|
||||
if !server.managementRoutesEnabled.Load() {
|
||||
t.Fatalf("expected managementRoutesEnabled to be true")
|
||||
}
|
||||
if server.cfg == nil {
|
||||
t.Fatalf("expected server cfg to be non-nil")
|
||||
}
|
||||
server.cfg.Home.Enabled = true
|
||||
redisqueue.SetEnabled(true)
|
||||
|
||||
addr, stop := startRedisMuxListener(t, server)
|
||||
t.Cleanup(stop)
|
||||
|
||||
conn, errDial := net.DialTimeout("tcp", addr, time.Second)
|
||||
if errDial != nil {
|
||||
t.Fatalf("failed to dial redis listener: %v", errDial)
|
||||
}
|
||||
t.Cleanup(func() { _ = conn.Close() })
|
||||
|
||||
_ = conn.SetDeadline(time.Now().Add(2 * time.Second))
|
||||
_ = writeTestRESPCommand(conn, "PING")
|
||||
|
||||
if msg, err := readTestRESPError(bufio.NewReader(conn)); err != nil {
|
||||
t.Fatalf("failed to read home-mode RESP error: %v", err)
|
||||
} else if msg != "ERR redis usage output disabled in home mode" {
|
||||
t.Fatalf("unexpected disabled RESP error: %q", msg)
|
||||
}
|
||||
|
||||
buf := make([]byte, 1)
|
||||
_, errRead := conn.Read(buf)
|
||||
if errRead == nil {
|
||||
t.Fatalf("expected connection to be closed after home-mode RESP error")
|
||||
}
|
||||
if ne, ok := errRead.(net.Error); ok && ne.Timeout() {
|
||||
t.Fatalf("expected connection to be closed after home-mode RESP error, got timeout: %v", errRead)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisProtocol_SUBSCRIBE_UsageSendsSupportRefresh(t *testing.T) {
|
||||
const managementPassword = "test-management-password"
|
||||
|
||||
t.Setenv("MANAGEMENT_PASSWORD", managementPassword)
|
||||
redisqueue.SetEnabled(false)
|
||||
t.Cleanup(func() { redisqueue.SetEnabled(false) })
|
||||
|
||||
server := newTestServer(t)
|
||||
if !server.managementRoutesEnabled.Load() {
|
||||
t.Fatalf("expected managementRoutesEnabled to be true")
|
||||
}
|
||||
|
||||
addr, stop := startRedisMuxListener(t, server)
|
||||
t.Cleanup(stop)
|
||||
|
||||
conn, errDial := net.DialTimeout("tcp", addr, time.Second)
|
||||
if errDial != nil {
|
||||
t.Fatalf("failed to dial redis listener: %v", errDial)
|
||||
}
|
||||
t.Cleanup(func() { _ = conn.Close() })
|
||||
|
||||
reader := bufio.NewReader(conn)
|
||||
_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "AUTH", managementPassword); errWrite != nil {
|
||||
t.Fatalf("failed to write AUTH command: %v", errWrite)
|
||||
}
|
||||
if msg, errRead := readTestRESPSimpleString(reader); errRead != nil {
|
||||
t.Fatalf("failed to read AUTH response: %v", errRead)
|
||||
} else if msg != "OK" {
|
||||
t.Fatalf("unexpected AUTH response: %q", msg)
|
||||
}
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "SUBSCRIBE", "usage"); errWrite != nil {
|
||||
t.Fatalf("failed to write SUBSCRIBE command: %v", errWrite)
|
||||
}
|
||||
channel, subscriptions, errSubscribe := readTestRESPPubSubSubscribe(reader)
|
||||
if errSubscribe != nil {
|
||||
t.Fatalf("failed to read subscribe response: %v", errSubscribe)
|
||||
}
|
||||
if channel != "usage" || subscriptions != 1 {
|
||||
t.Fatalf("unexpected subscribe response channel=%q subscriptions=%d", channel, subscriptions)
|
||||
}
|
||||
|
||||
channel, payload, errMessage := readTestRESPPubSubMessage(reader)
|
||||
if errMessage != nil {
|
||||
t.Fatalf("failed to read support refresh message: %v", errMessage)
|
||||
}
|
||||
if channel != "usage" || string(payload) != `{"support_refresh":true}` {
|
||||
t.Fatalf("unexpected support refresh message channel=%q payload=%q", channel, string(payload))
|
||||
}
|
||||
|
||||
redisqueue.Enqueue([]byte(`{"id":1}`))
|
||||
channel, payload, errMessage = readTestRESPPubSubMessage(reader)
|
||||
if errMessage != nil {
|
||||
t.Fatalf("failed to read usage message: %v", errMessage)
|
||||
}
|
||||
if channel != "usage" || string(payload) != `{"id":1}` {
|
||||
t.Fatalf("unexpected usage message channel=%q payload=%q", channel, string(payload))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisProtocol_SUBSCRIBE_ErrorsReceivesErrorEvents(t *testing.T) {
|
||||
const managementPassword = "test-management-password"
|
||||
|
||||
t.Setenv("MANAGEMENT_PASSWORD", managementPassword)
|
||||
redisqueue.SetEnabled(false)
|
||||
t.Cleanup(func() { redisqueue.SetEnabled(false) })
|
||||
|
||||
server := newTestServer(t)
|
||||
if !server.managementRoutesEnabled.Load() {
|
||||
t.Fatalf("expected managementRoutesEnabled to be true")
|
||||
}
|
||||
|
||||
addr, stop := startRedisMuxListener(t, server)
|
||||
t.Cleanup(stop)
|
||||
|
||||
conn, errDial := net.DialTimeout("tcp", addr, time.Second)
|
||||
if errDial != nil {
|
||||
t.Fatalf("failed to dial redis listener: %v", errDial)
|
||||
}
|
||||
t.Cleanup(func() { _ = conn.Close() })
|
||||
|
||||
reader := bufio.NewReader(conn)
|
||||
_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "AUTH", managementPassword); errWrite != nil {
|
||||
t.Fatalf("failed to write AUTH command: %v", errWrite)
|
||||
}
|
||||
if msg, errRead := readTestRESPSimpleString(reader); errRead != nil {
|
||||
t.Fatalf("failed to read AUTH response: %v", errRead)
|
||||
} else if msg != "OK" {
|
||||
t.Fatalf("unexpected AUTH response: %q", msg)
|
||||
}
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "SUBSCRIBE", "errors"); errWrite != nil {
|
||||
t.Fatalf("failed to write SUBSCRIBE command: %v", errWrite)
|
||||
}
|
||||
channel, subscriptions, errSubscribe := readTestRESPPubSubSubscribe(reader)
|
||||
if errSubscribe != nil {
|
||||
t.Fatalf("failed to read subscribe response: %v", errSubscribe)
|
||||
}
|
||||
if channel != "errors" || subscriptions != 1 {
|
||||
t.Fatalf("unexpected subscribe response channel=%q subscriptions=%d", channel, subscriptions)
|
||||
}
|
||||
|
||||
redisqueue.EnqueueError([]byte(`{"auth_index":"auth-1","status_code":401}`))
|
||||
channel, payload, errMessage := readTestRESPPubSubMessage(reader)
|
||||
if errMessage != nil {
|
||||
t.Fatalf("failed to read error message: %v", errMessage)
|
||||
}
|
||||
if channel != "errors" || string(payload) != `{"auth_index":"auth-1","status_code":401}` {
|
||||
t.Fatalf("unexpected error message channel=%q payload=%q", channel, string(payload))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisProtocol_AUTH_And_PopContracts(t *testing.T) {
|
||||
const managementPassword = "test-management-password"
|
||||
|
||||
t.Setenv("MANAGEMENT_PASSWORD", managementPassword)
|
||||
redisqueue.SetEnabled(false)
|
||||
t.Cleanup(func() { redisqueue.SetEnabled(false) })
|
||||
|
||||
server := newTestServer(t)
|
||||
if !server.managementRoutesEnabled.Load() {
|
||||
t.Fatalf("expected managementRoutesEnabled to be true")
|
||||
}
|
||||
|
||||
addr, stop := startRedisMuxListener(t, server)
|
||||
t.Cleanup(stop)
|
||||
|
||||
conn, errDial := net.DialTimeout("tcp", addr, time.Second)
|
||||
if errDial != nil {
|
||||
t.Fatalf("failed to dial redis listener: %v", errDial)
|
||||
}
|
||||
t.Cleanup(func() { _ = conn.Close() })
|
||||
|
||||
reader := bufio.NewReader(conn)
|
||||
|
||||
_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "AUTH", managementPassword); errWrite != nil {
|
||||
t.Fatalf("failed to write AUTH command: %v", errWrite)
|
||||
}
|
||||
if msg, errRead := readTestRESPSimpleString(reader); errRead != nil {
|
||||
t.Fatalf("failed to read AUTH response: %v", errRead)
|
||||
} else if msg != "OK" {
|
||||
t.Fatalf("unexpected AUTH response: %q", msg)
|
||||
}
|
||||
|
||||
if !redisqueue.Enabled() {
|
||||
t.Fatalf("expected redisqueue to be enabled")
|
||||
}
|
||||
redisqueue.Enqueue([]byte("a"))
|
||||
redisqueue.Enqueue([]byte("b"))
|
||||
redisqueue.Enqueue([]byte("c"))
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "RPOP", "usage"); errWrite != nil {
|
||||
t.Fatalf("failed to write RPOP command: %v", errWrite)
|
||||
}
|
||||
if item, errRead := readTestRESPBulkString(reader); errRead != nil {
|
||||
t.Fatalf("failed to read RPOP response: %v", errRead)
|
||||
} else if string(item) != "a" {
|
||||
t.Fatalf("unexpected RPOP item: %q", string(item))
|
||||
}
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "LPOP", "usage"); errWrite != nil {
|
||||
t.Fatalf("failed to write LPOP command: %v", errWrite)
|
||||
}
|
||||
if item, errRead := readTestRESPBulkString(reader); errRead != nil {
|
||||
t.Fatalf("failed to read LPOP response: %v", errRead)
|
||||
} else if string(item) != "b" {
|
||||
t.Fatalf("unexpected LPOP item: %q", string(item))
|
||||
}
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "RPOP", "usage", "10"); errWrite != nil {
|
||||
t.Fatalf("failed to write RPOP count command: %v", errWrite)
|
||||
}
|
||||
items, errItems := readRESPArrayOfBulkStrings(reader)
|
||||
if errItems != nil {
|
||||
t.Fatalf("failed to read RPOP count response: %v", errItems)
|
||||
}
|
||||
if len(items) != 1 || string(items[0]) != "c" {
|
||||
t.Fatalf("unexpected RPOP count items: %#v", items)
|
||||
}
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "LPOP", "usage"); errWrite != nil {
|
||||
t.Fatalf("failed to write LPOP empty command: %v", errWrite)
|
||||
}
|
||||
item, errItem := readTestRESPBulkString(reader)
|
||||
if errItem != nil {
|
||||
t.Fatalf("failed to read LPOP empty response: %v", errItem)
|
||||
}
|
||||
if item != nil {
|
||||
t.Fatalf("expected nil bulk string for empty queue, got %q", string(item))
|
||||
}
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "RPOP", "usage", "2"); errWrite != nil {
|
||||
t.Fatalf("failed to write RPOP empty count command: %v", errWrite)
|
||||
}
|
||||
emptyItems, errEmpty := readRESPArrayOfBulkStrings(reader)
|
||||
if errEmpty != nil {
|
||||
t.Fatalf("failed to read RPOP empty count response: %v", errEmpty)
|
||||
}
|
||||
if len(emptyItems) != 0 {
|
||||
t.Fatalf("expected empty array for empty queue with count, got %#v", emptyItems)
|
||||
}
|
||||
|
||||
if errWrite := writeTestRESPCommand(conn, "RPOP", "errors", "2"); errWrite != nil {
|
||||
t.Fatalf("failed to write RPOP errors count command: %v", errWrite)
|
||||
}
|
||||
if msg, errRead := readTestRESPError(reader); errRead != nil {
|
||||
t.Fatalf("failed to read RPOP errors response: %v", errRead)
|
||||
} else if msg != "ERR unsupported channel 'errors'" {
|
||||
t.Fatalf("unexpected RPOP errors response: %q", msg)
|
||||
}
|
||||
}
|
||||
397
backend/internal/api/server.go
Normal file
397
backend/internal/api/server.go
Normal file
|
|
@ -0,0 +1,397 @@
|
|||
// Package api provides the HTTP API server implementation for the CLI Proxy API.
|
||||
// It includes the main server struct, routing setup, middleware for CORS and authentication,
|
||||
// and integration with various AI API handlers (OpenAI, Claude, Gemini).
|
||||
// The server supports hot-reloading of clients and configuration.
|
||||
package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
managementHandlers "github.com/router-for-me/CLIProxyAPI/v7/internal/api/handlers/management"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/api/middleware"
|
||||
codexlive "github.com/router-for-me/CLIProxyAPI/v7/internal/client/codex/live"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue"
|
||||
sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/net/http2"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// Server represents the main API server.
|
||||
// It encapsulates the Gin engine, HTTP server, handlers, and configuration.
|
||||
type Server struct {
|
||||
// engine is the Gin web framework engine instance.
|
||||
engine *gin.Engine
|
||||
|
||||
// server is the underlying HTTP server.
|
||||
server *http.Server
|
||||
|
||||
// muxBaseListener is the shared TCP listener used to serve both HTTP and Redis protocol traffic.
|
||||
muxBaseListener net.Listener
|
||||
|
||||
// muxHTTPListener receives HTTP connections selected by the multiplexer.
|
||||
muxHTTPListener *muxListener
|
||||
|
||||
// handlers contains the API handlers for processing requests.
|
||||
handlers *handlers.BaseAPIHandler
|
||||
codexLiveHandler *codexlive.Handler
|
||||
|
||||
// cfg holds the current server configuration.
|
||||
cfg *config.Config
|
||||
|
||||
// oldConfigYaml stores a YAML snapshot of the previous configuration for change detection.
|
||||
// This prevents issues when the config object is modified in place by Management API.
|
||||
oldConfigYaml []byte
|
||||
|
||||
// accessManager handles request authentication providers.
|
||||
accessManager *sdkaccess.Manager
|
||||
|
||||
// requestLogger is the request logger instance for dynamic configuration updates.
|
||||
requestLogger logging.RequestLogger
|
||||
loggerToggle func(bool)
|
||||
|
||||
// configFilePath is the absolute path to the YAML config file for persistence.
|
||||
configFilePath string
|
||||
|
||||
// currentPath is the absolute path to the current working directory.
|
||||
currentPath string
|
||||
|
||||
// wsRoutes tracks registered websocket upgrade paths.
|
||||
wsRouteMu sync.Mutex
|
||||
wsRoutes map[string]struct{}
|
||||
wsAuthChanged func(bool, bool)
|
||||
wsAuthEnabled atomic.Bool
|
||||
|
||||
// management handler
|
||||
mgmt *managementHandlers.Handler
|
||||
|
||||
// pluginHost owns dynamic plugin Management API route dispatch.
|
||||
pluginHost *pluginhost.Host
|
||||
|
||||
// managementRoutesRegistered tracks whether the management routes have been attached to the engine.
|
||||
managementRoutesRegistered atomic.Bool
|
||||
// managementRoutesEnabled controls whether management endpoints serve real handlers.
|
||||
managementRoutesEnabled atomic.Bool
|
||||
|
||||
// envManagementSecret indicates whether MANAGEMENT_PASSWORD is configured.
|
||||
envManagementSecret bool
|
||||
|
||||
localPassword string
|
||||
|
||||
keepAliveEnabled bool
|
||||
keepAliveTimeout time.Duration
|
||||
keepAliveOnTimeout func()
|
||||
keepAliveHeartbeat chan struct{}
|
||||
keepAliveStop chan struct{}
|
||||
|
||||
exampleAPIKeySafeModeEnabled bool
|
||||
exampleAPIKeySafeModeActive atomic.Bool
|
||||
}
|
||||
|
||||
// NewServer creates and initializes a new API server instance.
|
||||
// It sets up the Gin engine, middleware, routes, and handlers.
|
||||
//
|
||||
// Parameters:
|
||||
// - cfg: The server configuration
|
||||
// - authManager: core runtime auth manager
|
||||
// - accessManager: request authentication manager
|
||||
//
|
||||
// Returns:
|
||||
// - *Server: A new server instance
|
||||
func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdkaccess.Manager, configFilePath string, opts ...ServerOption) *Server {
|
||||
optionState := &serverOptionConfig{
|
||||
requestLoggerFactory: defaultRequestLoggerFactory,
|
||||
}
|
||||
for i := range opts {
|
||||
opts[i](optionState)
|
||||
}
|
||||
// Set gin mode
|
||||
if !cfg.Debug {
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
}
|
||||
|
||||
// Create gin engine
|
||||
engine := gin.New()
|
||||
if optionState.engineConfigurator != nil {
|
||||
optionState.engineConfigurator(engine)
|
||||
}
|
||||
|
||||
// Add middleware
|
||||
engine.Use(logging.GinLogrusLogger())
|
||||
engine.Use(logging.GinLogrusRecovery())
|
||||
engine.Use(logging.CPATraceIDMiddleware())
|
||||
for _, mw := range optionState.extraMiddleware {
|
||||
engine.Use(mw)
|
||||
}
|
||||
|
||||
// Add request logging middleware (positioned after recovery, before auth)
|
||||
// Resolve logs directory relative to the configuration file directory.
|
||||
var requestLogger logging.RequestLogger
|
||||
var toggle func(bool)
|
||||
if !cfg.CommercialMode {
|
||||
if optionState.requestLoggerFactory != nil {
|
||||
requestLogger = optionState.requestLoggerFactory(cfg, configFilePath)
|
||||
}
|
||||
if requestLogger != nil {
|
||||
engine.Use(middleware.RequestLoggingMiddleware(requestLogger))
|
||||
if setter, ok := requestLogger.(interface{ SetEnabled(bool) }); ok {
|
||||
toggle = setter.SetEnabled
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
engine.Use(corsMiddleware())
|
||||
wd, err := os.Getwd()
|
||||
if err != nil {
|
||||
wd = configFilePath
|
||||
}
|
||||
|
||||
envAdminPassword, envAdminPasswordSet := os.LookupEnv("MANAGEMENT_PASSWORD")
|
||||
envAdminPassword = strings.TrimSpace(envAdminPassword)
|
||||
envManagementSecret := envAdminPasswordSet && envAdminPassword != ""
|
||||
|
||||
// Create server instance
|
||||
s := &Server{
|
||||
engine: engine,
|
||||
handlers: handlers.NewBaseAPIHandlers(effectiveSDKConfig(cfg), authManager),
|
||||
cfg: cfg,
|
||||
accessManager: accessManager,
|
||||
requestLogger: requestLogger,
|
||||
loggerToggle: toggle,
|
||||
configFilePath: configFilePath,
|
||||
currentPath: wd,
|
||||
envManagementSecret: envManagementSecret,
|
||||
wsRoutes: make(map[string]struct{}),
|
||||
pluginHost: optionState.pluginHost,
|
||||
|
||||
exampleAPIKeySafeModeEnabled: optionState.exampleAPIKeySafeMode,
|
||||
}
|
||||
s.wsAuthEnabled.Store(cfg.WebsocketAuth)
|
||||
s.exampleAPIKeySafeModeActive.Store(s.exampleAPIKeySafeModeRequired(cfg))
|
||||
s.handlers.SetPluginHost(optionState.pluginHost)
|
||||
if optionState.pluginHost != nil {
|
||||
optionState.pluginHost.SetModelExecutor(s.handlers)
|
||||
optionState.pluginHost.SetAuthManager(authManager)
|
||||
}
|
||||
// Save initial YAML snapshot
|
||||
s.oldConfigYaml, _ = yaml.Marshal(cfg)
|
||||
s.applyAccessConfig(nil, cfg)
|
||||
if authManager != nil {
|
||||
authManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second, cfg.MaxRetryCredentials)
|
||||
}
|
||||
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
|
||||
auth.SetTransientErrorCooldownSeconds(cfg.TransientErrorCooldownSeconds)
|
||||
applySignatureCacheConfig(nil, cfg)
|
||||
// Initialize management handler
|
||||
s.mgmt = managementHandlers.NewHandler(cfg, configFilePath, authManager)
|
||||
s.mgmt.SetPluginHost(optionState.pluginHost)
|
||||
s.mgmt.SetConfigReloadHook(optionState.configReloadHook)
|
||||
if optionState.localPassword != "" {
|
||||
s.mgmt.SetLocalPassword(optionState.localPassword)
|
||||
}
|
||||
logDir := logging.ResolveLogDirectory(cfg)
|
||||
s.mgmt.SetLogDirectory(logDir)
|
||||
if optionState.postAuthHook != nil {
|
||||
s.mgmt.SetPostAuthHook(optionState.postAuthHook)
|
||||
}
|
||||
if optionState.postAuthPersistHook != nil {
|
||||
s.mgmt.SetPostAuthPersistHook(optionState.postAuthPersistHook)
|
||||
}
|
||||
s.localPassword = optionState.localPassword
|
||||
|
||||
// Home heartbeat gate: when home is enabled, block all endpoints with 503 until the
|
||||
// subscribe-config heartbeat connection is healthy.
|
||||
engine.Use(s.homeHeartbeatMiddleware())
|
||||
engine.Use(s.exampleAPIKeySafeModeMiddleware())
|
||||
|
||||
// Setup routes
|
||||
s.setupRoutes()
|
||||
|
||||
// Apply additional router configurators from options
|
||||
if optionState.routerConfigurator != nil {
|
||||
optionState.routerConfigurator(engine, s.handlers, cfg)
|
||||
}
|
||||
|
||||
// Register management routes when configuration or environment secrets are available,
|
||||
// or when a local management password is provided (e.g. TUI mode).
|
||||
hasManagementSecret := cfg.RemoteManagement.SecretKey != "" || envManagementSecret || s.localPassword != ""
|
||||
s.managementRoutesEnabled.Store(hasManagementSecret)
|
||||
redisqueue.SetEnabled(hasManagementSecret || (cfg != nil && cfg.Home.Enabled))
|
||||
if hasManagementSecret {
|
||||
s.registerManagementRoutes()
|
||||
}
|
||||
s.refreshPluginManagementRoutes()
|
||||
engine.NoRoute(s.pluginManagementNoRoute)
|
||||
|
||||
if optionState.keepAliveEnabled {
|
||||
s.enableKeepAlive(optionState.keepAliveTimeout, optionState.keepAliveOnTimeout)
|
||||
}
|
||||
|
||||
// Create HTTP server
|
||||
s.server = &http.Server{
|
||||
Addr: fmt.Sprintf("%s:%d", cfg.Host, cfg.Port),
|
||||
Handler: engine,
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
// Start begins listening for and serving HTTP or HTTPS requests.
|
||||
// It's a blocking call and will only return on an unrecoverable error.
|
||||
//
|
||||
// Returns:
|
||||
// - error: An error if the server fails to start
|
||||
func (s *Server) Start() error {
|
||||
if s == nil || s.server == nil {
|
||||
return fmt.Errorf("failed to start HTTP server: server not initialized")
|
||||
}
|
||||
|
||||
addr := s.server.Addr
|
||||
listener, errListen := net.Listen("tcp", addr)
|
||||
if errListen != nil {
|
||||
return fmt.Errorf("failed to start HTTP server: %v", errListen)
|
||||
}
|
||||
|
||||
useTLS := s.cfg != nil && s.cfg.TLS.Enable
|
||||
if useTLS {
|
||||
certPath := strings.TrimSpace(s.cfg.TLS.Cert)
|
||||
keyPath := strings.TrimSpace(s.cfg.TLS.Key)
|
||||
if certPath == "" || keyPath == "" {
|
||||
if errClose := listener.Close(); errClose != nil {
|
||||
log.Errorf("failed to close listener after TLS validation failure: %v", errClose)
|
||||
}
|
||||
return fmt.Errorf("failed to start HTTPS server: tls.cert or tls.key is empty")
|
||||
}
|
||||
certPair, errLoad := tls.LoadX509KeyPair(certPath, keyPath)
|
||||
if errLoad != nil {
|
||||
if errClose := listener.Close(); errClose != nil {
|
||||
log.Errorf("failed to close listener after TLS key pair load failure: %v", errClose)
|
||||
}
|
||||
return fmt.Errorf("failed to start HTTPS server: %v", errLoad)
|
||||
}
|
||||
|
||||
tlsConfig := &tls.Config{
|
||||
Certificates: []tls.Certificate{certPair},
|
||||
NextProtos: []string{"h2", "http/1.1"},
|
||||
}
|
||||
s.server.TLSConfig = tlsConfig
|
||||
if errHTTP2 := http2.ConfigureServer(s.server, &http2.Server{}); errHTTP2 != nil {
|
||||
log.Warnf("failed to configure HTTP/2: %v", errHTTP2)
|
||||
}
|
||||
listener = tls.NewListener(listener, tlsConfig)
|
||||
log.Debugf("Starting API server on %s with TLS", addr)
|
||||
} else {
|
||||
log.Debugf("Starting API server on %s", addr)
|
||||
}
|
||||
|
||||
httpListener := newMuxListener(listener.Addr(), 1024)
|
||||
s.muxBaseListener = listener
|
||||
s.muxHTTPListener = httpListener
|
||||
|
||||
httpErrCh := make(chan error, 1)
|
||||
acceptErrCh := make(chan error, 1)
|
||||
|
||||
go func() {
|
||||
httpErrCh <- s.server.Serve(httpListener)
|
||||
}()
|
||||
go func() {
|
||||
acceptErrCh <- s.acceptMuxConnections(listener, httpListener)
|
||||
}()
|
||||
|
||||
select {
|
||||
case errServe := <-httpErrCh:
|
||||
if s.muxBaseListener != nil {
|
||||
if errClose := s.muxBaseListener.Close(); errClose != nil && !errors.Is(errClose, net.ErrClosed) {
|
||||
log.Debugf("failed to close shared listener after HTTP serve exit: %v", errClose)
|
||||
}
|
||||
}
|
||||
if s.muxHTTPListener != nil {
|
||||
_ = s.muxHTTPListener.Close()
|
||||
}
|
||||
errAccept := <-acceptErrCh
|
||||
errServe = normalizeHTTPServeError(errServe)
|
||||
errAccept = normalizeListenerError(errAccept)
|
||||
if errServe != nil {
|
||||
return fmt.Errorf("failed to start HTTP server: %v", errServe)
|
||||
}
|
||||
if errAccept != nil {
|
||||
return fmt.Errorf("failed to start HTTP server: %v", errAccept)
|
||||
}
|
||||
return nil
|
||||
case errAccept := <-acceptErrCh:
|
||||
if s.muxHTTPListener != nil {
|
||||
_ = s.muxHTTPListener.Close()
|
||||
}
|
||||
if s.muxBaseListener != nil {
|
||||
if errClose := s.muxBaseListener.Close(); errClose != nil && !errors.Is(errClose, net.ErrClosed) {
|
||||
log.Debugf("failed to close shared listener after accept loop exit: %v", errClose)
|
||||
}
|
||||
}
|
||||
errServe := <-httpErrCh
|
||||
errServe = normalizeHTTPServeError(errServe)
|
||||
errAccept = normalizeListenerError(errAccept)
|
||||
if errAccept != nil {
|
||||
return fmt.Errorf("failed to start HTTP server: %v", errAccept)
|
||||
}
|
||||
if errServe != nil {
|
||||
return fmt.Errorf("failed to start HTTP server: %v", errServe)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// Stop gracefully shuts down the API server without interrupting any
|
||||
// active connections.
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: The context for graceful shutdown
|
||||
//
|
||||
// Returns:
|
||||
// - error: An error if the server fails to stop
|
||||
func (s *Server) Stop(ctx context.Context) error {
|
||||
log.Debug("Stopping API server...")
|
||||
|
||||
if s.keepAliveEnabled {
|
||||
select {
|
||||
case s.keepAliveStop <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
if s.muxHTTPListener != nil {
|
||||
_ = s.muxHTTPListener.Close()
|
||||
}
|
||||
if s.muxBaseListener != nil {
|
||||
if errClose := s.muxBaseListener.Close(); errClose != nil && !errors.Is(errClose, net.ErrClosed) {
|
||||
log.Debugf("failed to close shared listener: %v", errClose)
|
||||
}
|
||||
}
|
||||
|
||||
// Shutdown the HTTP server.
|
||||
errShutdown := s.server.Shutdown(ctx)
|
||||
if s.codexLiveHandler != nil {
|
||||
s.codexLiveHandler.Close()
|
||||
}
|
||||
if errShutdown != nil {
|
||||
return fmt.Errorf("failed to shutdown HTTP server: %v", errShutdown)
|
||||
}
|
||||
|
||||
log.Debug("API server stopped")
|
||||
return nil
|
||||
}
|
||||
218
backend/internal/api/server_grok_models_test.go
Normal file
218
backend/internal/api/server_grok_models_test.go
Normal file
|
|
@ -0,0 +1,218 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/client/grokbuild"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/home"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
)
|
||||
|
||||
func TestModelsDispatchByGrokShellUserAgent(t *testing.T) {
|
||||
modelRegistry := registry.GetGlobalRegistry()
|
||||
clientID := "test-grok-shell-model-list"
|
||||
modelRegistry.RegisterClient(clientID, "openai", []*registry.ModelInfo{
|
||||
{ID: "grok-shell-openai-model", DisplayName: "Grok Shell Model", ContextLength: 256000, Thinking: ®istry.ThinkingSupport{Levels: []string{"high"}}},
|
||||
})
|
||||
modelRegistry.RegisterClient(clientID+"-claude", "claude", []*registry.ModelInfo{
|
||||
{ID: "grok-shell-claude-model", DisplayName: "Claude Catalog Model", ContextLength: 200000},
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
modelRegistry.UnregisterClient(clientID)
|
||||
modelRegistry.UnregisterClient(clientID + "-claude")
|
||||
})
|
||||
|
||||
server := newTestServer(t)
|
||||
for _, userAgent := range []string{
|
||||
"grok-shell/0.2.119 (macos; aarch64)",
|
||||
"grok-pager/0.2.119 grok-shell/0.2.119 (macos; aarch64)",
|
||||
} {
|
||||
t.Run(userAgent, func(t *testing.T) {
|
||||
req := httptest.NewRequest(http.MethodGet, "https://proxy.example.test/v1/models?client_version", nil)
|
||||
req.Header.Set("Authorization", "Bearer test-key")
|
||||
req.Header.Set("User-Agent", userAgent)
|
||||
recorder := httptest.NewRecorder()
|
||||
server.engine.ServeHTTP(recorder, req)
|
||||
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body = %s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
var response struct {
|
||||
Object string `json:"object"`
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
Model string `json:"model"`
|
||||
Name string `json:"name"`
|
||||
ContextWindow int `json:"context_window"`
|
||||
APIBackend string `json:"api_backend"`
|
||||
SupportedInAPI bool `json:"supported_in_api"`
|
||||
ReasoningEfforts []struct {
|
||||
Value string `json:"value"`
|
||||
} `json:"reasoning_efforts"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
|
||||
t.Fatalf("decode response: %v; body=%s", err, recorder.Body.String())
|
||||
}
|
||||
if response.Object != "list" {
|
||||
t.Fatalf("object = %q, want list", response.Object)
|
||||
}
|
||||
var foundOpenAI, foundClaude bool
|
||||
for _, model := range response.Data {
|
||||
switch model.ID {
|
||||
case "grok-shell-openai-model":
|
||||
foundOpenAI = true
|
||||
if model.Model != model.ID || model.Name != "Grok Shell Model" || model.ContextWindow != 256000 {
|
||||
t.Fatalf("OpenAI model mapping = %#v", model)
|
||||
}
|
||||
if model.APIBackend != "responses" || !model.SupportedInAPI {
|
||||
t.Fatalf("OpenAI model routing fields = %#v", model)
|
||||
}
|
||||
if len(model.ReasoningEfforts) != 1 || model.ReasoningEfforts[0].Value != "high" {
|
||||
t.Fatalf("OpenAI reasoning efforts = %#v", model.ReasoningEfforts)
|
||||
}
|
||||
case "grok-shell-claude-model":
|
||||
foundClaude = true
|
||||
if model.Model != model.ID || model.Name != "Claude Catalog Model" || model.ContextWindow != 200000 {
|
||||
t.Fatalf("Claude model mapping = %#v", model)
|
||||
}
|
||||
if len(model.ReasoningEfforts) != 0 {
|
||||
t.Fatalf("Claude reasoning efforts = %#v, want none", model.ReasoningEfforts)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !foundOpenAI {
|
||||
t.Fatalf("registered OpenAI Grok model missing: %s", recorder.Body.String())
|
||||
}
|
||||
if !foundClaude {
|
||||
t.Fatalf("registered Claude Grok model missing: %s", recorder.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelsDispatchKeepsOrdinaryOpenAIResponse(t *testing.T) {
|
||||
modelRegistry := registry.GetGlobalRegistry()
|
||||
clientID := "test-ordinary-model-list-after-grok"
|
||||
modelRegistry.RegisterClient(clientID, "openai", []*registry.ModelInfo{{ID: "ordinary-model"}})
|
||||
t.Cleanup(func() { modelRegistry.UnregisterClient(clientID) })
|
||||
|
||||
server := newTestServer(t)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
req.Header.Set("Authorization", "Bearer test-key")
|
||||
req.Header.Set("User-Agent", "curl/8.7.1")
|
||||
recorder := httptest.NewRecorder()
|
||||
server.engine.ServeHTTP(recorder, req)
|
||||
|
||||
var response struct {
|
||||
Object string `json:"object"`
|
||||
Data []map[string]any `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
|
||||
t.Fatalf("decode response: %v", err)
|
||||
}
|
||||
if response.Object != "list" {
|
||||
t.Fatalf("object = %q, want list", response.Object)
|
||||
}
|
||||
found := false
|
||||
for _, model := range response.Data {
|
||||
if _, exists := model["api_backend"]; exists {
|
||||
t.Fatalf("ordinary response contains Grok field: %#v", model)
|
||||
}
|
||||
if id, ok := model["id"].(string); ok && id == "ordinary-model" {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("registered ordinary model missing: %s", recorder.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokHomeModelAdapterOmitsReasoning(t *testing.T) {
|
||||
models := grokModelsFromHomeEntries([]homeModelEntry{
|
||||
{id: "home-model", displayName: "Home Model", contextLength: 1234},
|
||||
{id: "home-model-without-context", displayName: "No Context Model"},
|
||||
})
|
||||
if len(models) != 2 {
|
||||
t.Fatalf("Home model count = %d, want 2", len(models))
|
||||
}
|
||||
if models[0].ID != "home-model" || models[0].DisplayName != "Home Model" || models[0].ContextLength != 1234 {
|
||||
t.Fatalf("Home model adapter = %#v", models[0])
|
||||
}
|
||||
if models[1].ID != "home-model-without-context" || models[1].DisplayName != "No Context Model" || models[1].ContextLength != 0 {
|
||||
t.Fatalf("Home zero-context adapter = %#v", models[1])
|
||||
}
|
||||
|
||||
response := grokbuild.BuildResponse(models)
|
||||
if len(response.Data) != 2 || response.Data[0].ReasoningEfforts != nil || response.Data[1].ReasoningEfforts != nil {
|
||||
t.Fatalf("Home reasoning efforts = %#v", response.Data)
|
||||
}
|
||||
|
||||
wire, errMarshal := json.Marshal(response)
|
||||
if errMarshal != nil {
|
||||
t.Fatalf("marshal Home response: %v", errMarshal)
|
||||
}
|
||||
var wireResponse struct {
|
||||
Data []map[string]json.RawMessage `json:"data"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(wire, &wireResponse); errUnmarshal != nil {
|
||||
t.Fatalf("decode Home response JSON: %v; body=%s", errUnmarshal, wire)
|
||||
}
|
||||
if len(wireResponse.Data) != 2 {
|
||||
t.Fatalf("wire Home model count = %d, want 2; body=%s", len(wireResponse.Data), wire)
|
||||
}
|
||||
contextWindow, exists := wireResponse.Data[0]["context_window"]
|
||||
if !exists {
|
||||
t.Fatalf("Home model context_window missing from wire response: %s", wire)
|
||||
}
|
||||
var gotContextWindow int
|
||||
if errDecode := json.Unmarshal(contextWindow, &gotContextWindow); errDecode != nil {
|
||||
t.Fatalf("decode Home context_window: %v", errDecode)
|
||||
}
|
||||
if gotContextWindow != 1234 {
|
||||
t.Fatalf("Home context_window = %d, want 1234", gotContextWindow)
|
||||
}
|
||||
if _, exists := wireResponse.Data[0]["reasoning_efforts"]; exists {
|
||||
t.Fatalf("Home model contains omitted reasoning_efforts: %s", wire)
|
||||
}
|
||||
if _, exists := wireResponse.Data[1]["context_window"]; exists {
|
||||
t.Fatalf("zero-context Home model contains omitted context_window: %s", wire)
|
||||
}
|
||||
if _, exists := wireResponse.Data[1]["reasoning_efforts"]; exists {
|
||||
t.Fatalf("zero-context Home model contains omitted reasoning_efforts: %s", wire)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokModelsPreferHomeOverRegistry(t *testing.T) {
|
||||
previousHome := home.Current()
|
||||
home.ClearCurrent()
|
||||
t.Cleanup(func() { home.SetCurrent(previousHome) })
|
||||
modelRegistry := registry.GetGlobalRegistry()
|
||||
clientID := "test-grok-home-source"
|
||||
modelRegistry.RegisterClient(clientID, "openai", []*registry.ModelInfo{{ID: "local-only-model"}})
|
||||
t.Cleanup(func() { modelRegistry.UnregisterClient(clientID) })
|
||||
|
||||
server := newTestServer(t)
|
||||
server.cfg.Home.Enabled = true
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
req.Header.Set("Authorization", "Bearer test-key")
|
||||
req.Header.Set("User-Agent", "grok-shell/0.2.119")
|
||||
recorder := httptest.NewRecorder()
|
||||
ginContext, _ := gin.CreateTestContext(recorder)
|
||||
ginContext.Request = req
|
||||
server.handleGrokModels(ginContext)
|
||||
if recorder.Code != http.StatusServiceUnavailable {
|
||||
t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusServiceUnavailable, recorder.Body.String())
|
||||
}
|
||||
if !strings.Contains(recorder.Body.String(), "home control center unavailable") {
|
||||
t.Fatalf("Home failure response missing expected error: %s", recorder.Body.String())
|
||||
}
|
||||
if strings.Contains(recorder.Body.String(), "local-only-model") {
|
||||
t.Fatalf("Home failure response leaked local registry model: %s", recorder.Body.String())
|
||||
}
|
||||
}
|
||||
89
backend/internal/api/server_keepalive.go
Normal file
89
backend/internal/api/server_keepalive.go
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"crypto/subtle"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
func (s *Server) enableKeepAlive(timeout time.Duration, onTimeout func()) {
|
||||
if timeout <= 0 || onTimeout == nil {
|
||||
return
|
||||
}
|
||||
|
||||
s.keepAliveEnabled = true
|
||||
s.keepAliveTimeout = timeout
|
||||
s.keepAliveOnTimeout = onTimeout
|
||||
s.keepAliveHeartbeat = make(chan struct{}, 1)
|
||||
s.keepAliveStop = make(chan struct{}, 1)
|
||||
|
||||
s.engine.GET("/keep-alive", s.handleKeepAlive)
|
||||
|
||||
go s.watchKeepAlive()
|
||||
}
|
||||
|
||||
func (s *Server) handleKeepAlive(c *gin.Context) {
|
||||
if s.localPassword != "" {
|
||||
provided := strings.TrimSpace(c.GetHeader("Authorization"))
|
||||
if provided != "" {
|
||||
parts := strings.SplitN(provided, " ", 2)
|
||||
if len(parts) == 2 && strings.EqualFold(parts[0], "bearer") {
|
||||
provided = parts[1]
|
||||
}
|
||||
}
|
||||
if provided == "" {
|
||||
provided = strings.TrimSpace(c.GetHeader("X-Local-Password"))
|
||||
}
|
||||
if subtle.ConstantTimeCompare([]byte(provided), []byte(s.localPassword)) != 1 {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "invalid password"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
s.signalKeepAlive()
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||||
}
|
||||
|
||||
func (s *Server) signalKeepAlive() {
|
||||
if !s.keepAliveEnabled {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case s.keepAliveHeartbeat <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) watchKeepAlive() {
|
||||
if !s.keepAliveEnabled {
|
||||
return
|
||||
}
|
||||
|
||||
timer := time.NewTimer(s.keepAliveTimeout)
|
||||
defer timer.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-timer.C:
|
||||
log.Warnf("keep-alive endpoint idle for %s, shutting down", s.keepAliveTimeout)
|
||||
if s.keepAliveOnTimeout != nil {
|
||||
s.keepAliveOnTimeout()
|
||||
}
|
||||
return
|
||||
case <-s.keepAliveHeartbeat:
|
||||
if !timer.Stop() {
|
||||
select {
|
||||
case <-timer.C:
|
||||
default:
|
||||
}
|
||||
}
|
||||
timer.Reset(s.keepAliveTimeout)
|
||||
case <-s.keepAliveStop:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
362
backend/internal/api/server_management.go
Normal file
362
backend/internal/api/server_management.go
Normal file
|
|
@ -0,0 +1,362 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io/fs"
|
||||
"net/http"
|
||||
"path"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/managementasset"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
var hashedManagementAssetPattern = regexp.MustCompile(`(^|[-.])[A-Za-z0-9_-]{8,}(\.[^.]+)+$`)
|
||||
|
||||
func (s *Server) registerManagementRoutes() {
|
||||
if s == nil || s.engine == nil || s.mgmt == nil {
|
||||
return
|
||||
}
|
||||
if !s.managementRoutesRegistered.CompareAndSwap(false, true) {
|
||||
return
|
||||
}
|
||||
|
||||
log.Info("management routes registered after secret key configuration")
|
||||
|
||||
s.engine.POST("/v0/management/oauth-callback", s.managementAvailabilityMiddleware(), s.mgmt.PostOAuthCallback)
|
||||
s.engine.GET("/v0/management/oauth-callback", s.managementAvailabilityMiddleware(), s.mgmt.GetOAuthCallback)
|
||||
|
||||
mgmt := s.engine.Group("/v0/management")
|
||||
mgmt.Use(s.managementAvailabilityMiddleware(), s.mgmt.Middleware())
|
||||
{
|
||||
mgmt.GET("/config", s.mgmt.GetConfig)
|
||||
mgmt.GET("/config.yaml", s.mgmt.GetConfigYAML)
|
||||
mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML)
|
||||
mgmt.GET("/latest-version", s.mgmt.GetLatestVersion)
|
||||
mgmt.GET("/plugins", s.mgmt.ListPlugins)
|
||||
mgmt.GET("/plugin-store", s.mgmt.ListPluginStore)
|
||||
mgmt.POST("/plugin-store/:id/install", s.mgmt.InstallPluginFromStore)
|
||||
mgmt.DELETE("/plugins/:id", s.mgmt.DeletePlugin)
|
||||
mgmt.PATCH("/plugins/:id/enabled", s.mgmt.PatchPluginEnabled)
|
||||
mgmt.GET("/plugins/:id/config", s.mgmt.GetPluginConfig)
|
||||
mgmt.PUT("/plugins/:id/config", s.mgmt.PutPluginConfig)
|
||||
mgmt.PATCH("/plugins/:id/config", s.mgmt.PatchPluginConfig)
|
||||
|
||||
mgmt.GET("/debug", s.mgmt.GetDebug)
|
||||
mgmt.PUT("/debug", s.mgmt.PutDebug)
|
||||
mgmt.PATCH("/debug", s.mgmt.PutDebug)
|
||||
|
||||
mgmt.GET("/logging-to-file", s.mgmt.GetLoggingToFile)
|
||||
mgmt.PUT("/logging-to-file", s.mgmt.PutLoggingToFile)
|
||||
mgmt.PATCH("/logging-to-file", s.mgmt.PutLoggingToFile)
|
||||
|
||||
mgmt.GET("/logs-max-total-size-mb", s.mgmt.GetLogsMaxTotalSizeMB)
|
||||
mgmt.PUT("/logs-max-total-size-mb", s.mgmt.PutLogsMaxTotalSizeMB)
|
||||
mgmt.PATCH("/logs-max-total-size-mb", s.mgmt.PutLogsMaxTotalSizeMB)
|
||||
|
||||
mgmt.GET("/error-logs-max-files", s.mgmt.GetErrorLogsMaxFiles)
|
||||
mgmt.PUT("/error-logs-max-files", s.mgmt.PutErrorLogsMaxFiles)
|
||||
mgmt.PATCH("/error-logs-max-files", s.mgmt.PutErrorLogsMaxFiles)
|
||||
|
||||
mgmt.GET("/usage-statistics-enabled", s.mgmt.GetUsageStatisticsEnabled)
|
||||
mgmt.PUT("/usage-statistics-enabled", s.mgmt.PutUsageStatisticsEnabled)
|
||||
mgmt.PATCH("/usage-statistics-enabled", s.mgmt.PutUsageStatisticsEnabled)
|
||||
|
||||
mgmt.GET("/proxy-url", s.mgmt.GetProxyURL)
|
||||
mgmt.PUT("/proxy-url", s.mgmt.PutProxyURL)
|
||||
mgmt.PATCH("/proxy-url", s.mgmt.PutProxyURL)
|
||||
mgmt.DELETE("/proxy-url", s.mgmt.DeleteProxyURL)
|
||||
|
||||
mgmt.POST("/api-call", s.mgmt.APICall)
|
||||
|
||||
mgmt.GET("/quota-exceeded/switch-project", s.mgmt.GetSwitchProject)
|
||||
mgmt.PUT("/quota-exceeded/switch-project", s.mgmt.PutSwitchProject)
|
||||
mgmt.PATCH("/quota-exceeded/switch-project", s.mgmt.PutSwitchProject)
|
||||
|
||||
mgmt.GET("/quota-exceeded/switch-preview-model", s.mgmt.GetSwitchPreviewModel)
|
||||
mgmt.PUT("/quota-exceeded/switch-preview-model", s.mgmt.PutSwitchPreviewModel)
|
||||
mgmt.PATCH("/quota-exceeded/switch-preview-model", s.mgmt.PutSwitchPreviewModel)
|
||||
mgmt.POST("/reset-quota", s.mgmt.ResetQuota)
|
||||
|
||||
mgmt.GET("/api-keys", s.mgmt.GetAPIKeys)
|
||||
mgmt.PUT("/api-keys", s.mgmt.PutAPIKeys)
|
||||
mgmt.PATCH("/api-keys", s.mgmt.PatchAPIKeys)
|
||||
mgmt.DELETE("/api-keys", s.mgmt.DeleteAPIKeys)
|
||||
mgmt.GET("/api-key-usage", s.mgmt.GetAPIKeyUsage)
|
||||
mgmt.GET("/usage-queue", s.mgmt.GetUsageQueue)
|
||||
|
||||
mgmt.GET("/gemini-api-key", s.mgmt.GetGeminiKeys)
|
||||
mgmt.PUT("/gemini-api-key", s.mgmt.PutGeminiKeys)
|
||||
mgmt.PATCH("/gemini-api-key", s.mgmt.PatchGeminiKey)
|
||||
mgmt.DELETE("/gemini-api-key", s.mgmt.DeleteGeminiKey)
|
||||
|
||||
mgmt.GET("/interactions-api-key", s.mgmt.GetInteractionsKeys)
|
||||
mgmt.PUT("/interactions-api-key", s.mgmt.PutInteractionsKeys)
|
||||
mgmt.PATCH("/interactions-api-key", s.mgmt.PatchInteractionsKey)
|
||||
mgmt.DELETE("/interactions-api-key", s.mgmt.DeleteInteractionsKey)
|
||||
|
||||
mgmt.GET("/logs", s.mgmt.GetLogs)
|
||||
mgmt.DELETE("/logs", s.mgmt.DeleteLogs)
|
||||
mgmt.GET("/request-error-logs", s.mgmt.GetRequestErrorLogs)
|
||||
mgmt.GET("/request-error-logs/:name", s.mgmt.DownloadRequestErrorLog)
|
||||
mgmt.GET("/request-log-by-id/:id", s.mgmt.GetRequestLogByID)
|
||||
mgmt.GET("/request-log", s.mgmt.GetRequestLog)
|
||||
mgmt.PUT("/request-log", s.mgmt.PutRequestLog)
|
||||
mgmt.PATCH("/request-log", s.mgmt.PutRequestLog)
|
||||
mgmt.GET("/ws-auth", s.mgmt.GetWebsocketAuth)
|
||||
mgmt.PUT("/ws-auth", s.mgmt.PutWebsocketAuth)
|
||||
mgmt.PATCH("/ws-auth", s.mgmt.PutWebsocketAuth)
|
||||
|
||||
mgmt.GET("/request-retry", s.mgmt.GetRequestRetry)
|
||||
mgmt.PUT("/request-retry", s.mgmt.PutRequestRetry)
|
||||
mgmt.PATCH("/request-retry", s.mgmt.PutRequestRetry)
|
||||
mgmt.GET("/max-retry-credentials", s.mgmt.GetMaxRetryCredentials)
|
||||
mgmt.PUT("/max-retry-credentials", s.mgmt.PutMaxRetryCredentials)
|
||||
mgmt.PATCH("/max-retry-credentials", s.mgmt.PutMaxRetryCredentials)
|
||||
mgmt.GET("/max-retry-interval", s.mgmt.GetMaxRetryInterval)
|
||||
mgmt.PUT("/max-retry-interval", s.mgmt.PutMaxRetryInterval)
|
||||
mgmt.PATCH("/max-retry-interval", s.mgmt.PutMaxRetryInterval)
|
||||
|
||||
mgmt.GET("/force-model-prefix", s.mgmt.GetForceModelPrefix)
|
||||
mgmt.PUT("/force-model-prefix", s.mgmt.PutForceModelPrefix)
|
||||
mgmt.PATCH("/force-model-prefix", s.mgmt.PutForceModelPrefix)
|
||||
|
||||
mgmt.GET("/routing/strategy", s.mgmt.GetRoutingStrategy)
|
||||
mgmt.PUT("/routing/strategy", s.mgmt.PutRoutingStrategy)
|
||||
mgmt.PATCH("/routing/strategy", s.mgmt.PutRoutingStrategy)
|
||||
|
||||
mgmt.GET("/claude-api-key", s.mgmt.GetClaudeKeys)
|
||||
mgmt.PUT("/claude-api-key", s.mgmt.PutClaudeKeys)
|
||||
mgmt.PATCH("/claude-api-key", s.mgmt.PatchClaudeKey)
|
||||
mgmt.DELETE("/claude-api-key", s.mgmt.DeleteClaudeKey)
|
||||
|
||||
mgmt.GET("/codex-api-key", s.mgmt.GetCodexKeys)
|
||||
mgmt.PUT("/codex-api-key", s.mgmt.PutCodexKeys)
|
||||
mgmt.PATCH("/codex-api-key", s.mgmt.PatchCodexKey)
|
||||
mgmt.DELETE("/codex-api-key", s.mgmt.DeleteCodexKey)
|
||||
|
||||
mgmt.GET("/xai-api-key", s.mgmt.GetXAIKeys)
|
||||
mgmt.PUT("/xai-api-key", s.mgmt.PutXAIKeys)
|
||||
mgmt.PATCH("/xai-api-key", s.mgmt.PatchXAIKey)
|
||||
mgmt.DELETE("/xai-api-key", s.mgmt.DeleteXAIKey)
|
||||
|
||||
mgmt.GET("/openai-compatibility", s.mgmt.GetOpenAICompat)
|
||||
mgmt.PUT("/openai-compatibility", s.mgmt.PutOpenAICompat)
|
||||
mgmt.PATCH("/openai-compatibility", s.mgmt.PatchOpenAICompat)
|
||||
mgmt.DELETE("/openai-compatibility", s.mgmt.DeleteOpenAICompat)
|
||||
|
||||
mgmt.GET("/vertex-api-key", s.mgmt.GetVertexCompatKeys)
|
||||
mgmt.PUT("/vertex-api-key", s.mgmt.PutVertexCompatKeys)
|
||||
mgmt.PATCH("/vertex-api-key", s.mgmt.PatchVertexCompatKey)
|
||||
mgmt.DELETE("/vertex-api-key", s.mgmt.DeleteVertexCompatKey)
|
||||
|
||||
mgmt.GET("/oauth-excluded-models", s.mgmt.GetOAuthExcludedModels)
|
||||
mgmt.PUT("/oauth-excluded-models", s.mgmt.PutOAuthExcludedModels)
|
||||
mgmt.PATCH("/oauth-excluded-models", s.mgmt.PatchOAuthExcludedModels)
|
||||
mgmt.DELETE("/oauth-excluded-models", s.mgmt.DeleteOAuthExcludedModels)
|
||||
|
||||
mgmt.GET("/oauth-model-alias", s.mgmt.GetOAuthModelAlias)
|
||||
mgmt.PUT("/oauth-model-alias", s.mgmt.PutOAuthModelAlias)
|
||||
mgmt.PATCH("/oauth-model-alias", s.mgmt.PatchOAuthModelAlias)
|
||||
mgmt.DELETE("/oauth-model-alias", s.mgmt.DeleteOAuthModelAlias)
|
||||
|
||||
mgmt.GET("/oauth-request-scoped-errors", s.mgmt.GetOAuthRequestScopedErrors)
|
||||
mgmt.PUT("/oauth-request-scoped-errors", s.mgmt.PutOAuthRequestScopedErrors)
|
||||
mgmt.PATCH("/oauth-request-scoped-errors", s.mgmt.PatchOAuthRequestScopedErrors)
|
||||
mgmt.DELETE("/oauth-request-scoped-errors", s.mgmt.DeleteOAuthRequestScopedErrors)
|
||||
|
||||
mgmt.GET("/auth-files", s.mgmt.ListAuthFiles)
|
||||
mgmt.GET("/auth-files/models", s.mgmt.GetAuthFileModels)
|
||||
mgmt.GET("/model-definitions/:channel", s.mgmt.GetStaticModelDefinitions)
|
||||
mgmt.GET("/auth-files/download", s.mgmt.DownloadAuthFile)
|
||||
mgmt.POST("/auth-files", s.mgmt.UploadAuthFile)
|
||||
mgmt.DELETE("/auth-files", s.mgmt.DeleteAuthFile)
|
||||
mgmt.PATCH("/auth-files/status", s.mgmt.PatchAuthFileStatus)
|
||||
mgmt.PATCH("/auth-files/fields", s.mgmt.PatchAuthFileFields)
|
||||
mgmt.POST("/vertex/import", s.mgmt.ImportVertexCredential)
|
||||
|
||||
mgmt.GET("/anthropic-auth-url", s.mgmt.RequestAnthropicToken)
|
||||
mgmt.GET("/codex-auth-url", s.mgmt.RequestCodexToken)
|
||||
mgmt.GET("/antigravity-auth-url", s.mgmt.RequestAntigravityToken)
|
||||
mgmt.GET("/kimi-auth-url", s.mgmt.RequestKimiToken)
|
||||
mgmt.GET("/xai-auth-url", s.mgmt.RequestXAIToken)
|
||||
mgmt.GET("/get-auth-status", s.mgmt.GetAuthStatus)
|
||||
mgmt.DELETE("/oauth-session", s.mgmt.CancelAuthSession)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) managementAvailabilityMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if !s.managementAvailable(c) {
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) managementAvailable(c *gin.Context) bool {
|
||||
if s == nil || s.cfg == nil {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return false
|
||||
}
|
||||
if s.cfg.Home.Enabled {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return false
|
||||
}
|
||||
if !s.managementRoutesEnabled.Load() {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *Server) refreshPluginManagementRoutes() {
|
||||
if s == nil || s.pluginHost == nil || s.engine == nil {
|
||||
return
|
||||
}
|
||||
s.pluginHost.RegisterManagementRoutes(context.Background(), s.registeredManagementRouteKeys())
|
||||
}
|
||||
|
||||
// RefreshPluginManagementRoutes rebuilds plugin-owned Management API routes.
|
||||
func (s *Server) RefreshPluginManagementRoutes() {
|
||||
s.refreshPluginManagementRoutes()
|
||||
}
|
||||
|
||||
func (s *Server) registeredManagementRouteKeys() map[string]struct{} {
|
||||
out := make(map[string]struct{})
|
||||
if s == nil || s.engine == nil {
|
||||
return out
|
||||
}
|
||||
for _, route := range s.engine.Routes() {
|
||||
if strings.HasPrefix(route.Path, "/v0/management/") || route.Path == "/v0/management" {
|
||||
out[strings.ToUpper(strings.TrimSpace(route.Method))+" "+route.Path] = struct{}{}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (s *Server) pluginManagementNoRoute(c *gin.Context) {
|
||||
if s == nil || c == nil || c.Request == nil || c.Request.URL == nil {
|
||||
if c != nil {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
}
|
||||
return
|
||||
}
|
||||
path := c.Request.URL.Path
|
||||
if strings.HasPrefix(path, "/v0/resource/plugins/") {
|
||||
s.pluginResourceNoRoute(c)
|
||||
return
|
||||
}
|
||||
if path != "/v0/management" && !strings.HasPrefix(path, "/v0/management/") {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
if s.pluginHost == nil || s.mgmt == nil {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
if !s.managementAvailable(c) {
|
||||
return
|
||||
}
|
||||
s.mgmt.Middleware()(c)
|
||||
if c.IsAborted() {
|
||||
return
|
||||
}
|
||||
if s.mgmt.ServePluginAuthURL(c) {
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if s.pluginHost.ServeManagementHTTP(c.Writer, c.Request) {
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
}
|
||||
|
||||
func (s *Server) pluginResourceNoRoute(c *gin.Context) {
|
||||
if s == nil || c == nil || c.Request == nil || c.Request.URL == nil {
|
||||
if c != nil {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
}
|
||||
return
|
||||
}
|
||||
if s.cfg == nil || s.cfg.Home.Enabled || s.pluginHost == nil {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
if s.pluginHost.ServeResourceHTTP(c.Writer, c.Request) {
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
}
|
||||
|
||||
func (s *Server) serveManagementControlPanel(c *gin.Context) {
|
||||
cfg := s.cfg
|
||||
if cfg == nil || cfg.Home.Enabled || cfg.RemoteManagement.DisableControlPanel {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
source, err := managementasset.Current()
|
||||
if err != nil {
|
||||
if !errors.Is(err, fs.ErrNotExist) {
|
||||
log.WithError(err).Error("failed to resolve management control panel assets")
|
||||
}
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
s.serveManagementFile(c, source, source.Entry, false)
|
||||
}
|
||||
|
||||
func (s *Server) serveManagementAsset(c *gin.Context) {
|
||||
cfg := s.cfg
|
||||
if cfg == nil || cfg.Home.Enabled || cfg.RemoteManagement.DisableControlPanel {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
assetPath := strings.TrimPrefix(c.Param("filepath"), "/")
|
||||
if assetPath == "" || strings.Contains(assetPath, `\`) || !fs.ValidPath(assetPath) {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
source, err := managementasset.Current()
|
||||
if err != nil {
|
||||
if !errors.Is(err, fs.ErrNotExist) {
|
||||
log.WithError(err).Error("failed to resolve management control panel assets")
|
||||
}
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
s.serveManagementFile(c, source, assetPath, true)
|
||||
}
|
||||
|
||||
func (s *Server) serveManagementFile(c *gin.Context, source *managementasset.Source, name string, cacheHashed bool) {
|
||||
file, err := source.FileSystem.Open(name)
|
||||
if err != nil {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if errClose := file.Close(); errClose != nil {
|
||||
log.WithError(errClose).Debug("failed to close management asset")
|
||||
}
|
||||
}()
|
||||
|
||||
info, err := file.Stat()
|
||||
if err != nil || info.IsDir() {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
if !cacheHashed {
|
||||
c.Header("Content-Type", "text/html; charset=utf-8")
|
||||
} else if hashedManagementAssetPattern.MatchString(path.Base(name)) {
|
||||
c.Header("Cache-Control", "public, max-age=31536000, immutable")
|
||||
}
|
||||
http.ServeContent(c.Writer, c.Request, info.Name(), info.ModTime(), file)
|
||||
}
|
||||
233
backend/internal/api/server_middleware.go
Normal file
233
backend/internal/api/server_middleware.go
Normal file
|
|
@ -0,0 +1,233 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
codexlive "github.com/router-for-me/CLIProxyAPI/v7/internal/client/codex/live"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/home"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/safemode"
|
||||
sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
var corsExposedResponseHeaders = []string{
|
||||
logging.CPATraceIDHeader,
|
||||
"X-CPA-VERSION",
|
||||
"X-CPA-COMMIT",
|
||||
"X-CPA-BUILD-DATE",
|
||||
"X-CPA-SUPPORT-PLUGIN",
|
||||
"X-CPA-HOME-VERSION",
|
||||
"X-CPA-HOME-BUILD-DATE",
|
||||
"X-SERVER-VERSION",
|
||||
"X-SERVER-BUILD-DATE",
|
||||
"Location",
|
||||
"Retry-After",
|
||||
"X-Request-Id",
|
||||
"OpenAI-Request-Id",
|
||||
}
|
||||
|
||||
var corsExposedResponseHeadersJoined = strings.Join(corsExposedResponseHeaders, ", ")
|
||||
|
||||
const (
|
||||
exampleAPIKeyManagementPath = "/management.html"
|
||||
exampleAPIKeyManagementURL = "/management.html?safe-mode=configure"
|
||||
)
|
||||
|
||||
func (s *Server) homeHeartbeatMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if s == nil || s.cfg == nil || !s.cfg.Home.Enabled {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if c != nil && c.Request != nil {
|
||||
path := c.Request.URL.Path
|
||||
if strings.HasPrefix(path, "/v0/management/") || path == "/v0/management" || strings.HasPrefix(path, "/v0/resource/plugins/") || path == "/management.html" {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
}
|
||||
client := home.Current()
|
||||
if client == nil || !client.HeartbeatOK() {
|
||||
c.AbortWithStatus(http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) exampleAPIKeySafeModeRequired(cfg *config.Config) bool {
|
||||
return s != nil && s.exampleAPIKeySafeModeEnabled && cfg != nil && safemode.HasExampleAPIKeys(cfg.APIKeys)
|
||||
}
|
||||
|
||||
func (s *Server) exampleAPIKeySafeModeMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if s == nil || !s.exampleAPIKeySafeModeActive.Load() || c == nil || c.Request == nil || c.Request.URL == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
path := c.Request.URL.Path
|
||||
if (path == exampleAPIKeyManagementPath && c.Query("safe-mode") == "configure") || strings.HasPrefix(path, "/management-assets/") {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if (path == "/" || path == exampleAPIKeyManagementPath) && (c.Request.Method == http.MethodGet || c.Request.Method == http.MethodHead) {
|
||||
s.serveExampleAPIKeyWarningPage(c)
|
||||
return
|
||||
}
|
||||
if !isExampleAPIKeySafeModeProxyPath(path) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
c.Header("X-CPA-SAFE-MODE", "example-api-key")
|
||||
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
|
||||
"error": "unsafe_example_api_key",
|
||||
"message": "Proxy API endpoints are disabled because api-keys contains template values. Open /management.html?safe-mode=configure, update api-keys in Management, then retry.",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) serveExampleAPIKeyWarningPage(c *gin.Context) {
|
||||
cfg := s.cfg
|
||||
var keys []string
|
||||
if cfg != nil {
|
||||
keys = safemode.ExampleAPIKeys(cfg.APIKeys)
|
||||
}
|
||||
c.Header("Content-Type", "text/html; charset=utf-8")
|
||||
c.Header("Cache-Control", "no-store")
|
||||
if c.Request.Method == http.MethodHead {
|
||||
c.Status(http.StatusOK)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
c.String(http.StatusOK, safemode.ExampleAPIKeyWarningPageHTML(keys, exampleAPIKeyManagementURL))
|
||||
c.Abort()
|
||||
}
|
||||
|
||||
func isExampleAPIKeySafeModeProxyPath(path string) bool {
|
||||
switch {
|
||||
case path == "/v1" || strings.HasPrefix(path, "/v1/"):
|
||||
return true
|
||||
case path == "/v1beta" || strings.HasPrefix(path, "/v1beta/"):
|
||||
return true
|
||||
case path == "/openai/v1" || strings.HasPrefix(path, "/openai/v1/"):
|
||||
return true
|
||||
case path == "/backend-api/codex" || strings.HasPrefix(path, "/backend-api/codex/"):
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// corsMiddleware returns a Gin middleware handler that adds CORS headers
|
||||
// to every response, allowing cross-origin requests.
|
||||
//
|
||||
// Returns:
|
||||
// - gin.HandlerFunc: The CORS middleware handler
|
||||
func corsMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS")
|
||||
c.Header("Access-Control-Allow-Headers", "*")
|
||||
c.Header("Access-Control-Expose-Headers", corsExposedResponseHeadersJoined)
|
||||
|
||||
if c.Request.Method == "OPTIONS" {
|
||||
c.AbortWithStatus(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// AuthMiddleware returns a Gin middleware handler that authenticates requests
|
||||
// using the configured authentication providers. When no providers are available,
|
||||
// it allows all requests (legacy behaviour).
|
||||
func AuthMiddleware(manager *sdkaccess.Manager) gin.HandlerFunc {
|
||||
return accessAuthMiddleware(manager, false)
|
||||
}
|
||||
|
||||
func realtimeStandardAuthMiddleware(manager *sdkaccess.Manager) gin.HandlerFunc {
|
||||
return accessAuthMiddleware(manager, true)
|
||||
}
|
||||
|
||||
func accessAuthMiddleware(manager *sdkaccess.Manager, realtimeError bool) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if manager == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
result, err := manager.Authenticate(c.Request.Context(), c.Request)
|
||||
if err == nil {
|
||||
if result != nil {
|
||||
c.Set("userApiKey", result.Principal)
|
||||
c.Set("accessProvider", result.Provider)
|
||||
if len(result.Metadata) > 0 {
|
||||
c.Set("accessMetadata", result.Metadata)
|
||||
}
|
||||
}
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
statusCode := err.HTTPStatusCode()
|
||||
if statusCode >= http.StatusInternalServerError {
|
||||
log.Errorf("authentication middleware error: %v", err)
|
||||
}
|
||||
if realtimeError {
|
||||
errorType := "authentication_error"
|
||||
code := "invalid_api_key"
|
||||
if statusCode >= http.StatusInternalServerError {
|
||||
errorType = "server_error"
|
||||
code = "authentication_service_error"
|
||||
}
|
||||
c.AbortWithStatusJSON(statusCode, gin.H{"error": gin.H{
|
||||
"message": err.Message,
|
||||
"type": errorType,
|
||||
"param": nil,
|
||||
"code": code,
|
||||
}})
|
||||
return
|
||||
}
|
||||
c.AbortWithStatusJSON(statusCode, gin.H{"error": err.Message})
|
||||
}
|
||||
}
|
||||
|
||||
func realtimeAuthMiddleware(manager *sdkaccess.Manager, handler *codexlive.Handler) gin.HandlerFunc {
|
||||
fallback := realtimeStandardAuthMiddleware(manager)
|
||||
return func(c *gin.Context) {
|
||||
authorization, matched, errAuthenticate := handler.AuthenticateClientSecret(c.Request)
|
||||
if !matched {
|
||||
fallback(c)
|
||||
return
|
||||
}
|
||||
if errAuthenticate != nil {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": gin.H{
|
||||
"message": errAuthenticate.Error(),
|
||||
"type": "invalid_request_error",
|
||||
"param": nil,
|
||||
"code": "invalid_realtime_client_secret",
|
||||
}})
|
||||
return
|
||||
}
|
||||
principal := authorization.IssuerPrincipal
|
||||
if principal == "" {
|
||||
principal = authorization.Principal
|
||||
}
|
||||
provider := authorization.IssuerProvider
|
||||
if provider == "" {
|
||||
provider = "realtime-client-secret"
|
||||
}
|
||||
c.Set("userApiKey", principal)
|
||||
c.Set("accessProvider", provider)
|
||||
c.Set(codexlive.ClientSecretSessionContextKey, authorization.Session)
|
||||
c.Set(codexlive.ClientSecretPrincipalContextKey, authorization.Principal)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
135
backend/internal/api/server_options.go
Normal file
135
backend/internal/api/server_options.go
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
type serverOptionConfig struct {
|
||||
extraMiddleware []gin.HandlerFunc
|
||||
engineConfigurator func(*gin.Engine)
|
||||
routerConfigurator func(*gin.Engine, *handlers.BaseAPIHandler, *config.Config)
|
||||
requestLoggerFactory func(*config.Config, string) logging.RequestLogger
|
||||
localPassword string
|
||||
keepAliveEnabled bool
|
||||
keepAliveTimeout time.Duration
|
||||
keepAliveOnTimeout func()
|
||||
postAuthHook auth.PostAuthHook
|
||||
postAuthPersistHook auth.PostAuthHook
|
||||
pluginHost *pluginhost.Host
|
||||
configReloadHook func(context.Context, *config.Config)
|
||||
exampleAPIKeySafeMode bool
|
||||
}
|
||||
|
||||
// ServerOption customises HTTP server construction.
|
||||
type ServerOption func(*serverOptionConfig)
|
||||
|
||||
func defaultRequestLoggerFactory(cfg *config.Config, configPath string) logging.RequestLogger {
|
||||
configDir := filepath.Dir(configPath)
|
||||
logsDir := logging.ResolveLogDirectory(cfg)
|
||||
logger := logging.NewFileRequestLogger(cfg.RequestLog, logsDir, configDir, cfg.ErrorLogsMaxFiles)
|
||||
logger.SetHomeEnabled(cfg != nil && cfg.Home.Enabled)
|
||||
return logger
|
||||
}
|
||||
|
||||
func effectiveSDKConfig(cfg *config.Config) *config.SDKConfig {
|
||||
if cfg == nil {
|
||||
return nil
|
||||
}
|
||||
sdkCfg := cfg.SDKConfig
|
||||
sdkCfg.CodexOptimizeMultiAgentV2 = cfg.Codex.OptimizeMultiAgentV2
|
||||
if cfg.CommercialMode {
|
||||
sdkCfg.RequestLog = false
|
||||
}
|
||||
return &sdkCfg
|
||||
}
|
||||
|
||||
// WithMiddleware appends additional Gin middleware during server construction.
|
||||
func WithMiddleware(mw ...gin.HandlerFunc) ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
cfg.extraMiddleware = append(cfg.extraMiddleware, mw...)
|
||||
}
|
||||
}
|
||||
|
||||
// WithEngineConfigurator allows callers to mutate the Gin engine prior to middleware setup.
|
||||
func WithEngineConfigurator(fn func(*gin.Engine)) ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
cfg.engineConfigurator = fn
|
||||
}
|
||||
}
|
||||
|
||||
// WithRouterConfigurator appends a callback after default routes are registered.
|
||||
func WithRouterConfigurator(fn func(*gin.Engine, *handlers.BaseAPIHandler, *config.Config)) ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
cfg.routerConfigurator = fn
|
||||
}
|
||||
}
|
||||
|
||||
// WithLocalManagementPassword stores a runtime-only management password accepted for localhost requests.
|
||||
func WithLocalManagementPassword(password string) ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
cfg.localPassword = password
|
||||
}
|
||||
}
|
||||
|
||||
// WithKeepAliveEndpoint enables a keep-alive endpoint with the provided timeout and callback.
|
||||
func WithKeepAliveEndpoint(timeout time.Duration, onTimeout func()) ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
if timeout <= 0 || onTimeout == nil {
|
||||
return
|
||||
}
|
||||
cfg.keepAliveEnabled = true
|
||||
cfg.keepAliveTimeout = timeout
|
||||
cfg.keepAliveOnTimeout = onTimeout
|
||||
}
|
||||
}
|
||||
|
||||
// WithRequestLoggerFactory customises request logger creation.
|
||||
func WithRequestLoggerFactory(factory func(*config.Config, string) logging.RequestLogger) ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
cfg.requestLoggerFactory = factory
|
||||
}
|
||||
}
|
||||
|
||||
// WithPostAuthHook registers a hook to be called after auth record creation.
|
||||
func WithPostAuthHook(hook auth.PostAuthHook) ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
cfg.postAuthHook = hook
|
||||
}
|
||||
}
|
||||
|
||||
// WithPostAuthPersistHook registers a hook to be called after auth persistence.
|
||||
func WithPostAuthPersistHook(hook auth.PostAuthHook) ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
cfg.postAuthPersistHook = hook
|
||||
}
|
||||
}
|
||||
|
||||
// WithPluginHost registers dynamic plugin HTTP adapters with the server.
|
||||
func WithPluginHost(host *pluginhost.Host) ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
cfg.pluginHost = host
|
||||
}
|
||||
}
|
||||
|
||||
// WithConfigReloadHook registers a callback used after management saves config changes.
|
||||
func WithConfigReloadHook(hook func(context.Context, *config.Config)) ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
cfg.configReloadHook = hook
|
||||
}
|
||||
}
|
||||
|
||||
// WithExampleAPIKeySafeMode blocks proxy API endpoints while template API keys remain configured.
|
||||
func WithExampleAPIKeySafeMode() ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
cfg.exampleAPIKeySafeMode = true
|
||||
}
|
||||
}
|
||||
274
backend/internal/api/server_reload.go
Normal file
274
backend/internal/api/server_reload.go
Normal file
|
|
@ -0,0 +1,274 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/access"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/cache"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
func (s *Server) applyAccessConfig(oldCfg, newCfg *config.Config) bool {
|
||||
if s == nil || s.accessManager == nil || newCfg == nil {
|
||||
return false
|
||||
}
|
||||
if _, err := access.ApplyAccessProviders(s.accessManager, oldCfg, newCfg); err != nil {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// UpdateClients updates the server's client list and configuration.
|
||||
// This method is called when the configuration or authentication tokens change.
|
||||
//
|
||||
// Parameters:
|
||||
// - clients: The new slice of AI service clients
|
||||
// - cfg: The new application configuration
|
||||
func (s *Server) UpdateClients(cfg *config.Config) {
|
||||
s.UpdateClientsContext(context.Background(), cfg)
|
||||
}
|
||||
|
||||
// UpdateClientsContext updates runtime clients while honoring cancellation between
|
||||
// short configuration and filesystem operations.
|
||||
func (s *Server) UpdateClientsContext(ctx context.Context, cfg *config.Config) bool {
|
||||
if s == nil || cfg == nil {
|
||||
return false
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if errContext := ctx.Err(); errContext != nil {
|
||||
return false
|
||||
}
|
||||
// Reconstruct old config from YAML snapshot to avoid reference sharing issues
|
||||
var oldCfg *config.Config
|
||||
if len(s.oldConfigYaml) > 0 {
|
||||
_ = yaml.Unmarshal(s.oldConfigYaml, &oldCfg)
|
||||
}
|
||||
|
||||
// Update request logger enabled state if it has changed
|
||||
previousRequestLog := false
|
||||
if oldCfg != nil {
|
||||
previousRequestLog = oldCfg.RequestLog
|
||||
}
|
||||
if s.requestLogger != nil && (oldCfg == nil || previousRequestLog != cfg.RequestLog) {
|
||||
if s.loggerToggle != nil {
|
||||
s.loggerToggle(cfg.RequestLog)
|
||||
} else if toggler, ok := s.requestLogger.(interface{ SetEnabled(bool) }); ok {
|
||||
toggler.SetEnabled(cfg.RequestLog)
|
||||
}
|
||||
}
|
||||
|
||||
if oldCfg == nil || oldCfg.Home.Enabled != cfg.Home.Enabled {
|
||||
if setter, ok := s.requestLogger.(interface{ SetHomeEnabled(bool) }); ok {
|
||||
setter.SetHomeEnabled(cfg.Home.Enabled)
|
||||
}
|
||||
}
|
||||
|
||||
if oldCfg == nil || oldCfg.LoggingToFile != cfg.LoggingToFile || oldCfg.LogsMaxTotalSizeMB != cfg.LogsMaxTotalSizeMB {
|
||||
if err := logging.ConfigureLogOutput(cfg); err != nil {
|
||||
log.Errorf("failed to reconfigure log output: %v", err)
|
||||
}
|
||||
if errContext := ctx.Err(); errContext != nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
if oldCfg == nil || oldCfg.UsageStatisticsEnabled != cfg.UsageStatisticsEnabled {
|
||||
redisqueue.SetUsageStatisticsEnabled(cfg.UsageStatisticsEnabled)
|
||||
}
|
||||
|
||||
if oldCfg == nil || oldCfg.RedisUsageQueueRetentionSeconds != cfg.RedisUsageQueueRetentionSeconds {
|
||||
redisqueue.SetRetentionSeconds(cfg.RedisUsageQueueRetentionSeconds)
|
||||
}
|
||||
|
||||
if s.requestLogger != nil && (oldCfg == nil || oldCfg.ErrorLogsMaxFiles != cfg.ErrorLogsMaxFiles) {
|
||||
if setter, ok := s.requestLogger.(interface{ SetErrorLogsMaxFiles(int) }); ok {
|
||||
setter.SetErrorLogsMaxFiles(cfg.ErrorLogsMaxFiles)
|
||||
}
|
||||
}
|
||||
|
||||
if oldCfg == nil || oldCfg.DisableCooling != cfg.DisableCooling {
|
||||
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
|
||||
}
|
||||
if oldCfg == nil || oldCfg.TransientErrorCooldownSeconds != cfg.TransientErrorCooldownSeconds {
|
||||
auth.SetTransientErrorCooldownSeconds(cfg.TransientErrorCooldownSeconds)
|
||||
}
|
||||
|
||||
if oldCfg != nil && oldCfg.DisableImageGeneration != cfg.DisableImageGeneration {
|
||||
log.Infof("disable-image-generation updated: %v -> %v", oldCfg.DisableImageGeneration, cfg.DisableImageGeneration)
|
||||
}
|
||||
|
||||
applySignatureCacheConfig(oldCfg, cfg)
|
||||
|
||||
if s.handlers != nil && s.handlers.AuthManager != nil {
|
||||
s.handlers.AuthManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second, cfg.MaxRetryCredentials)
|
||||
}
|
||||
|
||||
// Update log level dynamically when debug flag changes
|
||||
if oldCfg == nil || oldCfg.Debug != cfg.Debug {
|
||||
util.SetLogLevel(cfg)
|
||||
}
|
||||
|
||||
prevSecretEmpty := true
|
||||
if oldCfg != nil {
|
||||
prevSecretEmpty = oldCfg.RemoteManagement.SecretKey == ""
|
||||
}
|
||||
newSecretEmpty := cfg.RemoteManagement.SecretKey == ""
|
||||
if s.envManagementSecret {
|
||||
s.registerManagementRoutes()
|
||||
if s.managementRoutesEnabled.CompareAndSwap(false, true) {
|
||||
log.Info("management routes enabled via MANAGEMENT_PASSWORD")
|
||||
} else {
|
||||
s.managementRoutesEnabled.Store(true)
|
||||
}
|
||||
} else {
|
||||
switch {
|
||||
case prevSecretEmpty && !newSecretEmpty:
|
||||
s.registerManagementRoutes()
|
||||
if s.managementRoutesEnabled.CompareAndSwap(false, true) {
|
||||
log.Info("management routes enabled after secret key update")
|
||||
} else {
|
||||
s.managementRoutesEnabled.Store(true)
|
||||
}
|
||||
case !prevSecretEmpty && newSecretEmpty:
|
||||
if s.managementRoutesEnabled.CompareAndSwap(true, false) {
|
||||
log.Info("management routes disabled after secret key removal")
|
||||
} else {
|
||||
s.managementRoutesEnabled.Store(false)
|
||||
}
|
||||
default:
|
||||
s.managementRoutesEnabled.Store(!newSecretEmpty)
|
||||
}
|
||||
}
|
||||
redisqueue.SetEnabled(s.managementRoutesEnabled.Load() || (cfg != nil && cfg.Home.Enabled))
|
||||
|
||||
exampleAPIKeySafeModeRequired := s.exampleAPIKeySafeModeRequired(cfg)
|
||||
if exampleAPIKeySafeModeRequired {
|
||||
s.exampleAPIKeySafeModeActive.Store(true)
|
||||
}
|
||||
accessConfigApplied := s.applyAccessConfig(oldCfg, cfg)
|
||||
if accessConfigApplied || exampleAPIKeySafeModeRequired {
|
||||
s.exampleAPIKeySafeModeActive.Store(exampleAPIKeySafeModeRequired)
|
||||
}
|
||||
s.cfg = cfg
|
||||
if s.codexLiveHandler != nil {
|
||||
if errUpdate := s.codexLiveHandler.UpdateConfig(cfg); errUpdate != nil {
|
||||
log.WithError(errUpdate).Error("failed to update Codex Live media relay configuration")
|
||||
}
|
||||
}
|
||||
s.wsAuthEnabled.Store(cfg.WebsocketAuth)
|
||||
if oldCfg != nil && s.wsAuthChanged != nil && oldCfg.WebsocketAuth != cfg.WebsocketAuth {
|
||||
s.wsAuthChanged(oldCfg.WebsocketAuth, cfg.WebsocketAuth)
|
||||
}
|
||||
if errContext := ctx.Err(); errContext != nil {
|
||||
return false
|
||||
}
|
||||
// Save YAML snapshot for next comparison
|
||||
s.oldConfigYaml, _ = yaml.Marshal(cfg)
|
||||
|
||||
s.handlers.UpdateClients(effectiveSDKConfig(cfg))
|
||||
s.handlers.SetPluginHost(s.pluginHost)
|
||||
if s.pluginHost != nil {
|
||||
s.pluginHost.SetModelExecutor(s.handlers)
|
||||
s.pluginHost.SetAuthManager(s.handlers.AuthManager)
|
||||
}
|
||||
|
||||
if s.mgmt != nil {
|
||||
s.mgmt.SetConfig(cfg)
|
||||
s.mgmt.SetAuthManager(s.handlers.AuthManager)
|
||||
s.mgmt.SetPluginHost(s.pluginHost)
|
||||
}
|
||||
s.refreshPluginManagementRoutes()
|
||||
|
||||
// Count client sources from configuration and auth store.
|
||||
authEntries := 0
|
||||
if cfg != nil && !cfg.Home.Enabled {
|
||||
tokenStore := sdkAuth.GetTokenStore()
|
||||
if dirSetter, ok := tokenStore.(interface{ SetBaseDir(string) }); ok {
|
||||
dirSetter.SetBaseDir(cfg.AuthDir)
|
||||
}
|
||||
authEntries = util.CountAuthFiles(ctx, tokenStore)
|
||||
if errContext := ctx.Err(); errContext != nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
geminiAPIKeyCount := len(cfg.GeminiKey)
|
||||
interactionsAPIKeyCount := len(cfg.InteractionsKey)
|
||||
claudeAPIKeyCount := len(cfg.ClaudeKey)
|
||||
codexAPIKeyCount := len(cfg.CodexKey)
|
||||
xaiAPIKeyCount := len(cfg.XAIKey)
|
||||
vertexAICompatCount := len(cfg.VertexCompatAPIKey)
|
||||
openAICompatCount := 0
|
||||
for i := range cfg.OpenAICompatibility {
|
||||
entry := cfg.OpenAICompatibility[i]
|
||||
if entry.Disabled {
|
||||
continue
|
||||
}
|
||||
openAICompatCount += len(entry.APIKeyEntries)
|
||||
}
|
||||
|
||||
total := authEntries + geminiAPIKeyCount + interactionsAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + xaiAPIKeyCount + vertexAICompatCount + openAICompatCount
|
||||
fmt.Printf("server clients and configuration updated: %d clients (%d auth entries + %d Gemini API keys + %d Interactions API keys + %d Claude API keys + %d Codex keys + %d xAI keys + %d Vertex-compat + %d OpenAI-compat)\n",
|
||||
total,
|
||||
authEntries,
|
||||
geminiAPIKeyCount,
|
||||
interactionsAPIKeyCount,
|
||||
claudeAPIKeyCount,
|
||||
codexAPIKeyCount,
|
||||
xaiAPIKeyCount,
|
||||
vertexAICompatCount,
|
||||
openAICompatCount,
|
||||
)
|
||||
return ctx.Err() == nil
|
||||
}
|
||||
|
||||
func (s *Server) SetWebsocketAuthChangeHandler(fn func(bool, bool)) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.wsAuthChanged = fn
|
||||
}
|
||||
|
||||
func configuredSignatureCacheEnabled(cfg *config.Config) bool {
|
||||
if cfg != nil && cfg.AntigravitySignatureCacheEnabled != nil {
|
||||
return *cfg.AntigravitySignatureCacheEnabled
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func applySignatureCacheConfig(oldCfg, cfg *config.Config) {
|
||||
newVal := configuredSignatureCacheEnabled(cfg)
|
||||
newStrict := configuredSignatureBypassStrict(cfg)
|
||||
if oldCfg == nil {
|
||||
cache.SetSignatureCacheEnabled(newVal)
|
||||
cache.SetSignatureBypassStrictMode(newStrict)
|
||||
return
|
||||
}
|
||||
|
||||
oldVal := configuredSignatureCacheEnabled(oldCfg)
|
||||
if oldVal != newVal {
|
||||
cache.SetSignatureCacheEnabled(newVal)
|
||||
}
|
||||
|
||||
oldStrict := configuredSignatureBypassStrict(oldCfg)
|
||||
if oldStrict != newStrict {
|
||||
cache.SetSignatureBypassStrictMode(newStrict)
|
||||
}
|
||||
}
|
||||
|
||||
func configuredSignatureBypassStrict(cfg *config.Config) bool {
|
||||
if cfg != nil && cfg.AntigravitySignatureBypassStrict != nil {
|
||||
return *cfg.AntigravitySignatureBypassStrict
|
||||
}
|
||||
return false
|
||||
}
|
||||
1053
backend/internal/api/server_routes.go
Normal file
1053
backend/internal/api/server_routes.go
Normal file
File diff suppressed because it is too large
Load diff
16
backend/internal/api/server_sdk_config_test.go
Normal file
16
backend/internal/api/server_sdk_config_test.go
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
func TestEffectiveSDKConfigCopiesCodexOptimizeMultiAgentV2(t *testing.T) {
|
||||
cfg := &config.Config{Codex: config.CodexConfig{OptimizeMultiAgentV2: true}}
|
||||
|
||||
sdkCfg := effectiveSDKConfig(cfg)
|
||||
if sdkCfg == nil || !sdkCfg.CodexOptimizeMultiAgentV2 {
|
||||
t.Fatalf("CodexOptimizeMultiAgentV2 = false, want true")
|
||||
}
|
||||
}
|
||||
2291
backend/internal/api/server_test.go
Normal file
2291
backend/internal/api/server_test.go
Normal file
File diff suppressed because it is too large
Load diff
Loading…
Reference in a new issue