87 lines
2.7 KiB
Go
87 lines
2.7 KiB
Go
package model
|
|
|
|
import (
|
|
"errors"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/shopspring/decimal"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
var (
|
|
ErrInvalidUserQuotaAdjustment = errors.New("invalid user quota adjustment")
|
|
ErrUserQuotaPermission = errors.New("cannot adjust quota for this user role")
|
|
)
|
|
|
|
// UserQuotaAdjustment is the immutable database snapshot of a committed manual
|
|
// adjustment. Pending relay deductions in the quota cache are not part of it.
|
|
type UserQuotaAdjustment struct {
|
|
UserID int
|
|
Username string
|
|
Before int
|
|
After int
|
|
}
|
|
|
|
func AdjustUserQuota(userID, operatorRole int, mode string, value int) (*UserQuotaAdjustment, error) {
|
|
if userID <= 0 || (mode != "add" && mode != "subtract" && mode != "override") {
|
|
return nil, ErrInvalidUserQuotaAdjustment
|
|
}
|
|
if mode != "override" && value <= 0 {
|
|
return nil, ErrInvalidUserQuotaAdjustment
|
|
}
|
|
if value > common.MaxWalletQuota || value < -common.MaxWalletQuota {
|
|
return nil, ErrWalletQuotaLimitExceeded
|
|
}
|
|
|
|
var adjustment UserQuotaAdjustment
|
|
err := DB.Transaction(func(tx *gorm.DB) error {
|
|
var user User
|
|
if err := lockForUpdate(tx).First(&user, userID).Error; err != nil {
|
|
return err
|
|
}
|
|
if operatorRole != common.RoleRootUser && operatorRole <= user.Role {
|
|
return ErrUserQuotaPermission
|
|
}
|
|
if user.Quota > common.MaxWalletQuota || user.Quota < -common.MaxWalletQuota {
|
|
return ErrWalletQuotaLimitExceeded
|
|
}
|
|
quota := decimal.NewFromInt(int64(value))
|
|
switch mode {
|
|
case "add":
|
|
quota = decimal.NewFromInt(int64(user.Quota)).Add(quota)
|
|
case "subtract":
|
|
quota = decimal.NewFromInt(int64(user.Quota)).Sub(quota)
|
|
}
|
|
after, err := common.WalletQuotaFromDecimalStrict(quota)
|
|
if err != nil {
|
|
return ErrWalletQuotaLimitExceeded
|
|
}
|
|
// An unchanged override is a successful operation, including on MySQL
|
|
// configurations that count only changed rows in RowsAffected.
|
|
if after != user.Quota {
|
|
result := tx.Model(&User{}).Where("id = ?", userID).Update("quota", after)
|
|
if result.Error != nil {
|
|
return result.Error
|
|
}
|
|
if result.RowsAffected != 1 {
|
|
return gorm.ErrRecordNotFound
|
|
}
|
|
}
|
|
adjustment = UserQuotaAdjustment{UserID: user.Id, Username: user.Username, Before: user.Quota, After: after}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Apply only the committed difference, preserving outstanding reservations.
|
|
// Both balances are bounded above, so their difference fits in int64.
|
|
delta := int64(adjustment.After) - int64(adjustment.Before)
|
|
if delta != 0 {
|
|
if err := cacheIncrUserQuota(userID, delta); err != nil {
|
|
common.SysError(common.LogText("failed to sync manual quota adjustment for user %d: %s", userID, err))
|
|
}
|
|
}
|
|
return &adjustment, nil
|
|
}
|