fix(accounting): bound upstream-controlled usage and counter keys

Passthrough accounting read the model id, the usage figures and the
response size from the upstream response without limits.

- Mint a model counter scope only for ids the provider serves. An
  upstream that echoed new ids created unbounded counter keys.
- Drop a reported token or cost figure above a fixed ceiling. One huge
  figure locked a user out of every budget for the whole window.
- Bill an echoed unpriced model at the price of the model that passed
  the cost-budget gate.
- Cap a body-size prompt estimate at one million tokens.
- Bound the total bytes that concurrent passthrough accounting tees
  hold. A tee that cannot reserve room falls back to the estimate.
- Add cache.accountHits. When true, a cache hit charges the cached usage
  to the user. The default stays false.

limits.go

@@ -1824,6 +1824,23 @@ func (l *limiter) accountWith(scopes []limitScope, u usage, costMicros int64, x
 	// misbehaving or hostile upstream is visible rather than silently
 	// rewriting a budget. Clamped per-component, so one bad field does not
 	// discard the other two, which may be perfectly valid.
+	// The positive direction needs a ceiling too: one absurd figure from an
+	// upstream (or a saturated cost) would otherwise put a counter past
+	// every budget until its window, up to 400 days on Redis, expires. No
+	// real request reports figures anywhere near these bounds, so an
+	// out-of-range figure is dropped like a negative one.
+	if u.prompt > maxUsageTokensPerRequest || u.completion > maxUsageTokensPerRequest || costMicros > maxCostMicrosPerRequest {
+		l.logf("%s", fmt.Sprintf("limits: dropping implausible usage reported by an upstream (prompt=%d completion=%d costMicros=%d)", u.prompt, u.completion, costMicros))
+		if u.prompt > maxUsageTokensPerRequest {
+			u.prompt = 0
+		}
+		if u.completion > maxUsageTokensPerRequest {
+			u.completion = 0
+		}
+		if costMicros > maxCostMicrosPerRequest {
+			costMicros = 0
+		}
+	}
 	if u.prompt < 0 || u.completion < 0 || costMicros < 0 {
 		l.logf("%s", fmt.Sprintf("limits: dropping negative usage reported by an upstream (prompt=%d completion=%d costMicros=%d); these counters are increment-only", u.prompt, u.completion, costMicros))
 		if u.prompt < 0 {
@@ -2250,6 +2267,15 @@ const (
 // scope can never collide with a user, group, provider or target scope.
 const kindModel = "model"
 
+// maxUsageTokensPerRequest and maxCostMicrosPerRequest bound what one
+// upstream response may add to a token or cost counter. 2^30 tokens is
+// far above any model context window, and 2^40 micro-USD is about
+// 1.1 million USD.
+const (
+	maxUsageTokensPerRequest = 1 << 30
+	maxCostMicrosPerRequest  = 1 << 40
+)
+
 // maxModelScopeIDLen bounds the "provider/model" id withModelScope will
 // accept as a counter-scope id. Generous next to any real model id —
 // the longest in the built-in table is well under 40 bytes — and far

llmgateway.go

@@ -447,6 +447,11 @@ type CacheConfig struct {
 	// value above maxCacheMaxBodyBytes is a construction error.
 	MaxBodyBytes int  `json:"maxBodyBytes,omitempty"`
 	Enabled      bool `json:"enabled,omitempty"`
+	// AccountHits charges the usage a cached response originally reported
+	// to the user and group that receive the hit, so token and cost
+	// budgets also bind repeated identical prompts. Defaults to false: a
+	// hit then adds only to the request counters.
+	AccountHits bool `json:"accountHits,omitempty"`
 }
 
 // BreakerConfig configures the per-provider discovery circuit breaker
@@ -761,6 +766,7 @@ type Gateway struct {
 	next                          http.Handler
 	events                        *eventLog
 	bodyAdmission                 chan struct{}
+	teeBudget                     *teeBudget
 	failoverHealth                *requestHealthTracker
 	targetHealth                  *targetHealthTracker
 	limiter                       *limiter
@@ -1049,6 +1055,7 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str
 		}
 	}
 	g.bodyAdmission = make(chan struct{}, bodyAdmissionCap)
+	g.teeBudget = newTeeBudget(maxAccountingTeeTotalBytes)
 
 	if config.Users != nil && config.Users.File != "" {
 		if err = attachUsersFile(auth, config.Users.File, g); err != nil {

registry.go

@@ -1555,6 +1555,13 @@ func (m *modelRegistry) bareWinner(id string, grp *group, extraIDs ...string) (s
 	return visible[0], true
 }
 
+// providerHasModel reports whether provider's explicit or discovered model
+// set contains id.
+func (m *modelRegistry) providerHasModel(provider, id string) bool {
+	st, ok := m.states[provider]
+	return ok && st.hasModel(id)
+}
+
 // resolve maps a client-requested model id to the adapter that serves it.
 // It returns the adapter, the upstream model id to send that adapter (the
 // id form the provider itself knows, with any gateway-side provider prefix

routes_passthrough.go

@@ -7,7 +7,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"math"
 	"net/http"
 	"net/url"
 	"strings"
@@ -29,6 +28,11 @@ const maxPassthroughBytes = 32 << 20
 // of any size — without ever waiting for the full body to buffer first.
 const maxAccountingTeeBytes = 4 << 20
 
+// maxAccountingTeeTotalBytes caps the bytes all concurrent accounting tees
+// may retain together. A response that cannot reserve room is accounted
+// from the truncated-body estimate instead of its parsed usage.
+const maxAccountingTeeTotalBytes = 256 << 20
+
 // sseAccountingHeadBytes/sseAccountingTailBytes cap sseAccountingBuffer
 // (verify-core fix, round 4): an SSE stream used to be teed through the
 // same 4MiB-capped cappedAccountingBuffer as a JSON body, which held
@@ -420,9 +424,43 @@ func hasTraversalSegment(rest string) bool {
 // caller knows a parse would only fail on cut-off JSON and skips it.
 type cappedAccountingBuffer struct {
 	buf       bytes.Buffer
+	budget    *teeBudget
+	reserved  int64
 	truncated bool
 }
 
+// teeBudget bounds the bytes every in-flight accounting tee may hold at
+// once, so N concurrent large responses cannot add up to N x
+// maxAccountingTeeBytes. Counters use the sync/atomic function forms, the
+// style this package already uses under Yaegi.
+type teeBudget struct {
+	used  int64
+	limit int64
+}
+
+func newTeeBudget(limit int64) *teeBudget { return &teeBudget{limit: limit} }
+
+// reserve claims n bytes and reports whether the budget had room.
+func (b *teeBudget) reserve(n int64) bool {
+	if atomic.AddInt64(&b.used, n) > b.limit {
+		atomic.AddInt64(&b.used, -n)
+		return false
+	}
+	return true
+}
+
+func (b *teeBudget) free(n int64) { atomic.AddInt64(&b.used, -n) }
+
+// release returns the tee's reserved bytes to the budget. It is safe on a
+// nil receiver and idempotent.
+func (c *cappedAccountingBuffer) release() {
+	if c == nil || c.budget == nil || c.reserved == 0 {
+		return
+	}
+	c.budget.free(c.reserved)
+	c.reserved = 0
+}
+
 // Write implements io.Writer. It always reports success for the full
 // input length — even once truncated is set and further bytes are being
 // discarded — because this type is driven by io.TeeReader, whose Read
@@ -444,6 +482,13 @@ func (c *cappedAccountingBuffer) Write(p []byte) (int, error) {
 		p = p[:remaining]
 		c.truncated = true
 	}
+	if c.budget != nil {
+		if !c.budget.reserve(int64(len(p))) {
+			c.truncated = true
+			return n, nil
+		}
+		c.reserved += int64(len(p))
+	}
 	c.buf.Write(p) //nolint:errcheck // bytes.Buffer.Write never errors
 	return n, nil
 }
@@ -1224,9 +1269,11 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
 	// was already restored via io.MultiReader after that first call) so
 	// this check runs independently of whether the group restricts
 	// models at all.
+	var gatedModel string
 	if !g.cfg.AllowUnpricedWithCostBudget && scopesHaveCostBudget(scopes) {
 		if model, hasModel, _, peekErr := peekPassthroughModel(r); peekErr == nil && hasModel {
 			canonical := providerName + "/" + model
+			gatedModel = model
 			if !priceKnown(canonical, model, g.cfg.Pricing, g.cfg.ModelMeta) {
 				g.logf("passthrough: refusing model %q: it has no configured price, and a cost budget applies to this caller that cannot be enforced without one", model)
 				g.recordUnpricedRefusalEvent(scopes, canonical, providerName, routePassthrough) // F3 hook 5 (item 5, this round)
@@ -1271,6 +1318,7 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
 	}
 
 	result, complete := g.proxyUpstream(w, r, upstreamURL, adapter.httpClient(), adapter.injectAuth, providerCredentialRetargetHeaders, true, "passthrough (provider "+providerName+")", adapter.requestTimeout())
+	defer result.tee.release()
 	g.recordProxyEvent(scopes, providerName, routePassthrough, result) // F3 hook 3 (v0.3 dashboard task)
 
 	// Finding 2 fix (review-routes.md): a build/dial failure that never
@@ -1335,7 +1383,7 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
 		if !canEstimate || result.requestBytes <= 0 {
 			return
 		}
-		respUsage = usage{prompt: int64(math.Ceil(float64(result.requestBytes) / 4)), estimated: true}
+		respUsage = usage{prompt: estimatePromptTokens(int64(result.requestBytes)), estimated: true}
 		if respModel == "" {
 			respModel = "unknown/" + providerName
 		}
@@ -1343,7 +1391,19 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
 	}
 
 	canonical := providerName + "/" + respModel
-	cost := unifiedCostMicros(canonical, respModel, respUsage, g.cfg.Pricing, g.cfg.ModelMeta, g.pricingWarn)
+	billModel := respModel
+	// Gate and bill on the same price: an upstream that echoes a model id
+	// the gate never saw must not turn a priced request into a free one.
+	if gatedModel != "" && respModel != gatedModel && !priceKnown(canonical, respModel, g.cfg.Pricing, g.cfg.ModelMeta) {
+		billModel = gatedModel
+	}
+	cost := unifiedCostMicros(providerName+"/"+billModel, billModel, respUsage, g.cfg.Pricing, g.cfg.ModelMeta, g.pricingWarn)
+	// Only ids the provider is known to serve get a model scope: the id is
+	// read from the upstream response body, so any other value would mint
+	// counter keys without bound.
+	if !g.registry.providerHasModel(providerName, respModel) {
+		canonical = ""
+	}
 	// withModelScope drops the model scope when respModel is empty (the
 	// upstream reported no model id), so a passthrough reply the gateway
 	// cannot attribute still accounts to user/group/total as before.
@@ -1714,7 +1774,7 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream
 	var reader io.Reader = resp.Body
 	switch {
 	case isJSON:
-		tee = &cappedAccountingBuffer{}
+		tee = &cappedAccountingBuffer{budget: g.teeBudget}
 		reader = io.TeeReader(resp.Body, tee)
 	case isSSE:
 		sseTee = &sseAccountingBuffer{}

routes_passthrough_test.go

@@ -3355,7 +3355,7 @@ func TestHandlePassthrough_ModelLatencyAndUModelSuppressed_ProviderHistogramsOnl
 
 	cfg := CreateConfig()
 	cfg.Providers = map[string]*ProviderConfig{
-		"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"},
+		"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-native-x"}},
 	}
 	cfg.Groups = map[string]*GroupConfig{"default": {}}
 	cfg.Users = &UsersConfig{Inline: []*UserConfig{

routes_unified.go

@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"math"
 	"net/http"
 	"strconv"
 	"strings"
@@ -501,6 +500,10 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop
 				// requires for a cache-enabled, admin-ENABLED-but-stats-off
 				// case. Gating here, not just inside recordCacheHit, is what
 				// actually saves the parse.
+				if g.cfg.Cache.AccountHits {
+					hitUsage := cachedUsage(cached.Body)
+					g.limiter.accountWith(scopes, hitUsage, unifiedCostMicros(cand.canonical, cand.upstreamModel, hitUsage, g.cfg.Pricing, g.cfg.ModelMeta, g.pricingWarn), accountExtras{provider: cand.providerName})
+				}
 				if g.limiter.statsAdmin {
 					saved := unifiedCostMicros(cand.canonical, cand.upstreamModel, cachedUsage(cached.Body), g.cfg.Pricing, g.cfg.ModelMeta, func(string) {})
 					g.limiter.recordCacheHit(cand.canonical, saved)
@@ -691,13 +694,13 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop
 				// never streamed at all, or a failover pair is double-
 				// billed: the failed candidate's phantom estimate plus
 				// the next candidate's real usage.
-				result.prompt = int64(math.Ceil(float64(len(body)) / 4))
+				result.prompt = estimatePromptTokens(int64(len(body)))
 				result.estimated = true
 				g.logf("%s: streaming response from model %q failed or was canceled after starting to write to the client with zero usage reported; accounting an estimate derived from request body size", logPrefix, cand.canonical)
 			case streaming:
 				g.logf("%s: zero usage reported for a streaming response from model %q; accounting the request only", logPrefix, cand.canonical)
 			case callErr == nil:
-				result.prompt = int64(math.Ceil(float64(len(body)) / 4))
+				result.prompt = estimatePromptTokens(int64(len(body)))
 				result.estimated = true
 			}
 		}
@@ -1435,3 +1438,18 @@ func (w *cacheCaptureWriter) Write(b []byte) (int, error) {
 	}
 	return w.statusTrackingWriter.Write(b)
 }
+
+// maxEstimatedPromptTokens caps a body-size estimate, so a padded request
+// cannot bill, or overshoot a budget by, more than a very large context
+// window would legitimately hold.
+const maxEstimatedPromptTokens = 1 << 20
+
+// estimatePromptTokens estimates prompt tokens as one per four body bytes,
+// rounded up and capped at maxEstimatedPromptTokens.
+func estimatePromptTokens(bodyBytes int64) int64 {
+	est := (bodyBytes + 3) / 4
+	if est > maxEstimatedPromptTokens {
+		return maxEstimatedPromptTokens
+	}
+	return est
+}

security_fixes_test.go

@@ -0,0 +1,180 @@
+package traefikllmgateway
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// newPassthroughUsageGateway builds a Gateway with one openai provider
+// pointed at an upstream that always answers upstreamBody, and one user
+// with the given limits and pricing.
+func newPassthroughUsageGateway(t *testing.T, upstreamBody string, models []string, limits *LimitsConfig, pricing map[string]*ModelPricing) *Gateway {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(upstreamBody))
+	}))
+	t.Cleanup(srv.Close)
+
+	cfg := CreateConfig()
+	cfg.Providers = map[string]*ProviderConfig{
+		"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: models},
+	}
+	cfg.Pricing = pricing
+	cfg.Groups = map[string]*GroupConfig{"default": {}}
+	cfg.Users = &UsersConfig{Inline: []*UserConfig{
+		{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: limits},
+	}}
+	h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
+	require.NoError(t, err)
+	gw, ok := h.(*Gateway)
+	require.True(t, ok)
+	gw.limiter.spawn = func(f func()) { f() }
+	return gw
+}
+
+func doPassthrough(t *testing.T, gw *Gateway) *httptest.ResponseRecorder {
+	t.Helper()
+	req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(`{"model":"gpt-1"}`))
+	req.Header.Set("Authorization", "Bearer sk-alice")
+	rec := httptest.NewRecorder()
+	gw.ServeHTTP(rec, req)
+	return rec
+}
+
+func TestPassthrough_EchoedUnknownModel_MintsNoModelScope(t *testing.T) {
+	t.Parallel()
+	gw := newPassthroughUsageGateway(t, `{"model":"echo-unknown","usage":{"prompt_tokens":4,"completion_tokens":2}}`, []string{"gpt-known"}, &LimitsConfig{}, nil)
+
+	require.Equal(t, http.StatusOK, doPassthrough(t, gw).Code)
+
+	reqs, _ := gw.limiter.getCounter(kindModel, "openai/echo-unknown", metricReq, windowDay, time.Now())
+	assert.Zero(t, reqs, "an id the provider does not serve must not get a model scope")
+	tokIn, _ := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
+	assert.EqualValues(t, 4, tokIn, "user accounting must be unchanged")
+}
+
+func TestPassthrough_EchoedKnownModel_KeepsModelScope(t *testing.T) {
+	t.Parallel()
+	gw := newPassthroughUsageGateway(t, `{"model":"gpt-known","usage":{"prompt_tokens":4,"completion_tokens":2}}`, []string{"gpt-known"}, &LimitsConfig{}, nil)
+
+	require.Equal(t, http.StatusOK, doPassthrough(t, gw).Code)
+
+	req, ok := gw.limiter.getCounter(kindModel, "openai/gpt-known", metricReq, windowDay, time.Now())
+	assert.True(t, ok)
+	assert.EqualValues(t, 1, req)
+}
+
+func TestPassthrough_ImplausibleUsage_DoesNotLockOutBudget(t *testing.T) {
+	t.Parallel()
+	gw := newPassthroughUsageGateway(t, `{"model":"gpt-known","usage":{"prompt_tokens":9000000000000000000,"completion_tokens":1}}`, []string{"gpt-known"}, &LimitsConfig{TokensPerDay: 1000}, nil)
+
+	require.Equal(t, http.StatusOK, doPassthrough(t, gw).Code)
+	assert.Equal(t, http.StatusOK, doPassthrough(t, gw).Code, "an absurd reported figure must not exhaust the budget")
+}
+
+func TestPassthrough_EchoedUnpricedModel_BilledAtRequestModelPrice(t *testing.T) {
+	t.Parallel()
+	pricing := map[string]*ModelPricing{"openai/gpt-1": {InputPerM: 10, OutputPerM: 10}}
+	gw := newPassthroughUsageGateway(t, `{"model":"unpriced-model","usage":{"prompt_tokens":1000,"completion_tokens":500}}`, nil, &LimitsConfig{CostPerDayUSD: 1}, pricing)
+
+	require.Equal(t, http.StatusOK, doPassthrough(t, gw).Code)
+
+	cost, ok := gw.limiter.getCounter("user", "alice", metricCost, windowDay, time.Now())
+	assert.True(t, ok, "cost must accrue even though the upstream echoed an unpriced model")
+	assert.EqualValues(t, 15000, cost, "1500 tokens at 10 USD per million = 15000 micro-USD")
+}
+
+func TestLimiterAccountWith_ImplausibleUsageDropped(t *testing.T) {
+	t.Parallel()
+	l := newLimiter(nil, true)
+	scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{TokensPerDay: 1000, CostPerDayUSD: 1}}}
+
+	l.account(scopes, usage{prompt: maxUsageTokensPerRequest + 1, completion: 7}, maxCostMicrosPerRequest+1)
+
+	assert.Nil(t, l.checkAndCount(scopes), "implausible figures must not push the budgets over")
+	out, _ := l.getCounter("user", "u", metricTokOut, windowDay, time.Now())
+	assert.EqualValues(t, 7, out, "a plausible component of the same request is kept")
+}
+
+func TestEstimatePromptTokens(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name  string
+		bytes int64
+		want  int64
+	}{
+		{"empty", 0, 0},
+		{"rounds up", 5, 2},
+		{"exact", 400, 100},
+		{"capped", 32 << 20, maxEstimatedPromptTokens},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, estimatePromptTokens(tt.bytes))
+		})
+	}
+}
+
+func TestCappedAccountingBuffer_SharedBudgetBoundsConcurrentTees(t *testing.T) {
+	t.Parallel()
+	budget := newTeeBudget(10)
+	a := &cappedAccountingBuffer{budget: budget}
+	b := &cappedAccountingBuffer{budget: budget}
+
+	_, _ = a.Write([]byte("0123456789"))
+	_, err := b.Write([]byte("abc"))
+	require.NoError(t, err)
+	assert.False(t, a.truncated)
+	assert.True(t, b.truncated, "a tee that cannot reserve budget must stop retaining bytes")
+	assert.Zero(t, b.buf.Len())
+
+	a.release()
+	c := &cappedAccountingBuffer{budget: budget}
+	_, _ = c.Write([]byte("abc"))
+	assert.False(t, c.truncated, "released bytes return to the budget")
+	a.release()
+	var nilTee *cappedAccountingBuffer
+	nilTee.release()
+}
+
+func TestCacheHit_AccountHits_ChargesServingUser(t *testing.T) {
+	const respBody = `{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`
+	for _, tt := range []struct {
+		name        string
+		accountHits bool
+		wantTokIn   int64
+	}{
+		{"off: a hit is free", false, 10},
+		{"on: a hit is charged", true, 20},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Content-Type", "application/json")
+				_, _ = w.Write([]byte(respBody))
+			}))
+			defer srv.Close()
+			gw := newCacheTestGateway(t, srv, nil, defaultCacheMaxBodyBytes)
+			gw.cfg.Cache.AccountHits = tt.accountHits
+			body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
+
+			for _, want := range []string{"miss", "hit"} {
+				rec := httptest.NewRecorder()
+				gw.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body))
+				require.Equal(t, want, rec.Header().Get("X-Llmgw-Cache"))
+			}
+
+			tokIn, _ := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
+			assert.Equal(t, tt.wantTokIn, tokIn)
+			reqs, _ := gw.limiter.getCounter("user", "alice", metricReq, windowDay, time.Now())
+			assert.EqualValues(t, 2, reqs, "the request counter must not double count a hit")
+		})
+	}
+}