From a073f74b38a33bb154821089c097658cbdcc0fbe Mon Sep 17 00:00:00 2001 From: Seefs <40468931+seefs001@users.noreply.github.com> Date: Wed, 26 Aug 2026 20:57:54 +0800 Subject: [PATCH] refactor: deprecate int32 (#7025) * refactor: deprecate int32 * fix(db): reject legacy user quota schemas at startup * fix(quota): enforce wallet bounds and saturating billing conversions * fix(rate-limit): keep count*duration from wrapping int64 * fix: error message --- .env.example | 2 ++ AGENTS.md | 2 +- common/quota_math.go | 53 +++++++++++++++++++--------- common/quota_math_test.go | 22 +++++++++++- controller/channel-test.go | 8 ++--- controller/redemption.go | 17 +++++++++ controller/token.go | 15 ++++++-- controller/topup.go | 12 ++++--- controller/topup_quota_limit_test.go | 34 ++++++++++-------- controller/user.go | 8 +++++ controller/user_manage_test.go | 21 +++++++++++ middleware/model-rate-limit.go | 27 ++++++++++++-- model/main.go | 49 +++++++++++++++++++++++++ model/payment_method_guard_test.go | 12 +++---- model/quota_reserve_test.go | 33 +++++++++++++++++ model/redemption.go | 14 +++++++- model/redemption_test.go | 29 +++++++++++++++ model/subscription.go | 2 +- model/topup.go | 33 ++++++++--------- model/user.go | 42 ++++++++++++++++------ model/utils.go | 19 ++++++++-- pkg/billingexpr/billingexpr_test.go | 5 +-- pkg/billingexpr/settle_clamp_test.go | 13 ++++--- pkg/billingexpr/types.go | 2 +- relay/common/relay_info.go | 2 +- relay/helper/price_test.go | 2 +- service/quota.go | 10 ++++-- service/quota_saturation_test.go | 25 +++++++++++++ service/text_quota.go | 4 +-- service/text_quota_test.go | 8 ++--- service/tiered_settle.go | 2 +- service/token_counter.go | 4 +-- service/violation_fee.go | 7 ++-- setting/rate_limit.go | 21 +++++++---- web/src/features/usage-logs/types.ts | 2 +- 35 files changed, 445 insertions(+), 116 deletions(-) diff --git a/.env.example b/.env.example index 3b8a2a5b..e2f28743 100644 --- a/.env.example +++ b/.env.example @@ -35,6 +35,8 @@ # SQL_MAX_LIFETIME=60 # 慢查询日志阈值(毫秒),0 表示关闭慢查询日志,超出 0-3600000 范围回退默认值 200 # SQL_SLOW_THRESHOLD_MS=200 +# 跳过用户额度列 64 位 schema 检查(仅在已确认数据库列可容纳 64 位时启用) +# SKIP_64BIT_QUOTA_SCHEMA_CHECK=true # 缓存相关配置 diff --git a/AGENTS.md b/AGENTS.md index fa942c71..9314a49b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -107,7 +107,7 @@ Do NOT directly import or call `encoding/json` in business code. `json.RawMessag - Every user-controlled quantity that becomes a billing multiplier (image `n`, video `seconds`/`duration`, resolution/quality ratios, batch counts) MUST be bounded before it reaches quota calculation. Reject out-of-range values at request validation with a 400. Existing bounds: `dto.MaxImageN` for image generation count, `relaycommon.MaxTaskDurationSeconds` for task video duration, `maxTokensLimit` (`relay/helper/valid_request.go`) for `max_tokens`-family fields on every relay format (OpenAI, Claude, Gemini, Responses). Reuse these constants instead of introducing new ad hoc limits for the same concepts. When adding a new relay format or request DTO, bound its max-tokens and count fields in its validator from day one. - Watch for validation bypass paths: passthrough fields (e.g. `Extra["parameters"]`), task `metadata` maps, and multipart form fields can carry the same quantities around the standard DTO validation. Any adaptor that reads a multiplier from such a path must enforce the same bound (or clamp) locally. - Durations parsed from media metadata are user/upstream-controlled too: audio file headers (transcription token counting, TTS response duration) and upstream deduction numbers (e.g. Kling `FinalUnitDeduction`) can claim absurd values. Convert them with saturation before they become token counts. -- Never convert a computed quota or token count to `int` with a bare cast like `int(float64(quota) * ratio)`, `int(math.Round(...))` on unbounded input, or `int(decimal.IntPart())`. All quota rounding/conversion is centralized in `common/quota_math.go`; use those helpers: `common.QuotaFromFloat` (truncating) for float products, `common.QuotaRound` (half-away-from-zero) where rounding is intended, and `common.QuotaFromDecimal` for decimal products. `billingexpr.QuotaRound` delegates to `common.QuotaRound`. Do not reintroduce local conversion helpers or bare casts. Saturation bounds are int32 because quota columns (user/token/log) are 32-bit integers in the database, and every clamp/NaN fallback is logged via `common.SysError` since a single request should never approach those bounds. +- Never convert a computed quota or token count to `int` with a bare cast like `int(float64(quota) * ratio)`, `int(math.Round(...))` on unbounded input, or `int(decimal.IntPart())`. All quota rounding/conversion is centralized in `common/quota_math.go`; use those helpers: `common.QuotaFromFloat` (truncating) for float products, `common.QuotaRound` (half-away-from-zero) where rounding is intended, and `common.QuotaFromDecimal` for decimal products. `billingexpr.QuotaRound` delegates to `common.QuotaRound`. Do not reintroduce local conversion helpers or bare casts. Single-request saturation stays at the int32 boundary so batch accumulation cannot approach 64-bit wraparound; wallet/top-up conversion uses `common.WalletQuotaFromDecimalStrict` with the JavaScript-safe `common.MaxWalletQuota` boundary. Every clamp/NaN fallback is logged via `common.SysError`. - Saturation events are also audited: each helper has a `*Checked` variant (`common.QuotaFromFloatChecked` / `QuotaRoundChecked` / `QuotaFromDecimalChecked`) that additionally returns a `*common.QuotaClamp` when clamping occurred. Billing paths that compute a charge capture that clamp onto `relayInfo.QuotaClamp` (or thread it into task settlement) and, right before writing the consume/task log, call `attachQuotaSaturation` (in `service/log_info_generate.go`) which nests the marker under the log's `other.admin_info.quota_saturation` and emits a request-correlated `logger.LogWarn`. Nesting under `admin_info` makes it admin-only for free (non-admin log views strip `admin_info`). When adding a new billing path, use the `*Checked` variant and surface the clamp the same way so the anomaly stays auditable in both the admin log UI and backend logs. - Multiplier maps go through `types.PriceData.AddOtherRatio`, which rejects non-positive, NaN, and +Inf ratios. Do not write to `PriceData.OtherRatios` directly, and do not weaken these guards. - Pre-consume (预扣费) and settle (结算/差额) must both be safe: a saturated oversized quota must fail pre-consume with insufficient-quota, never silently wrap. When adding a new billing path (new relay format, new task platform, new adjustment hook), trace the full chain — validation → EstimateBilling/OtherRatios → quota conversion → pre-consume → settle/refund — and confirm each step preserves these invariants. diff --git a/common/quota_math.go b/common/quota_math.go index 66d62093..fb0e71ca 100644 --- a/common/quota_math.go +++ b/common/quota_math.go @@ -8,14 +8,24 @@ import ( ) // Quota conversions are centralized here so every billing path shares one -// saturation + logging policy. Quota columns (user/token/log) are 32-bit -// integers in the database, so an oversized product must clamp to the int32 -// range instead of wrapping around and turning a charge into a credit. +// saturation + logging policy. Single-request charges stay bounded to int32; +// top-ups and wallet-priced purchases use a JavaScript-safe 64-bit domain. const ( - MaxQuota = math.MaxInt32 - MinQuota = math.MinInt32 + MaxQuota = math.MaxInt32 + MinQuota = math.MinInt32 + MaxWalletQuota = 1<<53 - 1 ) +// ValidateWalletQuota enforces the upper bound shared by wallet mutations. +// Negative balances remain valid because billing can temporarily overdraw a +// wallet; callers that accept credits must apply their own positive check. +func ValidateWalletQuota(quota int) error { + if quota > MaxWalletQuota { + return fmt.Errorf("wallet quota exceeds %d", MaxWalletQuota) + } + return nil +} + // QuotaClampKind identifies why a quota conversion had to be saturated. type QuotaClampKind string @@ -27,11 +37,11 @@ const ( ) // QuotaClamp describes a single saturation event: a quota conversion whose -// input fell outside the representable int32 range (or was NaN) and was +// input fell outside its supported range (or was NaN) and was // therefore clamped. It is surfaced to billing callers so the event can be // recorded on the related consume/task log for admin auditing. type QuotaClamp struct { - Op string `json:"op"` // "QuotaFromFloat" | "QuotaRound" | "QuotaFromDecimal" + Op string `json:"op"` // "QuotaFromFloat" | "QuotaRound" | "QuotaFromDecimal" | "WalletQuotaFromDecimal" Kind QuotaClampKind `json:"kind"` // "overflow" | "underflow" | "nan" Original float64 `json:"original"` // best-effort pre-clamp value (decimal -> float64 approx) Clamped int `json:"clamped"` // the saturated result actually used @@ -61,23 +71,27 @@ func (c *QuotaClamp) AuditMap() map[string]interface{} { } } -// saturateQuota converts an already-rounded quota value to int, clamping to -// the int32 range. Whenever clamping (what would otherwise be an integer -// wraparound) or a NaN fallback is triggered it logs a warning, because in +// saturateQuota converts an already-rounded single-request quota to int. +// Whenever clamping (what would otherwise be an integer wraparound) or a NaN +// fallback is triggered it logs a warning, because in // normal operation a single request never approaches these bounds — hitting // them signals a bug or an abusive request. `op` names the caller. When a // clamp occurs it returns a non-nil *QuotaClamp so callers can additionally // record the event (e.g. on the consume log); the returned pointer is nil for // in-range values. func saturateQuota(value float64, op string) (int, *QuotaClamp) { + return saturateQuotaBounded(value, op, MaxQuota, MinQuota) +} + +func saturateQuotaBounded(value float64, op string, maxQuota int, minQuota int) (int, *QuotaClamp) { var clamp *QuotaClamp switch { case math.IsNaN(value): clamp = &QuotaClamp{Op: op, Kind: QuotaClampNaN, Original: value, Clamped: 0} - case value >= MaxQuota: - clamp = &QuotaClamp{Op: op, Kind: QuotaClampOverflow, Original: value, Clamped: MaxQuota} - case value <= MinQuota: - clamp = &QuotaClamp{Op: op, Kind: QuotaClampUnderflow, Original: value, Clamped: MinQuota} + case value > float64(maxQuota): + clamp = &QuotaClamp{Op: op, Kind: QuotaClampOverflow, Original: value, Clamped: maxQuota} + case value < float64(minQuota): + clamp = &QuotaClamp{Op: op, Kind: QuotaClampUnderflow, Original: value, Clamped: minQuota} default: return int(value), nil } @@ -147,8 +161,15 @@ func QuotaFromDecimalChecked(d decimal.Decimal) (int, *QuotaClamp) { return saturateQuota(f, "QuotaFromDecimal") } -// QuotaFromDecimalStrict converts an in-range decimal quota and rejects a -// value that would otherwise be saturated at the database's int32 boundary. +// QuotaFromDecimalStrict converts an in-range single-request quota and rejects +// a value that would otherwise be saturated at the int32 boundary. func QuotaFromDecimalStrict(d decimal.Decimal) (int, error) { return strictQuota(QuotaFromDecimalChecked(d)) } + +// WalletQuotaFromDecimalStrict converts wallet and top-up values within the +// JavaScript-safe integer range, which is also exactly representable by float64. +func WalletQuotaFromDecimalStrict(d decimal.Decimal) (int, error) { + f, _ := d.Round(0).Float64() + return strictQuota(saturateQuotaBounded(f, "WalletQuotaFromDecimal", MaxWalletQuota, -MaxWalletQuota)) +} diff --git a/common/quota_math_test.go b/common/quota_math_test.go index 2d8742e6..efa8b9eb 100644 --- a/common/quota_math_test.go +++ b/common/quota_math_test.go @@ -1,6 +1,7 @@ package common import ( + "fmt" "math" "testing" @@ -21,6 +22,7 @@ func TestQuotaFromFloat(t *testing.T) { assert.Equal(t, 42, QuotaFromFloat(42.4)) assert.Equal(t, 42, QuotaFromFloat(42.9)) assert.Equal(t, -42, QuotaFromFloat(-42.9)) + assert.Equal(t, MaxQuota, QuotaFromFloat(float64(math.MaxInt32)+42)) assert.Equal(t, MaxQuota, QuotaFromFloat(overflowingProduct)) assert.Equal(t, MinQuota, QuotaFromFloat(-overflowingProduct)) assert.Equal(t, MaxQuota, QuotaFromFloat(math.Inf(1))) @@ -34,6 +36,7 @@ func TestQuotaRound(t *testing.T) { assert.Equal(t, 42, QuotaRound(41.5)) assert.Equal(t, 43, QuotaRound(42.5)) assert.Equal(t, -43, QuotaRound(-42.5)) + assert.Equal(t, MaxQuota, QuotaRound(float64(math.MaxInt32)+0.5)) assert.Equal(t, MaxQuota, QuotaRound(overflowingProduct)) assert.Equal(t, MinQuota, QuotaRound(-overflowingProduct)) assert.Equal(t, 0, QuotaRound(math.NaN())) @@ -93,7 +96,7 @@ func TestQuotaFromFloatStrictReturnsTypedClampError(t *testing.T) { assert.ErrorContains(t, err, "QuotaFromFloat") assert.ErrorContains(t, err, "overflow") assert.ErrorContains(t, err, "original=") - assert.ErrorContains(t, err, "clamped=2147483647") + assert.ErrorContains(t, err, fmt.Sprintf("clamped=%d", MaxQuota)) } // TestQuotaRoundChecked verifies the rounding entry point reports clamps the @@ -124,3 +127,20 @@ func TestQuotaFromDecimalChecked(t *testing.T) { assert.Equal(t, QuotaClampOverflow, clamp.Kind) } } + +func TestWalletQuotaFromDecimalStrict(t *testing.T) { + quota, err := WalletQuotaFromDecimalStrict(decimal.NewFromInt(4_294_500_000)) + require.NoError(t, err) + assert.Equal(t, 4_294_500_000, quota) + + quota, err = WalletQuotaFromDecimalStrict(decimal.NewFromInt(MaxWalletQuota)) + require.NoError(t, err) + assert.Equal(t, MaxWalletQuota, quota) + + quota, err = WalletQuotaFromDecimalStrict(decimal.NewFromInt(MaxWalletQuota + 1)) + assert.Zero(t, quota) + var clamp *QuotaClamp + require.ErrorAs(t, err, &clamp) + assert.Equal(t, "WalletQuotaFromDecimal", clamp.Op) + assert.Equal(t, QuotaClampOverflow, clamp.Kind) +} diff --git a/controller/channel-test.go b/controller/channel-test.go index b294979d..e1535d26 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -7,7 +7,6 @@ import ( "errors" "fmt" "io" - "math" "net/http" "net/http/httptest" "strconv" @@ -540,15 +539,16 @@ func settleTestQuota(info *relaycommon.RelayInfo, priceData hosttypes.PriceData, quota := 0 if !priceData.UsePrice { - quota = usage.PromptTokens + int(math.Round(float64(usage.CompletionTokens)*priceData.CompletionRatio)) - quota = int(math.Round(float64(quota) * priceData.ModelRatio)) + completionQuota := common.QuotaRound(float64(usage.CompletionTokens) * priceData.CompletionRatio) + quota = common.QuotaRound(float64(usage.PromptTokens) + float64(completionQuota)) + quota = common.QuotaRound(float64(quota) * priceData.ModelRatio) if priceData.ModelRatio != 0 && quota <= 0 { quota = 1 } return quota, nil } - return int(priceData.ModelPrice * common.QuotaPerUnit), nil + return common.QuotaFromFloat(priceData.ModelPrice * common.QuotaPerUnit), nil } func buildTestLogOther(c *gin.Context, info *relaycommon.RelayInfo, priceData hosttypes.PriceData, usage *dto.Usage, tieredResult *billingexpr.TieredResult) map[string]interface{} { diff --git a/controller/redemption.go b/controller/redemption.go index 838746e7..86289f8a 100644 --- a/controller/redemption.go +++ b/controller/redemption.go @@ -1,6 +1,7 @@ package controller import ( + "errors" "net/http" "strconv" "unicode/utf8" @@ -85,6 +86,14 @@ func AddRedemption(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgRedemptionCountMax) return } + if redemption.Quota <= 0 { + common.ApiError(c, errors.New("redemption quota must be positive")) + return + } + if err := common.ValidateWalletQuota(redemption.Quota); err != nil { + common.ApiError(c, err) + return + } if valid, msg := validateExpiredTime(c, redemption.ExpiredTime); !valid { c.JSON(http.StatusOK, gin.H{"success": false, "message": msg}) return @@ -153,6 +162,14 @@ func UpdateRedemption(c *gin.Context) { return } if statusOnly == "" { + if redemption.Quota <= 0 { + common.ApiError(c, errors.New("redemption quota must be positive")) + return + } + if err := common.ValidateWalletQuota(redemption.Quota); err != nil { + common.ApiError(c, err) + return + } if valid, msg := validateExpiredTime(c, redemption.ExpiredTime); !valid { c.JSON(http.StatusOK, gin.H{"success": false, "message": msg}) return diff --git a/controller/token.go b/controller/token.go index ff2aca8f..09d14fca 100644 --- a/controller/token.go +++ b/controller/token.go @@ -15,6 +15,7 @@ import ( "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/gin-gonic/gin" + "github.com/shopspring/decimal" ) type tokenAutoGroupsInput struct { @@ -41,6 +42,16 @@ type tokenResponse struct { AutoGroups []string `json:"auto_groups"` } +func maxTokenQuota() int { + quota, err := common.WalletQuotaFromDecimalStrict( + decimal.NewFromInt(1_000_000_000).Mul(decimal.NewFromFloat(common.QuotaPerUnit)), + ) + if err != nil { + return common.MaxWalletQuota + } + return quota +} + func buildMaskedTokenResponse(token *model.Token) *tokenResponse { if token == nil { return nil @@ -279,7 +290,7 @@ func AddToken(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative) return } - maxQuotaValue := common.QuotaFromFloat(1000000000 * common.QuotaPerUnit) + maxQuotaValue := maxTokenQuota() if token.RemainQuota > maxQuotaValue { common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue}) return @@ -373,7 +384,7 @@ func UpdateToken(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative) return } - maxQuotaValue := common.QuotaFromFloat(1000000000 * common.QuotaPerUnit) + maxQuotaValue := maxTokenQuota() if token.RemainQuota > maxQuotaValue { common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue}) return diff --git a/controller/topup.go b/controller/topup.go index 08aab813..30f8d221 100644 --- a/controller/topup.go +++ b/controller/topup.go @@ -182,7 +182,11 @@ func getMinTopup() int64 { if operation_setting.GetQuotaDisplayType() == operation_setting.QuotaDisplayTypeTokens { dMinTopup := decimal.NewFromInt(int64(minTopup)) dQuotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) - minTopup = common.QuotaFromDecimal(dMinTopup.Mul(dQuotaPerUnit)) + quota, err := common.WalletQuotaFromDecimalStrict(dMinTopup.Mul(dQuotaPerUnit)) + if err != nil { + return common.MaxWalletQuota + } + minTopup = quota } return int64(minTopup) } @@ -195,7 +199,7 @@ func getTopUpQuota(amount int64) (int, error) { } else { quota = quota.Mul(decimal.NewFromFloat(common.QuotaPerUnit)) } - return common.QuotaFromDecimalStrict(quota) + return common.WalletQuotaFromDecimalStrict(quota) } func getMaxTopUpAmount() int64 { @@ -203,7 +207,7 @@ func getMaxTopUpAmount() int64 { return 0 } quotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) - maxStoredAmount := decimal.NewFromInt(common.MaxQuota - 1). + maxStoredAmount := decimal.NewFromInt(common.MaxWalletQuota). Div(quotaPerUnit). Floor() if operation_setting.GetQuotaDisplayType() == operation_setting.QuotaDisplayTypeTokens { @@ -217,7 +221,7 @@ func getMaxTopUpAmount() int64 { } func validateCreditedQuota(quota decimal.Decimal) (int, error) { - value, err := common.QuotaFromDecimalStrict(quota) + value, err := common.WalletQuotaFromDecimalStrict(quota) if err != nil { return 0, errors.New("充值额度超出系统可表示范围") } diff --git a/controller/topup_quota_limit_test.go b/controller/topup_quota_limit_test.go index 5c291dea..710e87a3 100644 --- a/controller/topup_quota_limit_test.go +++ b/controller/topup_quota_limit_test.go @@ -1,6 +1,7 @@ package controller import ( + "fmt" "net/http" "net/http/httptest" "strings" @@ -43,19 +44,19 @@ func TestTopUpQuotaValidation(t *testing.T) { name: "currency amount above limit", displayType: operation_setting.QuotaDisplayTypeUSD, amount: 4295, - wantErr: true, + wantQuota: 2_147_500_000, }, { name: "token amount preserves settlement truncation", displayType: operation_setting.QuotaDisplayTypeTokens, - amount: common.MaxQuota, - wantQuota: 2_147_000_000, + amount: 2_147_500_000, + wantQuota: 2_147_500_000, }, { - name: "token amount above settlement limit", + name: "token amount above legacy int32 range", displayType: operation_setting.QuotaDisplayTypeTokens, - amount: 2_147_500_000, - wantErr: true, + amount: 4_294_500_000, + wantQuota: 4_294_500_000, }, } @@ -83,14 +84,14 @@ func TestValidateTopUpQuotaReturnsMaximumAmount(t *testing.T) { operation_setting.GetGeneralSetting().QuotaDisplayType = oldDisplayType }) - maxAmount := decimal.NewFromInt(common.MaxQuota - 1). + maxAmount := decimal.NewFromInt(common.MaxWalletQuota). Div(decimal.NewFromFloat(common.QuotaPerUnit)). Floor().IntPart() _, err := validateTopUpQuota(maxAmount) require.NoError(t, err) _, err = validateTopUpQuota(maxAmount + 1) - require.EqualError(t, err, "单笔充值数量不能大于 4294") + require.EqualError(t, err, fmt.Sprintf("单笔充值数量不能大于 %d", maxAmount)) } func TestRequestAmountRejectsTopUpThatCannotBeSettled(t *testing.T) { @@ -106,17 +107,20 @@ func TestRequestAmountRejectsTopUpThatCannotBeSettled(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() ctx, _ := gin.CreateTestContext(recorder) + maxAmount := decimal.NewFromInt(common.MaxWalletQuota). + Div(decimal.NewFromFloat(common.QuotaPerUnit)). + Floor().IntPart() ctx.Request = httptest.NewRequest( http.MethodPost, "/api/user/amount", - strings.NewReader(`{"amount":4295}`), + strings.NewReader(fmt.Sprintf(`{"amount":%d}`, maxAmount+1)), ) ctx.Request.Header.Set("Content-Type", "application/json") RequestAmount(ctx) assert.Equal(t, http.StatusOK, recorder.Code) - assert.JSONEq(t, `{"message":"error","data":"单笔充值数量不能大于 4294"}`, recorder.Body.String()) + assert.JSONEq(t, fmt.Sprintf(`{"message":"error","data":"单笔充值数量不能大于 %d"}`, maxAmount), recorder.Body.String()) } func TestRequestAmountRejectsTopUpThatWouldOverflowWallet(t *testing.T) { @@ -143,7 +147,7 @@ func TestRequestAmountRejectsTopUpThatWouldOverflowWallet(t *testing.T) { require.NoError(t, model.DB.Create(&model.User{ Id: 42, Username: "topup_capacity_user", - Quota: 1_000_000, + Quota: common.MaxWalletQuota - 100_000, Status: common.UserStatusEnabled, }).Error) @@ -154,7 +158,7 @@ func TestRequestAmountRejectsTopUpThatWouldOverflowWallet(t *testing.T) { ctx.Request = httptest.NewRequest( http.MethodPost, "/api/user/amount", - strings.NewReader(`{"amount":4294}`), + strings.NewReader(`{"amount":1}`), ) ctx.Request.Header.Set("Content-Type", "application/json") @@ -165,11 +169,11 @@ func TestRequestAmountRejectsTopUpThatWouldOverflowWallet(t *testing.T) { } func TestValidateCreditedQuotaRejectsOverflow(t *testing.T) { - _, err := validateCreditedQuota(decimal.NewFromInt(common.MaxQuota - 1)) + _, err := validateCreditedQuota(decimal.NewFromInt(int64(common.MaxWalletQuota / 2))) require.NoError(t, err) _, err = validateCreditedQuota(decimal.Zero) require.EqualError(t, err, "充值额度必须大于 0") - _, err = validateCreditedQuota(decimal.NewFromInt(common.MaxQuota)) + _, err = validateCreditedQuota(decimal.NewFromInt(common.MaxWalletQuota + 1)) require.EqualError( t, err, @@ -190,6 +194,8 @@ func TestStripeCreditedQuotaIncludesGroupRatio(t *testing.T) { _, err := validateCreditedQuota(getStripeCreditedQuota(2147, "vip")) require.NoError(t, err) _, err = validateCreditedQuota(getStripeCreditedQuota(2148, "vip")) + require.NoError(t, err) + _, err = validateCreditedQuota(getStripeCreditedQuota(int64(common.MaxWalletQuota), "vip")) require.Error(t, err) require.NoError(t, common.UpdateTopupGroupRatioByJSONString(`{"free":0}`)) diff --git a/controller/user.go b/controller/user.go index 9b8d931e..7020f1c6 100644 --- a/controller/user.go +++ b/controller/user.go @@ -1167,6 +1167,10 @@ func ManageUser(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgUserQuotaChangeZero) return } + if err := common.ValidateWalletQuota(req.Value); err != nil { + common.ApiError(c, err) + return + } if err := model.IncreaseUserQuota(user.Id, req.Value, true); err != nil { common.ApiError(c, err) return @@ -1187,6 +1191,10 @@ func ManageUser(c *gin.Context) { "quota": logger.LogQuota(req.Value), }) case "override": + if err := common.ValidateWalletQuota(req.Value); err != nil { + common.ApiError(c, err) + return + } oldQuota := user.Quota if err := model.DB.Model(&model.User{}).Where("id = ?", user.Id).Update("quota", req.Value).Error; err != nil { common.ApiError(c, err) diff --git a/controller/user_manage_test.go b/controller/user_manage_test.go index 1b52ece0..98564078 100644 --- a/controller/user_manage_test.go +++ b/controller/user_manage_test.go @@ -159,3 +159,24 @@ func TestManageUserDeleteReturnsImmediatelyAndUnknownActionFails(t *testing.T) { assert.EqualValues(t, 1, unchanged.AuthVersion) assert.Equal(t, common.UserStatusEnabled, unchanged.Status) } + +func TestManageUserQuotaRespectsWalletCeiling(t *testing.T) { + db := setupManageUserTestDB(t) + user := model.User{ + Username: "managed-quota-user", Password: "password", Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, Group: "default", Quota: common.MaxWalletQuota - 1, + } + require.NoError(t, db.Create(&user).Error) + + recorder := performManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"add_quota","mode":"add","value":2}`, user.Id)) + assert.Contains(t, recorder.Body.String(), `"success":false`) + + var updated model.User + require.NoError(t, db.First(&updated, user.Id).Error) + assert.Equal(t, common.MaxWalletQuota-1, updated.Quota) + + recorder = performManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"add_quota","mode":"override","value":%d}`, user.Id, common.MaxWalletQuota+1)) + assert.Contains(t, recorder.Body.String(), `"success":false`) + require.NoError(t, db.First(&updated, user.Id).Error) + assert.Equal(t, common.MaxWalletQuota-1, updated.Quota) +} diff --git a/middleware/model-rate-limit.go b/middleware/model-rate-limit.go index 9f1d9403..8cf0ea45 100644 --- a/middleware/model-rate-limit.go +++ b/middleware/model-rate-limit.go @@ -3,6 +3,7 @@ package middleware import ( "context" "fmt" + "math" "net/http" "strconv" "time" @@ -103,7 +104,7 @@ func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) g allowed, err = tb.Allow( ctx, totalKey, - limiter.WithCapacity(int64(totalMaxCount)*duration), + limiter.WithCapacity(rateLimitCapacity(totalMaxCount, duration)), limiter.WithRate(int64(totalMaxCount)), limiter.WithRequested(duration), ) @@ -174,7 +175,7 @@ func ModelRequestRateLimit() func(c *gin.Context) { } // 计算限流参数 - duration := int64(setting.ModelRequestRateLimitDurationMinutes * 60) + duration := rateLimitDurationSeconds(setting.ModelRequestRateLimitDurationMinutes) totalMaxCount := setting.ModelRequestRateLimitCount successMaxCount := setting.ModelRequestRateLimitSuccessCount @@ -199,3 +200,25 @@ func ModelRequestRateLimit() func(c *gin.Context) { } } } + +func rateLimitDurationSeconds(durationMinutes int) int64 { + if durationMinutes <= 0 { + return 0 + } + minutes := int64(durationMinutes) + if minutes > math.MaxInt64/60 { + return math.MaxInt64 + } + return minutes * 60 +} + +func rateLimitCapacity(count int, durationSeconds int64) int64 { + if count <= 0 || durationSeconds <= 0 { + return 0 + } + c := int64(count) + if c > math.MaxInt64/durationSeconds { + return math.MaxInt64 + } + return c * durationSeconds +} diff --git a/model/main.go b/model/main.go index 21445593..cd1569db 100644 --- a/model/main.go +++ b/model/main.go @@ -186,6 +186,9 @@ func InitDB() (err error) { panic(err) } } + if err := ensureUserQuotaColumns(DB, common.MainDatabaseType()); err != nil { + return err + } sqlDB, err := DB.DB() if err != nil { return err @@ -250,6 +253,52 @@ func InitLogDB() (err error) { return err } +var userQuotaColumns = []string{"quota", "used_quota", "aff_quota", "aff_history"} + +// ensureUserQuotaColumns rejects a legacy 32-bit wallet schema before any +// migrations run. The 64-bit-only build intentionally does not auto-upgrade +// an existing wallet; operators must migrate it explicitly before starting. +func ensureUserQuotaColumns(db *gorm.DB, dbType common.DatabaseType) error { + if common.GetEnvOrDefaultBool("SKIP_64BIT_QUOTA_SCHEMA_CHECK", false) { + common.SysLog("SKIP_64BIT_QUOTA_SCHEMA_CHECK=true; skipping user quota schema check") + return nil + } + if db == nil || dbType == common.DatabaseTypeSQLite { + return nil + } + if !db.Migrator().HasTable(&User{}) { + return nil + } + columnTypes, err := db.Migrator().ColumnTypes(&User{}) + if err != nil { + return fmt.Errorf("failed to inspect users schema: %w", err) + } + for _, expected := range userQuotaColumns { + for _, actual := range columnTypes { + if !strings.EqualFold(actual.Name(), expected) { + continue + } + dataType := actual.DatabaseTypeName() + if !is64BitIntegerType(dbType, dataType) { + return fmt.Errorf("users.%s uses %s; 32-bit is not supported", expected, dataType) + } + } + } + return nil +} + +func is64BitIntegerType(dbType common.DatabaseType, dataType string) bool { + normalized := strings.ToLower(strings.TrimSpace(dataType)) + switch dbType { + case common.DatabaseTypeMySQL: + return normalized == "bigint" || normalized == "unsigned bigint" || normalized == "bigint unsigned" + case common.DatabaseTypePostgreSQL: + return normalized == "bigint" || normalized == "int8" + default: + return false + } +} + func migrateDB() error { // Migrate price_amount column from float/double to decimal for existing tables migrateSubscriptionPlanPriceAmount() diff --git a/model/payment_method_guard_test.go b/model/payment_method_guard_test.go index 33da6cfa..80da9e3e 100644 --- a/model/payment_method_guard_test.go +++ b/model/payment_method_guard_test.go @@ -297,7 +297,7 @@ func TestRechargeEpayRejectsQuotaOverflowBeforeCompletingOrder(t *testing.T) { truncateTables(t) oldQuotaPerUnit := common.QuotaPerUnit - common.QuotaPerUnit = float64(common.MaxQuota) + common.QuotaPerUnit = float64(common.MaxWalletQuota + 1) t.Cleanup(func() { common.QuotaPerUnit = oldQuotaPerUnit }) user := insertUserForPaymentGuardTest(t, 505, 3) @@ -323,15 +323,15 @@ func TestRechargeEpayEnforcesFinalWalletQuotaLimit(t *testing.T) { }{ { name: "allows exact highest representable wallet balance", - currentQuota: common.MaxQuota - 1 - 1_000_000, - wantQuota: common.MaxQuota - 1, + currentQuota: common.MaxWalletQuota - 1_000_000, + wantQuota: common.MaxWalletQuota, wantStatus: common.TopUpStatusSuccess, }, { - name: "rejects balance above int32 quota domain", - currentQuota: common.MaxQuota - 1_000_000, + name: "rejects balance above wallet quota domain", + currentQuota: common.MaxWalletQuota - 999_999, wantErr: true, - wantQuota: common.MaxQuota - 1_000_000, + wantQuota: common.MaxWalletQuota - 999_999, wantStatus: common.TopUpStatusPending, }, } diff --git a/model/quota_reserve_test.go b/model/quota_reserve_test.go index 76eab83a..865f73ed 100644 --- a/model/quota_reserve_test.go +++ b/model/quota_reserve_test.go @@ -1,6 +1,7 @@ package model import ( + "math" "testing" "time" @@ -137,6 +138,38 @@ func TestRedisBatchReserveNeverFallsBackToStaleDatabaseBalance(t *testing.T) { assert.Equal(t, 7, reloadedToken.UsedQuota) } +func TestBatchUpdateAccumulatesTwoMaximumRequestCharges(t *testing.T) { + truncateTables(t) + resetBatchUpdateTestState(t) + common.BatchUpdateEnabled = true + + user := createReserveTestUser(t, common.MaxQuota*2+100) + require.NoError(t, DecreaseUserQuota(user.Id, common.MaxQuota, false)) + require.NoError(t, DecreaseUserQuota(user.Id, common.MaxQuota, false)) + + batchUpdate() + assert.Equal(t, 100, getUserQuotaFromDB(t, user.Id)) +} + +func TestBatchUpdateAccumulatorSaturatesOverflow(t *testing.T) { + resetBatchUpdateTestState(t) + + addNewRecord(BatchUpdateTypeUserQuota, 1, math.MaxInt) + addNewRecord(BatchUpdateTypeUserQuota, 1, 1) + batchUpdateLocks[BatchUpdateTypeUserQuota].Lock() + assert.Equal(t, math.MaxInt, batchUpdateStores[BatchUpdateTypeUserQuota][1]) + batchUpdateLocks[BatchUpdateTypeUserQuota].Unlock() + + batchUpdateLocks[BatchUpdateTypeUserQuota].Lock() + batchUpdateStores[BatchUpdateTypeUserQuota] = make(map[int]int) + batchUpdateLocks[BatchUpdateTypeUserQuota].Unlock() + addNewRecord(BatchUpdateTypeUserQuota, 1, math.MinInt) + addNewRecord(BatchUpdateTypeUserQuota, 1, -1) + batchUpdateLocks[BatchUpdateTypeUserQuota].Lock() + assert.Equal(t, math.MinInt, batchUpdateStores[BatchUpdateTypeUserQuota][1]) + batchUpdateLocks[BatchUpdateTypeUserQuota].Unlock() +} + func TestReserveFallsBackToDatabaseWhenRedisIsUnavailable(t *testing.T) { truncateTables(t) resetBatchUpdateTestState(t) diff --git a/model/redemption.go b/model/redemption.go index a7751d90..54750d50 100644 --- a/model/redemption.go +++ b/model/redemption.go @@ -175,7 +175,7 @@ func Redeem(key string, userId int) (quota int, err error) { if result.RowsAffected == 0 { return errors.New("该兑换码已被使用") } - return tx.Model(&User{}).Where("id = ?", userId).Update("quota", gorm.Expr("quota + ?", redemption.Quota)).Error + return creditTopUpQuota(tx, userId, redemption.Quota, nil) }) if err != nil { common.SysError("redemption failed: " + err.Error()) @@ -187,6 +187,12 @@ func Redeem(key string, userId int) (quota int, err error) { } func (redemption *Redemption) Insert() error { + if redemption.Quota <= 0 { + return errors.New("redemption quota must be positive") + } + if err := common.ValidateWalletQuota(redemption.Quota); err != nil { + return err + } var err error err = DB.Create(redemption).Error return err @@ -199,6 +205,12 @@ func (redemption *Redemption) SelectUpdate() error { // Update Make sure your token's fields is completed, because this will update non-zero values func (redemption *Redemption) Update() error { + if redemption.Quota <= 0 { + return errors.New("redemption quota must be positive") + } + if err := common.ValidateWalletQuota(redemption.Quota); err != nil { + return err + } var err error err = DB.Model(redemption).Select("name", "status", "quota", "redeemed_time", "expired_time").Updates(redemption).Error return err diff --git a/model/redemption_test.go b/model/redemption_test.go index 0ba2e8e8..0150fc19 100644 --- a/model/redemption_test.go +++ b/model/redemption_test.go @@ -148,6 +148,35 @@ func TestRedeemCreditsQuotaExactlyOnce(t *testing.T) { assert.Equal(t, 500, user.Quota) } +func TestRedeemRejectsWalletOverflow(t *testing.T) { + userId, key := setupRedeemFixture(t, 11) + require.NoError(t, DB.Model(&User{}).Where("id = ?", userId).Update("quota", common.MaxWalletQuota-10).Error) + + _, err := Redeem(key, userId) + require.ErrorIs(t, err, ErrRedeemFailed) + + var user User + require.NoError(t, DB.First(&user, "id = ?", userId).Error) + assert.Equal(t, common.MaxWalletQuota-10, user.Quota) + + var redemption Redemption + require.NoError(t, DB.First(&redemption, "key = ?", key).Error) + assert.Equal(t, common.RedemptionCodeStatusEnabled, redemption.Status) +} + +func TestRedemptionQuotaRejectsWalletOverflow(t *testing.T) { + setupRedeemFixture(t, 500) + + redemption := &Redemption{ + Name: "overflow-redemption", + Key: "10000000000000000000000000000002", + Status: common.RedemptionCodeStatusEnabled, + Quota: common.MaxWalletQuota + 1, + CreatedTime: common.GetTimestamp(), + } + require.Error(t, redemption.Insert()) +} + // Exactly one of several concurrent redeems of the same code may win, and // quota must be credited exactly once. func TestRedeemConcurrentSingleSuccess(t *testing.T) { diff --git a/model/subscription.go b/model/subscription.go index c89e63bf..769d7b94 100644 --- a/model/subscription.go +++ b/model/subscription.go @@ -749,7 +749,7 @@ func calcSubscriptionBalanceQuota(priceAmount float64) (int, error) { quota := decimal.NewFromFloat(priceAmount). Mul(decimal.NewFromFloat(common.QuotaPerUnit)). Ceil() - return common.QuotaFromDecimalStrict(quota) + return common.WalletQuotaFromDecimalStrict(quota) } // PurchaseSubscriptionWithBalance creates a subscription by deducting the user's wallet quota. diff --git a/model/topup.go b/model/topup.go index d837ced2..a8b91b77 100644 --- a/model/topup.go +++ b/model/topup.go @@ -42,11 +42,12 @@ const ( ) var ( - ErrPaymentMethodMismatch = errors.New("payment method mismatch") - ErrTopUpNotFound = errors.New("topup not found") - ErrTopUpStatusInvalid = errors.New("topup status invalid") - ErrInvalidTopUpQuota = errors.New("invalid top-up quota") - ErrTopUpQuotaLimitExceeded = errors.New("top-up quota limit exceeded") + ErrPaymentMethodMismatch = errors.New("payment method mismatch") + ErrTopUpNotFound = errors.New("topup not found") + ErrTopUpStatusInvalid = errors.New("topup status invalid") + ErrInvalidTopUpQuota = errors.New("invalid top-up quota") + ErrTopUpQuotaLimitExceeded = errors.New("top-up quota limit exceeded") + ErrWalletQuotaLimitExceeded = errors.New("wallet quota limit exceeded") ) func (topUp *TopUp) Insert() error { @@ -56,10 +57,10 @@ func (topUp *TopUp) Insert() error { } func topUpQuotaMaxCurrent(creditedQuota int) (int, error) { - if creditedQuota <= 0 || creditedQuota >= common.MaxQuota { + if creditedQuota <= 0 || creditedQuota > common.MaxWalletQuota { return 0, ErrInvalidTopUpQuota } - return common.MaxQuota - 1 - creditedQuota, nil + return common.MaxWalletQuota - creditedQuota, nil } // ValidateTopUpQuotaCapacity performs the user-facing pre-payment check. The @@ -81,8 +82,8 @@ func ValidateTopUpQuotaCapacity(userId int, creditedQuota int) error { return nil } -// creditTopUpQuota atomically enforces the int32 wallet ceiling while adding -// quota. Keeping the predicate and increment in one UPDATE prevents two +// creditTopUpQuota atomically enforces the wallet ceiling while adding quota. +// Keeping the predicate and increment in one UPDATE prevents two // concurrent callbacks from both passing a separate read/check. func creditTopUpQuota(tx *gorm.DB, userId int, creditedQuota int, updates map[string]interface{}) error { maxCurrentQuota, err := topUpQuotaMaxCurrent(creditedQuota) @@ -203,7 +204,7 @@ func RechargeEpay(tradeNo string, actualPaymentMethod string, callerIp string) ( topUp.PaymentMethod = actualPaymentMethod } var quotaErr error - quotaToAdd, quotaErr = common.QuotaFromDecimalStrict( + quotaToAdd, quotaErr = common.WalletQuotaFromDecimalStrict( decimal.NewFromInt(topUp.Amount).Mul(decimal.NewFromFloat(common.QuotaPerUnit)), ) if quotaErr != nil || quotaToAdd <= 0 { @@ -266,7 +267,7 @@ func Recharge(referenceId string, customerId string, callerIp string) (err error return err } - quota, err = common.QuotaFromDecimalStrict( + quota, err = common.WalletQuotaFromDecimalStrict( decimal.NewFromFloat(topUp.Money).Mul(decimal.NewFromFloat(common.QuotaPerUnit)), ) if err != nil || quota <= 0 { @@ -482,11 +483,11 @@ func ManualCompleteTopUp(tradeNo string, callerIp string) error { // - 其他订单(如易支付):Amount 为美元数量,* QuotaPerUnit var quotaErr error if topUp.PaymentProvider == PaymentProviderStripe { - quotaToAdd, quotaErr = common.QuotaFromDecimalStrict( + quotaToAdd, quotaErr = common.WalletQuotaFromDecimalStrict( decimal.NewFromFloat(topUp.Money).Mul(decimal.NewFromFloat(common.QuotaPerUnit)), ) } else { - quotaToAdd, quotaErr = common.QuotaFromDecimalStrict( + quotaToAdd, quotaErr = common.WalletQuotaFromDecimalStrict( decimal.NewFromInt(topUp.Amount).Mul(decimal.NewFromFloat(common.QuotaPerUnit)), ) } @@ -556,7 +557,7 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string } // Creem 直接使用 Amount 作为充值额度(整数) - quota, err = common.QuotaFromDecimalStrict(decimal.NewFromInt(topUp.Amount)) + quota, err = common.WalletQuotaFromDecimalStrict(decimal.NewFromInt(topUp.Amount)) if err != nil || quota <= 0 { return ErrInvalidTopUpQuota } @@ -624,7 +625,7 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) { return errors.New("充值订单状态错误") } - quotaToAdd, err = common.QuotaFromDecimalStrict( + quotaToAdd, err = common.WalletQuotaFromDecimalStrict( decimal.NewFromInt(topUp.Amount).Mul(decimal.NewFromFloat(common.QuotaPerUnit)), ) if err != nil || quotaToAdd <= 0 { @@ -684,7 +685,7 @@ func RechargeWaffoPancake(tradeNo string) (err error) { return errors.New("充值订单状态错误") } - quotaToAdd, err = common.QuotaFromDecimalStrict( + quotaToAdd, err = common.WalletQuotaFromDecimalStrict( decimal.NewFromInt(topUp.Amount).Mul(decimal.NewFromFloat(common.QuotaPerUnit)), ) if err != nil || quotaToAdd <= 0 { diff --git a/model/user.go b/model/user.go index 7bc060ad..1fbaa34c 100644 --- a/model/user.go +++ b/model/user.go @@ -1271,25 +1271,47 @@ func IncreaseUserQuota(id int, quota int, db bool) (err error) { if quota < 0 { return errors.New("quota 不能为负数!") } + if err := common.ValidateWalletQuota(quota); err != nil { + return err + } + if !db && common.BatchUpdateEnabled { + addNewRecord(BatchUpdateTypeUserQuota, id, quota) + gopool.Go(func() { + if err := cacheIncrUserQuota(id, int64(quota)); err != nil { + common.SysLog("failed to increase user quota: " + err.Error()) + } + }) + return nil + } + if err := increaseUserQuota(id, quota); err != nil { + return err + } gopool.Go(func() { - err := cacheIncrUserQuota(id, int64(quota)) - if err != nil { + if err := cacheIncrUserQuota(id, int64(quota)); err != nil { common.SysLog("failed to increase user quota: " + err.Error()) } }) - if !db && common.BatchUpdateEnabled { - addNewRecord(BatchUpdateTypeUserQuota, id, quota) - return nil - } - return increaseUserQuota(id, quota) + return nil } func increaseUserQuota(id int, quota int) (err error) { - err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error - if err != nil { + result := DB.Model(&User{}). + Where("id = ? AND quota <= ?", id, common.MaxWalletQuota-quota). + Update("quota", gorm.Expr("quota + ?", quota)) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 1 { + return nil + } + var count int64 + if err := DB.Model(&User{}).Where("id = ?", id).Count(&count).Error; err != nil { return err } - return err + if count == 0 { + return gorm.ErrRecordNotFound + } + return ErrWalletQuotaLimitExceeded } func DecreaseUserQuota(id int, quota int, db bool) (err error) { diff --git a/model/utils.go b/model/utils.go index b1793706..63d51c4b 100644 --- a/model/utils.go +++ b/model/utils.go @@ -2,6 +2,8 @@ package model import ( "errors" + "fmt" + "math" "sync" "time" @@ -42,11 +44,22 @@ func InitBatchUpdater() { func addNewRecord(type_ int, id int, value int) { batchUpdateLocks[type_].Lock() defer batchUpdateLocks[type_].Unlock() - if _, ok := batchUpdateStores[type_][id]; !ok { + old, ok := batchUpdateStores[type_][id] + if !ok { batchUpdateStores[type_][id] = value - } else { - batchUpdateStores[type_][id] += value + return } + + sum := old + value + if (value > 0 && sum < old) || (value < 0 && sum > old) { + common.SysError(fmt.Sprintf("batch update overflow: type=%d id=%d old=%d value=%d", type_, id, old, value)) + if value > 0 { + sum = math.MaxInt + } else { + sum = math.MinInt + } + } + batchUpdateStores[type_][id] = sum } func batchUpdate() { diff --git a/pkg/billingexpr/billingexpr_test.go b/pkg/billingexpr/billingexpr_test.go index 90485571..d34bdcbf 100644 --- a/pkg/billingexpr/billingexpr_test.go +++ b/pkg/billingexpr/billingexpr_test.go @@ -4,6 +4,7 @@ import ( "math" "testing" + "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/pkg/billingexpr" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -344,9 +345,9 @@ func TestQuotaRound(t *testing.T) { {999.4999, 999}, {999.5, 1000}, {1e9 + 0.5, 1e9 + 1}, - // Oversized expression results saturate at int32 (delegated to + // Oversized expression results saturate at the single-request limit (delegated to // common.QuotaRound); full saturation coverage lives in common. - {3.6893488147419103e19, math.MaxInt32}, + {3.6893488147419103e19, common.MaxQuota}, } for _, tt := range tests { got := billingexpr.QuotaRound(tt.in) diff --git a/pkg/billingexpr/settle_clamp_test.go b/pkg/billingexpr/settle_clamp_test.go index 4d765b23..2d082311 100644 --- a/pkg/billingexpr/settle_clamp_test.go +++ b/pkg/billingexpr/settle_clamp_test.go @@ -1,7 +1,6 @@ package billingexpr_test import ( - "math" "testing" "github.com/QuantumNous/new-api/common" @@ -11,13 +10,13 @@ import ( ) // TestComputeTieredQuota_ClampOnOverflow guards the billing-safety invariant -// that an oversized tiered settlement clamps to the int32 max instead of +// that an oversized tiered settlement clamps to the single-request max instead of // wrapping into a credit, and that the saturation event is surfaced on the // result so callers can record it for admin auditing. func TestComputeTieredQuota_ClampOnOverflow(t *testing.T) { - // exprOutput = p * 1e9 = 1e18; quotaBeforeGroup = 1e18 / 1e6 * 5e5 = 5e17, - // which far exceeds MaxInt32 and must saturate. - exprStr := `tier("base", p * 1000000000)` + // exprOutput = p * 1e12 = 1e21; quotaBeforeGroup = 1e21 / 1e6 * 5e5 = 5e20, + // which far exceeds the supported single-request range and must saturate. + exprStr := `tier("base", p * 1000000000000)` snap := &billingexpr.BillingSnapshot{ BillingMode: "tiered_expr", ExprString: exprStr, @@ -29,10 +28,10 @@ func TestComputeTieredQuota_ClampOnOverflow(t *testing.T) { result, err := billingexpr.ComputeTieredQuota(snap, billingexpr.TokenParams{P: 1_000_000_000}) require.NoError(t, err) - assert.Equal(t, math.MaxInt32, result.ActualQuotaAfterGroup, "oversized quota must clamp to int32 max, never wrap negative") + assert.Equal(t, common.MaxQuota, result.ActualQuotaAfterGroup, "oversized quota must clamp, never wrap negative") require.NotNil(t, result.Clamp, "clamp event must be surfaced so it can be audited") assert.Equal(t, common.QuotaClampOverflow, result.Clamp.Kind) - assert.Equal(t, math.MaxInt32, result.Clamp.Clamped) + assert.Equal(t, common.MaxQuota, result.Clamp.Clamped) } // TestComputeTieredQuota_NoClampInRange confirms an in-range settlement leaves diff --git a/pkg/billingexpr/types.go b/pkg/billingexpr/types.go index 1bdf6834..a6711903 100644 --- a/pkg/billingexpr/types.go +++ b/pkg/billingexpr/types.go @@ -68,7 +68,7 @@ type TieredResult struct { MatchedTier string `json:"matched_tier"` RequestRules []RequestRuleTrace `json:"request_rules,omitempty"` CrossedTier bool `json:"crossed_tier"` - // Clamp records an int32 saturation event during quota conversion so the + // Clamp records a single-request saturation event during quota conversion so the // caller can surface it on the consume log for admin auditing. Nil when no // clamping occurred. Not serialized: the marker is attached separately via // the shared quota-saturation audit path. diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index b0bb19bd..56f57234 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -152,7 +152,7 @@ type RelayInfo struct { PriceData hosttypes.PriceData // QuotaClamp is set (non-nil) when a quota conversion saturated at the - // int32 bound (or NaN fallback) while computing this request's charge. + // supported single-request bound (or NaN fallback) while computing this request's charge. // It is surfaced onto the consume/task log's admin_info for auditing. QuotaClamp *common.QuotaClamp diff --git a/relay/helper/price_test.go b/relay/helper/price_test.go index 0f28b5a4..ca38b548 100644 --- a/relay/helper/price_test.go +++ b/relay/helper/price_test.go @@ -156,7 +156,7 @@ func TestModelPriceHelperTieredRejectsPreConsumeOverflow(t *testing.T) { require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{ "billing_setting.billing_mode": `{"tiered-overflow-model":"tiered_expr"}`, - "billing_setting.billing_expr": `{"tiered-overflow-model":"tier(\"overflow\", p * 1000000000000000)"}`, + "billing_setting.billing_expr": `{"tiered-overflow-model":"tier(\"overflow\", p * 100000000000000000)"}`, "group_ratio_setting.group_ratio": `{"default":1}`, })) diff --git a/service/quota.go b/service/quota.go index 359956a5..3639ee5f 100644 --- a/service/quota.go +++ b/service/quota.go @@ -3,7 +3,6 @@ package service import ( "errors" "fmt" - "math" "strings" "time" @@ -272,11 +271,16 @@ func CalcOpenRouterCacheCreateTokens(usage dto.Usage, priceData types.PriceData) completionTokens := float64(usage.CompletionTokens) promptCacheReadTokens := float64(usage.PromptTokensDetails.CachedTokens) - return int(math.Round((cost - + value := (cost - totalPromptTokens*quotaPrice + promptCacheReadTokens*(quotaPrice-promptCacheReadPrice) - completionTokens*completionPrice) / - (promptCacheCreatePrice - quotaPrice))) + (promptCacheCreatePrice - quotaPrice) + quota, clamp := common.QuotaRoundChecked(value) + if clamp != nil { + return -1 + } + return quota } func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent string) { diff --git a/service/quota_saturation_test.go b/service/quota_saturation_test.go index e8cd55c2..da4ff350 100644 --- a/service/quota_saturation_test.go +++ b/service/quota_saturation_test.go @@ -1,12 +1,15 @@ package service import ( + "math" "net/http" "testing" "github.com/QuantumNous/new-api/common" relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/types" + hosttypes "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -42,6 +45,28 @@ func TestAttachQuotaSaturationNestsUnderAdminInfo(t *testing.T) { require.Equal(t, common.MaxQuota, sat["clamped"]) } +func TestCalcViolationFeeQuotaSaturates(t *testing.T) { + oldQuotaPerUnit := common.QuotaPerUnit + common.QuotaPerUnit = 500_000 + t.Cleanup(func() { common.QuotaPerUnit = oldQuotaPerUnit }) + + require.Equal(t, common.MaxQuota, calcViolationFeeQuota(1e20, 1)) +} + +func TestCalcOpenRouterCacheCreateTokensDoesNotWrap(t *testing.T) { + oldQuotaPerUnit := common.QuotaPerUnit + common.QuotaPerUnit = 500_000 + t.Cleanup(func() { common.QuotaPerUnit = oldQuotaPerUnit }) + + got := CalcOpenRouterCacheCreateTokens(dto.Usage{Cost: math.Inf(1)}, hosttypes.PriceData{ + ModelRatio: 1, + CacheCreationRatio: 2, + CacheRatio: 1, + CompletionRatio: 1, + }) + require.Equal(t, -1, got) +} + // TestAttachQuotaSaturationPreservesExistingAdminInfo verifies the marker is // merged into a pre-existing admin_info map without clobbering it. func TestAttachQuotaSaturationPreservesExistingAdminInfo(t *testing.T) { diff --git a/service/text_quota.go b/service/text_quota.go index b7578f73..19f0e946 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -216,8 +216,8 @@ func composeTieredTextQuota(relayInfo *relaycommon.RelayInfo, summary textQuotaS } // Saturate the final sum, not just the surcharge: tieredQuota can be near - // MaxQuota and adding the surcharge could push the total past the int32 - // quota policy bound (persisted quota columns are 32-bit). + // MaxQuota and adding the surcharge could push the total past the + // single-request quota policy bound. total, clamp := common.QuotaFromDecimalChecked( decimal.NewFromInt(int64(tieredQuota)).Add(summary.ToolCallSurchargeQuota), ) diff --git a/service/text_quota_test.go b/service/text_quota_test.go index e4a1ed68..9e935b40 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -774,9 +774,9 @@ func TestComposeTieredTextQuotaErrorFallbackUsesPreConsumedQuota(t *testing.T) { // settlement both saturates the quota and records the clamp on RelayInfo, so // every consume path (text, audio, WSS) can surface it under admin_info. func TestTryTieredSettleRecordsClampOnOverflow(t *testing.T) { - // exprOutput = p * 1e9; quotaBeforeGroup = p*1e9 / 1e6 * 5e5 far exceeds - // MaxInt32 and must saturate. - exprStr := `tier("base", p * 1000000000)` + // exprOutput = p * 1e12; quotaBeforeGroup = p*1e12 / 1e6 * 5e5 far exceeds + // the supported single-request range and must saturate. + exprStr := `tier("base", p * 1000000000000)` relayInfo := &relaycommon.RelayInfo{ OriginModelName: "overflow-model", TieredBillingSnapshot: &billingexpr.BillingSnapshot{ @@ -792,7 +792,7 @@ func TestTryTieredSettleRecordsClampOnOverflow(t *testing.T) { require.True(t, ok) require.NotNil(t, result) - require.Equal(t, math.MaxInt32, quota, "oversized settlement must clamp, never wrap negative") + require.Equal(t, common.MaxQuota, quota, "oversized settlement must clamp, never wrap negative") require.NotNil(t, relayInfo.QuotaClamp, "clamp must be recorded on RelayInfo for admin auditing") require.Equal(t, common.QuotaClampOverflow, relayInfo.QuotaClamp.Kind) } diff --git a/service/tiered_settle.go b/service/tiered_settle.go index 1f3f58fe..0e9618ea 100644 --- a/service/tiered_settle.go +++ b/service/tiered_settle.go @@ -180,7 +180,7 @@ func TryTieredSettle(relayInfo *relaycommon.RelayInfo, params billingexpr.TokenP return true, quota, nil } - // Surface any int32 saturation from settlement onto RelayInfo so the + // Surface any single-request saturation from settlement onto RelayInfo so the // consume log records it under admin_info, regardless of which caller // (text, audio, WSS) consumes the returned quota. First non-nil wins. noteQuotaClamp(relayInfo, tr.Clamp) diff --git a/service/token_counter.go b/service/token_counter.go index aad320a1..3b0b5cd1 100644 --- a/service/token_counter.go +++ b/service/token_counter.go @@ -140,11 +140,11 @@ func getImageToken(c *gin.Context, fileMeta *types.FileMeta, model string, strea if imageTokens > 1536 { imageTokens = 1536 } - return int(math.Round(float64(imageTokens) * multiplier)), nil + return common.QuotaRound(float64(imageTokens) * multiplier), nil } // below cap imageTokens := rawPatches - return int(math.Round(float64(imageTokens) * multiplier)), nil + return common.QuotaRound(float64(imageTokens) * multiplier), nil } // Tile-based calculation for 4o/4.1/4.5/o1/o3/etc. diff --git a/service/violation_fee.go b/service/violation_fee.go index f5153362..e063d4d8 100644 --- a/service/violation_fee.go +++ b/service/violation_fee.go @@ -88,15 +88,14 @@ func calcViolationFeeQuota(amount, groupRatio float64) int { if groupRatio <= 0 { return 0 } - quota := decimal.NewFromFloat(amount). + quota := common.QuotaFromDecimal(decimal.NewFromFloat(amount). Mul(decimal.NewFromFloat(common.QuotaPerUnit)). Mul(decimal.NewFromFloat(groupRatio)). - Round(0). - IntPart() + Round(0)) if quota <= 0 { return 0 } - return int(quota) + return quota } // ChargeViolationFeeIfNeeded charges an additional fee after the normal flow finishes (including refund). diff --git a/setting/rate_limit.go b/setting/rate_limit.go index 413f3958..d046ae18 100644 --- a/setting/rate_limit.go +++ b/setting/rate_limit.go @@ -1,7 +1,6 @@ package setting import ( - "encoding/json" "fmt" "math" "sync" @@ -9,6 +8,16 @@ import ( "github.com/QuantumNous/new-api/common" ) +// maxRateLimitDurationSeconds is the largest window the count cap is computed +// against (24h). Token-bucket capacity is count*duration; this keeps that +// product inside int64 when the window is at most a day. +const maxRateLimitDurationSeconds = 24 * 60 * 60 + +// maxModelRequestRateLimitCount is math.MaxInt64 / maxRateLimitDurationSeconds. +// It is the largest count that cannot overflow int64(count)*duration for a +// window of at most 24 hours. +const maxModelRequestRateLimitCount int64 = math.MaxInt64 / maxRateLimitDurationSeconds + var ModelRequestRateLimitEnabled = false var ModelRequestRateLimitDurationMinutes = 1 var ModelRequestRateLimitCount = 0 @@ -20,7 +29,7 @@ func ModelRequestRateLimitGroup2JSONString() string { ModelRequestRateLimitMutex.RLock() defer ModelRequestRateLimitMutex.RUnlock() - jsonBytes, err := json.Marshal(ModelRequestRateLimitGroup) + jsonBytes, err := common.Marshal(ModelRequestRateLimitGroup) if err != nil { common.SysLog("error marshalling model ratio: " + err.Error()) } @@ -32,7 +41,7 @@ func UpdateModelRequestRateLimitGroupByJSONString(jsonStr string) error { defer ModelRequestRateLimitMutex.RUnlock() ModelRequestRateLimitGroup = make(map[string][2]int) - return json.Unmarshal([]byte(jsonStr), &ModelRequestRateLimitGroup) + return common.Unmarshal([]byte(jsonStr), &ModelRequestRateLimitGroup) } func GetGroupRateLimit(group string) (totalCount, successCount int, found bool) { @@ -52,7 +61,7 @@ func GetGroupRateLimit(group string) (totalCount, successCount int, found bool) func CheckModelRequestRateLimitGroup(jsonStr string) error { checkModelRequestRateLimitGroup := make(map[string][2]int) - err := json.Unmarshal([]byte(jsonStr), &checkModelRequestRateLimitGroup) + err := common.Unmarshal([]byte(jsonStr), &checkModelRequestRateLimitGroup) if err != nil { return err } @@ -60,8 +69,8 @@ func CheckModelRequestRateLimitGroup(jsonStr string) error { if limits[0] < 0 || limits[1] < 1 { return fmt.Errorf("group %s has negative rate limit values: [%d, %d]", group, limits[0], limits[1]) } - if limits[0] > math.MaxInt32 || limits[1] > math.MaxInt32 { - return fmt.Errorf("group %s [%d, %d] has max rate limits value 2147483647", group, limits[0], limits[1]) + if int64(limits[0]) > maxModelRequestRateLimitCount || int64(limits[1]) > maxModelRequestRateLimitCount { + return fmt.Errorf("group %s [%d, %d] exceeds max rate limit %d", group, limits[0], limits[1], maxModelRequestRateLimitCount) } } diff --git a/web/src/features/usage-logs/types.ts b/web/src/features/usage-logs/types.ts index 3e3789b9..f2af3155 100644 --- a/web/src/features/usage-logs/types.ts +++ b/web/src/features/usage-logs/types.ts @@ -134,7 +134,7 @@ export interface LogOtherData { admin_role?: number auth_method?: 'session' | 'access_token' | string // Quota saturation marker: set when a quota conversion clamped at the - // int32 bound (overflow/underflow) or hit a NaN fallback while computing + // supported single-request bound (overflow/underflow) or hit a NaN fallback while computing // this request's charge. Admin-only (nested under admin_info). quota_saturation?: { op: string