583 lines
13 KiB
Go
583 lines
13 KiB
Go
package fqapi
|
|
|
|
import (
|
|
"fanqie/internal/common"
|
|
"fanqie/internal/config"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
poolLogger = common.NewLogger("device_key_pool", "INFO", "")
|
|
globalPool *DeviceKeyPool
|
|
poolOnce sync.Once
|
|
)
|
|
|
|
// DeviceKeyItem 设备密钥项
|
|
type DeviceKeyItem struct {
|
|
ID int `json:"id"`
|
|
DeviceHeader string `json:"device_header"`
|
|
Key string `json:"key"`
|
|
KeyVer int `json:"key_ver"`
|
|
IsValid bool `json:"is_valid"`
|
|
CreatedAt time.Time `json:"created_at"`
|
|
LastUsedAt time.Time `json:"last_used_at"`
|
|
}
|
|
|
|
// DeviceKeyPool 设备密钥池
|
|
type DeviceKeyPool struct {
|
|
mu sync.RWMutex
|
|
currentIndex int
|
|
items []*DeviceKeyItem
|
|
enabled bool
|
|
}
|
|
|
|
// GetDeviceKeyPool 获取全局设备密钥池实例
|
|
func GetDeviceKeyPool() *DeviceKeyPool {
|
|
poolOnce.Do(func() {
|
|
globalPool = &DeviceKeyPool{
|
|
items: make([]*DeviceKeyItem, 0),
|
|
enabled: false,
|
|
}
|
|
})
|
|
return globalPool
|
|
}
|
|
|
|
// Initialize 初始化设备密钥池
|
|
func (p *DeviceKeyPool) Initialize() error {
|
|
cfg := config.GetConfig()
|
|
if cfg == nil || !cfg.DeviceKeyPool.Enabled {
|
|
poolLogger.Info("设备密钥池未启用")
|
|
p.enabled = false
|
|
return nil
|
|
}
|
|
|
|
p.enabled = true
|
|
poolLogger.Info("开始初始化设备密钥池...")
|
|
|
|
// 从数据库加载已有的设备
|
|
if err := p.loadFromDB(); err != nil {
|
|
poolLogger.Warning("从数据库加载设备失败: " + err.Error())
|
|
}
|
|
|
|
// 检查并补充设备到初始化数量
|
|
validCount := p.getValidCount()
|
|
initialSize := cfg.DeviceKeyPool.InitialSize
|
|
|
|
if validCount < initialSize {
|
|
needCount := initialSize - validCount
|
|
poolLogger.Info(fmt.Sprintf("当前有效设备数: %d, 需要补充: %d", validCount, needCount))
|
|
|
|
for i := 0; i < needCount; i++ {
|
|
if err := p.addNewDevice(); err != nil {
|
|
poolLogger.Error(fmt.Sprintf("添加设备失败: %v", err))
|
|
}
|
|
}
|
|
}
|
|
|
|
poolLogger.Info(fmt.Sprintf("设备密钥池初始化完成，当前有效设备数: %d", p.getValidCount()))
|
|
return nil
|
|
}
|
|
|
|
// IsEnabled 检查设备密钥池是否启用
|
|
func (p *DeviceKeyPool) IsEnabled() bool {
|
|
return p.enabled
|
|
}
|
|
|
|
// GetNext 获取下一个可用的设备密钥（轮询方式）
|
|
func (p *DeviceKeyPool) GetNext() (*DeviceKeyItem, error) {
|
|
if !p.enabled {
|
|
return nil, fmt.Errorf("设备密钥池未启用")
|
|
}
|
|
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
|
|
// 过滤有效设备
|
|
validItems := make([]*DeviceKeyItem, 0)
|
|
for _, item := range p.items {
|
|
if item.IsValid {
|
|
validItems = append(validItems, item)
|
|
}
|
|
}
|
|
|
|
if len(validItems) == 0 {
|
|
return nil, fmt.Errorf("没有可用的设备")
|
|
}
|
|
|
|
// 轮询获取下一个设备
|
|
item := validItems[p.currentIndex%len(validItems)]
|
|
p.currentIndex++
|
|
|
|
// 更新最后使用时间
|
|
item.LastUsedAt = time.Now()
|
|
p.updateDeviceInDB(item)
|
|
|
|
// poolLogger.Debug(fmt.Sprintf("使用设备 ID=%d", item.ID))
|
|
|
|
// 异步检查是否需要补充
|
|
go p.checkAndRefill()
|
|
|
|
return item, nil
|
|
}
|
|
|
|
// MarkInvalid 标记设备为无效
|
|
func (p *DeviceKeyPool) MarkInvalid(deviceHeader string) error {
|
|
if !p.enabled {
|
|
return nil
|
|
}
|
|
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
|
|
for _, item := range p.items {
|
|
if item.DeviceHeader == deviceHeader {
|
|
item.IsValid = false
|
|
// poolLogger.Warning(fmt.Sprintf("设备 ID=%d 已标记为无效", item.ID))
|
|
|
|
// 更新数据库
|
|
if err := p.updateDeviceInDB(item); err != nil {
|
|
return err
|
|
}
|
|
|
|
// 异步补充新设备
|
|
go p.checkAndRefill()
|
|
|
|
return nil
|
|
}
|
|
}
|
|
|
|
return fmt.Errorf("未找到设备")
|
|
}
|
|
|
|
// RemoveInvalid 移除所有无效设备并自动补充
|
|
func (p *DeviceKeyPool) RemoveInvalid() error {
|
|
if !p.enabled {
|
|
return nil
|
|
}
|
|
|
|
p.mu.Lock()
|
|
|
|
// 从数据库中删除无效设备
|
|
if err := p.removeInvalidFromDB(); err != nil {
|
|
p.mu.Unlock()
|
|
return err
|
|
}
|
|
|
|
// 从内存中移除无效设备
|
|
validItems := make([]*DeviceKeyItem, 0)
|
|
for _, item := range p.items {
|
|
if item.IsValid {
|
|
validItems = append(validItems, item)
|
|
}
|
|
}
|
|
|
|
removedCount := len(p.items) - len(validItems)
|
|
p.items = validItems
|
|
|
|
if removedCount > 0 {
|
|
poolLogger.Info(fmt.Sprintf("已移除 %d 个无效设备", removedCount))
|
|
|
|
// 立即检查并补充设备（使用已持有锁的版本）
|
|
p.checkAndRefillLocked()
|
|
}
|
|
|
|
p.mu.Unlock()
|
|
return nil
|
|
}
|
|
|
|
// GetStats 获取池统计信息
|
|
func (p *DeviceKeyPool) GetStats() map[string]interface{} {
|
|
p.mu.RLock()
|
|
defer p.mu.RUnlock()
|
|
|
|
validCount := 0
|
|
invalidCount := 0
|
|
|
|
for _, item := range p.items {
|
|
if item.IsValid {
|
|
validCount++
|
|
} else {
|
|
invalidCount++
|
|
}
|
|
}
|
|
|
|
return map[string]interface{}{
|
|
"enabled": p.enabled,
|
|
"total": len(p.items),
|
|
"valid": validCount,
|
|
"invalid": invalidCount,
|
|
"current_index": p.currentIndex,
|
|
}
|
|
}
|
|
|
|
// checkAndRefill 检查并补充设备（私有方法，会自动加锁）
|
|
func (p *DeviceKeyPool) checkAndRefill() {
|
|
if !p.enabled {
|
|
return
|
|
}
|
|
|
|
cfg := config.GetConfig()
|
|
if cfg == nil {
|
|
return
|
|
}
|
|
|
|
p.mu.Lock()
|
|
validCount := p.getValidCount()
|
|
p.mu.Unlock()
|
|
|
|
minSize := cfg.DeviceKeyPool.MinSize
|
|
maxSize := cfg.DeviceKeyPool.MaxSize
|
|
|
|
if validCount < minSize {
|
|
needCount := minSize - validCount
|
|
// 限制不超过最大值
|
|
if validCount+needCount > maxSize {
|
|
needCount = maxSize - validCount
|
|
}
|
|
|
|
if needCount > 0 {
|
|
poolLogger.Info(fmt.Sprintf("有效设备数 %d 低于最小值 %d，开始补充 %d 个设备", validCount, minSize, needCount))
|
|
|
|
for i := 0; i < needCount; i++ {
|
|
if err := p.addNewDevice(); err != nil {
|
|
poolLogger.Error(fmt.Sprintf("补充设备失败: %v", err))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// checkAndRefillLocked 检查并补充设备（已持有锁的内部版本）
|
|
func (p *DeviceKeyPool) checkAndRefillLocked() {
|
|
if !p.enabled {
|
|
return
|
|
}
|
|
|
|
cfg := config.GetConfig()
|
|
if cfg == nil {
|
|
return
|
|
}
|
|
|
|
validCount := p.getValidCount()
|
|
minSize := cfg.DeviceKeyPool.MinSize
|
|
maxSize := cfg.DeviceKeyPool.MaxSize
|
|
|
|
if validCount < minSize {
|
|
needCount := minSize - validCount
|
|
// 限制不超过最大值
|
|
if validCount+needCount > maxSize {
|
|
needCount = maxSize - validCount
|
|
}
|
|
|
|
if needCount > 0 {
|
|
poolLogger.Info(fmt.Sprintf("有效设备数 %d 低于最小值 %d，开始补充 %d 个设备", validCount, minSize, needCount))
|
|
|
|
// 临时释放锁，因为addNewDevice需要加锁
|
|
p.mu.Unlock()
|
|
for i := 0; i < needCount; i++ {
|
|
if err := p.addNewDevice(); err != nil {
|
|
poolLogger.Error(fmt.Sprintf("补充设备失败: %v", err))
|
|
}
|
|
}
|
|
p.mu.Lock()
|
|
}
|
|
}
|
|
}
|
|
|
|
// getValidCount 获取有效设备数量（调用前需已加锁或在RLock内）
|
|
func (p *DeviceKeyPool) getValidCount() int {
|
|
count := 0
|
|
for _, item := range p.items {
|
|
if item.IsValid {
|
|
count++
|
|
}
|
|
}
|
|
return count
|
|
}
|
|
|
|
// addNewDevice 添加新设备（会自动加锁）
|
|
func (p *DeviceKeyPool) addNewDevice() error {
|
|
// 注册新设备
|
|
deviceHeader, err := Register()
|
|
if err != nil {
|
|
return fmt.Errorf("注册设备失败: %w", err)
|
|
}
|
|
|
|
// 获取密钥
|
|
keyInfo, err := GetKey(deviceHeader)
|
|
if err != nil {
|
|
return fmt.Errorf("获取密钥失败: %w", err)
|
|
}
|
|
|
|
// 创建设备项
|
|
item := &DeviceKeyItem{
|
|
DeviceHeader: deviceHeader,
|
|
Key: keyInfo.Key,
|
|
KeyVer: keyInfo.Ver,
|
|
IsValid: true,
|
|
CreatedAt: time.Now(),
|
|
LastUsedAt: time.Now(),
|
|
}
|
|
|
|
// 保存到数据库
|
|
if err := p.saveDeviceToDB(item); err != nil {
|
|
return fmt.Errorf("保存设备到数据库失败: %w", err)
|
|
}
|
|
|
|
// 添加到内存
|
|
p.mu.Lock()
|
|
p.items = append(p.items, item)
|
|
p.mu.Unlock()
|
|
|
|
// poolLogger.Info(fmt.Sprintf("已添加新设备 ID=%d", item.ID))
|
|
return nil
|
|
}
|
|
|
|
// ===== 数据库操作方法 =====
|
|
|
|
// loadFromDB 从数据库加载设备
|
|
func (p *DeviceKeyPool) loadFromDB() error {
|
|
if cacheDB == nil {
|
|
return fmt.Errorf("数据库未初始化")
|
|
}
|
|
|
|
rows, err := cacheDB.Query(`
|
|
SELECT id, device_header, key, key_ver, is_valid, created_at, last_used_at
|
|
FROM device_pool
|
|
ORDER BY id
|
|
`)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer rows.Close()
|
|
|
|
items := make([]*DeviceKeyItem, 0)
|
|
for rows.Next() {
|
|
var item DeviceKeyItem
|
|
var createdAt, lastUsedAt int64
|
|
|
|
err := rows.Scan(
|
|
&item.ID,
|
|
&item.DeviceHeader,
|
|
&item.Key,
|
|
&item.KeyVer,
|
|
&item.IsValid,
|
|
&createdAt,
|
|
&lastUsedAt,
|
|
)
|
|
if err != nil {
|
|
poolLogger.Error("扫描行失败: " + err.Error())
|
|
continue
|
|
}
|
|
|
|
item.CreatedAt = time.Unix(createdAt, 0)
|
|
item.LastUsedAt = time.Unix(lastUsedAt, 0)
|
|
items = append(items, &item)
|
|
}
|
|
|
|
p.mu.Lock()
|
|
p.items = items
|
|
p.mu.Unlock()
|
|
|
|
poolLogger.Info(fmt.Sprintf("从数据库加载了 %d 个设备", len(items)))
|
|
return nil
|
|
}
|
|
|
|
// saveDeviceToDB 保存设备到数据库
|
|
func (p *DeviceKeyPool) saveDeviceToDB(item *DeviceKeyItem) error {
|
|
if cacheDB == nil {
|
|
return fmt.Errorf("数据库未初始化")
|
|
}
|
|
|
|
cacheMutex.Lock()
|
|
defer cacheMutex.Unlock()
|
|
|
|
result, err := cacheDB.Exec(`
|
|
INSERT INTO device_pool (device_header, key, key_ver, is_valid, created_at, last_used_at)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
`, item.DeviceHeader, item.Key, item.KeyVer, item.IsValid,
|
|
item.CreatedAt.Unix(), item.LastUsedAt.Unix())
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// 获取插入的ID
|
|
id, err := result.LastInsertId()
|
|
if err == nil {
|
|
item.ID = int(id)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// updateDeviceInDB 更新设备信息
|
|
func (p *DeviceKeyPool) updateDeviceInDB(item *DeviceKeyItem) error {
|
|
if cacheDB == nil {
|
|
return fmt.Errorf("数据库未初始化")
|
|
}
|
|
|
|
cacheMutex.Lock()
|
|
defer cacheMutex.Unlock()
|
|
|
|
_, err := cacheDB.Exec(`
|
|
UPDATE device_pool
|
|
SET device_header = ?, key = ?, key_ver = ?, is_valid = ?, last_used_at = ?
|
|
WHERE id = ?
|
|
`, item.DeviceHeader, item.Key, item.KeyVer, item.IsValid,
|
|
item.LastUsedAt.Unix(), item.ID)
|
|
|
|
return err
|
|
}
|
|
|
|
// removeInvalidFromDB 从数据库移除无效设备
|
|
func (p *DeviceKeyPool) removeInvalidFromDB() error {
|
|
if cacheDB == nil {
|
|
return fmt.Errorf("数据库未初始化")
|
|
}
|
|
|
|
cacheMutex.Lock()
|
|
defer cacheMutex.Unlock()
|
|
|
|
result, err := cacheDB.Exec(`DELETE FROM device_pool WHERE is_valid = 0`)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
affected, _ := result.RowsAffected()
|
|
if affected > 0 {
|
|
poolLogger.Info(fmt.Sprintf("从数据库删除了 %d 个无效设备", affected))
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// InitDevicePoolTable 初始化设备池表
|
|
func InitDevicePoolTable() error {
|
|
if cacheDB == nil {
|
|
return fmt.Errorf("数据库未初始化")
|
|
}
|
|
|
|
_, err := cacheDB.Exec(`CREATE TABLE IF NOT EXISTS device_pool (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
device_header TEXT NOT NULL,
|
|
key TEXT NOT NULL,
|
|
key_ver INTEGER NOT NULL,
|
|
is_valid INTEGER DEFAULT 1,
|
|
created_at INTEGER NOT NULL,
|
|
last_used_at INTEGER NOT NULL
|
|
)`)
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// 创建索引
|
|
_, err = cacheDB.Exec(`CREATE INDEX IF NOT EXISTS idx_device_pool_valid ON device_pool(is_valid)`)
|
|
if err != nil {
|
|
poolLogger.Warning("创建索引失败: " + err.Error())
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// GetDeviceHeaderFromPool 从池中获取设备header（兼容旧接口）
|
|
func GetDeviceHeaderFromPool() string {
|
|
pool := GetDeviceKeyPool()
|
|
if !pool.IsEnabled() {
|
|
return ""
|
|
}
|
|
|
|
item, err := pool.GetNext()
|
|
if err != nil {
|
|
poolLogger.Error("从池中获取设备失败: " + err.Error())
|
|
return ""
|
|
}
|
|
|
|
return item.DeviceHeader
|
|
}
|
|
|
|
// GetKeyInfoFromPool 从池中获取密钥信息（兼容旧接口）
|
|
func GetKeyInfoFromPool(deviceHeader string) *KeyInfo {
|
|
pool := GetDeviceKeyPool()
|
|
if !pool.IsEnabled() {
|
|
return nil
|
|
}
|
|
|
|
pool.mu.RLock()
|
|
defer pool.mu.RUnlock()
|
|
|
|
for _, item := range pool.items {
|
|
if item.DeviceHeader == deviceHeader && item.IsValid {
|
|
return &KeyInfo{
|
|
Key: item.Key,
|
|
Ver: item.KeyVer,
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// UpdateKeyInPool 更新池中设备的密钥
|
|
func UpdateKeyInPool(deviceHeader string, key string, ver int) error {
|
|
pool := GetDeviceKeyPool()
|
|
if !pool.IsEnabled() {
|
|
return nil
|
|
}
|
|
|
|
pool.mu.Lock()
|
|
defer pool.mu.Unlock()
|
|
|
|
for _, item := range pool.items {
|
|
if item.DeviceHeader == deviceHeader {
|
|
item.Key = key
|
|
item.KeyVer = ver
|
|
return pool.updateDeviceInDB(item)
|
|
}
|
|
}
|
|
|
|
return fmt.Errorf("未找到设备")
|
|
}
|
|
|
|
// GetAnyValidDeviceHeader 获取任意一个有效设备的header（用于辅助功能）
|
|
func GetAnyValidDeviceHeader() string {
|
|
pool := GetDeviceKeyPool()
|
|
if !pool.IsEnabled() {
|
|
return ""
|
|
}
|
|
|
|
pool.mu.RLock()
|
|
defer pool.mu.RUnlock()
|
|
|
|
for _, item := range pool.items {
|
|
if item.IsValid {
|
|
return item.DeviceHeader
|
|
}
|
|
}
|
|
|
|
return ""
|
|
}
|
|
|
|
// GetValidDevice 获取一个有效的设备
|
|
func GetValidDevice() (*DeviceKeyItem, error) {
|
|
pool := GetDeviceKeyPool()
|
|
if !pool.IsEnabled() {
|
|
return nil, fmt.Errorf("设备池未启用")
|
|
}
|
|
return pool.GetNext()
|
|
}
|
|
|
|
// GetCurrentDevice 获取当前使用的设备（和GetValidDevice一样，返回下一个可用设备）
|
|
func GetCurrentDevice() *DeviceKeyItem {
|
|
pool := GetDeviceKeyPool()
|
|
if !pool.IsEnabled() {
|
|
return nil
|
|
}
|
|
item, _ := pool.GetNext()
|
|
return item
|
|
}
|
|
|