feat: 密钥管理

2025-07-04 21:19:15 +08:00
parent 7c10474d19
commit 01b86f7e30
23 changed files with 1427 additions and 250 deletions
--- a/internal/handler/group_handler.go
+++ b/internal/handler/group_handler.go
@@ -3,6 +3,7 @@ package handler

 import (
 	"encoding/json"
+	"fmt"
 	app_errors "gpt-load/internal/errors"
 	"gpt-load/internal/models"
 	"gpt-load/internal/response"
@@ -10,6 +11,7 @@ import (
 	"strconv"

 	"github.com/gin-gonic/gin"
+	"gorm.io/datatypes"
 )

 // isValidGroupName checks if the group name is valid.
@@ -22,6 +24,50 @@ func isValidGroupName(name string) bool {
 	return match
 }

+// validateAndCleanConfig validates the group config against the GroupConfig struct.
+func validateAndCleanConfig(configMap map[string]interface{}) (map[string]interface{}, error) {
+	if configMap == nil {
+		return nil, nil
+	}
+
+	configBytes, err := json.Marshal(configMap)
+	if err != nil {
+		return nil, err
+	}
+
+	var validatedConfig models.GroupConfig
+	if err := json.Unmarshal(configBytes, &validatedConfig); err != nil {
+		return nil, err
+	}
+
+	// 验证配置项的合理范围
+	if validatedConfig.BlacklistThreshold != nil && *validatedConfig.BlacklistThreshold < 0 {
+		return nil, fmt.Errorf("blacklist_threshold must be >= 0")
+	}
+	if validatedConfig.MaxRetries != nil && (*validatedConfig.MaxRetries < 0 || *validatedConfig.MaxRetries > 10) {
+		return nil, fmt.Errorf("max_retries must be between 0 and 10")
+	}
+	if validatedConfig.RequestTimeout != nil && (*validatedConfig.RequestTimeout < 1 || *validatedConfig.RequestTimeout > 3600) {
+		return nil, fmt.Errorf("request_timeout must be between 1 and 3600 seconds")
+	}
+	if validatedConfig.KeyValidationIntervalMinutes != nil && (*validatedConfig.KeyValidationIntervalMinutes < 5 || *validatedConfig.KeyValidationIntervalMinutes > 1440) {
+		return nil, fmt.Errorf("key_validation_interval_minutes must be between 5 and 1440 minutes")
+	}
+
+	// Marshal back to a map to remove any fields not in GroupConfig
+	validatedBytes, err := json.Marshal(validatedConfig)
+	if err != nil {
+		return nil, err
+	}
+
+	var cleanedMap map[string]interface{}
+	if err := json.Unmarshal(validatedBytes, &cleanedMap); err != nil {
+		return nil, err
+	}
+
+	return cleanedMap, nil
+}
+
 // CreateGroup handles the creation of a new group.
 func (s *Server) CreateGroup(c *gin.Context) {
 	var group models.Group
@@ -43,6 +89,17 @@ func (s *Server) CreateGroup(c *gin.Context) {
 		response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, "Channel type is required"))
 		return
 	}
+	if group.TestModel == "" {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, "Test model is required"))
+		return
+	}
+
+	cleanedConfig, err := validateAndCleanConfig(group.Config)
+	if err != nil {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, "Invalid config format"))
+		return
+	}
+	group.Config = cleanedConfig

 	if err := s.DB.Create(&group).Error; err != nil {
 		response.Error(c, app_errors.ParseDBError(err))
@@ -62,6 +119,20 @@ func (s *Server) ListGroups(c *gin.Context) {
 	response.Success(c, groups)
 }

+// GroupUpdateRequest defines the payload for updating a group.
+// Using a dedicated struct avoids issues with zero values being ignored by GORM's Update.
+type GroupUpdateRequest struct {
+	Name           string                 `json:"name"`
+	DisplayName    string                 `json:"display_name"`
+	Description    string                 `json:"description"`
+	Upstreams      json.RawMessage        `json:"upstreams"`
+	ChannelType    string                 `json:"channel_type"`
+	Sort           *int                   `json:"sort"`
+	TestModel      string                 `json:"test_model"`
+	ParamOverrides map[string]interface{} `json:"param_overrides"`
+	Config         map[string]interface{} `json:"config"`
+}
+
 // UpdateGroup handles updating an existing group.
 func (s *Server) UpdateGroup(c *gin.Context) {
 	id, err := strconv.Atoi(c.Param("id"))
@@ -76,84 +147,70 @@ func (s *Server) UpdateGroup(c *gin.Context) {
 		return
 	}

-	var updateData models.Group
-	if err := c.ShouldBindJSON(&updateData); err != nil {
+	var req GroupUpdateRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
 		response.Error(c, app_errors.NewAPIError(app_errors.ErrInvalidJSON, err.Error()))
 		return
 	}

-	// Validate group name if it's being updated
-	if updateData.Name != "" && !isValidGroupName(updateData.Name) {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, "Invalid group name format. Use 3-30 lowercase letters, numbers, and underscores."))
-		return
-	}
-
-	// Use a transaction to ensure atomicity
+	// Start a transaction
 	tx := s.DB.Begin()
 	if tx.Error != nil {
 		response.Error(c, app_errors.ErrDatabase)
 		return
 	}
