fork-it/gpt-load
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: 调整key验证结构

Browse Source
This commit is contained in:
tbphp
2025-07-07 20:37:45 +08:00
parent 231708896e
commit e461f1af5f
6 changed files with 25 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
internal/services/key_manual_validation_service.go
View File

@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"gpt-load/internal/config"
"gpt-load/internal/keypool"
"gpt-load/internal/models"
"gpt-load/internal/types"
"sync"
@@ -23,14 +24,14 @@ type ManualValidationResult struct {
// KeyManualValidationService handles user-initiated key validation for a group.
type KeyManualValidationService struct {
DB *gorm.DB
Validator *KeyValidatorService
Validator *keypool.KeyValidator
TaskService *TaskService
SettingsManager *config.SystemSettingsManager
ConfigManager types.ConfigManager
}
// NewKeyManualValidationService creates a new KeyManualValidationService.
func NewKeyManualValidationService(db *gorm.DB, validator *KeyValidatorService, taskService *TaskService, settingsManager *config.SystemSettingsManager, configManager types.ConfigManager) *KeyManualValidationService {
func NewKeyManualValidationService(db *gorm.DB, validator *keypool.KeyValidator, taskService *TaskService, settingsManager *config.SystemSettingsManager, configManager types.ConfigManager) *KeyManualValidationService {
return &KeyManualValidationService{
DB: db,
Validator: validator,

5
internal/services/key_validation_concurrency.go
View File

@@ -2,6 +2,7 @@ package services
import (
"context"
"gpt-load/internal/keypool"
"gpt-load/internal/models"
"gpt-load/internal/types"
"sync"
@@ -27,7 +28,7 @@ type ValidationResult struct {
// KeyValidationPool manages a global worker pool for key validation.
type KeyValidationPool struct {
validator *KeyValidatorService
validator *keypool.KeyValidator
configManager types.ConfigManager
jobs chan ValidationJob
results chan ValidationResult // 定时任务结果
@@ -36,7 +37,7 @@ type KeyValidationPool struct {
}
// NewKeyValidationPool creates a new KeyValidationPool.
func NewKeyValidationPool(validator *KeyValidatorService, configManager types.ConfigManager) *KeyValidationPool {
func NewKeyValidationPool(validator *keypool.KeyValidator, configManager types.ConfigManager) *KeyValidationPool {
return &KeyValidationPool{
validator: validator,
configManager: configManager,

118
internal/services/key_validator_service.go
View File

@@ -1,118 +0,0 @@
package services
import (
"context"
"fmt"
"gpt-load/internal/channel"
"gpt-load/internal/config"
"gpt-load/internal/keypool"
"gpt-load/internal/models"
"github.com/sirupsen/logrus"
"go.uber.org/dig"
"gorm.io/gorm"
)
// KeyTestResult holds the validation result for a single key.
type KeyTestResult struct {
KeyValue string `json:"key_value"`
IsValid bool `json:"is_valid"`
Error string `json:"error,omitempty"`
}
// KeyValidatorService provides methods to validate API keys.
type KeyValidatorService struct {
DB *gorm.DB
channelFactory *channel.Factory
SettingsManager *config.SystemSettingsManager
keypoolProvider *keypool.KeyProvider
}
type KeyValidatorServiceParams struct {
dig.In
DB *gorm.DB
ChannelFactory *channel.Factory
SettingsManager *config.SystemSettingsManager
KeypoolProvider *keypool.KeyProvider
}
// NewKeyValidatorService creates a new KeyValidatorService.
func NewKeyValidatorService(params KeyValidatorServiceParams) *KeyValidatorService {
return &KeyValidatorService{
DB: params.DB,
channelFactory: params.ChannelFactory,
SettingsManager: params.SettingsManager,
keypoolProvider: params.KeypoolProvider,
}
}
// ValidateSingleKey performs a validation check on a single API key.
func (s *KeyValidatorService) ValidateSingleKey(ctx context.Context, key *models.APIKey, group *models.Group) (bool, error) {
if ctx.Err() != nil {
return false, fmt.Errorf("context cancelled or timed out: %w", ctx.Err())
}
ch, err := s.channelFactory.GetChannel(group)
if err != nil {
return false, fmt.Errorf("failed to get channel for group %s: %w", group.Name, err)
}
isValid, validationErr := ch.ValidateKey(ctx, key.KeyValue)
s.keypoolProvider.UpdateStatus(key.ID, group.ID, isValid)
if !isValid {
logrus.WithFields(logrus.Fields{
"error": validationErr,
"key_id": key.ID,
"group_id": group.ID,
}).Debug("Key validation failed")
return false, validationErr
}
logrus.WithFields(logrus.Fields{
"key_id": key.ID,
"is_valid": isValid,
}).Debug("Key validation successful")
return true, nil
}
// TestMultipleKeys performs a synchronous validation for a list of key values within a specific group.
func (s *KeyValidatorService) TestMultipleKeys(ctx context.Context, group *models.Group, keyValues []string) ([]KeyTestResult, error) {
results := make([]KeyTestResult, len(keyValues))
// Find which of the provided keys actually exist in the database for this group
var existingKeys []models.APIKey
if err := s.DB.Where("group_id = ? AND key_value IN ?", group.ID, keyValues).Find(&existingKeys).Error; err != nil {
return nil, fmt.Errorf("failed to query keys from DB: %w", err)
}
existingKeyMap := make(map[string]models.APIKey)
for _, k := range existingKeys {
existingKeyMap[k.KeyValue] = k
}
for i, kv := range keyValues {
apiKey, exists := existingKeyMap[kv]
if !exists {
results[i] = KeyTestResult{
KeyValue: kv,
IsValid: false,
Error: "Key does not exist in this group or has been removed.",
}
continue
}
isValid, validationErr := s.ValidateSingleKey(ctx, &apiKey, group)
results[i] = KeyTestResult{
KeyValue: kv,
IsValid: isValid,
Error: "",
}
if validationErr != nil {
results[i].Error = validationErr.Error()
}
}
return results, nil
}