From 3ce3715c7916c8675c670113dde20cd29ea57a32 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 20 Aug 2026 11:45:53 +0100 Subject: [PATCH] fix(passthrough): decompression-safe accounting, escaped-path proxying, streaming tee --- llmgateway.go | 2 +- routes_passthrough.go | 208 ++++++++++++++++++++----- routes_passthrough_test.go | 304 ++++++++++++++++++++++++++++++++++++- 3 files changed, 472 insertions(+), 42 deletions(-) diff --git a/llmgateway.go b/llmgateway.go index d766fd2..a05aa15 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -268,7 +268,7 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } - if providerName, rest, ok := passthroughRoute(r.URL.Path); ok { + if providerName, rest, ok := passthroughRoute(r.URL.EscapedPath()); ok { if _, known := g.adapters[providerName]; known { u, grp, ok := g.auth.identify(r) if !ok { diff --git a/routes_passthrough.go b/routes_passthrough.go index eb13a22..eb60a8d 100644 --- a/routes_passthrough.go +++ b/routes_passthrough.go @@ -1,11 +1,13 @@ package traefikllmgateway import ( + "bytes" "context" "encoding/json" "errors" "io" "net/http" + "net/url" "strings" ) @@ -16,6 +18,13 @@ import ( // or malicious body. const maxPassthroughBytes = 32 << 20 +// maxAccountingTeeBytes caps how much of a non-streaming JSON response +// body handlePassthrough tees off for best-effort usage accounting: 4MiB, +// generous for a usage object even inside a very large chat response, +// while the client-facing io.Copy streams every byte of the response — +// of any size — without ever waiting for the full body to buffer first. +const maxAccountingTeeBytes = 4 << 20 + // hopByHopHeaders lists the RFC 7230 §6.1 hop-by-hop headers, stripped // from both the outgoing upstream request and the response copied back // to the client — they describe this one TCP hop, not the end-to-end @@ -41,6 +50,21 @@ var gatewayCredentialHeaders = map[string]bool{ "X-Api-Key": true, } +// clientNegotiationHeaders are content-negotiation headers stripped from +// the outgoing upstream request only, never from the response: +// Accept-Encoding, so Go's http.Transport negotiates compression itself +// and transparently decompresses the reply — per its documented contract, +// that automatic decompression only happens when Transport itself added +// the header, not when the caller (this gateway, forwarding the client's +// own Accept-Encoding) set one explicitly. Without this strip, a client +// that sends "Accept-Encoding: gzip" (every mainstream OpenAI SDK and +// curl --compressed does) gets its exact header forwarded verbatim, +// Transport hands back raw gzip bytes, and the accounting JSON parse in +// handlePassthrough silently fails on binary it cannot decode. +var clientNegotiationHeaders = map[string]bool{ + "Accept-Encoding": true, +} + // copyHeadersExcept copies every header in src to dst, skipping any key // present in any of excepts. Header keys from both an *http.Request // parsed off the wire and an *http.Response parsed by the client's @@ -64,12 +88,18 @@ func copyHeadersExcept(dst, src http.Header, excepts ...map[string]bool) { // passthroughRoute splits path into a candidate provider name and the // remaining rest path, per the native passthrough convention -// "/{providerName}/{rest...}". ok is false for an empty or "/"-only path, -// which carries no candidate provider name at all — the caller falls -// through to its next-route/404 handling in that case. passthroughRoute -// makes no claim the returned providerName is actually configured; the -// caller checks that against g.adapters before treating the request as -// passthrough. +// "/{providerName}/{rest...}". Callers must pass path as +// r.URL.EscapedPath(), never the decoded r.URL.Path: a decoded path +// collapses a client's percent-encoded "%2f" into a literal "/" before it +// ever reaches this split, silently turning what the client sent as one +// opaque path segment into extra routing segments (and, downstream, into +// a "#" that would truncate the rest of the upstream URL at a fragment +// boundary if the encoded byte were "%23"). ok is false for an empty or +// "/"-only path, which carries no candidate provider name at all — the +// caller falls through to its next-route/404 handling in that case. +// passthroughRoute makes no claim the returned providerName is actually +// configured; the caller checks that against g.adapters before treating +// the request as passthrough. func passthroughRoute(path string) (providerName, rest string, ok bool) { trimmed := strings.TrimPrefix(path, "/") if trimmed == "" { @@ -82,6 +112,71 @@ func passthroughRoute(path string) (providerName, rest string, ok bool) { return providerName, rest, true } +// hasTraversalSegment reports whether rest, once percent-decoded, contains +// a path segment exactly equal to "..". rest is still in its escaped form +// here (see passthroughRoute) — decoding it is a validation-only step, +// never used to build the outgoing upstream URL, so a legitimate +// percent-encoded segment (e.g. "a%2Fb" naming a literal "a/b" resource) +// still reaches the upstream exactly as the client sent it. Checking the +// decoded form, not the escaped one, is what catches an encoded traversal +// attempt like "..%2f..%2fsecret": as a single escaped segment it has no +// literal "/" for a naive split to catch, but a permissive upstream +// server decoding that same "%2f" itself would read it as "../../secret" +// — this check rejects it here instead. A rest that fails to +// url.PathUnescape at all is rejected too: a malformed percent-encoding +// is not a path this gateway can reason about safely. +func hasTraversalSegment(rest string) bool { + decoded, err := url.PathUnescape(rest) + if err != nil { + return true + } + for _, seg := range strings.Split(decoded, "/") { + if seg == ".." { + return true + } + } + return false +} + +// cappedAccountingBuffer is an io.Writer that accumulates up to +// maxAccountingTeeBytes of a response body, then silently discards the +// rest while still reporting every byte as written. It is the accounting +// side of an io.TeeReader wrapped around a non-streaming JSON response +// body in handlePassthrough: the client-facing io.Copy reads the real +// response (of any size) through the tee, and this buffer only ever +// retains the first maxAccountingTeeBytes of it for a best-effort usage +// parse afterward. truncated records whether the cap was hit, so the +// caller knows a parse would only fail on cut-off JSON and skips it. +type cappedAccountingBuffer struct { + buf bytes.Buffer + truncated bool +} + +// 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 +// aborts the underlying copy with an error the moment a tee Write returns +// anything short of len(p); silently dropping bytes past the cap must +// never disrupt the real, client-facing copy this buffer is only +// observing. +func (c *cappedAccountingBuffer) Write(p []byte) (int, error) { + n := len(p) + if c.truncated { + return n, nil + } + remaining := maxAccountingTeeBytes - c.buf.Len() + if remaining <= 0 { + c.truncated = true + return n, nil + } + if len(p) > remaining { + p = p[:remaining] + c.truncated = true + } + c.buf.Write(p) //nolint:errcheck // bytes.Buffer.Write never errors + return n, nil +} + // passthroughUsagePayload captures every shape the three provider types' // native non-streaming JSON responses report token usage in, plus the // optional top-level "model" field OpenAI- and Anthropic-shaped responses @@ -104,15 +199,19 @@ type passthroughUsagePayload struct { // extractPassthroughUsage best-effort parses body as one of the three // provider types' native non-streaming JSON response shapes, returning // its reported token usage and model id, keyed by typeName (one of -// providerTypeOpenAI/Anthropic/Gemini). A body that fails to decode, or -// that carries none of the recognized usage shapes, returns a zero usage -// — the caller's account call is a no-op for zero usage, not a wrong -// charge. The model id falls back to "unknown/"+providerName when the -// response carries no "model" field of its own, which is always true for -// Gemini's native response. -func extractPassthroughUsage(typeName, providerName string, body []byte) (usage, string) { +// providerTypeOpenAI/Anthropic/Gemini). The returned error is body's own +// json.Unmarshal error, when decoding failed — the caller logs it rather +// than discarding it, since a body that should have been JSON (its +// Content-Type said so) but did not decode as JSON is worth knowing +// about, not silently swallowing. Usage and model still come back as +// their best-effort zero/fallback values in that case: a zero usage is a +// no-op for the caller's account call, not a wrong charge. The model id +// falls back to "unknown/"+providerName when the response carries no +// "model" field of its own, which is always true for Gemini's native +// response. +func extractPassthroughUsage(typeName, providerName string, body []byte) (usage, string, error) { var payload passthroughUsagePayload - _ = json.Unmarshal(body, &payload) // best-effort; zero usage/empty model on decode failure + err := json.Unmarshal(body, &payload) model := payload.Model if model == "" { @@ -121,13 +220,13 @@ func extractPassthroughUsage(typeName, providerName string, body []byte) (usage, switch typeName { case providerTypeOpenAI: - return usage{prompt: payload.Usage.PromptTokens, completion: payload.Usage.CompletionTokens}, model + return usage{prompt: payload.Usage.PromptTokens, completion: payload.Usage.CompletionTokens}, model, err case providerTypeAnthropic: - return usage{prompt: payload.Usage.InputTokens, completion: payload.Usage.OutputTokens}, model + return usage{prompt: payload.Usage.InputTokens, completion: payload.Usage.OutputTokens}, model, err case providerTypeGemini: - return usage{prompt: payload.UsageMetadata.PromptTokenCount, completion: payload.UsageMetadata.CandidatesTokenCount}, model + return usage{prompt: payload.UsageMetadata.PromptTokenCount, completion: payload.UsageMetadata.CandidatesTokenCount}, model, err default: - return usage{}, model + return usage{}, model, err } } @@ -140,6 +239,11 @@ func extractPassthroughUsage(typeName, providerName string, body []byte) (usage, // in the gateway's own error envelope the way the unified routes wrap a // providerHTTPError. // +// Checks run in this order: group authorization (403) before capability +// checks (Upgrade→501, path validity→400) before rate limits (429/503) — +// a caller who cannot use providerName at all learns that first, rather +// than learning something about how they tried to use it. +// // For a Gemini provider, rest is appended to base() exactly as the client // sent it: there is no model extraction or URL rewriting here, so a // Gemini passthrough client must address it with Gemini's own native URL @@ -147,13 +251,18 @@ func extractPassthroughUsage(typeName, providerName string, body []byte) (usage, // paths — injectAuth still sets the same x-goog-api-key header it sets // for every other Gemini request. func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *user, grp *group, providerName, rest string) { + if !grp.allowsProvider(providerName) { + writeOAIError(w, http.StatusForbidden, "invalid_request_error", "provider access denied") + return + } + if r.Header.Get("Upgrade") != "" { writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported") return } - if !grp.allowsProvider(providerName) { - writeOAIError(w, http.StatusForbidden, "invalid_request_error", "provider access denied") + if hasTraversalSegment(rest) { + writeOAIError(w, http.StatusBadRequest, "invalid_request_error", "invalid path") return } @@ -181,7 +290,14 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u // upstreamURL is built as base()+"/"+rest — string concatenation, not a // second url.Parse of rest alone — so no client-supplied value can // change the scheme or host url.Parse resolves for the final string. - upstreamReq, err := http.NewRequestWithContext(r.Context(), r.Method, upstreamURL, bodyReader) //nolint:gosec // path/query passthrough onto an operator-fixed host; see comment above + // rest is additionally traversal-checked above (hasTraversalSegment + // rejects any ".." segment in its decoded form before this point is + // reached) and is forwarded here in its original escaped form — see + // passthroughRoute's EscapedPath()-based split in llmgateway.go — so a + // client cannot smuggle an encoded "/" (%2f) past this gateway's own + // routing only to have a permissive upstream reinterpret it as a real + // separator. + upstreamReq, err := http.NewRequestWithContext(r.Context(), r.Method, upstreamURL, bodyReader) //nolint:gosec // operator-fixed host, rest is traversal-checked and forwarded escaped; see comment above if err != nil { g.errorf("passthrough: build upstream request (provider %q): %v", providerName, err) writeOAIError(w, http.StatusBadGateway, "server_error", "upstream connection error") @@ -194,10 +310,10 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u // bodyReader above may not actually deliver. upstreamReq.ContentLength = -1 } - copyHeadersExcept(upstreamReq.Header, r.Header, hopByHopHeaders, gatewayCredentialHeaders) + copyHeadersExcept(upstreamReq.Header, r.Header, hopByHopHeaders, gatewayCredentialHeaders, clientNegotiationHeaders) adapter.injectAuth(upstreamReq) - resp, err := adapter.httpClient().Do(upstreamReq) //nolint:gosec // same upstreamReq built above; host is operator-fixed, see its construction comment + resp, err := adapter.httpClient().Do(upstreamReq) //nolint:gosec // same upstreamReq built above; operator-fixed host, traversal-checked, see its construction comment if err != nil { if errors.Is(err, context.Canceled) { g.logf("passthrough: client canceled request to provider %q: %v", providerName, err) @@ -213,25 +329,37 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u w.WriteHeader(resp.StatusCode) fw := newFlushWriter(w) - if strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "application/json") { - body, readErr := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes)) - if readErr != nil { - g.errorf("passthrough: read response body (provider %q): %v", providerName, readErr) - return - } - if _, err := fw.Write(body); err != nil { - g.errorf("passthrough: write response body (provider %q): %v", providerName, err) - return - } - respUsage, model := extractPassthroughUsage(adapter.typeName(), providerName, body) - cost := unifiedCostMicros(providerName+"/"+model, model, respUsage, g.cfg.Pricing) - g.limiter.account(scopes, respUsage, cost) + // Only a non-streaming JSON response gets teed off for best-effort + // accounting; every other response streams straight through fw with no + // extra buffering — a large or slow upstream body must never sit + // waiting for full receipt before the client sees its first byte. + isJSON := strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "application/json") + var tee *cappedAccountingBuffer + var reader io.Reader = resp.Body + if isJSON { + tee = &cappedAccountingBuffer{} + reader = io.TeeReader(resp.Body, tee) + } + + if _, err := io.Copy(fw, reader); err != nil { + g.errorf("passthrough: stream response body (provider %q): %v", providerName, err) + return + } + if !isJSON { + // Streaming or non-JSON: usage is unknowable from here, so only the + // request itself gets accounted — already done by checkAndCount + // above. + return + } + if tee.truncated { + g.logf("passthrough: response body exceeded %d bytes; skipping usage accounting (provider %q)", maxAccountingTeeBytes, providerName) return } - // Streaming or non-JSON: usage is unknowable from here, so only the - // request itself gets accounted — already done by checkAndCount above. - if _, err := io.Copy(fw, resp.Body); err != nil { - g.errorf("passthrough: stream response body (provider %q): %v", providerName, err) + respUsage, model, unmarshalErr := extractPassthroughUsage(adapter.typeName(), providerName, tee.buf.Bytes()) + if unmarshalErr != nil { + g.logf("passthrough: response body did not decode as JSON for usage accounting (provider %q): %v", providerName, unmarshalErr) } + cost := unifiedCostMicros(providerName+"/"+model, model, respUsage, g.cfg.Pricing) + g.limiter.account(scopes, respUsage, cost) } diff --git a/routes_passthrough_test.go b/routes_passthrough_test.go index 4769578..5cd55ea 100644 --- a/routes_passthrough_test.go +++ b/routes_passthrough_test.go @@ -2,6 +2,7 @@ package traefikllmgateway import ( "bytes" + "compress/gzip" "context" "io" "net/http" @@ -49,6 +50,7 @@ func TestExtractPassthroughUsage(t *testing.T) { body string wantModel string wantUsage usage + wantErr bool }{ { name: "openai shape", @@ -81,17 +83,46 @@ func TestExtractPassthroughUsage(t *testing.T) { body: `not json`, wantUsage: usage{}, wantModel: "unknown/openai", + wantErr: true, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - gotUsage, gotModel := extractPassthroughUsage(tt.typeName, tt.providerName, []byte(tt.body)) + gotUsage, gotModel, gotErr := extractPassthroughUsage(tt.typeName, tt.providerName, []byte(tt.body)) if gotUsage != tt.wantUsage { t.Errorf("usage = %+v, want %+v", gotUsage, tt.wantUsage) } if gotModel != tt.wantModel { t.Errorf("model = %q, want %q", gotModel, tt.wantModel) } + if (gotErr != nil) != tt.wantErr { + t.Errorf("err = %v, want err present = %v", gotErr, tt.wantErr) + } + }) + } +} + +// --- unit tests: hasTraversalSegment --- + +func TestHasTraversalSegment(t *testing.T) { + tests := []struct { + rest string + want bool + }{ + {"v1/audio/speech", false}, + {"v1beta/models/gemini-pro:generateContent", false}, + {"a%2Fb", false}, + {"../secret", true}, + {"v1/../secret", true}, + {"..%2f..%2fsecret", true}, + {"..%2F..%2Fsecret", true}, + {"%zz", true}, // invalid percent-encoding + } + for _, tt := range tests { + t.Run(tt.rest, func(t *testing.T) { + if got := hasTraversalSegment(tt.rest); got != tt.want { + t.Errorf("hasTraversalSegment(%q) = %v, want %v", tt.rest, got, tt.want) + } }) } } @@ -492,3 +523,274 @@ func TestHandlePassthrough_DeadUpstream_Returns502(t *testing.T) { t.Fatalf("status = %d, want 502, body=%s", rec.Code, rec.Body.String()) } } + +// TestHandlePassthrough_AcceptEncodingStripped_AccountsUsageThroughGzip +// proves a client's own "Accept-Encoding: gzip" (sent by every mainstream +// OpenAI SDK, Node fetch, and curl --compressed) never reaches the +// upstream verbatim: stripping it lets Go's Transport negotiate and +// transparently decompress gzip itself, so the accounting JSON parse +// still sees plaintext instead of silently failing on raw gzip bytes. +func TestHandlePassthrough_AcceptEncodingStripped_AccountsUsageThroughGzip(t *testing.T) { + const respBody = `{"model":"gpt-native-x","usage":{"prompt_tokens":7,"completion_tokens":3}}` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") { + w.Header().Set("Content-Encoding", "gzip") + w.WriteHeader(http.StatusOK) + gz := gzip.NewWriter(w) + _, _ = gz.Write([]byte(respBody)) + _ = gz.Close() + return + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{ + {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}, + }} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + + req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(`{}`)) + req.Header.Set("Authorization", "Bearer sk-alice") + req.Header.Set("Accept-Encoding", "gzip") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if rec.Body.String() != respBody { + t.Errorf("body = %q, want decompressed upstream body %q", rec.Body.String(), respBody) + } + + tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now()) + if !ok || tok != 10 { + t.Errorf("user token/day counter = %d (ok=%v), want 10 — usage must still be accounted through gzip", tok, ok) + } +} + +// TestHandlePassthrough_TraversalPath_Returns400 covers both a literal +// and a percent-encoded ".." segment in rest; neither must ever reach the +// upstream. +func TestHandlePassthrough_TraversalPath_Returns400(t *testing.T) { + upstreamCalled := false + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalled = true + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + for _, path := range []string{ + "/openai/../secret", + "/openai/..%2f..%2fsecret", + } { + t.Run(path, func(t *testing.T) { + req := httptest.NewRequest(http.MethodGet, path, nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400, body=%s", rec.Code, rec.Body.String()) + } + }) + } + if upstreamCalled { + t.Error("upstream must never be called for a rejected traversal path") + } +} + +// TestHandlePassthrough_EncodedSlashSegment_PreservedAtUpstream proves a +// legitimately percent-encoded "/" within one rest segment (e.g. a +// resource id that itself contains a slash) reaches the upstream in its +// original encoded form, not silently decoded into an extra path +// separator. +func TestHandlePassthrough_EncodedSlashSegment_PreservedAtUpstream(t *testing.T) { + var gotRequestURI string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotRequestURI = r.RequestURI + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "/openai/v1/files/a%2Fb", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if !strings.Contains(gotRequestURI, "a%2Fb") { + t.Errorf("upstream RequestURI = %q, want the encoded \"a%%2Fb\" preserved, not decoded", gotRequestURI) + } +} + +// TestHandlePassthrough_OversizedJSONBody_ClientGetsFullBody_RequestOnlyAccounting +// proves a JSON response body larger than maxAccountingTeeBytes still +// reaches the client in full, while the accounting parse is skipped — +// only the request itself was already accounted, by checkAndCount. +func TestHandlePassthrough_OversizedJSONBody_ClientGetsFullBody_RequestOnlyAccounting(t *testing.T) { + // Pad well past maxAccountingTeeBytes (4MiB) with an oversized field + // ahead of "usage", so the tee's cap is hit before "usage" is reached. + padding := strings.Repeat("x", maxAccountingTeeBytes+1024) + respBody := `{"model":"gpt-native-x","padding":"` + padding + `","usage":{"prompt_tokens":7,"completion_tokens":3}}` + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{ + {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}, + }} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + + req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(`{}`)) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200", rec.Code) + } + if rec.Body.String() != respBody { + t.Errorf("client body = %d bytes, want the full %d bytes (client copy must never truncate)", rec.Body.Len(), len(respBody)) + } + + tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now()) + if !ok || tok != 0 { + t.Errorf("user token/day counter = %d (ok=%v), want 0 — usage accounting must be skipped over the tee cap", tok, ok) + } + reqCount, ok := gw.limiter.getCounter("user", "alice", metricReq, windowMin, time.Now()) + if !ok || reqCount != 1 { + t.Errorf("user request/min counter = %d (ok=%v), want 1 — the request itself is still accounted", reqCount, ok) + } +} + +// TestHandlePassthrough_QueryStringForwarded proves the client's query +// string reaches the upstream via RawQuery, unmodified. +func TestHandlePassthrough_QueryStringForwarded(t *testing.T) { + var gotQuery string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotQuery = r.URL.RawQuery + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "/openai/v1/thing?a=1&b=2", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if gotQuery != "a=1&b=2" { + t.Errorf("upstream RawQuery = %q, want %q", gotQuery, "a=1&b=2") + } +} + +// TestHandlePassthrough_AnthropicClientAPIKeyStripped_ProviderKeyInjected +// proves an anthropic passthrough client's own x-api-key — the gateway's +// own auth header, which happens to share anthropic's own upstream +// credential header name — never reaches the upstream: the provider's own +// key is injected in its place. +func TestHandlePassthrough_AnthropicClientAPIKeyStripped_ProviderKeyInjected(t *testing.T) { + var gotAPIKeyHeader string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAPIKeyHeader = r.Header.Get("x-api-key") + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "anthropic": {Type: "anthropic", BaseURL: srv.URL, APIKey: "sk-provider-real"}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "client-key-should-not-leak"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodPost, "/anthropic/v1/messages", strings.NewReader(`{}`)) + req.Header.Set("x-api-key", "client-key-should-not-leak") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if gotAPIKeyHeader != "sk-provider-real" { + t.Errorf("x-api-key = %q, want the provider's own key injected", gotAPIKeyHeader) + } +}