fix(ratelimit): enforce negative token validation before quota lookup
Some checks failed
Security Scan / security (push) Has been cancelled
Test / test (push) Has been cancelled

This commit is contained in:
Virgil 2026-04-03 07:37:23 +00:00
parent 2ad4870bd0
commit 781e0ee3d6

View file

@ -249,7 +249,12 @@ func (rl *RateLimiter) Load() error {
return err return err
} }
return yaml.Unmarshal([]byte(content), rl) if err := yaml.Unmarshal([]byte(content), rl); err != nil {
return err
}
ensureMaps(rl)
return nil
} }
// loadSQLite reads quotas and state from the SQLite backend. // loadSQLite reads quotas and state from the SQLite backend.
@ -596,6 +601,14 @@ func (rl *RateLimiter) Decide(model string, estimatedTokens int) Decision {
now := time.Now() now := time.Now()
decision := Decision{} decision := Decision{}
if estimatedTokens < 0 {
decision.Allowed = false
decision.Code = DecisionInvalidTokens
decision.Reason = "estimated tokens must be non-negative"
decision.Stats = rl.snapshotLocked(model)
return decision
}
quota, ok := rl.Quotas[model] quota, ok := rl.Quotas[model]
if !ok { if !ok {
decision.Allowed = true decision.Allowed = true
@ -606,13 +619,6 @@ func (rl *RateLimiter) Decide(model string, estimatedTokens int) Decision {
} }
if quota.MaxRPM == 0 && quota.MaxTPM == 0 && quota.MaxRPD == 0 { if quota.MaxRPM == 0 && quota.MaxTPM == 0 && quota.MaxRPD == 0 {
if estimatedTokens < 0 {
decision.Allowed = false
decision.Code = DecisionInvalidTokens
decision.Reason = "estimated tokens must be non-negative"
decision.Stats = rl.snapshotLocked(model)
return decision
}
decision.Allowed = true decision.Allowed = true
decision.Code = DecisionUnlimited decision.Code = DecisionUnlimited
decision.Reason = "all limits are unlimited" decision.Reason = "all limits are unlimited"
@ -627,14 +633,6 @@ func (rl *RateLimiter) Decide(model string, estimatedTokens int) Decision {
rl.State[model] = stats rl.State[model] = stats
} }
if estimatedTokens < 0 {
decision.Allowed = false
decision.Code = DecisionInvalidTokens
decision.Reason = "estimated tokens must be non-negative"
decision.Stats = rl.snapshotLocked(model)
return decision
}
decision.Stats = rl.snapshotLocked(model) decision.Stats = rl.snapshotLocked(model)
if quota.MaxRPD > 0 && stats.DayCount >= quota.MaxRPD { if quota.MaxRPD > 0 && stats.DayCount >= quota.MaxRPD {
@ -836,6 +834,15 @@ func newConfiguredRateLimiter(cfg Config) *RateLimiter {
return rl return rl
} }
func ensureMaps(rl *RateLimiter) {
if rl.Quotas == nil {
rl.Quotas = make(map[string]ModelQuota)
}
if rl.State == nil {
rl.State = make(map[string]*UsageStats)
}
}
func applyConfig(rl *RateLimiter, cfg Config) { func applyConfig(rl *RateLimiter, cfg Config) {
profiles := DefaultProfiles() profiles := DefaultProfiles()
providers := cfg.Providers providers := cfg.Providers