+	defer tx.Rollback() // Rollback on panic

-	// Convert updateData to a map to ensure zero values (like Sort: 0) are updated
-	var updateMap map[string]interface{}
-	updateBytes, _ := json.Marshal(updateData)
-	if err := json.Unmarshal(updateBytes, &updateMap); err != nil {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Failed to process update data"))
-		return
-	}
-
-	// If config is being updated, it needs to be marshalled to JSON string for GORM
-	if config, ok := updateMap["config"]; ok {
-		if configMap, isMap := config.(map[string]interface{}); isMap {
-			configJSON, err := json.Marshal(configMap)
-			if err != nil {
-				response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Failed to process config data"))
-				return
-			}
-			updateMap["config"] = string(configJSON)
+	// Apply updates from the request
+	if req.Name != "" {
+		if !isValidGroupName(req.Name) {
+			response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, "Invalid group name format."))
+			return
 		}
+		group.Name = req.Name
 	}
-
-	// Handle upstreams field specifically
-	if upstreams, ok := updateMap["upstreams"]; ok {
-		if upstreamsSlice, isSlice := upstreams.([]interface{}); isSlice {
-			upstreamsJSON, err := json.Marshal(upstreamsSlice)
-			if err != nil {
-				response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Failed to process upstreams data"))
-				return
-			}
-			updateMap["upstreams"] = string(upstreamsJSON)
+	if req.DisplayName != "" {
+		group.DisplayName = req.DisplayName
+	}
+	if req.Description != "" {
+		group.Description = req.Description
+	}
+	if req.Upstreams != nil {
+		group.Upstreams = datatypes.JSON(req.Upstreams)
+	}
+	if req.ChannelType != "" {
+		group.ChannelType = req.ChannelType
+	}
+	if req.Sort != nil {
+		group.Sort = *req.Sort
+	}
+	if req.TestModel != "" {
+		group.TestModel = req.TestModel
+	}
+	if req.ParamOverrides != nil {
+		group.ParamOverrides = req.ParamOverrides
+	}
+	if req.Config != nil {
+		cleanedConfig, err := validateAndCleanConfig(req.Config)
+		if err != nil {
+			response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, "Invalid config format"))
+			return
 		}
+		group.Config = cleanedConfig
 	}

-	// Remove fields that are not actual columns or should not be updated from the map
-	delete(updateMap, "id")
-	delete(updateMap, "api_keys")
-	delete(updateMap, "created_at")
-	delete(updateMap, "updated_at")
-
-	// Use Updates with a map to only update provided fields, including zero values
-	if err := tx.Model(&group).Updates(updateMap).Error; err != nil {
-		tx.Rollback()
+	// Save the updated group object
+	if err := tx.Save(&group).Error; err != nil {
 		response.Error(c, app_errors.ParseDBError(err))
 		return
 	}

 	if err := tx.Commit().Error; err != nil {
-		tx.Rollback()
 		response.Error(c, app_errors.ErrDatabase)
 		return
 	}

-	// Re-fetch the group to return the updated data
-	var updatedGroup models.Group
-	if err := s.DB.First(&updatedGroup, id).Error; err != nil {
-		response.Error(c, app_errors.ParseDBError(err))
-		return
-	}
-
-	response.Success(c, updatedGroup)
+	response.Success(c, group)
 }

 // DeleteGroup handles deleting a group.
