feat: 算力引擎 Go 后端核心（new-api）
- 统一 OpenAI 兼容 /v1 中继 + 渠道/额度/令牌/流水 - 各领域包：relay 模型网关、model 数据层、controller 管理 API
This commit is contained in:
+423
@@ -0,0 +1,423 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// TwoFA 用户2FA设置表
|
||||
type TwoFA struct {
|
||||
Id int `json:"id" gorm:"primaryKey"`
|
||||
UserId int `json:"user_id" gorm:"unique;not null;index"`
|
||||
Secret string `json:"-" gorm:"type:varchar(255);not null"` // TOTP密钥，不返回给前端
|
||||
IsEnabled bool `json:"is_enabled"`
|
||||
FailedAttempts int `json:"failed_attempts" gorm:"default:0"`
|
||||
LockedUntil *time.Time `json:"locked_until,omitempty"`
|
||||
LastUsedAt *time.Time `json:"last_used_at,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
|
||||
}
|
||||
|
||||
// TwoFABackupCode 备用码使用记录表
|
||||
type TwoFABackupCode struct {
|
||||
Id int `json:"id" gorm:"primaryKey"`
|
||||
UserId int `json:"user_id" gorm:"not null;index"`
|
||||
CodeHash string `json:"-" gorm:"type:varchar(255);not null"` // 备用码哈希
|
||||
IsUsed bool `json:"is_used"`
|
||||
UsedAt *time.Time `json:"used_at,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
|
||||
}
|
||||
|
||||
// GetTwoFAByUserId 根据用户ID获取2FA设置
|
||||
func GetTwoFAByUserId(userId int) (*TwoFA, error) {
|
||||
if userId == 0 {
|
||||
return nil, errors.New("用户ID不能为空")
|
||||
}
|
||||
|
||||
var twoFA TwoFA
|
||||
err := DB.Where("user_id = ?", userId).First(&twoFA).Error
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil // 返回nil表示未设置2FA
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &twoFA, nil
|
||||
}
|
||||
|
||||
// IsTwoFAEnabled 检查用户是否启用了2FA
|
||||
func IsTwoFAEnabled(userId int) (bool, error) {
|
||||
twoFA, err := GetTwoFAByUserId(userId)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return twoFA != nil && twoFA.IsEnabled, nil
|
||||
}
|
||||
|
||||
// CreatePendingTwoFASetup stores a disabled factor while the user completes
|
||||
// enrollment. Enabling a factor must use EnableWithAuthVersion.
|
||||
func (t *TwoFA) CreatePendingTwoFASetup() error {
|
||||
if t == nil || t.UserId <= 0 || t.IsEnabled {
|
||||
return errors.New("无效的2FA待验证设置")
|
||||
}
|
||||
// 检查用户是否已存在2FA设置
|
||||
existing, err := GetTwoFAByUserId(t.UserId)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if existing != nil {
|
||||
return errors.New("用户已存在2FA设置")
|
||||
}
|
||||
|
||||
// 验证用户存在
|
||||
var user User
|
||||
if err := DB.First(&user, t.UserId).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return errors.New("用户不存在")
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
return DB.Create(t).Error
|
||||
}
|
||||
|
||||
func (t *TwoFA) updateUsageState() error {
|
||||
if t.Id == 0 {
|
||||
return errors.New("2FA记录ID不能为空")
|
||||
}
|
||||
return DB.Model(&TwoFA{}).Where("id = ?", t.Id).Updates(map[string]interface{}{
|
||||
"failed_attempts": t.FailedAttempts,
|
||||
"locked_until": t.LockedUntil,
|
||||
"last_used_at": t.LastUsedAt,
|
||||
}).Error
|
||||
}
|
||||
|
||||
// DeletePendingTwoFASetup removes only an unverified setup. Enabled factors
|
||||
// must use DisableTwoFAWithAuthVersion.
|
||||
func (t *TwoFA) DeletePendingTwoFASetup() error {
|
||||
if t == nil || t.Id == 0 || t.UserId <= 0 {
|
||||
return errors.New("2FA记录ID不能为空")
|
||||
}
|
||||
|
||||
return DB.Transaction(func(tx *gorm.DB) error {
|
||||
var pending TwoFA
|
||||
if err := lockForUpdate(tx).
|
||||
Where("id = ? AND user_id = ? AND is_enabled = ?", t.Id, t.UserId, false).
|
||||
First(&pending).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Unscoped().Where("user_id = ?", t.UserId).Delete(&TwoFABackupCode{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Unscoped().Delete(&pending).Error
|
||||
})
|
||||
}
|
||||
|
||||
// ResetFailedAttempts 重置失败尝试次数
|
||||
func (t *TwoFA) ResetFailedAttempts() error {
|
||||
t.FailedAttempts = 0
|
||||
t.LockedUntil = nil
|
||||
return t.updateUsageState()
|
||||
}
|
||||
|
||||
// IncrementFailedAttempts 增加失败尝试次数
|
||||
func (t *TwoFA) IncrementFailedAttempts() error {
|
||||
if t.Id == 0 {
|
||||
return errors.New("2FA记录ID不能为空")
|
||||
}
|
||||
|
||||
const maxUpdateRetries = 5
|
||||
for range maxUpdateRetries {
|
||||
var current TwoFA
|
||||
if err := DB.Select("id", "failed_attempts", "locked_until").First(¤t, t.Id).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
if current.LockedUntil != nil && now.Before(*current.LockedUntil) {
|
||||
t.FailedAttempts = current.FailedAttempts
|
||||
t.LockedUntil = current.LockedUntil
|
||||
return nil
|
||||
}
|
||||
|
||||
nextFailedAttempts := current.FailedAttempts + 1
|
||||
nextLockedUntil := current.LockedUntil
|
||||
if nextFailedAttempts >= common.MaxFailAttempts {
|
||||
lockUntil := now.Add(time.Duration(common.LockoutDuration) * time.Second)
|
||||
nextLockedUntil = &lockUntil
|
||||
}
|
||||
|
||||
result := DB.Model(&TwoFA{}).
|
||||
Where("id = ? AND failed_attempts = ? AND (locked_until IS NULL OR locked_until <= ?)", current.Id, current.FailedAttempts, now).
|
||||
Updates(map[string]interface{}{
|
||||
"failed_attempts": nextFailedAttempts,
|
||||
"locked_until": nextLockedUntil,
|
||||
})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
t.FailedAttempts = nextFailedAttempts
|
||||
t.LockedUntil = nextLockedUntil
|
||||
return nil
|
||||
}
|
||||
|
||||
return errors.New("更新2FA失败次数冲突，请重试")
|
||||
}
|
||||
|
||||
// IsLocked 检查账户是否被锁定
|
||||
func (t *TwoFA) IsLocked() bool {
|
||||
if t.LockedUntil == nil {
|
||||
return false
|
||||
}
|
||||
return time.Now().Before(*t.LockedUntil)
|
||||
}
|
||||
|
||||
// CreatePendingTwoFASetupBackupCodes stores recovery codes for an unverified
|
||||
// setup. Regeneration for an enabled factor must advance auth_version.
|
||||
func CreatePendingTwoFASetupBackupCodes(userId int, codes []string) error {
|
||||
return DB.Transaction(func(tx *gorm.DB) error {
|
||||
var pending TwoFA
|
||||
if err := lockForUpdate(tx).Where("user_id = ? AND is_enabled = ?", userId, false).First(&pending).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return replaceBackupCodesWithTx(tx, userId, codes)
|
||||
})
|
||||
}
|
||||
|
||||
func replaceBackupCodesWithTx(tx *gorm.DB, userId int, codes []string) error {
|
||||
if err := tx.Where("user_id = ?", userId).Delete(&TwoFABackupCode{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, code := range codes {
|
||||
hashedCode, err := common.HashBackupCode(code)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Create(&TwoFABackupCode{UserId: userId, CodeHash: hashedCode, IsUsed: false}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReplaceBackupCodesWithAuthVersion atomically replaces the factor's recovery
|
||||
// credentials and advances the user's authentication version.
|
||||
func ReplaceBackupCodesWithAuthVersion(userId int, codes []string) error {
|
||||
if err := DB.Transaction(func(tx *gorm.DB) error {
|
||||
var enabled TwoFA
|
||||
if err := lockForUpdate(tx).Where("user_id = ? AND is_enabled = ?", userId, true).First(&enabled).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return ErrTwoFANotEnabled
|
||||
}
|
||||
return err
|
||||
}
|
||||
if _, err := IncrementUserAuthVersionWithTx(tx, userId); err != nil {
|
||||
return err
|
||||
}
|
||||
return replaceBackupCodesWithTx(tx, userId, codes)
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
return PublishUserAuthCache(userId)
|
||||
}
|
||||
|
||||
// ValidateBackupCode 验证并使用备用码
|
||||
func ValidateBackupCode(userId int, code string) (bool, error) {
|
||||
if !common.ValidateBackupCode(code) {
|
||||
return false, errors.New("验证码或备用码不正确")
|
||||
}
|
||||
|
||||
normalizedCode := common.NormalizeBackupCode(code)
|
||||
|
||||
// 查找未使用的备用码
|
||||
var backupCodes []TwoFABackupCode
|
||||
if err := DB.Where("user_id = ? AND is_used = false", userId).Find(&backupCodes).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
// 验证备用码
|
||||
for _, bc := range backupCodes {
|
||||
if common.ValidatePasswordAndHash(normalizedCode, bc.CodeHash) {
|
||||
now := time.Now()
|
||||
result := DB.Model(&TwoFABackupCode{}).
|
||||
Where("id = ? AND is_used = ?", bc.Id, false).
|
||||
Updates(map[string]interface{}{
|
||||
"is_used": true,
|
||||
"used_at": now,
|
||||
})
|
||||
if result.Error != nil {
|
||||
return false, result.Error
|
||||
}
|
||||
return result.RowsAffected == 1, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// GetUnusedBackupCodeCount 获取未使用的备用码数量
|
||||
func GetUnusedBackupCodeCount(userId int) (int, error) {
|
||||
var count int64
|
||||
err := DB.Model(&TwoFABackupCode{}).Where("user_id = ? AND is_used = false", userId).Count(&count).Error
|
||||
return int(count), err
|
||||
}
|
||||
|
||||
// DisableTwoFAWithAuthVersion atomically removes the factor and invalidates
|
||||
// every access token issued against the previous security configuration.
|
||||
func DisableTwoFAWithAuthVersion(userId int) error {
|
||||
if err := DB.Transaction(func(tx *gorm.DB) error {
|
||||
var twoFA TwoFA
|
||||
if err := lockForUpdate(tx).Where("user_id = ? AND is_enabled = ?", userId, true).First(&twoFA).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return ErrTwoFANotEnabled
|
||||
}
|
||||
return err
|
||||
}
|
||||
if _, err := IncrementUserAuthVersionWithTx(tx, userId); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Unscoped().Where("user_id = ?", userId).Delete(&TwoFABackupCode{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Unscoped().Delete(&twoFA).Error
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
return PublishUserAuthCache(userId)
|
||||
}
|
||||
|
||||
// EnableWithAuthVersion atomically enables this factor and advances the user
|
||||
// authentication version so pre-enrollment sessions cannot remain valid.
|
||||
func (t *TwoFA) EnableWithAuthVersion() error {
|
||||
if t == nil || t.Id == 0 || t.UserId == 0 {
|
||||
return errors.New("2FA记录ID不能为空")
|
||||
}
|
||||
if err := DB.Transaction(func(tx *gorm.DB) error {
|
||||
var pending TwoFA
|
||||
if err := lockForUpdate(tx).Where("id = ? AND user_id = ? AND is_enabled = ?", t.Id, t.UserId, false).First(&pending).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return ErrTwoFAAlreadyEnabled
|
||||
}
|
||||
return err
|
||||
}
|
||||
if _, err := IncrementUserAuthVersionWithTx(tx, t.UserId); err != nil {
|
||||
return err
|
||||
}
|
||||
result := tx.Model(&pending).
|
||||
Updates(map[string]interface{}{"is_enabled": true, "failed_attempts": 0, "locked_until": nil})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected != 1 {
|
||||
return ErrTwoFAAlreadyEnabled
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
t.IsEnabled = true
|
||||
t.FailedAttempts = 0
|
||||
t.LockedUntil = nil
|
||||
return PublishUserAuthCache(t.UserId)
|
||||
}
|
||||
|
||||
// ValidateTOTPAndUpdateUsage 验证TOTP并更新使用记录
|
||||
func (t *TwoFA) ValidateTOTPAndUpdateUsage(code string) (bool, error) {
|
||||
// 检查是否被锁定
|
||||
if t.IsLocked() {
|
||||
return false, fmt.Errorf("账户已被锁定，请在%v后重试", t.LockedUntil.Format("2006-01-02 15:04:05"))
|
||||
}
|
||||
|
||||
// 验证TOTP码
|
||||
if !common.ValidateTOTPCode(t.Secret, code) {
|
||||
// 增加失败次数
|
||||
if err := t.IncrementFailedAttempts(); err != nil {
|
||||
common.SysLog("更新2FA失败次数失败: " + err.Error())
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// 验证成功，重置失败次数并更新最后使用时间
|
||||
now := time.Now()
|
||||
t.FailedAttempts = 0
|
||||
t.LockedUntil = nil
|
||||
t.LastUsedAt = &now
|
||||
|
||||
if err := t.updateUsageState(); err != nil {
|
||||
common.SysLog("更新2FA使用记录失败: " + err.Error())
|
||||
}
|
||||
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// ValidateBackupCodeAndUpdateUsage 验证备用码并更新使用记录
|
||||
func (t *TwoFA) ValidateBackupCodeAndUpdateUsage(code string) (bool, error) {
|
||||
// 检查是否被锁定
|
||||
if t.IsLocked() {
|
||||
return false, fmt.Errorf("账户已被锁定，请在%v后重试", t.LockedUntil.Format("2006-01-02 15:04:05"))
|
||||
}
|
||||
|
||||
// 验证备用码
|
||||
valid, err := ValidateBackupCode(t.UserId, code)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
if !valid {
|
||||
// 增加失败次数
|
||||
if err := t.IncrementFailedAttempts(); err != nil {
|
||||
common.SysLog("更新2FA失败次数失败: " + err.Error())
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// 验证成功，重置失败次数并更新最后使用时间
|
||||
now := time.Now()
|
||||
t.FailedAttempts = 0
|
||||
t.LockedUntil = nil
|
||||
t.LastUsedAt = &now
|
||||
|
||||
if err := t.updateUsageState(); err != nil {
|
||||
common.SysLog("更新2FA使用记录失败: " + err.Error())
|
||||
}
|
||||
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// GetTwoFAStats 获取2FA统计信息（管理员使用）
|
||||
func GetTwoFAStats() (map[string]interface{}, error) {
|
||||
var totalUsers, enabledUsers int64
|
||||
|
||||
// 总用户数
|
||||
if err := DB.Model(&User{}).Count(&totalUsers).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 启用2FA的用户数
|
||||
if err := DB.Model(&TwoFA{}).Where("is_enabled = true").Count(&enabledUsers).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
enabledRate := float64(0)
|
||||
if totalUsers > 0 {
|
||||
enabledRate = float64(enabledUsers) / float64(totalUsers) * 100
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"total_users": totalUsers,
|
||||
"enabled_users": enabledUsers,
|
||||
"enabled_rate": fmt.Sprintf("%.1f%%", enabledRate),
|
||||
}, nil
|
||||
}
|
||||
Reference in New Issue
Block a user