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.
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
@@ -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
|
||||
|
||||
+7
-6
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+2
-2
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user