--- a/internal/handler/handler.go
+++ b/internal/handler/handler.go
@@ -6,6 +6,7 @@ import (
 	"time"

 	"gpt-load/internal/models"
+	"gpt-load/internal/services"
 	"gpt-load/internal/types"

 	"github.com/gin-gonic/gin"
@@ -14,15 +15,30 @@ import (

 // Server contains dependencies for HTTP handlers
 type Server struct {
-	DB     *gorm.DB
-	config types.ConfigManager
+	DB                         *gorm.DB
+	config                     types.ConfigManager
+	KeyValidatorService        *services.KeyValidatorService
+	KeyManualValidationService *services.KeyManualValidationService
+	TaskService                *services.TaskService
+	KeyService                 *services.KeyService
 }

 // NewServer creates a new handler instance
-func NewServer(db *gorm.DB, config types.ConfigManager) *Server {
+func NewServer(
+	db *gorm.DB,
+	config types.ConfigManager,
+	keyValidatorService *services.KeyValidatorService,
+	keyManualValidationService *services.KeyManualValidationService,
+	taskService *services.TaskService,
+	keyService *services.KeyService,
+) *Server {
 	return &Server{
-		DB:     db,
-		config: config,
+		DB:                         db,
+		config:                     config,
+		KeyValidatorService:        keyValidatorService,
+		KeyManualValidationService: keyManualValidationService,
+		TaskService:                taskService,
+		KeyService:                 keyService,
 	}
 }

--- a/internal/handler/key_handler.go
+++ b/internal/handler/key_handler.go
@@ -1,60 +1,123 @@
-// Package handler provides HTTP handlers for the application
 package handler

 import (
+	"fmt"
 	app_errors "gpt-load/internal/errors"
 	"gpt-load/internal/models"
 	"gpt-load/internal/response"
 	"strconv"
+	"strings"

 	"github.com/gin-gonic/gin"
+	"gorm.io/gorm"
 )

-type CreateKeysRequest struct {
-	Keys []string `json:"keys" binding:"required"`
+// validateGroupID validates and parses group ID from request parameter
+func validateGroupID(c *gin.Context) (uint, error) {
+	groupIDStr := c.Param("id")
+	if groupIDStr == "" {
+		return 0, fmt.Errorf("group ID is required")
+	}
+
+	groupID, err := strconv.Atoi(groupIDStr)
+	if err != nil || groupID <= 0 {
+		return 0, fmt.Errorf("invalid group ID format")
+	}
+
+	return uint(groupID), nil
 }

-// CreateKeysInGroup handles creating new keys within a specific group.
-func (s *Server) CreateKeysInGroup(c *gin.Context) {
-	groupID, err := strconv.Atoi(c.Param("id"))
+// validateKeyID validates and parses key ID from request parameter
+func validateKeyID(c *gin.Context) (uint, error) {
+	keyIDStr := c.Param("key_id")
+	if keyIDStr == "" {
+		return 0, fmt.Errorf("key ID is required")
+	}
+
+	keyID, err := strconv.Atoi(keyIDStr)
+	if err != nil || keyID <= 0 {
+		return 0, fmt.Errorf("invalid key ID format")
+	}
+
+	return uint(keyID), nil
+}
+
+// validateKeysText validates the keys text input
+func validateKeysText(keysText string) error {
+	if strings.TrimSpace(keysText) == "" {
+		return fmt.Errorf("keys text cannot be empty")
+	}
+
+	if len(keysText) > 1024*1024 { // 1MB limit
+		return fmt.Errorf("keys text is too large (max 1MB)")
+	}
+
+	return nil
+}
+
+// findGroupByID is a helper function to find a group by its ID.
+func (s *Server) findGroupByID(c *gin.Context, groupID int) (*models.Group, bool) {
+	var group models.Group
+	if err := s.DB.First(&group, groupID).Error; err != nil {
+		if err == gorm.ErrRecordNotFound {
+			response.Error(c, app_errors.ErrResourceNotFound)
+		} else {
+			response.Error(c, app_errors.ParseDBError(err))
+		}
+		return nil, false
+	}
+	return &group, true
+}
+
+// AddMultipleKeysRequest defines the payload for adding multiple keys from a text block.
+type AddMultipleKeysRequest struct {
+	KeysText string `json:"keys_text" binding:"required"`
+}
+
+// AddMultipleKeys handles creating new keys from a text block within a specific group.
+func (s *Server) AddMultipleKeys(c *gin.Context) {
+	groupID, err := validateGroupID(c)
 	if err != nil {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid group ID format"))
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, err.Error()))
 		return
 	}

-	var req CreateKeysRequest
+	var req AddMultipleKeysRequest
 	if err := c.ShouldBindJSON(&req); err != nil {
 		response.Error(c, app_errors.NewAPIError(app_errors.ErrInvalidJSON, err.Error()))
 		return
 	}

-	var newKeys []models.APIKey
-	for _, keyVal := range req.Keys {
-		newKeys = append(newKeys, models.APIKey{
-			GroupID:  uint(groupID),
-			KeyValue: keyVal,
-			Status:   "active",
-		})
+	if err := validateKeysText(req.KeysText); err != nil {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
+		return
 	}

-	if err := s.DB.Create(&newKeys).Error; err != nil {
+	result, err := s.KeyService.AddMultipleKeys(groupID, req.KeysText)
+	if err != nil {
 		response.Error(c, app_errors.ParseDBError(err))
 		return
 	}

-	response.Success(c, newKeys)
+	response.Success(c, result)
 }

 // ListKeysInGroup handles listing all keys within a specific group.
 func (s *Server) ListKeysInGroup(c *gin.Context) {
 	groupID, err := strconv.Atoi(c.Param("id"))
 	if err != nil {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid group ID format"))
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid group ID"))
 		return
 	}

-	var keys []models.APIKey
-	if err := s.DB.Where("group_id = ?", groupID).Find(&keys).Error; err != nil {
+	statusFilter := c.Query("status")
+	if statusFilter != "" && statusFilter != "active" && statusFilter != "inactive" {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, "Invalid status filter"))
+		return
+	}
+
+	keys, err := s.KeyService.ListKeysInGroup(uint(groupID), statusFilter)
+	if err != nil {
 		response.Error(c, app_errors.ParseDBError(err))
 		return
 	}
@@ -62,90 +125,124 @@ func (s *Server) ListKeysInGroup(c *gin.Context) {
 	response.Success(c, keys)
 }

-// UpdateKey handles updating a specific key.
-func (s *Server) UpdateKey(c *gin.Context) {
+// DeleteSingleKey handles deleting a specific key.
+func (s *Server) DeleteSingleKey(c *gin.Context) {
 	groupID, err := strconv.Atoi(c.Param("id"))
 	if err != nil {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid group ID format"))
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid group ID"))
 		return
 	}

 	keyID, err := strconv.Atoi(c.Param("key_id"))
 	if err != nil {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid key ID format"))
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid key ID"))
 		return
 	}

-	var key models.APIKey
-	if err := s.DB.Where("group_id = ? AND id = ?", groupID, keyID).First(&key).Error; err != nil {
+	rowsAffected, err := s.KeyService.DeleteSingleKey(uint(groupID), uint(keyID))
+	if err != nil {
 		response.Error(c, app_errors.ParseDBError(err))
 		return
 	}
-
-	var updateData struct {
-		Status string `json:"status"`
-	}
-	if err := c.ShouldBindJSON(&updateData); err != nil {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrInvalidJSON, err.Error()))
+	if rowsAffected == 0 {
+		response.Error(c, app_errors.ErrResourceNotFound)
 		return
 	}

-	key.Status = updateData.Status
-	if err := s.DB.Save(&key).Error; err != nil {
-		response.Error(c, app_errors.ParseDBError(err))
+	response.Success(c, gin.H{"message": "Key deleted successfully"})
+}
+
+// TestSingleKey handles a one-off validation test for a single key.
+func (s *Server) TestSingleKey(c *gin.Context) {
+	keyID, err := validateKeyID(c)
+	if err != nil {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, err.Error()))
 		return
 	}

-	response.Success(c, key)
+	isValid, validationErr := s.KeyValidatorService.TestSingleKeyByID(c.Request.Context(), keyID)
+	if validationErr != nil {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadGateway, validationErr.Error()))
+		return
+	}
+
+	if isValid {
+		response.Success(c, gin.H{"success": true, "message": "Key is valid."})
+	} else {
+		response.Success(c, gin.H{"success": false, "message": "Key is invalid or has insufficient quota."})
+	}
 }

-type DeleteKeysRequest struct {
-	KeyIDs []uint `json:"key_ids" binding:"required"`
-}
-
-// DeleteKeys handles deleting one or more keys.
-func (s *Server) DeleteKeys(c *gin.Context) {
+// ValidateGroupKeys initiates a manual validation task for all keys in a group.
+func (s *Server) ValidateGroupKeys(c *gin.Context) {
 	groupID, err := strconv.Atoi(c.Param("id"))
 	if err != nil {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid group ID format"))
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid group ID"))
 		return
 	}

-	var req DeleteKeysRequest
-	if err := c.ShouldBindJSON(&req); err != nil {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrInvalidJSON, err.Error()))
+	group, ok := s.findGroupByID(c, groupID)
+	if !ok {
 		return
 	}

-	if len(req.KeyIDs) == 0 {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, "No key IDs provided"))
+	taskStatus, err := s.KeyManualValidationService.StartValidationTask(group)
+	if err != nil {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrTaskInProgress, err.Error()))
 		return
 	}

-	// 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, app_errors.ParseDBError(err))
-		return
-	}
-
-	if count != int64(len(req.KeyIDs)) {
-		tx.Rollback()
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrForbidden, "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, app_errors.ParseDBError(err))
-		return
-	}
-
-	tx.Commit()
-	response.Success(c, gin.H{"message": "Keys deleted successfully"})
+	response.Success(c, taskStatus)
+}
+
+// RestoreAllInvalidKeys sets the status of all 'inactive' keys in a group to 'active'.
+func (s *Server) RestoreAllInvalidKeys(c *gin.Context) {
+	groupID, err := strconv.Atoi(c.Param("id"))
+	if err != nil {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid group ID"))
+		return
+	}
+
+	rowsAffected, err := s.KeyService.RestoreAllInvalidKeys(uint(groupID))
+	if err != nil {
+		response.Error(c, app_errors.ParseDBError(err))
+		return
+	}
+
+	response.Success(c, gin.H{"message": fmt.Sprintf("%d keys restored.", rowsAffected)})
+}
+
+// ClearAllInvalidKeys deletes all 'inactive' keys from a group.
+func (s *Server) ClearAllInvalidKeys(c *gin.Context) {
+	groupID, err := strconv.Atoi(c.Param("id"))
+	if err != nil {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid group ID"))
+		return
+	}
+
+	rowsAffected, err := s.KeyService.ClearAllInvalidKeys(uint(groupID))
+	if err != nil {
+		response.Error(c, app_errors.ParseDBError(err))
+		return
+	}
+
+	response.Success(c, gin.H{"message": fmt.Sprintf("%d invalid keys cleared.", rowsAffected)})
+}
+
+// ExportKeys returns a list of keys for a group, filtered by status.
+func (s *Server) ExportKeys(c *gin.Context) {
+	groupID, err := strconv.Atoi(c.Param("id"))
+	if err != nil {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Invalid group ID"))
+		return
+	}
+
+	filter := c.DefaultQuery("filter", "all")
+	keys, err := s.KeyService.ExportKeys(uint(groupID), filter)
+	if err != nil {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
+		return
+	}
+
+	response.Success(c, gin.H{"keys": keys})
 }
--- a/internal/handler/task_handler.go
+++ b/internal/handler/task_handler.go
@@ -0,0 +1,31 @@
+package handler
+
+import (
+	"gpt-load/internal/response"
+	app_errors "gpt-load/internal/errors"
+
+	"github.com/gin-gonic/gin"
+)
+
+// GetTaskStatus handles requests for the status of the global long-running task.
+func (s *Server) GetTaskStatus(c *gin.Context) {
+	taskStatus := s.TaskService.GetTaskStatus()
+	response.Success(c, taskStatus)
+}
+
+// GetTaskResult handles requests for the result of a finished task.
+func (s *Server) GetTaskResult(c *gin.Context) {
+	taskID := c.Param("task_id")
+	if taskID == "" {
+		response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Task ID is required"))
+		return
+	}
+
+	result, found := s.TaskService.GetResult(taskID)
+	if !found {
+		response.Error(c, app_errors.ErrResourceNotFound)
+		return
+	}
+
+	response.Success(c, result)
+}