From c9943d37ad93477dd937fc4901cc3c4e0fd8aaab Mon Sep 17 00:00:00 2001 From: CaIon Date: Tue, 7 Jul 2026 01:08:52 +0800 Subject: [PATCH] fix(billing): extend quantity validation and saturating conversions to remaining paths Bound max-tokens fields across all relay format validators, saturate tiered-expression rounding and audio/tool/task token conversions, and route legacy remix ratios through the guarded setter. --- AGENTS.md | 5 +- pkg/billingexpr/billingexpr_test.go | 6 +++ pkg/billingexpr/round.go | 16 +++++- relay/channel/openai/audio.go | 5 +- relay/channel/task/kling/adaptor.go | 3 +- relay/helper/max_tokens_bounds_test.go | 72 ++++++++++++++++++++++++++ relay/helper/price.go | 8 +-- relay/helper/valid_request.go | 25 ++++++++- relay/relay_task.go | 13 ++--- service/quota.go | 4 +- service/text_quota.go | 8 ++- service/token_counter.go | 16 ++++-- service/tool_billing.go | 4 +- 13 files changed, 155 insertions(+), 30 deletions(-) create mode 100644 relay/helper/max_tokens_bounds_test.go diff --git a/AGENTS.md b/AGENTS.md index 0a1534b5..aea1a633 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -100,9 +100,10 @@ Do NOT directly import or call `encoding/json` in business code. `json.RawMessag **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. +- 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. -- 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. +- 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())`. Use the saturating converters: `common.QuotaFromFloat` for float products, `decimalToQuota` in `service` for decimal products, `billingexpr.QuotaRound` (saturating) for tiered-expression results. 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. diff --git a/pkg/billingexpr/billingexpr_test.go b/pkg/billingexpr/billingexpr_test.go index aa979e37..ee9135bc 100644 --- a/pkg/billingexpr/billingexpr_test.go +++ b/pkg/billingexpr/billingexpr_test.go @@ -291,6 +291,12 @@ func TestQuotaRound(t *testing.T) { {999.4999, 999}, {999.5, 1000}, {1e9 + 0.5, 1e9 + 1}, + // Oversized expression results must saturate at int32 bounds + // instead of wrapping into a negative charge. + {3.6893488147419103e19, math.MaxInt32}, + {-3.6893488147419103e19, math.MinInt32}, + {math.Inf(1), math.MaxInt32}, + {math.NaN(), 0}, } for _, tt := range tests { got := billingexpr.QuotaRound(tt.in) diff --git a/pkg/billingexpr/round.go b/pkg/billingexpr/round.go index 35a5534a..be6c8f1a 100644 --- a/pkg/billingexpr/round.go +++ b/pkg/billingexpr/round.go @@ -5,6 +5,20 @@ import "math" // QuotaRound converts a float64 quota value to int using half-away-from-zero // rounding. Every tiered billing path (pre-consume, settlement, breakdown // validation, log fields) MUST use this function to avoid +-1 discrepancies. +// +// The result saturates at int32 bounds: quota columns are 32-bit integers in +// the database, and an oversized expression result must never wrap around +// and turn a charge into a credit. func QuotaRound(f float64) int { - return int(math.Round(f)) + r := math.Round(f) + if math.IsNaN(r) { + return 0 + } + if r >= math.MaxInt32 { + return math.MaxInt32 + } + if r <= math.MinInt32 { + return math.MinInt32 + } + return int(r) } diff --git a/relay/channel/openai/audio.go b/relay/channel/openai/audio.go index 6a87d89f..ec1cbe72 100644 --- a/relay/channel/openai/audio.go +++ b/relay/channel/openai/audio.go @@ -103,8 +103,9 @@ func OpenaiTTSHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel usage.CompletionTokens = estimatedTokens usage.CompletionTokenDetails.AudioTokens = estimatedTokens } else if duration > 0 { - // 计算 token: ceil(duration) / 60.0 * 1000，即每分钟 1000 tokens - completionTokens := int(math.Round(math.Ceil(duration) / 60.0 * 1000)) + // 计算 token: ceil(duration) / 60.0 * 1000，即每分钟 1000 tokens。 + // duration 解析自上游返回的音频元数据，饱和转换防止 int 回绕。 + completionTokens := common.QuotaFromFloat(math.Round(math.Ceil(duration) / 60.0 * 1000)) usage.CompletionTokens = completionTokens usage.CompletionTokenDetails.AudioTokens = completionTokens } diff --git a/relay/channel/task/kling/adaptor.go b/relay/channel/task/kling/adaptor.go index 413ade04..c7a492b9 100644 --- a/relay/channel/task/kling/adaptor.go +++ b/relay/channel/task/kling/adaptor.go @@ -358,7 +358,8 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e taskInfo.Url = video.Url } if tokens, err := strconv.ParseFloat(resPayload.Data.FinalUnitDeduction, 64); err == nil { - rounded := int(math.Ceil(tokens)) + // 上游返回的扣费数值，饱和转换防止超大数值回绕成负数 + rounded := common.QuotaFromFloat(math.Ceil(tokens)) if rounded > 0 { taskInfo.CompletionTokens = rounded taskInfo.TotalTokens = rounded diff --git a/relay/helper/max_tokens_bounds_test.go b/relay/helper/max_tokens_bounds_test.go new file mode 100644 index 00000000..7cdcba87 --- /dev/null +++ b/relay/helper/max_tokens_bounds_test.go @@ -0,0 +1,72 @@ +package helper + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// TestMaxTokensBounds guards the billing invariant that user-supplied max +// token fields are bounded on every relay format. These values feed +// pre-consume quota math (preConsumedTokens * ratio); a huge or +// wrapped-negative value (e.g. 18446744073686646784 parsed into *uint) must +// be rejected at validation instead of corrupting the pre-charge. +func TestMaxTokensBounds(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, "/relay", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + return c + } + + const hugeN = "18446744073686646784" + + t.Run("openai max_tokens overflow rejected", func(t *testing.T) { + c := newJSONContext(t, `{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"max_tokens":`+hugeN+`}`) + _, err := GetAndValidateTextRequest(c, relayconstant.RelayModeChatCompletions) + require.Error(t, err) + require.Contains(t, err.Error(), "max_tokens is invalid") + }) + + t.Run("openai max_completion_tokens overflow rejected", func(t *testing.T) { + c := newJSONContext(t, `{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"max_completion_tokens":`+hugeN+`}`) + _, err := GetAndValidateTextRequest(c, relayconstant.RelayModeChatCompletions) + require.Error(t, err) + require.Contains(t, err.Error(), "max_tokens is invalid") + }) + + t.Run("claude max_tokens overflow rejected", func(t *testing.T) { + c := newJSONContext(t, `{"model":"claude-sonnet-4","messages":[{"role":"user","content":"hi"}],"max_tokens":`+hugeN+`}`) + _, err := GetAndValidateClaudeRequest(c) + require.Error(t, err) + require.Contains(t, err.Error(), "max_tokens is invalid") + }) + + t.Run("claude normal max_tokens accepted", func(t *testing.T) { + c := newJSONContext(t, `{"model":"claude-sonnet-4","messages":[{"role":"user","content":"hi"}],"max_tokens":8192}`) + req, err := GetAndValidateClaudeRequest(c) + require.NoError(t, err) + require.EqualValues(t, 8192, *req.MaxTokens) + }) + + t.Run("gemini maxOutputTokens overflow rejected", func(t *testing.T) { + c := newJSONContext(t, `{"contents":[{"parts":[{"text":"hi"}]}],"generationConfig":{"maxOutputTokens":`+hugeN+`}}`) + _, err := GetAndValidateGeminiRequest(c) + require.Error(t, err) + require.Contains(t, err.Error(), "maxOutputTokens is invalid") + }) + + t.Run("responses max_output_tokens overflow rejected", func(t *testing.T) { + c := newJSONContext(t, `{"model":"gpt-4o","input":"hi","max_output_tokens":`+hugeN+`}`) + _, err := GetAndValidateResponsesRequest(c) + require.Error(t, err) + require.Contains(t, err.Error(), "max_output_tokens is invalid") + }) +} diff --git a/relay/helper/price.go b/relay/helper/price.go index d1e16bca..313424fb 100644 --- a/relay/helper/price.go +++ b/relay/helper/price.go @@ -112,12 +112,12 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens audioRatio = ratio_setting.GetAudioRatio(info.OriginModelName) audioCompletionRatio = ratio_setting.GetAudioCompletionRatio(info.OriginModelName) ratio := modelRatio * groupRatioInfo.GroupRatio - preConsumedQuota = int(float64(preConsumedTokens) * ratio) + preConsumedQuota = common.QuotaFromFloat(float64(preConsumedTokens) * ratio) } else { if meta.ImagePriceRatio != 0 { modelPrice = modelPrice * meta.ImagePriceRatio } - preConsumedQuota = int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + preConsumedQuota = common.QuotaFromFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) } // check if free model pre-consume is disabled @@ -194,7 +194,7 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types freeModel := false if usePrice { - quota = int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + quota = common.QuotaFromFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume { if groupRatioInfo.GroupRatio == 0 || modelPrice == 0 { quota = 0 @@ -203,7 +203,7 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types } } else { // 按量计费：以模型倍率的一半作为预扣额度 - quota = int(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + quota = common.QuotaFromFloat(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio) modelPrice = -1 if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume { if groupRatioInfo.GroupRatio == 0 || modelRatio == 0 { diff --git a/relay/helper/valid_request.go b/relay/helper/valid_request.go index 480f2cf7..afbd59ff 100644 --- a/relay/helper/valid_request.go +++ b/relay/helper/valid_request.go @@ -114,6 +114,20 @@ func GetAndValidateEmbeddingRequest(c *gin.Context, relayMode int) (*dto.Embeddi return embeddingRequest, nil } +// maxTokensLimit bounds user-supplied max token fields. These values feed +// pre-consume quota math (preConsumedTokens * ratio); an unbounded value can +// overflow the conversion and corrupt billing. +const maxTokensLimit = math.MaxInt32 / 2 + +func exceedsMaxTokensLimit(values ...*uint) bool { + for _, v := range values { + if lo.FromPtrOr(v, uint(0)) > maxTokensLimit { + return true + } + } + return false +} + func GetAndValidateResponsesRequest(c *gin.Context) (*dto.OpenAIResponsesRequest, error) { request := &dto.OpenAIResponsesRequest{} err := common.UnmarshalBodyReusable(c, request) @@ -126,6 +140,9 @@ func GetAndValidateResponsesRequest(c *gin.Context) (*dto.OpenAIResponsesRequest if request.Input == nil { return nil, errors.New("input is required") } + if exceedsMaxTokensLimit(request.MaxOutputTokens) { + return nil, errors.New("max_output_tokens is invalid") + } return request, nil } @@ -259,6 +276,9 @@ func GetAndValidateClaudeRequest(c *gin.Context) (textRequest *dto.ClaudeRequest if textRequest.Model == "" { return nil, errors.New("field model is required") } + if exceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxTokensToSample) { + return nil, errors.New("max_tokens is invalid") + } //if textRequest.Stream { // relayInfo.IsStream = true @@ -281,7 +301,7 @@ func GetAndValidateTextRequest(c *gin.Context, relayMode int) (*dto.GeneralOpenA textRequest.Model = c.Param("model") } - if lo.FromPtrOr(textRequest.MaxTokens, uint(0)) > math.MaxInt32/2 { + if exceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxCompletionTokens) { return nil, errors.New("max_tokens is invalid") } if textRequest.Model == "" { @@ -334,6 +354,9 @@ func GetAndValidateGeminiRequest(c *gin.Context) (*dto.GeminiChatRequest, error) if len(request.Contents) == 0 && len(request.Requests) == 0 { return nil, errors.New("contents is required") } + if exceedsMaxTokensLimit(request.GenerationConfig.MaxOutputTokens) { + return nil, errors.New("maxOutputTokens is invalid") + } //if c.Query("alt") == "sse" { // relayInfo.IsStream = true diff --git a/relay/relay_task.go b/relay/relay_task.go index 2f4df214..037e4076 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -121,14 +121,15 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr if seconds <= 0 { seconds = 4 } - sizeStr, _ := taskData["size"].(string) - if info.PriceData.OtherRatios == nil { - info.PriceData.OtherRatios = map[string]float64{} + // 历史任务数据可能包含未经校验的时长，作为计费乘数前必须钳制 + if seconds > relaycommon.MaxTaskDurationSeconds { + seconds = relaycommon.MaxTaskDurationSeconds } - info.PriceData.OtherRatios["seconds"] = float64(seconds) - info.PriceData.OtherRatios["size"] = 1 + sizeStr, _ := taskData["size"].(string) + info.PriceData.AddOtherRatio("seconds", float64(seconds)) + info.PriceData.AddOtherRatio("size", 1) if sizeStr == "1792x1024" || sizeStr == "1024x1792" { - info.PriceData.OtherRatios["size"] = 1.666667 + info.PriceData.AddOtherRatio("size", 1.666667) } } } diff --git a/service/quota.go b/service/quota.go index 862805d7..d3aab3dc 100644 --- a/service/quota.go +++ b/service/quota.go @@ -54,7 +54,7 @@ func calculateAudioQuota(info QuotaInfo) int { groupRatio := decimal.NewFromFloat(info.GroupRatio) quota := modelPrice.Mul(quotaPerUnit).Mul(groupRatio) - return int(quota.IntPart()) + return decimalToQuota(quota) } completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(info.ModelName)) @@ -83,7 +83,7 @@ func calculateAudioQuota(info QuotaInfo) int { quota = decimal.NewFromInt(1) } - return int(quota.Round(0).IntPart()) + return decimalToQuota(quota) } func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.RealtimeUsage) error { diff --git a/service/text_quota.go b/service/text_quota.go index 241e3d9d..e30e0b3d 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -145,15 +145,13 @@ func composeTieredTextQuota(relayInfo *relaycommon.RelayInfo, summary textQuotaS if tieredResult != nil { if snap := relayInfo.TieredBillingSnapshot; snap != nil { - return int(decimal.NewFromFloat(tieredResult.ActualQuotaBeforeGroup). + return decimalToQuota(decimal.NewFromFloat(tieredResult.ActualQuotaBeforeGroup). Mul(decimal.NewFromFloat(snap.GroupRatio)). - Add(summary.ToolCallSurchargeQuota). - Round(0). - IntPart()) + Add(summary.ToolCallSurchargeQuota)) } } - return tieredQuota + int(summary.ToolCallSurchargeQuota.Round(0).IntPart()) + return tieredQuota + decimalToQuota(summary.ToolCallSurchargeQuota) } func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage) textQuotaSummary { diff --git a/service/token_counter.go b/service/token_counter.go index 933d9fd1..e045c407 100644 --- a/service/token_counter.go +++ b/service/token_counter.go @@ -208,8 +208,14 @@ func EstimateRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *rela if err != nil { return 0, fmt.Errorf("error getting audio duration: %v", err) } - // 一分钟 1000 token，与 $price / minute 对齐 - totalAudioToken += int(math.Round(math.Ceil(duration) / 60.0 * 1000)) + // 一分钟 1000 token，与 $price / minute 对齐。 + // duration 来自用户上传文件的元数据，可被伪造成天文数字， + // 必须饱和转换防止 int 回绕成负数 token。 + audioTokens := common.QuotaFromFloat(math.Round(math.Ceil(duration) / 60.0 * 1000)) + if audioTokens < 0 { + audioTokens = 0 + } + totalAudioToken += audioTokens } return totalAudioToken, nil } @@ -377,7 +383,8 @@ func CountAudioTokenInput(audioBase64 string, audioFormat string) (int, error) { if err != nil { return 0, err } - return int(duration / 60 * 100 / 0.06), nil + // duration 来自用户提供的音频元数据，饱和转换防止 int 回绕 + return common.QuotaFromFloat(duration / 60 * 100 / 0.06), nil } func CountAudioTokenOutput(audioBase64 string, audioFormat string) (int, error) { @@ -388,7 +395,8 @@ func CountAudioTokenOutput(audioBase64 string, audioFormat string) (int, error) if err != nil { return 0, err } - return int(duration / 60 * 200 / 0.24), nil + // duration 来自上游返回的音频元数据，饱和转换防止 int 回绕 + return common.QuotaFromFloat(duration / 60 * 200 / 0.24), nil } // CountTextToken 统计文本的token数量，仅OpenAI模型使用tokenizer，其余模型使用估算 diff --git a/service/tool_billing.go b/service/tool_billing.go index fd28fddb..15993067 100644 --- a/service/tool_billing.go +++ b/service/tool_billing.go @@ -49,7 +49,7 @@ func ComputeToolCallQuota(usage ToolCallUsage, groupRatio float64) ToolCallResul return } totalPrice := pricePer1K * float64(count) / 1000 - quota := int(math.Round(totalPrice * common.QuotaPerUnit * groupRatio)) + quota := common.QuotaFromFloat(math.Round(totalPrice * common.QuotaPerUnit * groupRatio)) items = append(items, ToolCallItem{ Name: toolName, CallCount: count, @@ -70,7 +70,7 @@ func ComputeToolCallQuota(usage ToolCallUsage, groupRatio float64) ToolCallResul if usage.ImageGenerationCall { price := operation_setting.GetGPTImage1PriceOnceCall(usage.ImageGenerationQuality, usage.ImageGenerationSize) - quota := int(math.Round(price * common.QuotaPerUnit * groupRatio)) + quota := common.QuotaFromFloat(math.Round(price * common.QuotaPerUnit * groupRatio)) items = append(items, ToolCallItem{ Name: "image_generation", CallCount: 1,