Luis Paterfirst commit827bd6e1 year agoHistory
176 lines · 3.7 KB · Go
RawDownload
package api import (	"context"	"errors"	"fmt"	"github.com/gin-gonic/gin"	"github.com/luispater/CLIProxyAPI/internal/client"	log "github.com/sirupsen/logrus"	"net/http"	"strings") // Server represents the API servertype Server struct {	engine   *gin.Engine	server   *http.Server	handlers *APIHandlers	cfg      *ServerConfig} // ServerConfig contains configuration for the API servertype ServerConfig struct {	Port    string	Debug   bool	ApiKeys []string} // NewServer creates a new API server instancefunc NewServer(config *ServerConfig, cliClients []*client.Client) *Server {	// Set gin mode	if !config.Debug {		gin.SetMode(gin.ReleaseMode)	} 	// Create handlers	handlers := NewAPIHandlers(cliClients, config.Debug) 	// Create gin engine	engine := gin.New() 	// Add middleware	engine.Use(gin.Logger())	engine.Use(gin.Recovery())	engine.Use(corsMiddleware()) 	// Create server instance	s := &Server{		engine:   engine,		handlers: handlers,		cfg:      config,	} 	// Setup routes	s.setupRoutes() 	// Create HTTP server	s.server = &http.Server{		Addr:    ":" + config.Port,		Handler: engine,	} 	return s} // setupRoutes configures the API routesfunc (s *Server) setupRoutes() {	// OpenAI compatible API routes	v1 := s.engine.Group("/v1")	v1.Use(AuthMiddleware(s.cfg))	{		v1.GET("/models", s.handlers.Models)		v1.POST("/chat/completions", s.handlers.ChatCompletions)	} 	// Root endpoint	s.engine.GET("/", func(c *gin.Context) {		c.JSON(http.StatusOK, gin.H{			"message": "CLI Proxy API Server",			"version": "1.0.0",			"endpoints": []string{				"POST /v1/chat/completions",				"GET /v1/models",			},		})	})} // Start starts the API serverfunc (s *Server) Start() error {	log.Debugf("Starting API server on %s", s.server.Addr) 	// Start the HTTP server	if err := s.server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {		return fmt.Errorf("failed to start HTTP server: %v", err)	} 	return nil} // Stop gracefully stops the API serverfunc (s *Server) Stop(ctx context.Context) error {	log.Debug("Stopping API server...") 	// Shutdown the HTTP server	if err := s.server.Shutdown(ctx); err != nil {		return fmt.Errorf("failed to shutdown HTTP server: %v", err)	} 	log.Debug("API server stopped")	return nil} // corsMiddleware adds CORS headersfunc corsMiddleware() gin.HandlerFunc {	return func(c *gin.Context) {		c.Header("Access-Control-Allow-Origin", "*")		c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")		c.Header("Access-Control-Allow-Headers", "Origin, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization") 		if c.Request.Method == "OPTIONS" {			c.AbortWithStatus(http.StatusNoContent)			return		} 		c.Next()	}} // AuthMiddleware authenticates requests using API keysfunc AuthMiddleware(cfg *ServerConfig) gin.HandlerFunc {	return func(c *gin.Context) {		if len(cfg.ApiKeys) == 0 {			c.Next()			return		} 		// Get the Authorization header		authHeader := c.GetHeader("Authorization")		if authHeader == "" {			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{				"error": "Missing API key",			})			return		} 		// Extract the API key		parts := strings.Split(authHeader, " ")		var apiKey string		if len(parts) == 2 && strings.ToLower(parts[0]) == "bearer" {			apiKey = parts[1]		} else {			apiKey = authHeader		} 		// Find the API key in the in-memory list		var foundKey string		for i := range cfg.ApiKeys {			if cfg.ApiKeys[i] == apiKey {				foundKey = cfg.ApiKeys[i]				break			}		}		if foundKey == "" {			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{				"error": "Invalid API key",			})			return		} 		// Store the API key and user in the context		c.Set("apiKey", foundKey) 		c.Next()	}}