1198 lines · 40.5 KB · Go
RawDownload
// 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/subtle"	"crypto/tls"	"errors"	"fmt"	"net"	"net/http"	"os"	"path/filepath"	"reflect"	"strings"	"sync"	"sync/atomic"	"time" 	"github.com/gin-gonic/gin"	"github.com/router-for-me/CLIProxyAPI/v6/internal/access"	managementHandlers "github.com/router-for-me/CLIProxyAPI/v6/internal/api/handlers/management"	"github.com/router-for-me/CLIProxyAPI/v6/internal/api/middleware"	"github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules"	ampmodule "github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules/amp"	"github.com/router-for-me/CLIProxyAPI/v6/internal/cache"	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"	"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"	"github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset"	"github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue"	"github.com/router-for-me/CLIProxyAPI/v6/internal/usage"	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"	sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"	"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"	"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/claude"	"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/gemini"	"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/openai"	sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"	"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"	log "github.com/sirupsen/logrus"	"golang.org/x/net/http2"	"gopkg.in/yaml.v3") const oauthCallbackSuccessHTML = `<html><head><meta charset="utf-8"><title>Authentication successful</title><script>setTimeout(function(){window.close();},5000);</script></head><body><h1>Authentication successful!</h1><p>You can close this window.</p><p>This window will close automatically in 5 seconds.</p></body></html>` 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} // 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)	return logging.NewFileRequestLogger(cfg.RequestLog, logsDir, configDir, cfg.ErrorLogsMaxFiles)} // 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	}} // 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 	// 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 	// ampModule is the Amp routing module for model mapping hot-reload	ampModule *ampmodule.AmpModule 	// 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{}} // 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 instancefunc 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())	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(&cfg.SDKConfig, authManager),		cfg:                 cfg,		accessManager:       accessManager,		requestLogger:       requestLogger,		loggerToggle:        toggle,		configFilePath:      configFilePath,		currentPath:         wd,		envManagementSecret: envManagementSecret,		wsRoutes:            make(map[string]struct{}),	}	s.wsAuthEnabled.Store(cfg.WebsocketAuth)	// 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)	}	managementasset.SetCurrentConfig(cfg)	auth.SetQuotaCooldownDisabled(cfg.DisableCooling)	applySignatureCacheConfig(nil, cfg)	// Initialize management handler	s.mgmt = managementHandlers.NewHandler(cfg, configFilePath, authManager)	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)	}	s.localPassword = optionState.localPassword 	// Setup routes	s.setupRoutes() 	// Register Amp module using V2 interface with Context	s.ampModule = ampmodule.NewLegacy(accessManager, AuthMiddleware(accessManager))	ctx := modules.Context{		Engine:         engine,		BaseHandler:    s.handlers,		Config:         cfg,		AuthMiddleware: AuthMiddleware(accessManager),	}	if err := modules.RegisterModule(ctx, s.ampModule); err != nil {		log.Errorf("Failed to register Amp module: %v", err)	} 	// 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)	if hasManagementSecret {		s.registerManagementRoutes()	} 	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} // setupRoutes configures the API routes for the server.// It defines the endpoints and associates them with their respective handlers.func (s *Server) setupRoutes() {	healthzHandler := func(c *gin.Context) {		if c.Request.Method == http.MethodHead {			c.Status(http.StatusOK)			return		} 		c.JSON(http.StatusOK, gin.H{"status": "ok"})	}	s.engine.GET("/healthz", healthzHandler)	s.engine.HEAD("/healthz", healthzHandler) 	s.engine.GET("/management.html", s.serveManagementControlPanel)	openaiHandlers := openai.NewOpenAIAPIHandler(s.handlers)	geminiHandlers := gemini.NewGeminiAPIHandler(s.handlers)	geminiCLIHandlers := gemini.NewGeminiCLIAPIHandler(s.handlers)	claudeCodeHandlers := claude.NewClaudeCodeAPIHandler(s.handlers)	openaiResponsesHandlers := openai.NewOpenAIResponsesAPIHandler(s.handlers) 	// OpenAI compatible API routes	v1 := s.engine.Group("/v1")	v1.Use(AuthMiddleware(s.accessManager))	{		v1.GET("/models", s.unifiedModelsHandler(openaiHandlers, claudeCodeHandlers))		v1.POST("/chat/completions", openaiHandlers.ChatCompletions)		v1.POST("/completions", openaiHandlers.Completions)		v1.POST("/images/generations", openaiHandlers.ImagesGenerations)		v1.POST("/images/edits", openaiHandlers.ImagesEdits)		v1.POST("/messages", claudeCodeHandlers.ClaudeMessages)		v1.POST("/messages/count_tokens", claudeCodeHandlers.ClaudeCountTokens)		v1.GET("/responses", openaiResponsesHandlers.ResponsesWebsocket)		v1.POST("/responses", openaiResponsesHandlers.Responses)		v1.POST("/responses/compact", openaiResponsesHandlers.Compact)	} 	// Codex CLI direct route aliases (chatgpt_base_url compatible)	codexDirect := s.engine.Group("/backend-api/codex")	codexDirect.Use(AuthMiddleware(s.accessManager))	{		codexDirect.GET("/responses", openaiResponsesHandlers.ResponsesWebsocket)		codexDirect.POST("/responses", openaiResponsesHandlers.Responses)		codexDirect.POST("/responses/compact", openaiResponsesHandlers.Compact)	} 	// Gemini compatible API routes	v1beta := s.engine.Group("/v1beta")	v1beta.Use(AuthMiddleware(s.accessManager))	{		v1beta.GET("/models", geminiHandlers.GeminiModels)		v1beta.POST("/models/*action", geminiHandlers.GeminiHandler)		v1beta.GET("/models/*action", geminiHandlers.GeminiGetHandler)	} 	// Root endpoint	s.engine.GET("/", func(c *gin.Context) {		c.JSON(http.StatusOK, gin.H{			"message": "CLI Proxy API Server",			"endpoints": []string{				"POST /v1/chat/completions",				"POST /v1/completions",				"GET /v1/models",			},		})	})	s.engine.POST("/v1internal:method", geminiCLIHandlers.CLIHandler) 	// OAuth callback endpoints (reuse main server port)	// These endpoints receive provider redirects and persist	// the short-lived code/state for the waiting goroutine.	s.engine.GET("/anthropic/callback", func(c *gin.Context) {		code := c.Query("code")		state := c.Query("state")		errStr := c.Query("error")		if errStr == "" {			errStr = c.Query("error_description")		}		if state != "" {			_, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "anthropic", state, code, errStr)		}		c.Header("Content-Type", "text/html; charset=utf-8")		c.String(http.StatusOK, oauthCallbackSuccessHTML)	}) 	s.engine.GET("/codex/callback", func(c *gin.Context) {		code := c.Query("code")		state := c.Query("state")		errStr := c.Query("error")		if errStr == "" {			errStr = c.Query("error_description")		}		if state != "" {			_, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "codex", state, code, errStr)		}		c.Header("Content-Type", "text/html; charset=utf-8")		c.String(http.StatusOK, oauthCallbackSuccessHTML)	}) 	s.engine.GET("/google/callback", func(c *gin.Context) {		code := c.Query("code")		state := c.Query("state")		errStr := c.Query("error")		if errStr == "" {			errStr = c.Query("error_description")		}		if state != "" {			_, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "gemini", state, code, errStr)		}		c.Header("Content-Type", "text/html; charset=utf-8")		c.String(http.StatusOK, oauthCallbackSuccessHTML)	}) 	s.engine.GET("/antigravity/callback", func(c *gin.Context) {		code := c.Query("code")		state := c.Query("state")		errStr := c.Query("error")		if errStr == "" {			errStr = c.Query("error_description")		}		if state != "" {			_, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "antigravity", state, code, errStr)		}		c.Header("Content-Type", "text/html; charset=utf-8")		c.String(http.StatusOK, oauthCallbackSuccessHTML)	}) 	// Management routes are registered lazily by registerManagementRoutes when a secret is configured.} // AttachWebsocketRoute registers a websocket upgrade handler on the primary Gin engine.// The handler is served as-is without additional middleware beyond the standard stack already configured.func (s *Server) AttachWebsocketRoute(path string, handler http.Handler) {	if s == nil || s.engine == nil || handler == nil {		return	}	trimmed := strings.TrimSpace(path)	if trimmed == "" {		trimmed = "/v1/ws"	}	if !strings.HasPrefix(trimmed, "/") {		trimmed = "/" + trimmed	}	s.wsRouteMu.Lock()	if _, exists := s.wsRoutes[trimmed]; exists {		s.wsRouteMu.Unlock()		return	}	s.wsRoutes[trimmed] = struct{}{}	s.wsRouteMu.Unlock() 	authMiddleware := AuthMiddleware(s.accessManager)	conditionalAuth := func(c *gin.Context) {		if !s.wsAuthEnabled.Load() {			c.Next()			return		}		authMiddleware(c)	}	finalHandler := func(c *gin.Context) {		handler.ServeHTTP(c.Writer, c.Request)		c.Abort()	} 	s.engine.GET(trimmed, conditionalAuth, finalHandler)} func (s *Server) registerManagementRoutes() {	if s == nil || s.engine == nil || s.mgmt == nil {		return	}

Lines 1–500 of 1198.