From 0b0c9f9d04c36e94539b01bfbc02fc82c7a0cb0f Mon Sep 17 00:00:00 2001 From: Yuzhong Zhang Date: Tue, 25 Aug 2026 04:15:18 +0000 Subject: [PATCH] fix: reject quota deductions that would go negative Guard decreaseUserQuota and decreaseTokenQuota with a WHERE remaining >= amount check so concurrent consumption cannot drive balances below zero. --- model/quota_test.go | 206 ++++++++++++++++++++++++++++++++++++++++++++ model/token.go | 12 ++- model/user.go | 10 ++- 3 files changed, 223 insertions(+), 5 deletions(-) create mode 100644 model/quota_test.go diff --git a/model/quota_test.go b/model/quota_test.go new file mode 100644 index 0000000000..662afc2ffb --- /dev/null +++ b/model/quota_test.go @@ -0,0 +1,206 @@ +package model + +import ( + "fmt" + "sync" + "sync/atomic" + "testing" + + "gorm.io/driver/sqlite" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +var quotaTestSeq atomic.Int64 + +func setupQuotaTestDB(t *testing.T) { + t.Helper() + dsn := fmt.Sprintf("file:quota-%s-%d?mode=memory&cache=shared&_busy_timeout=5000", t.Name(), quotaTestSeq.Add(1)) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{ + Logger: logger.Default.LogMode(logger.Silent), + }) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.AutoMigrate(&User{}, &Token{}); err != nil { + t.Fatalf("migrate: %v", err) + } + DB = db +} + +func createTestUser(t *testing.T, quota int64) *User { + t.Helper() + seq := quotaTestSeq.Add(1) + user := &User{ + Username: fmt.Sprintf("u%d", seq), + Password: "password1234", + Quota: quota, + Status: UserStatusEnabled, + Role: RoleCommonUser, + AccessToken: fmt.Sprintf("%032d", seq), + AffCode: fmt.Sprintf("a%d", seq), + } + if err := DB.Create(user).Error; err != nil { + t.Fatalf("create user: %v", err) + } + return user +} + +func createTestToken(t *testing.T, userId int, remain int64) *Token { + t.Helper() + seq := quotaTestSeq.Add(1) + token := &Token{ + UserId: userId, + Key: fmt.Sprintf("%048d", seq), + Name: "quota-test", + Status: TokenStatusEnabled, + RemainQuota: remain, + ExpiredTime: -1, + } + if err := DB.Create(token).Error; err != nil { + t.Fatalf("create token: %v", err) + } + return token +} + +func TestDecreaseUserQuotaSequentialOversellDoesNotGoNegative(t *testing.T) { + setupQuotaTestDB(t) + user := createTestUser(t, 100) + + if err := decreaseUserQuota(user.Id, 60); err != nil { + t.Fatalf("first decrease: %v", err) + } + err := decreaseUserQuota(user.Id, 60) + remaining, qErr := GetUserQuota(user.Id) + if qErr != nil { + t.Fatalf("GetUserQuota: %v", qErr) + } + if remaining < 0 { + t.Fatalf("user quota went negative: %d", remaining) + } + if err == nil { + t.Fatalf("expected oversell to fail, remaining=%d", remaining) + } + if remaining != 40 { + t.Fatalf("remaining quota = %d, want 40", remaining) + } +} + +func TestDecreaseTokenQuotaSequentialOversellDoesNotGoNegative(t *testing.T) { + setupQuotaTestDB(t) + user := createTestUser(t, 1000) + token := createTestToken(t, user.Id, 100) + + if err := decreaseTokenQuota(token.Id, 60); err != nil { + t.Fatalf("first decrease: %v", err) + } + err := decreaseTokenQuota(token.Id, 60) + got, gErr := GetTokenById(token.Id) + if gErr != nil { + t.Fatalf("GetTokenById: %v", gErr) + } + if got.RemainQuota < 0 { + t.Fatalf("token remain_quota went negative: %d", got.RemainQuota) + } + if err == nil { + t.Fatalf("expected oversell to fail, remain_quota=%d", got.RemainQuota) + } + if got.RemainQuota != 40 { + t.Fatalf("remain_quota = %d, want 40", got.RemainQuota) + } + if got.UsedQuota != 60 { + t.Fatalf("used_quota = %d, want 60", got.UsedQuota) + } +} + +func TestDecreaseUserQuotaExactDrainThenReject(t *testing.T) { + setupQuotaTestDB(t) + user := createTestUser(t, 100) + + if err := decreaseUserQuota(user.Id, 100); err != nil { + t.Fatalf("exact drain: %v", err) + } + err := decreaseUserQuota(user.Id, 1) + remaining, qErr := GetUserQuota(user.Id) + if qErr != nil { + t.Fatalf("GetUserQuota: %v", qErr) + } + if remaining != 0 { + t.Fatalf("remaining quota = %d, want 0", remaining) + } + if err == nil { + t.Fatal("expected decrease below zero to fail") + } +} + +func TestDecreaseUserQuotaConcurrentOversellDoesNotGoNegative(t *testing.T) { + setupQuotaTestDB(t) + user := createTestUser(t, 100) + + const workers = 10 + const each = int64(60) + var success atomic.Int64 + var wg sync.WaitGroup + wg.Add(workers) + for i := 0; i < workers; i++ { + go func() { + defer wg.Done() + if err := decreaseUserQuota(user.Id, each); err == nil { + success.Add(1) + } + }() + } + wg.Wait() + + remaining, err := GetUserQuota(user.Id) + if err != nil { + t.Fatalf("GetUserQuota: %v", err) + } + if remaining < 0 { + t.Fatalf("user quota went negative under concurrent decrease: %d (successes=%d)", remaining, success.Load()) + } + if success.Load() != 1 { + t.Fatalf("successful decreases = %d, want 1; remaining=%d", success.Load(), remaining) + } + if remaining != 40 { + t.Fatalf("remaining quota = %d, want 40", remaining) + } +} + +func TestDecreaseTokenQuotaConcurrentOversellDoesNotGoNegative(t *testing.T) { + setupQuotaTestDB(t) + user := createTestUser(t, 1000) + token := createTestToken(t, user.Id, 100) + + const workers = 10 + const each = int64(60) + var success atomic.Int64 + var wg sync.WaitGroup + wg.Add(workers) + for i := 0; i < workers; i++ { + go func() { + defer wg.Done() + if err := decreaseTokenQuota(token.Id, each); err == nil { + success.Add(1) + } + }() + } + wg.Wait() + + got, err := GetTokenById(token.Id) + if err != nil { + t.Fatalf("GetTokenById: %v", err) + } + if got.RemainQuota < 0 { + t.Fatalf("token remain_quota went negative under concurrent decrease: %d (successes=%d)", got.RemainQuota, success.Load()) + } + if success.Load() != 1 { + t.Fatalf("successful decreases = %d, want 1; remain_quota=%d", success.Load(), got.RemainQuota) + } + if got.RemainQuota != 40 { + t.Fatalf("remain_quota = %d, want 40", got.RemainQuota) + } + if got.UsedQuota != 60 { + t.Fatalf("used_quota = %d, want 60", got.UsedQuota) + } +} diff --git a/model/token.go b/model/token.go index 52ee63ef5b..d54ae44d34 100644 --- a/model/token.go +++ b/model/token.go @@ -204,14 +204,20 @@ func DecreaseTokenQuota(id int, quota int64) (err error) { } func decreaseTokenQuota(id int, quota int64) (err error) { - err = DB.Model(&Token{}).Where("id = ?", id).Updates( + result := DB.Model(&Token{}).Where("id = ? AND remain_quota >= ?", id, quota).Updates( map[string]interface{}{ "remain_quota": gorm.Expr("remain_quota - ?", quota), "used_quota": gorm.Expr("used_quota + ?", quota), "accessed_time": helper.GetTimestamp(), }, - ).Error - return err + ) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return errors.New("令牌额度不足") + } + return nil } func PreConsumeTokenQuota(tokenId int, quota int64) (err error) { diff --git a/model/user.go b/model/user.go index 021810c0f1..400b879bd0 100644 --- a/model/user.go +++ b/model/user.go @@ -399,8 +399,14 @@ func DecreaseUserQuota(id int, quota int64) (err error) { } func decreaseUserQuota(id int, quota int64) (err error) { - err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error - return err + result := DB.Model(&User{}).Where("id = ? AND quota >= ?", id, quota).Update("quota", gorm.Expr("quota - ?", quota)) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return errors.New("用户额度不足") + } + return nil } func GetRootUserEmail() (email string) {