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

@@ -127,16 +127,6 @@ func (h *Handler) ResetKeys(c *gin.Context) {
 	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),
-	})
-}
-
 // MethodNotAllowed handles 405 requests
 func (h *Handler) MethodNotAllowed(c *gin.Context) {
 	c.JSON(http.StatusMethodNotAllowed, gin.H{

internal/keymanager/manager.go

@@ -204,8 +204,9 @@ 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 {
@@ -213,6 +214,7 @@ func (km *Manager) RecordFailure(key string, err error) {
 			logrus.Debugf("Key blacklisted after %d failures", newFailCount)
 		}
 	}
+}
 
 // isPermanentError checks if an error is permanent
 func (km *Manager) isPermanentError(err error) bool {
@@ -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"`