From 9562f0f6dd4ed3446952f98fa0deb84b2e481f47 Mon Sep 17 00:00:00 2001 From: tbphp Date: Mon, 9 Jun 2025 22:03:15 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BC=98=E5=8C=96=E4=BB=A3=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/gpt-load/main.go | 5 ++- internal/handler/handler.go | 52 +++++++++++++------------------ internal/keymanager/manager.go | 18 ++++++----- internal/middleware/middleware.go | 2 +- pkg/types/interfaces.go | 17 ++++++++-- 5 files changed, 48 insertions(+), 46 deletions(-) diff --git a/cmd/gpt-load/main.go b/cmd/gpt-load/main.go index e19df2c..883c110 100644 --- a/cmd/gpt-load/main.go +++ b/cmd/gpt-load/main.go @@ -125,11 +125,10 @@ func setupRoutes(handlers *handler.Handler, proxyServer *proxy.ProxyServer, conf router.GET("/reset-keys", handlers.ResetKeys) router.GET("/config", handlers.GetConfig) // Debug endpoint - // Handle 404 and 405 - router.NoRoute(handlers.NotFound) + // Handle 405 Method Not Allowed router.NoMethod(handlers.MethodNotAllowed) - // Proxy all other requests + // Proxy all other requests (this handles 404 as well) router.NoRoute(proxyServer.HandleProxy) return router diff --git a/internal/handler/handler.go b/internal/handler/handler.go index 0e6d61c..574d7fb 100644 --- a/internal/handler/handler.go +++ b/internal/handler/handler.go @@ -29,10 +29,10 @@ func NewHandler(keyManager types.KeyManager, config types.ConfigManager) *Handle // Health handles health check requests func (h *Handler) Health(c *gin.Context) { stats := h.keyManager.GetStats() - + status := "healthy" httpStatus := http.StatusOK - + // Check if there are any healthy keys if stats.HealthyKeys == 0 { status = "unhealthy" @@ -51,16 +51,16 @@ func (h *Handler) Health(c *gin.Context) { // Stats handles statistics requests func (h *Handler) Stats(c *gin.Context) { stats := h.keyManager.GetStats() - + // Add additional system information var m runtime.MemStats runtime.ReadMemStats(&m) - + response := gin.H{ "keys": gin.H{ - "total": stats.TotalKeys, - "healthy": stats.HealthyKeys, - "blacklisted": stats.BlacklistedKeys, + "total": stats.TotalKeys, + "healthy": stats.HealthyKeys, + "blacklisted": stats.BlacklistedKeys, "current_index": stats.CurrentIndex, }, "requests": gin.H{ @@ -77,9 +77,9 @@ func (h *Handler) Stats(c *gin.Context) { "next_gc_mb": bToMb(m.NextGC), }, "system": gin.H{ - "goroutines": runtime.NumGoroutine(), - "cpu_count": runtime.NumCPU(), - "go_version": runtime.Version(), + "goroutines": runtime.NumGoroutine(), + "cpu_count": runtime.NumCPU(), + "go_version": runtime.Version(), }, "timestamp": time.Now().UTC().Format(time.RFC3339), } @@ -90,11 +90,11 @@ func (h *Handler) Stats(c *gin.Context) { // Blacklist handles blacklist requests func (h *Handler) Blacklist(c *gin.Context) { blacklist := h.keyManager.GetBlacklist() - + response := gin.H{ "blacklisted_keys": blacklist, - "count": len(blacklist), - "timestamp": time.Now().UTC().Format(time.RFC3339), + "count": len(blacklist), + "timestamp": time.Now().UTC().Format(time.RFC3339), } c.JSON(http.StatusOK, response) @@ -104,7 +104,7 @@ func (h *Handler) Blacklist(c *gin.Context) { func (h *Handler) ResetKeys(c *gin.Context) { // Reset blacklist h.keyManager.ResetBlacklist() - + // Reload keys from file if err := h.keyManager.LoadKeys(); err != nil { logrus.Errorf("Failed to reload keys: %v", err) @@ -116,25 +116,15 @@ func (h *Handler) ResetKeys(c *gin.Context) { } stats := h.keyManager.GetStats() - - c.JSON(http.StatusOK, gin.H{ - "message": "Keys reset and reloaded successfully", - "total_keys": stats.TotalKeys, - "healthy_keys": stats.HealthyKeys, - "timestamp": time.Now().UTC().Format(time.RFC3339), - }) - - logrus.Info("Keys reset and reloaded successfully") -} -// NotFound handles 404 requests -func (h *Handler) NotFound(c *gin.Context) { - c.JSON(http.StatusNotFound, gin.H{ - "error": "Endpoint not found", - "path": c.Request.URL.Path, - "method": c.Request.Method, - "timestamp": time.Now().UTC().Format(time.RFC3339), + c.JSON(http.StatusOK, gin.H{ + "message": "Keys reset and reloaded successfully", + "total_keys": stats.TotalKeys, + "healthy_keys": stats.HealthyKeys, + "timestamp": time.Now().UTC().Format(time.RFC3339), }) + + logrus.Info("Keys reset and reloaded successfully") } // MethodNotAllowed handles 405 requests diff --git a/internal/keymanager/manager.go b/internal/keymanager/manager.go index 7fe1376..827e27e 100644 --- a/internal/keymanager/manager.go +++ b/internal/keymanager/manager.go @@ -204,13 +204,15 @@ func (km *Manager) RecordFailure(key string, err error) { } // Increment failure count - failCount, _ := km.keyFailureCounts.LoadOrStore(key, int64(0)) - newFailCount := atomic.AddInt64(failCount.(*int64), 1) + failCount, _ := km.keyFailureCounts.LoadOrStore(key, new(int64)) + if counter, ok := failCount.(*int64); ok { + newFailCount := atomic.AddInt64(counter, 1) - // Blacklist if threshold exceeded - if int(newFailCount) >= km.config.BlacklistThreshold { - km.blacklistedKeys.Store(key, time.Now()) - logrus.Debugf("Key blacklisted after %d failures", newFailCount) + // Blacklist if threshold exceeded + if int(newFailCount) >= km.config.BlacklistThreshold { + km.blacklistedKeys.Store(key, time.Now()) + logrus.Debugf("Key blacklisted after %d failures", newFailCount) + } } } @@ -236,7 +238,7 @@ func (km *Manager) GetStats() types.Stats { km.keysMutex.RUnlock() blacklistedCount := 0 - km.blacklistedKeys.Range(func(key, value interface{}) bool { + km.blacklistedKeys.Range(func(key, value any) bool { blacklistedCount++ return true }) @@ -273,7 +275,7 @@ func (km *Manager) ResetBlacklist() { func (km *Manager) GetBlacklist() []types.BlacklistEntry { var blacklist []types.BlacklistEntry - km.blacklistedKeys.Range(func(key, value interface{}) bool { + km.blacklistedKeys.Range(func(key, value any) bool { keyStr := key.(string) blacklistTime := value.(time.Time) diff --git a/internal/middleware/middleware.go b/internal/middleware/middleware.go index ead038e..367f94b 100644 --- a/internal/middleware/middleware.go +++ b/internal/middleware/middleware.go @@ -173,7 +173,7 @@ func Auth(config types.AuthConfig) gin.HandlerFunc { // Recovery creates a recovery middleware with custom error handling func Recovery() gin.HandlerFunc { - return gin.CustomRecovery(func(c *gin.Context, recovered interface{}) { + return gin.CustomRecovery(func(c *gin.Context, recovered any) { if err, ok := recovered.(string); ok { logrus.Errorf("Panic recovered: %s", err) c.JSON(500, gin.H{ diff --git a/pkg/types/interfaces.go b/pkg/types/interfaces.go index 1181114..2eefcac 100644 --- a/pkg/types/interfaces.go +++ b/pkg/types/interfaces.go @@ -1,7 +1,11 @@ // Package types defines common interfaces and types used across the application package types -import "time" +import ( + "time" + + "github.com/gin-gonic/gin" +) // ConfigManager defines the interface for configuration management type ConfigManager interface { @@ -13,6 +17,7 @@ type ConfigManager interface { GetPerformanceConfig() PerformanceConfig GetLogConfig() LogConfig Validate() error + DisplayConfig() } // KeyManager defines the interface for API key management @@ -27,6 +32,12 @@ type KeyManager interface { Close() } +// ProxyServer defines the interface for proxy server +type ProxyServer interface { + HandleProxy(c *gin.Context) + Close() +} + // ServerConfig represents server configuration type ServerConfig struct { Port int `json:"port"` @@ -64,8 +75,8 @@ type CORSConfig struct { // PerformanceConfig represents performance configuration type PerformanceConfig struct { - MaxConcurrentRequests int `json:"maxConcurrentRequests"` - RequestTimeout int `json:"requestTimeout"` + MaxConcurrentRequests int `json:"maxConcurrentRequests"` + RequestTimeout int `json:"requestTimeout"` EnableGzip bool `json:"enableGzip"` }