diff --git a/limits.go b/limits.go index 3749854..a51b898 100644 --- a/limits.go +++ b/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 diff --git a/llmgateway.go b/llmgateway.go index 8901039..9677f9d 100644 --- a/llmgateway.go +++ b/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 { diff --git a/registry.go b/registry.go index 71eeb47..178a8c0 100644 --- a/registry.go +++ b/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 diff --git a/routes_passthrough.go b/routes_passthrough.go index 0ec5281..2522b90 100644 --- a/routes_passthrough.go +++ b/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{} diff --git a/routes_passthrough_test.go b/routes_passthrough_test.go index 41d1be1..96410fb 100644 --- a/routes_passthrough_test.go +++ b/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{ diff --git a/routes_unified.go b/routes_unified.go index a681885..86468bd 100644 --- a/routes_unified.go +++ b/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 +} diff --git a/security_fixes_test.go b/security_fixes_test.go new file mode 100644 index 0000000..b180372 --- /dev/null +++ b/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") + }) + } +}