// Filename: internal/handlers/apikey_handler.go package handlers import ( "gemini-balancer/internal/errors" "gemini-balancer/internal/models" "gemini-balancer/internal/response" "gemini-balancer/internal/service" "gemini-balancer/internal/task" "strconv" "github.com/gin-gonic/gin" "gorm.io/gorm" ) type APIKeyHandler struct { apiKeyService *service.APIKeyService db *gorm.DB keyImportService *service.KeyImportService keyValidationService *service.KeyValidationService } func NewAPIKeyHandler(apiKeyService *service.APIKeyService, db *gorm.DB, keyImportService *service.KeyImportService, keyValidationService *service.KeyValidationService) *APIKeyHandler { return &APIKeyHandler{ apiKeyService: apiKeyService, db: db, keyImportService: keyImportService, keyValidationService: keyValidationService, } } // DTOs for API requests type BulkAddKeysToGroupRequest struct { KeyGroupID uint `json:"key_group_id" binding:"required"` Keys string `json:"keys" binding:"required"` ValidateOnImport bool `json:"validate_on_import"` // OmitEmpty/default is false } type BulkUnlinkKeysFromGroupRequest struct { KeyGroupID uint `json:"key_group_id" binding:"required"` Keys string `json:"keys" binding:"required"` } type BulkHardDeleteKeysRequest struct { Keys string `json:"keys" binding:"required"` } type BulkRestoreKeysRequest struct { Keys string `json:"keys" binding:"required"` } type UpdateAPIKeyRequest struct { Status *string `json:"status" binding:"omitempty,oneof=ACTIVE,PENDING_VALIDATION,COOLDOWN,DISABLED,BANNED"` } type UpdateMappingRequest struct { Status models.APIKeyStatus `json:"status" binding:"required,oneof=ACTIVE PENDING_VALIDATION COOLDOWN DISABLED BANNED"` } type BulkTestKeysRequest struct { KeyGroupID uint `json:"key_group_id" binding:"required"` Keys string `json:"keys" binding:"required"` } type RestoreKeysRequest struct { KeyIDs []uint `json:"key_ids" binding:"required,gt=0"` } type BulkTestKeysForGroupRequest struct { Keys string `json:"keys" binding:"required"` } type BulkActionFilter struct { Status []string `json:"status"` // Changed to slice to accept multiple statuses } type BulkActionRequest struct { Action string `json:"action" binding:"required,oneof=revalidate set_status delete"` NewStatus string `json:"new_status" binding:"omitempty,oneof=active disabled cooldown banned"` // For 'set_status' action Filter BulkActionFilter `json:"filter" binding:"required"` } // --- Handler Methods --- // AddMultipleKeysToGroup handles adding/linking multiple keys to a specific group. func (h *APIKeyHandler) AddMultipleKeysToGroup(c *gin.Context) { var req BulkAddKeysToGroupRequest if err := c.ShouldBindJSON(&req); err != nil { response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error())) return } taskStatus, err := h.keyImportService.StartAddKeysTask(req.KeyGroupID, req.Keys, req.ValidateOnImport) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error())) return } response.Success(c, taskStatus) } // UnlinkMultipleKeysFromGroup handles unlinking multiple keys from a specific group. func (h *APIKeyHandler) UnlinkMultipleKeysFromGroup(c *gin.Context) { var req BulkUnlinkKeysFromGroupRequest if err := c.ShouldBindJSON(&req); err != nil { response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error())) return } taskStatus, err := h.keyImportService.StartUnlinkKeysTask(req.KeyGroupID, req.Keys) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error())) return } response.Success(c, taskStatus) } // HardDeleteMultipleKeys handles globally deleting multiple key entities. func (h *APIKeyHandler) HardDeleteMultipleKeys(c *gin.Context) { var req BulkHardDeleteKeysRequest if err := c.ShouldBindJSON(&req); err != nil { response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error())) return } taskStatus, err := h.keyImportService.StartHardDeleteKeysTask(req.Keys) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error())) return } response.Success(c, taskStatus) } // RestoreMultipleKeys handles restoring multiple keys to ACTIVE status globally. func (h *APIKeyHandler) RestoreMultipleKeys(c *gin.Context) { var req BulkRestoreKeysRequest if err := c.ShouldBindJSON(&req); err != nil { response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error())) return } taskStatus, err := h.keyImportService.StartRestoreKeysTask(req.Keys) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error())) return } response.Success(c, taskStatus) } func (h *APIKeyHandler) TestMultipleKeys(c *gin.Context) { var req BulkTestKeysRequest if err := c.ShouldBindJSON(&req); err != nil { response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error())) return } taskStatus, err := h.keyValidationService.StartTestKeysTask(req.KeyGroupID, req.Keys) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error())) return } response.Success(c, taskStatus) } func (h *APIKeyHandler) ListAPIKeys(c *gin.Context) { var params models.APIKeyQueryParams if err := c.ShouldBindQuery(¶ms); err != nil { response.Error(c, errors.NewAPIError(errors.ErrBadRequest, err.Error())) return } if params.Page <= 0 { params.Page = 1 } if params.PageSize <= 0 { params.PageSize = 20 } result, err := h.apiKeyService.ListAPIKeys(¶ms) if err != nil { response.Error(c, errors.ParseDBError(err)) return } response.Success(c, result) } // ListKeysForGroup handles the GET /keygroups/:id/keys request. func (h *APIKeyHandler) ListKeysForGroup(c *gin.Context) { // 1. Manually handle the path parameter. groupID, err := strconv.ParseUint(c.Param("id"), 10, 32) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid group ID format")) return } // 2. Bind query parameters using the correctly tagged struct. var params models.APIKeyQueryParams if err := c.ShouldBindQuery(¶ms); err != nil { response.Error(c, errors.NewAPIError(errors.ErrBadRequest, err.Error())) return } // 3. Set server-side defaults and the path parameter. if params.Page <= 0 { params.Page = 1 } if params.PageSize <= 0 { params.PageSize = 20 } params.KeyGroupID = uint(groupID) // 4. Call the service layer. paginatedResult, err := h.apiKeyService.ListAPIKeys(¶ms) if err != nil { response.Error(c, errors.ParseDBError(err)) return } // 5. [THE FIX] Return a successful response using the standard `response.Success` // and a gin.H map, as confirmed to exist in your project. response.Success(c, gin.H{ "items": paginatedResult.Items, "total": paginatedResult.Total, "page": paginatedResult.Page, "pages": paginatedResult.TotalPages, }) } func (h *APIKeyHandler) TestKeysForGroup(c *gin.Context) { // Group ID is now correctly sourced from the URL path. groupID, err := strconv.ParseUint(c.Param("id"), 10, 32) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid group ID format")) return } // The request body is now simpler, only needing the keys. var req BulkTestKeysForGroupRequest if err := c.ShouldBindJSON(&req); err != nil { response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error())) return } // Call the same underlying service, but with unambiguous context. taskStatus, err := h.keyValidationService.StartTestKeysTask(uint(groupID), req.Keys) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrTaskInProgress, err.Error())) return } response.Success(c, taskStatus) } // UpdateAPIKey is DEPRECATED. Status is now contextual to a group. func (h *APIKeyHandler) UpdateAPIKey(c *gin.Context) { err := errors.NewAPIError(errors.ErrBadRequest, "This endpoint is deprecated. Use 'PUT /keygroups/:id/apikeys/:keyId' to update key status within a group context.") response.Error(c, err) } // UpdateGroupAPIKeyMapping handles updating a key's status within a specific group. // Route: PUT /keygroups/:id/apikeys/:keyId func (h *APIKeyHandler) UpdateGroupAPIKeyMapping(c *gin.Context) { groupID, err := strconv.ParseUint(c.Param("id"), 10, 32) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Group ID format")) return } keyID, err := strconv.ParseUint(c.Param("keyId"), 10, 32) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Key ID format")) return } var req UpdateMappingRequest if err := c.ShouldBindJSON(&req); err != nil { response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error())) return } // Directly use the service to handle the logic updatedMapping, err := h.apiKeyService.UpdateMappingStatus(uint(groupID), uint(keyID), req.Status) if err != nil { var apiErr *errors.APIError if errors.As(err, &apiErr) { response.Error(c, apiErr) } else { response.Error(c, errors.ParseDBError(err)) } return } response.Success(c, updatedMapping) } // HardDeleteAPIKey handles globally deleting a single key entity. func (h *APIKeyHandler) HardDeleteAPIKey(c *gin.Context) { id, err := strconv.Atoi(c.Param("id")) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid ID format")) return } if err := h.apiKeyService.HardDeleteAPIKeyByID(uint(id)); err != nil { response.Error(c, errors.ParseDBError(err)) return } response.Success(c, gin.H{"message": "API key globally deleted successfully"}) } // RestoreKeysInGroup 恢复指定Key的接口 // POST /keygroups/:id/apikeys/restore func (h *APIKeyHandler) RestoreKeysInGroup(c *gin.Context) { groupID, err := strconv.ParseUint(c.Param("id"), 10, 32) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Group ID format")) return } var req RestoreKeysRequest if err := c.ShouldBindJSON(&req); err != nil { response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error())) return } taskStatus, err := h.apiKeyService.StartRestoreKeysTask(uint(groupID), req.KeyIDs) if err != nil { var apiErr *errors.APIError if errors.As(err, &apiErr) { response.Error(c, apiErr) } else { response.Error(c, errors.NewAPIError(errors.ErrInternalServer, err.Error())) } return } response.Success(c, taskStatus) } // RestoreAllBannedInGroup 一键恢复所有Banned Key的接口 // POST /keygroups/:id/apikeys/restore-all-banned func (h *APIKeyHandler) RestoreAllBannedInGroup(c *gin.Context) { groupID, err := strconv.ParseUint(c.Param("groupId"), 10, 32) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Group ID format")) return } taskStatus, err := h.apiKeyService.StartRestoreAllBannedTask(uint(groupID)) if err != nil { var apiErr *errors.APIError if errors.As(err, &apiErr) { response.Error(c, apiErr) } else { response.Error(c, errors.NewAPIError(errors.ErrInternalServer, err.Error())) } return } response.Success(c, taskStatus) } // HandleBulkAction handles generic bulk actions on a key group based on server-side filters. // Route: POST /keygroups/:id/bulk-actions func (h *APIKeyHandler) HandleBulkAction(c *gin.Context) { // 1. Parse GroupID from URL groupID, err := strconv.ParseUint(c.Param("id"), 10, 32) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Group ID format")) return } // 2. Bind the JSON payload to our new DTO var req BulkActionRequest if err := c.ShouldBindJSON(&req); err != nil { response.Error(c, errors.NewAPIError(errors.ErrInvalidJSON, err.Error())) return } // 3. Central logic: based on the action, call the appropriate service method. var task *task.Status var apiErr *errors.APIError switch req.Action { case "revalidate": // Assume keyValidationService has a method that accepts a filter task, err = h.keyValidationService.StartTestKeysByFilterTask(uint(groupID), req.Filter.Status) case "set_status": if req.NewStatus == "" { apiErr = errors.NewAPIError(errors.ErrBadRequest, "new_status is required for set_status action") break } // Assume apiKeyService has a method to update status by filter targetStatus := models.APIKeyStatus(req.NewStatus) // Convert string to your model's type task, err = h.apiKeyService.StartUpdateStatusByFilterTask(uint(groupID), req.Filter.Status, targetStatus) case "delete": // Assume keyImportService has a method to unlink by filter task, err = h.keyImportService.StartUnlinkKeysByFilterTask(uint(groupID), req.Filter.Status) default: apiErr = errors.NewAPIError(errors.ErrBadRequest, "Unsupported action: "+req.Action) } // 4. Handle errors from the switch block if apiErr != nil { response.Error(c, apiErr) return } if err != nil { // Attempt to parse it as a known APIError, otherwise, wrap it. var parsedErr *errors.APIError if errors.As(err, &parsedErr) { response.Error(c, parsedErr) } else { response.Error(c, errors.NewAPIError(errors.ErrInternalServer, err.Error())) } return } // 5. Return the task status on success response.Success(c, task) } // ExportKeysForGroup handles requests to export all keys for a group based on status filters. // Route: GET /keygroups/:id/apikeys/export func (h *APIKeyHandler) ExportKeysForGroup(c *gin.Context) { groupID, err := strconv.ParseUint(c.Param("id"), 10, 32) if err != nil { response.Error(c, errors.NewAPIError(errors.ErrBadRequest, "Invalid Group ID format")) return } // Use QueryArray to correctly parse `status[]=active&status[]=cooldown` statuses := c.QueryArray("status") keyStrings, err := h.apiKeyService.GetAPIKeyStringsForExport(uint(groupID), statuses) if err != nil { response.Error(c, errors.ParseDBError(err)) return } response.Success(c, keyStrings) }