mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-19 05:47:48 +00:00
fix(gateway): close billing, limit, routing and translation gaps
Billing: - Passthrough SSE streams now bill the usage the upstream reports. The gateway falls back to a request-size estimate only on a 2xx stream or a truncated JSON body. - A unified stream that fails after it starts is estimated, not billed at zero. - Billing uses the built-in LiteLLM price table and honours modelMeta free: true. - Pricing entries are validated. Cost arithmetic rounds the total and saturates on overflow. - Gemini thinking tokens count as completion tokens. Limits and storage: - A rejected request rolls back its counts on the other scopes, so one user cannot lock out a group. - Redis EXPIRE replies are checked. Keys recreated without a TTL are repaired at once. - A mutating pipeline is not retried after reply bytes arrive. - History, model totals and target usage report a store outage. Routing and translation: - Discovered ids with a provider-like prefix are listed only in owner/id form, so listing and routing agree. - Anthropic and Gemini discovery follows pagination. - The developer role maps to system. - stream_options.include_usage is honoured on translated streams. - tool_choice none keeps tools when history holds tool blocks. - refusal maps to content_filter. Data URI parameters are stripped. Streaming, retry and failover: - SSE accepts CR, LF and CRLF line ends. This closes a CR injection path. - Retry-After is never jittered. A cancelled wait returns the context error. - The watchdog timer stops at end of body. - Oversized bodies return 413. - A failed stream ends with an error event. - Failover restores pre-existing headers. Client cancels do not count as provider successes. Also: "admin" is a reserved provider name. metrics.path must start with "/". Pricing warnings go to the right gateway instance. The admin targets view sends access as null or [] so the two states differ. Upstream organisation headers are stripped. The users-file reload error is logged once per error. Dead code is removed, and test-only helpers move to test files. Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
This commit is contained in:
+162
@@ -4,6 +4,7 @@ import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"reflect"
|
||||
"strings"
|
||||
@@ -253,6 +254,167 @@ func TestReadSSE_StripsLeadingBOM(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestReadSSE_LoneCRLineEndings proves a stream using ONLY "\r" as its
|
||||
// line terminator (finding 3, review-routes.md) is parsed into the same
|
||||
// events a "\n"-terminated stream would produce. bufio.ScanLines, the
|
||||
// previous split function, left a lone "\r" embedded in the line's own
|
||||
// content instead of ending it, so a CR-only stream parsed as one giant
|
||||
// line instead of the two events here.
|
||||
func TestReadSSE_LoneCRLineEndings(t *testing.T) {
|
||||
raw := "event: message\rdata: first\r\rdata: second\r\r"
|
||||
|
||||
var got []sseEvent
|
||||
err := readSSE(strings.NewReader(raw), func(ev sseEvent) error {
|
||||
got = append(got, ev)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("readSSE: %v", err)
|
||||
}
|
||||
|
||||
want := []sseEvent{
|
||||
{event: "message", data: []byte("first")},
|
||||
{event: "", data: []byte("second")},
|
||||
}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Errorf("got %#v, want %#v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestReadSSE_MixedLineTerminators proves CRLF, a lone LF, and a lone CR
|
||||
// are all recognized as terminators within the SAME stream, matching the
|
||||
// WHATWG EventSource specification's three-way definition.
|
||||
func TestReadSSE_MixedLineTerminators(t *testing.T) {
|
||||
raw := "data: crlf\r\ndata: lf\ndata: cr\r\n\n"
|
||||
|
||||
var got []sseEvent
|
||||
err := readSSE(strings.NewReader(raw), func(ev sseEvent) error {
|
||||
got = append(got, ev)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("readSSE: %v", err)
|
||||
}
|
||||
|
||||
want := []sseEvent{{event: "", data: []byte("crlf\nlf\ncr")}}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Errorf("got %#v, want %#v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// oneByteReader returns exactly one byte per Read call, forcing
|
||||
// bufio.Scanner to invoke scanSSELines with a buffer that can end
|
||||
// mid-terminator — the shape TestReadSSE_CRLFSplitAcrossReads needs to
|
||||
// prove a CRLF pair split across two underlying Read calls is never
|
||||
// misread as a lone CR followed by a spurious empty line.
|
||||
type oneByteReader struct{ s string }
|
||||
|
||||
func (r *oneByteReader) Read(p []byte) (int, error) {
|
||||
if len(r.s) == 0 {
|
||||
return 0, io.EOF
|
||||
}
|
||||
p[0] = r.s[0]
|
||||
r.s = r.s[1:]
|
||||
return 1, nil
|
||||
}
|
||||
|
||||
// TestReadSSE_CRLFSplitAcrossReads proves scanSSELines' own "wait for
|
||||
// more data before deciding" branch: when a CRLF pair straddles two
|
||||
// separate underlying Read calls (the "\r" arrives alone, with the "\n"
|
||||
// only available on the NEXT Read), it must still be recognized as one
|
||||
// CRLF terminator, not a lone CR followed by an extra, spurious empty
|
||||
// data line.
|
||||
func TestReadSSE_CRLFSplitAcrossReads(t *testing.T) {
|
||||
raw := "data: first\r\ndata: second\r\n\r\n"
|
||||
|
||||
var got []sseEvent
|
||||
err := readSSE(&oneByteReader{s: raw}, func(ev sseEvent) error {
|
||||
got = append(got, ev)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("readSSE: %v", err)
|
||||
}
|
||||
|
||||
want := []sseEvent{{event: "", data: []byte("first\nsecond")}}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Errorf("got %#v, want %#v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSSEWriter_WriteData_EncodesLoneCR proves a payload containing a raw
|
||||
// "\r" — including the exact upstream shape finding 3 (review-routes.md)
|
||||
// named, `{"a":1}\r\revent: evil\rdata: injected` — is re-encoded as a
|
||||
// sequence of clean "data: "-prefixed lines joined by "\n" alone, with no
|
||||
// raw "\r" reaching the wire: a spec-conforming EventSource client treats
|
||||
// a lone "\r" as its own line terminator, so any "\r" left unencoded here
|
||||
// could be reinterpreted by that client as extra, attacker-chosen
|
||||
// "event:"/"data:" lines this gateway never intended to send.
|
||||
func TestSSEWriter_WriteData_EncodesLoneCR(t *testing.T) {
|
||||
w := newRecordingResponseWriter()
|
||||
s := newSSEWriter(w)
|
||||
|
||||
if err := s.writeData([]byte("{\"a\":1}\r\revent: evil\rdata: injected")); err != nil {
|
||||
t.Fatalf("writeData: %v", err)
|
||||
}
|
||||
|
||||
got := w.body.String()
|
||||
if strings.Contains(got, "\r") {
|
||||
t.Fatalf("body = %q, must not contain a raw carriage return", got)
|
||||
}
|
||||
want := "data: {\"a\":1}\ndata: \ndata: event: evil\ndata: data: injected\n\n"
|
||||
if got != want {
|
||||
t.Errorf("body = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSSEWriter_WriteData_CRLFCollapsesToOneTerminator proves a CRLF pair
|
||||
// inside a payload produces exactly one "data:" line, not two (with a
|
||||
// spurious empty line between them) — CRLF is one terminator, never two.
|
||||
func TestSSEWriter_WriteData_CRLFCollapsesToOneTerminator(t *testing.T) {
|
||||
w := newRecordingResponseWriter()
|
||||
s := newSSEWriter(w)
|
||||
|
||||
if err := s.writeData([]byte("line one\r\nline two")); err != nil {
|
||||
t.Fatalf("writeData: %v", err)
|
||||
}
|
||||
|
||||
want := "data: line one\ndata: line two\n\n"
|
||||
if got := w.body.String(); got != want {
|
||||
t.Errorf("body = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestReadSSE_WriteData_Roundtrip_NoRawCRInjectionSurvives is the
|
||||
// end-to-end regression for finding 3 (review-routes.md): the exact
|
||||
// malicious upstream shape from the review's own probe —
|
||||
// `data: {"a":1}\r\revent: evil\rdata: injected\n\n` — must neither
|
||||
// collapse into one giant misread line on the read side, nor leak a raw
|
||||
// "\r" back out on the write side, once both readSSE (scanSSELines) and
|
||||
// writeData (splitSSELines) apply their own finding-3 fixes.
|
||||
// openaiAdapter.forwardStream is exactly this readSSE-then-writeData
|
||||
// pipeline in production.
|
||||
func TestReadSSE_WriteData_Roundtrip_NoRawCRInjectionSurvives(t *testing.T) {
|
||||
raw := "data: {\"a\":1}\r\revent: evil\rdata: injected\n\n"
|
||||
w := newRecordingResponseWriter()
|
||||
s := newSSEWriter(w)
|
||||
|
||||
var events int
|
||||
err := readSSE(strings.NewReader(raw), func(ev sseEvent) error {
|
||||
events++
|
||||
return s.writeData(ev.data)
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("readSSE: %v", err)
|
||||
}
|
||||
if events != 2 {
|
||||
t.Fatalf("events = %d, want 2 (readSSE must split the CR-terminated fields into two events, not one giant misread line)", events)
|
||||
}
|
||||
if got := w.body.String(); strings.Contains(got, "\r") {
|
||||
t.Errorf("wire body = %q, must not contain a raw carriage return", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewSSEWriter_SetsHeaders proves newSSEWriter sets the three SSE
|
||||
// response headers on construction, before any body write.
|
||||
func TestNewSSEWriter_SetsHeaders(t *testing.T) {
|
||||
|
||||
Reference in New Issue
Block a user