From d5ebbf10f556d395280142e0a000b1bb04d6427f Mon Sep 17 00:00:00 2001 From: tbphp Date: Tue, 1 Jul 2025 00:13:01 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E8=B7=AF=E7=94=B1=E4=BC=98=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Dockerfile | 2 +- cmd/gpt-load/embed.go | 19 +++++- cmd/gpt-load/main.go | 92 ++------------------------ internal/db/database.go | 19 +++--- internal/handler/handler.go | 15 +---- internal/handler/key_handler.go | 38 ++++++++++- internal/router/router.go | 88 ++++++++++++++++++++++++ web/src/api/dashboard.ts | 2 +- web/src/api/index.ts | 2 +- web/src/api/logs.ts | 2 +- web/src/api/settings.ts | 4 +- web/src/components/GroupConfigForm.vue | 65 +++++++++++------- 12 files changed, 206 insertions(+), 142 deletions(-) create mode 100644 internal/router/router.go diff --git a/Dockerfile b/Dockerfile index ac2132e..e8dae89 100644 --- a/Dockerfile +++ b/Dockerfile @@ -33,7 +33,7 @@ RUN go mod download COPY . . # Copy the built frontend from the previous stage -COPY --from=frontend-builder /app/web/dist ./web/dist +COPY --from=frontend-builder /app/web/dist ./cmd/gpt-load/dist # Build the Go application # We use CGO_ENABLED=0 to create a static binary diff --git a/cmd/gpt-load/embed.go b/cmd/gpt-load/embed.go index 241b8cf..ac32778 100644 --- a/cmd/gpt-load/embed.go +++ b/cmd/gpt-load/embed.go @@ -1,6 +1,21 @@ package main -import "embed" +import ( + "embed" + "io/fs" + "log" +) //go:embed all:dist -var WebUI embed.FS \ No newline at end of file +var content embed.FS + +// WebUI is the filesystem for the embedded web UI. +var WebUI fs.FS + +func init() { + var err error + WebUI, err = fs.Sub(content, "dist") + if err != nil { + log.Fatalf("Failed to create sub filesystem for UI: %v", err) + } +} \ No newline at end of file diff --git a/cmd/gpt-load/main.go b/cmd/gpt-load/main.go index 386a08a..b8d22a7 100644 --- a/cmd/gpt-load/main.go +++ b/cmd/gpt-load/main.go @@ -5,13 +5,10 @@ import ( "context" "fmt" "io" - "io/fs" "net/http" "os" "os/signal" - "path" "path/filepath" - "strings" "sync" "syscall" "time" @@ -19,12 +16,11 @@ import ( "gpt-load/internal/config" "gpt-load/internal/db" "gpt-load/internal/handler" - "gpt-load/internal/middleware" "gpt-load/internal/models" "gpt-load/internal/proxy" + "gpt-load/internal/router" // <-- 引入新的 router 包 "gpt-load/internal/types" - "github.com/gin-gonic/gin" "github.com/sirupsen/logrus" "gorm.io/gorm" ) @@ -48,7 +44,6 @@ func main() { // Display startup information displayStartupInfo(configManager) - // --- Asynchronous Request Logging Setup --- requestLogChan := make(chan models.RequestLog, 1000) var wg sync.WaitGroup @@ -66,14 +61,14 @@ func main() { // Create handlers serverHandler := handler.NewServer(database, configManager) - // Setup routes - router := setupRoutes(serverHandler, proxyServer, configManager) + // Setup routes using the new router package + appRouter := router.New(serverHandler, proxyServer, configManager, WebUI) // Create HTTP server with optimized timeout configuration serverConfig := configManager.GetServerConfig() server := &http.Server{ Addr: fmt.Sprintf("%s:%d", serverConfig.Host, serverConfig.Port), - Handler: router, + Handler: appRouter, ReadTimeout: time.Duration(serverConfig.ReadTimeout) * time.Second, WriteTimeout: time.Duration(serverConfig.WriteTimeout) * time.Second, IdleTimeout: time.Duration(serverConfig.IdleTimeout) * time.Second, @@ -117,83 +112,8 @@ func main() { logrus.Info("Server exited gracefully") } -// setupRoutes configures the HTTP routes -func setupRoutes(serverHandler *handler.Server, proxyServer *proxy.ProxyServer, configManager types.ConfigManager) *gin.Engine { - // Set Gin mode - gin.SetMode(gin.ReleaseMode) - - router := gin.New() - - // Add server start time middleware for uptime calculation - startTime := time.Now() - router.Use(func(c *gin.Context) { - c.Set("serverStartTime", startTime) - c.Next() - }) - - // Add middleware - router.Use(middleware.Recovery()) - router.Use(middleware.ErrorHandler()) - router.Use(middleware.Logger(configManager.GetLogConfig())) - router.Use(middleware.CORS(configManager.GetCORSConfig())) - router.Use(middleware.RateLimiter(configManager.GetPerformanceConfig())) - - // Add authentication middleware if enabled - if configManager.GetAuthConfig().Enabled { - router.Use(middleware.Auth(configManager.GetAuthConfig())) - } - - // Management endpoints - router.GET("/health", serverHandler.Health) - router.GET("/stats", serverHandler.Stats) - router.GET("/config", serverHandler.GetConfig) // Debug endpoint - - // Register API routes for group and key management - api := router.Group("/api") - serverHandler.RegisterAPIRoutes(api) - - // Register the main proxy route - proxy := router.Group("/proxy") - proxyServer.RegisterProxyRoutes(proxy) - - // Handle 405 Method Not Allowed - router.NoMethod(serverHandler.MethodNotAllowed) - - // Serve the frontend UI for all other requests - router.NoRoute(ServeUI()) - - return router -} - -// ServeUI returns a gin.HandlerFunc to serve the embedded frontend UI. -func ServeUI() gin.HandlerFunc { - subFS, err := fs.Sub(WebUI, "dist") - if err != nil { - // This should not happen at runtime if embed is correct. - // Panic is acceptable here as it's a startup failure. - panic(fmt.Sprintf("Failed to create sub filesystem for UI: %v", err)) - } - fileServer := http.FileServer(http.FS(subFS)) - - return func(c *gin.Context) { - // Clean the path to prevent directory traversal attacks. - upath := path.Clean(c.Request.URL.Path) - if !strings.HasPrefix(upath, "/") { - upath = "/" + upath - } - - // Check if the file exists in the embedded filesystem. - _, err := subFS.Open(strings.TrimPrefix(upath, "/")) - if os.IsNotExist(err) { - // The file does not exist, so we serve index.html for SPA routing. - // This allows the Vue router to handle the path. - c.Request.URL.Path = "/" - } - - // Let the http.FileServer handle the request. - fileServer.ServeHTTP(c.Writer, c.Request) - } -} +// setupLogger, displayStartupInfo, and startRequestLogger functions remain unchanged. +// The old setupRoutes and ServeUI functions are now removed from this file. // setupLogger configures the logging system func setupLogger(configManager types.ConfigManager) { diff --git a/internal/db/database.go b/internal/db/database.go index 584a608..1fde309 100644 --- a/internal/db/database.go +++ b/internal/db/database.go @@ -48,15 +48,16 @@ func InitDB() (*gorm.DB, error) { sqlDB.SetMaxOpenConns(100) sqlDB.SetConnMaxLifetime(time.Hour) - // Auto-migrate models - err = DB.AutoMigrate( - &models.SystemSetting{}, - &models.Group{}, - &models.APIKey{}, - &models.RequestLog{}, - ) - if err != nil { - return nil, fmt.Errorf("failed to auto-migrate database: %w", err) + if os.Getenv("DB_AUTO_MIGRATE") != "false" { + err = DB.AutoMigrate( + &models.SystemSetting{}, + &models.Group{}, + &models.APIKey{}, + &models.RequestLog{}, + ) + if err != nil { + return nil, fmt.Errorf("failed to auto-migrate database: %w", err) + } } fmt.Println("Database connection initialized and models migrated.") diff --git a/internal/handler/handler.go b/internal/handler/handler.go index 87aee43..5b9edb7 100644 --- a/internal/handler/handler.go +++ b/internal/handler/handler.go @@ -42,13 +42,11 @@ func (s *Server) RegisterAPIRoutes(api *gin.RouterGroup) { { keys.POST("", s.CreateKeysInGroup) keys.GET("", s.ListKeysInGroup) + keys.PUT("/:key_id", s.UpdateKey) + keys.DELETE("", s.DeleteKeys) } } - // Key management routes - api.PUT("/keys/:key_id", s.UpdateKey) - api.DELETE("/keys", s.DeleteKeys) - // Dashboard and logs routes dashboard := api.Group("/dashboard") { @@ -100,15 +98,6 @@ func (s *Server) Health(c *gin.Context) { }) } -// MethodNotAllowed handles 405 requests -func (s *Server) MethodNotAllowed(c *gin.Context) { - c.JSON(http.StatusMethodNotAllowed, gin.H{ - "error": "Method not allowed", - "path": c.Request.URL.Path, - "method": c.Request.Method, - "timestamp": time.Now().UTC().Format(time.RFC3339), - }) -} // GetConfig returns configuration information (for debugging) func (s *Server) GetConfig(c *gin.Context) { diff --git a/internal/handler/key_handler.go b/internal/handler/key_handler.go index f1e4be2..1795bd4 100644 --- a/internal/handler/key_handler.go +++ b/internal/handler/key_handler.go @@ -64,6 +64,12 @@ func (s *Server) ListKeysInGroup(c *gin.Context) { // UpdateKey handles updating a specific key. func (s *Server) UpdateKey(c *gin.Context) { + groupID, err := strconv.Atoi(c.Param("id")) + if err != nil { + response.Error(c, http.StatusBadRequest, "Invalid group ID") + return + } + keyID, err := strconv.Atoi(c.Param("key_id")) if err != nil { response.Error(c, http.StatusBadRequest, "Invalid key ID") @@ -71,8 +77,8 @@ func (s *Server) UpdateKey(c *gin.Context) { } var key models.APIKey - if err := s.DB.First(&key, keyID).Error; err != nil { - response.Error(c, http.StatusNotFound, "Key not found") + if err := s.DB.Where("group_id = ? AND id = ?", groupID, keyID).First(&key).Error; err != nil { + response.Error(c, http.StatusNotFound, "Key not found in this group") return } @@ -99,6 +105,12 @@ type DeleteKeysRequest struct { // DeleteKeys handles deleting one or more keys. func (s *Server) DeleteKeys(c *gin.Context) { + groupID, err := strconv.Atoi(c.Param("id")) + if err != nil { + response.Error(c, http.StatusBadRequest, "Invalid group ID") + return + } + var req DeleteKeysRequest if err := c.ShouldBindJSON(&req); err != nil { response.Error(c, http.StatusBadRequest, "Invalid request body") @@ -110,10 +122,30 @@ func (s *Server) DeleteKeys(c *gin.Context) { return } - if err := s.DB.Delete(&models.APIKey{}, req.KeyIDs).Error; err != nil { + // Start a transaction + tx := s.DB.Begin() + + // Verify all keys belong to the specified group + var count int64 + if err := tx.Model(&models.APIKey{}).Where("id IN ? AND group_id = ?", req.KeyIDs, groupID).Count(&count).Error; err != nil { + tx.Rollback() + response.Error(c, http.StatusInternalServerError, "Failed to verify keys") + return + } + + if count != int64(len(req.KeyIDs)) { + tx.Rollback() + response.Error(c, http.StatusForbidden, "One or more keys do not belong to the specified group") + return + } + + // Delete the keys + if err := tx.Where("id IN ?", req.KeyIDs).Delete(&models.APIKey{}).Error; err != nil { + tx.Rollback() response.Error(c, http.StatusInternalServerError, "Failed to delete keys") return } + tx.Commit() response.Success(c, gin.H{"message": "Keys deleted successfully"}) } \ No newline at end of file diff --git a/internal/router/router.go b/internal/router/router.go new file mode 100644 index 0000000..cb66653 --- /dev/null +++ b/internal/router/router.go @@ -0,0 +1,88 @@ +package router + +import ( + "gpt-load/internal/handler" + "gpt-load/internal/middleware" + "gpt-load/internal/proxy" + "gpt-load/internal/types" + "io/fs" + "net/http" + "strings" + "time" + + "github.com/gin-gonic/gin" +) + +// New 创建并配置一个完整的 gin.Engine 实例 +func New( + serverHandler *handler.Server, + proxyServer *proxy.ProxyServer, + configManager types.ConfigManager, + webUI fs.FS, +) *gin.Engine { + // 设置 Gin 模式 + gin.SetMode(gin.ReleaseMode) + + router := gin.New() + + // 注册全局中间件 + router.Use(middleware.Recovery()) + router.Use(middleware.ErrorHandler()) + router.Use(middleware.Logger(configManager.GetLogConfig())) + router.Use(middleware.CORS(configManager.GetCORSConfig())) + router.Use(middleware.RateLimiter(configManager.GetPerformanceConfig())) + + // 添加服务器启动时间中间件 + startTime := time.Now() + router.Use(func(c *gin.Context) { + c.Set("serverStartTime", startTime) + c.Next() + }) + + // 注册 Web UI 和通用端点 + router.GET("/health", serverHandler.Health) + router.GET("/stats", serverHandler.Stats) + router.GET("/config", serverHandler.GetConfig) // Debug endpoint + + // 注册管理 API 路由 + api := router.Group("/api") + authConfig := configManager.GetAuthConfig() + if authConfig.Enabled { + api.Use(middleware.Auth(authConfig)) + } + serverHandler.RegisterAPIRoutes(api) + + // 注册代理路由 + proxyGroup := router.Group("/proxy") + if authConfig.Enabled { + proxyGroup.Use(middleware.Auth(authConfig)) + } + proxyServer.RegisterProxyRoutes(proxyGroup) + + // 处理 405 Method Not Allowed + router.NoMethod(func(c *gin.Context) { + c.JSON(http.StatusMethodNotAllowed, gin.H{"error": "Method not allowed"}) + }) + + // 其他所有路由都交给前端 UI 处理 + router.NoRoute(ServeUI(webUI)) + + return router +} + +// ServeUI 返回一个 gin.HandlerFunc 来服务嵌入式前端 UI +func ServeUI(webUI fs.FS) gin.HandlerFunc { + fileServer := http.FileServer(http.FS(webUI)) + + return func(c *gin.Context) { + // 检查文件是否存在于嵌入的文件系统中 + if _, err := webUI.Open(strings.TrimPrefix(c.Request.URL.Path, "/")); err != nil { + // 如果文件不存在，并且不是API或代理请求，则将请求重写为 / + // 这将提供 index.html，以支持 SPA 的前端路由 + if !strings.HasPrefix(c.Request.URL.Path, "/api/") && !strings.HasPrefix(c.Request.URL.Path, "/proxy/") { + c.Request.URL.Path = "/" + } + } + fileServer.ServeHTTP(c.Writer, c.Request) + } +} diff --git a/web/src/api/dashboard.ts b/web/src/api/dashboard.ts index 731e1c5..98da43f 100644 --- a/web/src/api/dashboard.ts +++ b/web/src/api/dashboard.ts @@ -2,5 +2,5 @@ import request from './index'; import type { DashboardStats } from '@/types/models'; export const getDashboardStats = (): Promise => { - return request.get('/api/dashboard/stats'); + return request.get('/dashboard/stats'); }; \ No newline at end of file diff --git a/web/src/api/index.ts b/web/src/api/index.ts index d019f36..96899c5 100644 --- a/web/src/api/index.ts +++ b/web/src/api/index.ts @@ -1,7 +1,7 @@ import axios from "axios"; const apiClient = axios.create({ - baseURL: "/", + baseURL: "/api", headers: { "Content-Type": "application/json", }, diff --git a/web/src/api/logs.ts b/web/src/api/logs.ts index 0a348ef..5d3dc3e 100644 --- a/web/src/api/logs.ts +++ b/web/src/api/logs.ts @@ -20,5 +20,5 @@ export interface PaginatedLogs { } export const getLogs = (query: LogQuery): Promise => { - return request.get('/api/logs', { params: query }); + return request.get('/logs', { params: query }); }; \ No newline at end of file diff --git a/web/src/api/settings.ts b/web/src/api/settings.ts index b925b95..0b3d344 100644 --- a/web/src/api/settings.ts +++ b/web/src/api/settings.ts @@ -2,9 +2,9 @@ import request from './index'; import type { Setting } from '@/types/models'; export function getSettings() { - return request.get('/api/settings'); + return request.get('/settings'); } export function updateSettings(settings: Setting[]) { - return request.put('/api/settings', settings); + return request.put('/settings', settings); } \ No newline at end of file diff --git a/web/src/components/GroupConfigForm.vue b/web/src/components/GroupConfigForm.vue index d169fbe..5c2708e 100644 --- a/web/src/components/GroupConfigForm.vue +++ b/web/src/components/GroupConfigForm.vue @@ -4,11 +4,17 @@ - + @@ -24,29 +30,42 @@