fix: 优化代码

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

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

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)
 

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{

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"`
 }