diff --git a/internal/handler/key_handler.go b/internal/handler/key_handler.go index a555ee1..0344829 100644 --- a/internal/handler/key_handler.go +++ b/internal/handler/key_handler.go @@ -157,6 +157,36 @@ func (s *Server) DeleteMultipleKeys(c *gin.Context) { response.Success(c, result) } +// RestoreMultipleKeys handles restoring keys from a text block within a specific group. +func (s *Server) RestoreMultipleKeys(c *gin.Context) { + var req KeyTextRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.Error(c, app_errors.NewAPIError(app_errors.ErrInvalidJSON, err.Error())) + return + } + + if _, ok := s.findGroupByID(c, req.GroupID); !ok { + return + } + + if err := validateKeysText(req.KeysText); err != nil { + response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error())) + return + } + + result, err := s.KeyService.RestoreMultipleKeys(req.GroupID, req.KeysText) + if err != nil { + if err.Error() == "no valid keys found in the input text" { + response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error())) + } else { + response.Error(c, app_errors.ParseDBError(err)) + } + return + } + + response.Success(c, result) +} + // TestMultipleKeys handles a one-off validation test for multiple keys. func (s *Server) TestMultipleKeys(c *gin.Context) { var req KeyTextRequest diff --git a/internal/keypool/provider.go b/internal/keypool/provider.go index 2b45e21..96a768c 100644 --- a/internal/keypool/provider.go +++ b/internal/keypool/provider.go @@ -362,6 +362,55 @@ func (p *KeyProvider) RestoreKeys(groupID uint) (int64, error) { return restoredCount, err } +// RestoreMultipleKeys 恢复指定的 Key。 +func (p *KeyProvider) RestoreMultipleKeys(groupID uint, keyValues []string) (int64, error) { + if len(keyValues) == 0 { + return 0, nil + } + + var keysToRestore []models.APIKey + var restoredCount int64 + + err := p.db.Transaction(func(tx *gorm.DB) error { + // 1. 查找要恢复的密钥 + if err := tx.Where("group_id = ? AND key_value IN ? AND status = ?", groupID, keyValues, models.KeyStatusInvalid).Find(&keysToRestore).Error; err != nil { + return err + } + + if len(keysToRestore) == 0 { + return nil + } + + keyIDsToRestore := pluckIDs(keysToRestore) + + // 2. 更新数据库中的状态 + updates := map[string]any{ + "status": models.KeyStatusActive, + "failure_count": 0, + } + result := tx.Model(&models.APIKey{}).Where("id IN ?", keyIDsToRestore).Updates(updates) + if result.Error != nil { + return result.Error + } + restoredCount = result.RowsAffected + + // 3. 将密钥添加回 Redis + for _, key := range keysToRestore { + key.Status = models.KeyStatusActive + key.FailureCount = 0 + if err := p.addKeyToStore(&key); err != nil { + // 在事务中,单个失败会回滚整个事务,但这里的日志记录仍然有用 + logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to restore key in store after DB update") + return err // 返回错误以回滚事务 + } + } + + return nil + }) + + return restoredCount, err +} + // RemoveInvalidKeys 移除组内所有无效的 Key。 func (p *KeyProvider) RemoveInvalidKeys(groupID uint) (int64, error) { var invalidKeys []models.APIKey diff --git a/internal/router/router.go b/internal/router/router.go index 62a7e0d..5901e55 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -117,6 +117,7 @@ func registerProtectedAPIRoutes(api *gin.RouterGroup, serverHandler *handler.Ser keys.GET("", serverHandler.ListKeysInGroup) keys.POST("/add-multiple", serverHandler.AddMultipleKeys) keys.POST("/delete-multiple", serverHandler.DeleteMultipleKeys) + keys.POST("/restore-multiple", serverHandler.RestoreMultipleKeys) keys.POST("/restore-all-invalid", serverHandler.RestoreAllInvalidKeys) keys.POST("/clear-all-invalid", serverHandler.ClearAllInvalidKeys) keys.POST("/validate-group", serverHandler.ValidateGroupKeys) diff --git a/internal/services/key_service.go b/internal/services/key_service.go index f737fea..d0ee85a 100644 --- a/internal/services/key_service.go +++ b/internal/services/key_service.go @@ -25,6 +25,13 @@ type DeleteKeysResult struct { TotalInGroup int64 `json:"total_in_group"` } +// RestoreKeysResult holds the result of restoring multiple keys. +type RestoreKeysResult struct { + RestoredCount int `json:"restored_count"` + IgnoredCount int `json:"ignored_count"` + TotalInGroup int64 `json:"total_in_group"` +} + // KeyService provides services related to API keys. type KeyService struct { DB *gorm.DB @@ -154,6 +161,32 @@ func (s *KeyService) isValidKeyFormat(key string) bool { return validChars.MatchString(key) } +// RestoreMultipleKeys handles the business logic of restoring keys from a text block. +func (s *KeyService) RestoreMultipleKeys(groupID uint, keysText string) (*RestoreKeysResult, error) { + keysToRestore := s.ParseKeysFromText(keysText) + if len(keysToRestore) == 0 { + return nil, fmt.Errorf("no valid keys found in the input text") + } + + restoredCount, err := s.KeyProvider.RestoreMultipleKeys(groupID, keysToRestore) + if err != nil { + return nil, err + } + + ignoredCount := len(keysToRestore) - int(restoredCount) + + var totalInGroup int64 + if err := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup).Error; err != nil { + return nil, err + } + + return &RestoreKeysResult{ + RestoredCount: int(restoredCount), + IgnoredCount: ignoredCount, + TotalInGroup: totalInGroup, + }, nil +} + // RestoreAllInvalidKeys sets the status of all 'inactive' keys in a group to 'active'. func (s *KeyService) RestoreAllInvalidKeys(groupID uint) (int64, error) { return s.KeyProvider.RestoreKeys(groupID)