From d0bd8aac742d1e160a5ca61743fe35f4fff880e8 Mon Sep 17 00:00:00 2001 From: CaIon Date: Tue, 7 Jul 2026 00:21:06 +0800 Subject: [PATCH] fix(billing): validate quantity parameters and harden quota calculations Bound user-supplied count/duration parameters at request validation, route ratio multipliers through guarded setters, and use saturating int conversions in all quota math paths. --- AGENTS.md | 10 +++ common/quota_math.go | 21 ++++++ common/quota_math_test.go | 22 +++++++ controller/task_video.go | 4 +- dto/openai_image.go | 4 ++ relay/channel/ali/image.go | 6 ++ relay/channel/task/ali/adaptor.go | 4 +- relay/channel/task/gemini/billing.go | 10 ++- relay/common/relay_utils.go | 24 +++++++ relay/common/relay_utils_test.go | 64 +++++++++++++++++++ relay/helper/openai_image_request_test.go | 78 +++++++++++++++++++++++ relay/helper/valid_request.go | 12 +++- relay/relay_task.go | 14 ++-- service/task_billing.go | 4 +- service/text_quota.go | 12 +++- service/text_quota_test.go | 14 ++++ types/price_data.go | 9 ++- 17 files changed, 293 insertions(+), 19 deletions(-) create mode 100644 common/quota_math.go create mode 100644 common/quota_math_test.go diff --git a/AGENTS.md b/AGENTS.md index ec6af979..0a1534b5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -98,6 +98,16 @@ Do NOT directly import or call `encoding/json` in business code. `json.RawMessag **Billing expression system:** When working on tiered/dynamic billing (expression-based pricing), MUST read `pkg/billingexpr/expr.md` first. It documents the design philosophy, expression language, full architecture, token normalization rules, quota conversion, and expression versioning. All billing expression changes must follow that document. +**Billing safety invariants:** Quota/billing code MUST never produce a negative charge (a credit) from arithmetic overflow or unvalidated input. Apply defense in depth: + +- 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. Reuse these constants instead of introducing new ad hoc limits for the same concepts. +- 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. +- Never convert a computed quota to `int` with a bare cast like `int(float64(quota) * ratio)` or `int(decimal.IntPart())`. Use the saturating converters: `common.QuotaFromFloat` for float products, `decimalToQuota` in `service` for decimal products. Saturation bounds are int32 because quota columns (user/token/log) are 32-bit integers in the database. +- 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. +- Fields parsed into unsigned types (`*uint`) accept huge positive JSON numbers (e.g. `18446744073686646784`, a wrapped negative); a `>= 0` check is not sufficient, an upper bound is mandatory. +- Regression tests for these invariants belong with the boundary they protect (request validators, converter helpers). See `relay/helper/openai_image_request_test.go`, `relay/common/relay_utils_test.go`, and `common/quota_math_test.go` for the expected style. + **Backend test quality:** Backend tests must protect real behavior, API contracts, billing/accounting invariants, data compatibility, or regression paths. - Do not add tests that only improve coverage numbers, prove that code happens to run, or lock in implementation details without a user-visible or cross-module contract. diff --git a/common/quota_math.go b/common/quota_math.go new file mode 100644 index 00000000..878b3526 --- /dev/null +++ b/common/quota_math.go @@ -0,0 +1,21 @@ +package common + +import "math" + +// QuotaFromFloat converts a computed quota value to int with saturation. +// Quota products can include user-controlled multipliers (image n, video +// seconds, resolution ratios); an oversized product must never wrap around +// and turn a charge into a credit. The bound is int32 because quota columns +// (user/token/log) are 32-bit integers in the database. +func QuotaFromFloat(value float64) int { + if math.IsNaN(value) { + return 0 + } + if value >= math.MaxInt32 { + return math.MaxInt32 + } + if value <= math.MinInt32 { + return math.MinInt32 + } + return int(value) +} diff --git a/common/quota_math_test.go b/common/quota_math_test.go new file mode 100644 index 00000000..847d3db0 --- /dev/null +++ b/common/quota_math_test.go @@ -0,0 +1,22 @@ +package common + +import ( + "math" + "testing" + + "github.com/stretchr/testify/assert" +) + +// TestQuotaFromFloat guards the billing invariant that oversized quota +// products (e.g. price multiplied by a huge user-supplied count) saturate +// instead of wrapping into a negative charge (credit). +func TestQuotaFromFloat(t *testing.T) { + assert.Equal(t, 42, QuotaFromFloat(42.4)) + assert.Equal(t, -42, QuotaFromFloat(-42.4)) + // 2000 quota per call * n=18446744073686646784 overflows int64. + assert.Equal(t, math.MaxInt32, QuotaFromFloat(2000*1.8446744073686647e19)) + assert.Equal(t, math.MinInt32, QuotaFromFloat(-2000*1.8446744073686647e19)) + assert.Equal(t, math.MaxInt32, QuotaFromFloat(math.Inf(1))) + assert.Equal(t, math.MinInt32, QuotaFromFloat(math.Inf(-1))) + assert.Equal(t, 0, QuotaFromFloat(math.NaN())) +} diff --git a/controller/task_video.go b/controller/task_video.go index 8be2e176..18c20c06 100644 --- a/controller/task_video.go +++ b/controller/task_video.go @@ -178,8 +178,8 @@ func updateVideoSingleTask(ctx context.Context, adaptor channel.TaskAdaptor, cha finalGroupRatio = groupRatio } - // 计算实际应扣费额度: totalTokens * modelRatio * groupRatio - actualQuota := int(float64(taskResult.TotalTokens) * modelRatio * finalGroupRatio) + // 计算实际应扣费额度: totalTokens * modelRatio * groupRatio（饱和转换，防止溢出成负数） + actualQuota := common.QuotaFromFloat(float64(taskResult.TotalTokens) * modelRatio * finalGroupRatio) // 计算差额 preConsumedQuota := task.Quota diff --git a/dto/openai_image.go b/dto/openai_image.go index 547b0d18..ef99c1c7 100644 --- a/dto/openai_image.go +++ b/dto/openai_image.go @@ -11,6 +11,10 @@ import ( "github.com/gin-gonic/gin" ) +// MaxImageN caps the image generation count. Without this bound a huge or +// wrapped-negative n overflows quota calculation into a negative charge. +const MaxImageN = 128 + type ImageRequest struct { Model string `json:"model"` Prompt string `json:"prompt" binding:"required"` diff --git a/relay/channel/ali/image.go b/relay/channel/ali/image.go index a391e40f..af0717a3 100644 --- a/relay/channel/ali/image.go +++ b/relay/channel/ali/image.go @@ -54,6 +54,12 @@ func oaiImage2AliImageRequest(info *relaycommon.RelayInfo, request dto.ImageRequ } } + // Parameters may come from Extra["parameters"], bypassing the standard + // top-level n validation; enforce the same bound before it becomes a + // billing multiplier. + if imageRequest.Parameters.N < 0 || imageRequest.Parameters.N > dto.MaxImageN { + return nil, fmt.Errorf("parameters.n must be an integer between 1 and %d", dto.MaxImageN) + } if imageRequest.Parameters.N != 0 { info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.N)) } diff --git a/relay/channel/task/ali/adaptor.go b/relay/channel/task/ali/adaptor.go index 2b85c435..51ffa743 100644 --- a/relay/channel/task/ali/adaptor.go +++ b/relay/channel/task/ali/adaptor.go @@ -456,8 +456,10 @@ func (a *TaskAdaptor) EstimateBilling(c *gin.Context, info *relaycommon.RelayInf return nil } + // metadata can override Duration past standard request validation; + // cap it because it is used as a billing multiplier. otherRatios := map[string]float64{ - "seconds": float64(aliReq.Parameters.Duration), + "seconds": float64(min(aliReq.Parameters.Duration, relaycommon.MaxTaskDurationSeconds)), } ratios, err := ProcessAliOtherRatios(aliReq) if err != nil { diff --git a/relay/channel/task/gemini/billing.go b/relay/channel/task/gemini/billing.go index b081eb2f..d3b981b9 100644 --- a/relay/channel/task/gemini/billing.go +++ b/relay/channel/task/gemini/billing.go @@ -3,6 +3,8 @@ package gemini import ( "strconv" "strings" + + relaycommon "github.com/QuantumNous/new-api/relay/common" ) // ParseVeoDurationSeconds extracts durationSeconds from metadata. @@ -46,19 +48,21 @@ func ParseVeoResolution(metadata map[string]any) string { // ResolveVeoDuration returns the effective duration in seconds. // Priority: metadata["durationSeconds"] > stdDuration > stdSeconds > default (8). +// The result is capped because it is used as a billing multiplier and the +// metadata path bypasses standard request validation. func ResolveVeoDuration(metadata map[string]any, stdDuration int, stdSeconds string) int { if metadata != nil { if _, exists := metadata["durationSeconds"]; exists { if d := ParseVeoDurationSeconds(metadata); d > 0 { - return d + return min(d, relaycommon.MaxTaskDurationSeconds) } } } if stdDuration > 0 { - return stdDuration + return min(stdDuration, relaycommon.MaxTaskDurationSeconds) } if s, err := strconv.Atoi(stdSeconds); err == nil && s > 0 { - return s + return min(s, relaycommon.MaxTaskDurationSeconds) } return 8 } diff --git a/relay/common/relay_utils.go b/relay/common/relay_utils.go index e3bb4383..eb027587 100644 --- a/relay/common/relay_utils.go +++ b/relay/common/relay_utils.go @@ -78,6 +78,22 @@ func validatePrompt(prompt string) *dto.TaskError { return nil } +// MaxTaskDurationSeconds caps user-supplied video duration. Duration is used +// as a billing multiplier (OtherRatio "seconds"); an unbounded value could +// overflow quota calculation into a negative charge. +const MaxTaskDurationSeconds = 3600 + +func validateTaskDurationBounds(req TaskSubmitReq) *dto.TaskError { + seconds := req.Duration + if seconds == 0 && req.Seconds != "" { + seconds, _ = strconv.Atoi(req.Seconds) + } + if seconds < 0 || seconds > MaxTaskDurationSeconds { + return createTaskError(fmt.Errorf("seconds must be between 1 and %d", MaxTaskDurationSeconds), "invalid_seconds", http.StatusBadRequest, true) + } + return nil +} + func validateMultipartTaskRequest(c *gin.Context, info *RelayInfo, action string) (TaskSubmitReq, error) { var req TaskSubmitReq if _, err := c.MultipartForm(); err != nil { @@ -156,6 +172,10 @@ func ValidateMultipartDirect(c *gin.Context, info *RelayInfo) *dto.TaskError { return taskErr } + if taskErr := validateTaskDurationBounds(req); taskErr != nil { + return taskErr + } + action := constant.TaskActionTextGenerate if hasInputReference { action = constant.TaskActionGenerate @@ -217,6 +237,10 @@ func ValidateBasicTaskRequest(c *gin.Context, info *RelayInfo, action string) *d return taskErr } + if taskErr := validateTaskDurationBounds(req); taskErr != nil { + return taskErr + } + if len(req.Images) == 0 && strings.TrimSpace(req.Image) != "" { // 兼容单图上传 req.Images = []string{req.Image} diff --git a/relay/common/relay_utils_test.go b/relay/common/relay_utils_test.go index 62390da0..5bd51e8e 100644 --- a/relay/common/relay_utils_test.go +++ b/relay/common/relay_utils_test.go @@ -31,3 +31,67 @@ func TestValidateMultipartDirectNormalizesImageField(t *testing.T) { require.Equal(t, []string{"https://example.com/first.png"}, storedReq.Images) require.Equal(t, constant.TaskActionGenerate, info.Action) } + +// TestTaskDurationBounds guards the billing invariant that user-supplied +// video duration (a quota multiplier via OtherRatio "seconds") is bounded, so +// it can never overflow quota calculation into a negative charge. +func TestTaskDurationBounds(t *testing.T) { + gin.SetMode(gin.TestMode) + + newContext := func(t *testing.T, body string) (*gin.Context, *RelayInfo) { + request := httptest.NewRequest(http.MethodPost, "/v1/video/generations", strings.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + context, _ := gin.CreateTestContext(httptest.NewRecorder()) + context.Request = request + return context, &RelayInfo{TaskRelayInfo: &TaskRelayInfo{}} + } + + tests := []struct { + name string + body string + wantErr bool + }{ + { + name: "huge duration is rejected", + body: `{"model":"sora-2","prompt":"a cat","duration":9999999999}`, + wantErr: true, + }, + { + name: "huge seconds string is rejected", + body: `{"model":"sora-2","prompt":"a cat","seconds":"9999999999"}`, + wantErr: true, + }, + { + name: "negative duration is rejected", + body: `{"model":"sora-2","prompt":"a cat","duration":-8}`, + wantErr: true, + }, + { + name: "normal duration is accepted", + body: `{"model":"sora-2","prompt":"a cat","seconds":"8"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name+" (multipart direct)", func(t *testing.T) { + context, info := newContext(t, tt.body) + taskErr := ValidateMultipartDirect(context, info) + if tt.wantErr { + require.NotNil(t, taskErr) + require.Equal(t, "invalid_seconds", taskErr.Code) + } else { + require.Nil(t, taskErr) + } + }) + t.Run(tt.name+" (basic task request)", func(t *testing.T) { + context, info := newContext(t, tt.body) + taskErr := ValidateBasicTaskRequest(context, info, constant.TaskActionGenerate) + if tt.wantErr { + require.NotNil(t, taskErr) + require.Equal(t, "invalid_seconds", taskErr.Code) + } else { + require.Nil(t, taskErr) + } + }) + } +} diff --git a/relay/helper/openai_image_request_test.go b/relay/helper/openai_image_request_test.go index 94c90872..7068b895 100644 --- a/relay/helper/openai_image_request_test.go +++ b/relay/helper/openai_image_request_test.go @@ -2,6 +2,7 @@ package helper import ( "bytes" + "fmt" "io" "mime/multipart" "net/http" @@ -10,6 +11,7 @@ import ( "testing" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -69,3 +71,79 @@ func TestGetAndValidOpenAIImageRequestMultipartStream(t *testing.T) { require.Contains(t, err.Error(), "invalid stream value") }) } + +// TestGetAndValidOpenAIImageRequestNBounds guards the billing invariant that +// the image generation count can never reach quota calculation with a value +// large enough to overflow int64 into a negative charge. +func TestGetAndValidOpenAIImageRequestNBounds(t *testing.T) { + gin.SetMode(gin.TestMode) + + newJSONContext := func(t *testing.T, body string) *gin.Context { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + return c + } + + boundErr := fmt.Sprintf("n must be an integer between 1 and %d", dto.MaxImageN) + + tests := []struct { + name string + body string + wantErr string + wantN uint + }{ + { + name: "overflowed uint64 n is rejected", + body: `{"model":"gpt-image-1","prompt":"a cat","n":18446744073686646784}`, + wantErr: boundErr, + }, + { + name: "n above max is rejected", + body: fmt.Sprintf(`{"model":"gpt-image-1","prompt":"a cat","n":%d}`, dto.MaxImageN+1), + wantErr: boundErr, + }, + { + name: "n at max is accepted", + body: fmt.Sprintf(`{"model":"gpt-image-1","prompt":"a cat","n":%d}`, dto.MaxImageN), + wantN: dto.MaxImageN, + }, + { + name: "absent n defaults to 1", + body: `{"model":"gpt-image-1","prompt":"a cat"}`, + wantN: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c := newJSONContext(t, tt.body) + req, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesGenerations) + if tt.wantErr != "" { + require.Error(t, err) + require.Contains(t, err.Error(), tt.wantErr) + return + } + require.NoError(t, err) + require.NotNil(t, req.N) + require.Equal(t, tt.wantN, *req.N) + }) + } + + t.Run("negative multipart n is rejected", func(t *testing.T) { + var body bytes.Buffer + writer := multipart.NewWriter(&body) + require.NoError(t, writer.WriteField("model", "gpt-image-1")) + require.NoError(t, writer.WriteField("prompt", "edit this image")) + require.NoError(t, writer.WriteField("n", "-22904832")) + require.NoError(t, writer.Close()) + + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body) + c.Request.Header.Set("Content-Type", writer.FormDataContentType()) + + _, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesEdits) + require.Error(t, err) + require.Contains(t, err.Error(), boundErr) + }) +} diff --git a/relay/helper/valid_request.go b/relay/helper/valid_request.go index dd866bcb..480f2cf7 100644 --- a/relay/helper/valid_request.go +++ b/relay/helper/valid_request.go @@ -155,7 +155,13 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq c.Request.PostForm = formData imageRequest.Prompt = formData.Get("prompt") imageRequest.Model = formData.Get("model") - imageRequest.N = common.GetPointer(uint(common.String2Int(formData.Get("n")))) + if nValue := strings.TrimSpace(formData.Get("n")); nValue != "" { + n, err := strconv.Atoi(nValue) + if err != nil || n < 0 || n > dto.MaxImageN { + return nil, fmt.Errorf("n must be an integer between 1 and %d", dto.MaxImageN) + } + imageRequest.N = common.GetPointer(uint(n)) + } imageRequest.Quality = formData.Get("quality") imageRequest.Size = formData.Get("size") if streamValue := strings.TrimSpace(formData.Get("stream")); streamValue != "" { @@ -201,6 +207,10 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq return nil, errors.New("size an unexpected error occurred in the parameter, please use 'x' instead of the multiplication sign '×'") } + if imageRequest.N != nil && *imageRequest.N > dto.MaxImageN { + return nil, fmt.Errorf("n must be an integer between 1 and %d", dto.MaxImageN) + } + // Not "256x256", "512x512", or "1024x1024" if imageRequest.Model == "dall-e-2" || imageRequest.Model == "dall-e" { if imageRequest.Size != "" && imageRequest.Size != "256x256" && imageRequest.Size != "512x512" && imageRequest.Size != "1024x1024" { diff --git a/relay/relay_task.go b/relay/relay_task.go index 098e2382..2f4df214 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -193,13 +193,15 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe } } - // 6. 将 OtherRatios 应用到基础额度 + // 6. 将 OtherRatios 应用到基础额度（饱和转换，防止溢出成负数） if !common.StringsContains(constant.TaskPricePatches, modelName) { + quotaWithRatios := float64(info.PriceData.Quota) for _, ra := range info.PriceData.OtherRatios { if ra != 1.0 { - info.PriceData.Quota = int(float64(info.PriceData.Quota) * ra) + quotaWithRatios *= ra } } + info.PriceData.Quota = common.QuotaFromFloat(quotaWithRatios) } // 7. 预扣费（仅首次 — 重试时 info.Billing 已存在，跳过） @@ -261,21 +263,21 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe // 公式: baseQuota × ∏(ratio) — 其中 baseQuota 是不含 OtherRatios 的基础额度。 func recalcQuotaFromRatios(info *relaycommon.RelayInfo, ratios map[string]float64) int { // 从 PriceData 获取不含 OtherRatios 的基础价格 - baseQuota := info.PriceData.Quota + baseQuota := float64(info.PriceData.Quota) // 先除掉原有的 OtherRatios 恢复基础额度 for _, ra := range info.PriceData.OtherRatios { if ra != 1.0 && ra > 0 { - baseQuota = int(float64(baseQuota) / ra) + baseQuota /= ra } } // 应用新的 ratios - result := float64(baseQuota) + result := baseQuota for _, ra := range ratios { if ra != 1.0 { result *= ra } } - return int(result) + return common.QuotaFromFloat(result) } var fetchRespBuilders = map[int]func(c *gin.Context) (respBody []byte, taskResp *dto.TaskError){ diff --git a/service/task_billing.go b/service/task_billing.go index 785a8675..2c75b06c 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -297,8 +297,8 @@ func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTo } } - // 计算实际应扣费额度: totalTokens * modelRatio * groupRatio * otherMultiplier - actualQuota := int(float64(totalTokens) * modelRatio * finalGroupRatio * otherMultiplier) + // 计算实际应扣费额度: totalTokens * modelRatio * groupRatio * otherMultiplier（饱和转换，防止溢出成负数） + actualQuota := common.QuotaFromFloat(float64(totalTokens) * modelRatio * finalGroupRatio * otherMultiplier) reason := fmt.Sprintf("token重算：tokens=%d, modelRatio=%.2f, groupRatio=%.2f, otherMultiplier=%.4f", totalTokens, modelRatio, finalGroupRatio, otherMultiplier) RecalculateTaskQuota(ctx, task, actualQuota, reason) diff --git a/service/text_quota.go b/service/text_quota.go index 3f344dc3..241e3d9d 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -287,7 +287,7 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf if !ratio.IsZero() && quotaCalculateDecimal.LessThanOrEqual(decimal.Zero) { quotaCalculateDecimal = decimal.NewFromInt(1) } - summary.Quota = int(quotaCalculateDecimal.Round(0).IntPart()) + summary.Quota = decimalToQuota(quotaCalculateDecimal) } else { quotaCalculateDecimal := dModelPrice.Mul(dQuotaPerUnit).Mul(dGroupRatio) quotaCalculateDecimal = quotaCalculateDecimal.Add(summary.ToolCallSurchargeQuota) @@ -297,7 +297,7 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf quotaCalculateDecimal = quotaCalculateDecimal.Mul(decimal.NewFromFloat(otherRatio)) } } - summary.Quota = int(quotaCalculateDecimal.Round(0).IntPart()) + summary.Quota = decimalToQuota(quotaCalculateDecimal) } if summary.TotalTokens == 0 { @@ -309,6 +309,14 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf return summary } +// decimalToQuota converts a computed quota decimal to int with saturation +// (see common.QuotaFromFloat). Oversized multipliers (e.g. an absurd image +// generation count) must never wrap around and turn a charge into a credit. +func decimalToQuota(d decimal.Decimal) int { + f, _ := d.Round(0).Float64() + return common.QuotaFromFloat(f) +} + func usageSemanticFromUsage(relayInfo *relaycommon.RelayInfo, usage *dto.Usage) string { if usage != nil && usage.UsageSemantic != "" { return usage.UsageSemantic diff --git a/service/text_quota_test.go b/service/text_quota_test.go index 37ce1877..2988a2ce 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -1,6 +1,7 @@ package service import ( + "math" "net/http/httptest" "testing" "time" @@ -12,9 +13,22 @@ import ( "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" + "github.com/shopspring/decimal" "github.com/stretchr/testify/require" ) +// TestDecimalToQuotaSaturation guards the billing invariant that an oversized +// quota product (e.g. per-call price multiplied by a huge image n ratio) must +// saturate instead of wrapping into a negative charge (credit). +func TestDecimalToQuotaSaturation(t *testing.T) { + // 2000 quota per call * n=18446744073686646784 overflows int64. + overflowing := decimal.NewFromInt(2000).Mul(decimal.NewFromFloat(1.8446744073686647e19)) + require.Equal(t, math.MaxInt32, decimalToQuota(overflowing)) + + require.Equal(t, math.MinInt32, decimalToQuota(overflowing.Neg())) + require.Equal(t, 42, decimalToQuota(decimal.NewFromFloat(41.7))) +} + func TestCalculateTextQuotaSummaryUnifiedForClaudeSemantic(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() diff --git a/types/price_data.go b/types/price_data.go index 93bc6ae8..b4fad00a 100644 --- a/types/price_data.go +++ b/types/price_data.go @@ -1,6 +1,9 @@ package types -import "fmt" +import ( + "fmt" + "math" +) type GroupRatioInfo struct { GroupRatio float64 @@ -31,7 +34,9 @@ func (p *PriceData) AddOtherRatio(key string, ratio float64) { if p.OtherRatios == nil { p.OtherRatios = make(map[string]float64) } - if ratio <= 0 { + // NaN/Inf would poison every downstream quota multiplication + // (int(NaN * quota) wraps to a negative charge). + if !(ratio > 0) || math.IsInf(ratio, 1) { return } p.OtherRatios[key] = ratio