mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-14 05:18:54 +00:00
fix(passthrough): decompression-safe accounting, escaped-path proxying, streaming tee
This commit is contained in:
+1
-1
@@ -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
@@ -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
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user