Commit c9943d37 by CaIon

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.
parent d0bd8aac
......@@ -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.
......
......@@ -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)
......
......@@ -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)
}
......@@ -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
}
......
......@@ -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
......
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")
})
}
......@@ -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 {
......
......@@ -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
......
......@@ -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)
}
}
}
......
......@@ -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 {
......
......@@ -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 {
......
......@@ -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,其余模型使用估算
......
......@@ -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,
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment