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
This commit is contained in:
Seefs
2026-08-26 20:57:54 +08:00
committed by GitHub
parent 2d8e50bf36
commit a073f74b38
35 changed files with 445 additions and 116 deletions
+2
View File
@@ -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
# 缓存相关配置
+1 -1
View File
@@ -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.
+37 -16
View File
@@ -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))
}
+21 -1
View File
@@ -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)
}
+4 -4
View File
@@ -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{} {
+17
View File
@@ -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
+13 -2
View File
@@ -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
+8 -4
View File
@@ -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("充值额度超出系统可表示范围")
}
+20 -14
View File
@@ -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}`))
+8
View File
@@ -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)
+21
View File
@@ -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)
}
+25 -2
View File
@@ -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
}
+49
View File
@@ -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()
+6 -6
View File
@@ -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,
},
}
+33
View File
@@ -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)
+13 -1
View File
@@ -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
+29
View File
@@ -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) {
+1 -1
View File
@@ -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.
+17 -16
View File
@@ -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 {
+32 -10
View File
@@ -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) {
+16 -3
View File
@@ -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() {
+3 -2
View File
@@ -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)
+6 -7
View File
@@ -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
+1 -1
View File
@@ -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.
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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}`,
}))
+7 -3
View File
@@ -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) {
+25
View File
@@ -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) {
+2 -2
View File
@@ -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),
)
+4 -4
View File
@@ -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)
}
+1 -1
View File
@@ -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)
+2 -2
View File
@@ -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.
+3 -4
View File
@@ -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).
+15 -6
View File
@@ -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)
}
}
+1 -1
View File
@@ -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