fix(passthrough): decompression-safe accounting, escaped-path proxying, streaming tee

This commit is contained in:
2026-08-20 11:45:53 +01:00
parent 8d470b173b
commit 3ce3715c79
3 changed files with 472 additions and 42 deletions
+1 -1
View File
@@ -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 {
+168 -40
View File
@@ -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)
}
+303 -1
View File
@@ -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)
}
}