mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
BLOCKER (v0.22 review round, round 2): matchesSentinel's hand-rolled
Unwrap walk (round 1's SHOULD-1 fix) stopped dead the moment it hit
providers.go's own multi-%w wrapping. fmt.Errorf with more than one %w
verb builds a *fmt.wrapErrors implementing Unwrap() []error, not
Unwrap() error — a comma-ok assertion against a single-error Unwrap
interface alone never reaches past it. All three of providers.go's own
wrap sites use multi-%w (upstreamJSON's encode-failure path, "%w:
encode request body: %w"; upstreamBytes' build-request path, "%w: %w:
build request: %w"; and upstreamBytes' client.Do path, "%w: %w" — the
one a real request-context deadline reaches through). Result: a hung
upstream on the unified chat/embeddings/media routes recorded as 100%
healthy; only native passthrough (whose error is never re-wrapped by
upstreamBytes) worked.
Round 2's first attempt added a second hand-rolled interface
(multiUnwrapper, Unwrap() []error) and recursed into it — verified
correct against every wrap shape under `go test`, INCLUDING the exact
double-/triple-%w shapes providers.go itself produces. It still did not
work. tools/yaegi-check's own SHOULD-A extension (below) caught it: a
type assertion of a COMPILED concrete value (*fmt.wrapErrors, built by
the real fmt package) against an INTERPRETED interface type
(multiUnwrapper, declared in this plugin's own source) silently returns
ok=false under Yaegi, even though %T correctly reports *fmt.wrapErrors
and that type genuinely implements Unwrap() []error in real Go — the
reverse direction of the already-documented errors.As-on-interpreted-
types trap, not previously known to this codebase.
Fix: matchesSentinel now calls real errors.Is instead of hand-rolling
anything. This is safe here specifically, despite the "avoid errors.Is/
As" guidance elsewhere in this codebase: the panic risk that guidance
targets is errors.As checking whether an INTERPRETED type implements
error, which never arises for a plain sentinel comparison against
compiled-origin values (context.DeadlineExceeded, context.Canceled, and
every err this function ever sees) — the identical safe shape retry.go's
isTransient already relies on for these same two sentinels, proven
correct under Yaegi across every prior round. Empirically verified: this
is what made the extended yaegi-check harness pass; the hand-rolled
version did not, despite passing every compiled `go test` row.
Also in this commit:
- Route-level regression tests (TestHandleChat_
ContextDeadlineExceeded_RecordsProviderFailure, routes_unified_test.go;
TestHandlePassthrough_ContextDeadlineExceeded_RecordsProviderFailure,
routes_passthrough_test.go) — a real httptest upstream sleeping past a
real short request-context deadline, adapted from a reviewer-provided
probe. These are what exposed the blocker in the first place; the
sentinel-level table test alone (hand-built error shapes) is exactly
what let it ship.
- Prod-shape rows (double-/triple-%w, byte-for-byte providers.go's own
format strings) folded into the existing classification table
(renamed TestIsDeadlineExceeded_Classification — its own doc comment
now records that this table alone is NOT proof under the interpreter,
after round 2's own history).
- tools/yaegi-check/main.go (SHOULD-A): the harness's claim to prove the
interpreted deadline path was never actually true — its only upstream
answered instantly, so isDeadlineExceeded was never exercised with a
non-nil error at all. Added a second, deliberately slow provider
("slow", provider-prefixed model id so it cannot collide with the
existing fast one) driven with a short request-context deadline,
asserting failuresDay==1 under the real interpreter. This is the
harness that caught the blocker above.
- FOLDED-1: recordProviderAttempt's fire-and-forget store write
(SHOULD-5, round 1) is now bounded — limiter.spawnTokens, a per-limiter
buffered-channel semaphore (providerAttemptSpawnCap=64) the production
spawn closure acquires non-blockingly; a write that cannot start
immediately is DROPPED, not queued. Bounded and lossy-under-pressure by
design: telemetry only, enforcement never goes through spawn.
Per-limiter, not package-level, so the test suite's many parallel
*limiter instances cannot starve each other.
- FOLDED-2: TestRecordProviderAttempt_ProductionSpawn_
WriteEventuallyLandsAsync leaves limiter.spawn at its real production
default (every other test in this file overrides it to run
synchronously) and polls (waitUntil, deadline-bounded, not a fixed
sleep) for the counter to land — closing the -race async-path gap the
synchronous overrides otherwise leave entirely untested.
990 lines
37 KiB
Go
990 lines
37 KiB
Go
package traefikllmgateway
|
|
|
|
import (
|
|
"bytes"
|
|
"compress/gzip"
|
|
"context"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// --- unit tests: passthroughRoute ---
|
|
|
|
func TestPassthroughRoute(t *testing.T) {
|
|
tests := []struct {
|
|
path string
|
|
wantProvider string
|
|
wantRest string
|
|
wantOK bool
|
|
}{
|
|
{"/openai/v1/audio/speech", "openai", "v1/audio/speech", true},
|
|
{"/gemini/v1beta/models/gemini-pro:generateContent", "gemini", "v1beta/models/gemini-pro:generateContent", true},
|
|
{"/anthropic", "anthropic", "", true},
|
|
{"/anthropic/", "anthropic", "", true},
|
|
{"/", "", "", false},
|
|
{"", "", "", false},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.path, func(t *testing.T) {
|
|
gotProvider, gotRest, gotOK := passthroughRoute(tt.path)
|
|
if gotOK != tt.wantOK || gotProvider != tt.wantProvider || gotRest != tt.wantRest {
|
|
t.Errorf("passthroughRoute(%q) = (%q, %q, %v), want (%q, %q, %v)",
|
|
tt.path, gotProvider, gotRest, gotOK, tt.wantProvider, tt.wantRest, tt.wantOK)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// --- unit tests: extractPassthroughUsage ---
|
|
|
|
func TestExtractPassthroughUsage(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
typeName string
|
|
providerName string
|
|
body string
|
|
wantModel string
|
|
wantUsage usage
|
|
wantErr bool
|
|
}{
|
|
{
|
|
name: "openai shape",
|
|
typeName: providerTypeOpenAI,
|
|
providerName: "openai",
|
|
body: `{"model":"gpt-native","usage":{"prompt_tokens":7,"completion_tokens":3}}`,
|
|
wantUsage: usage{prompt: 7, completion: 3},
|
|
wantModel: "gpt-native",
|
|
},
|
|
{
|
|
name: "anthropic shape",
|
|
typeName: providerTypeAnthropic,
|
|
providerName: "anthropic",
|
|
body: `{"model":"claude-native","usage":{"input_tokens":11,"output_tokens":4}}`,
|
|
wantUsage: usage{prompt: 11, completion: 4},
|
|
wantModel: "claude-native",
|
|
},
|
|
{
|
|
name: "gemini shape, no model field",
|
|
typeName: providerTypeGemini,
|
|
providerName: "gemini",
|
|
body: `{"usageMetadata":{"promptTokenCount":20,"candidatesTokenCount":6}}`,
|
|
wantUsage: usage{prompt: 20, completion: 6},
|
|
wantModel: "unknown/gemini",
|
|
},
|
|
{
|
|
name: "malformed body",
|
|
typeName: providerTypeOpenAI,
|
|
providerName: "openai",
|
|
body: `not json`,
|
|
wantUsage: usage{},
|
|
wantModel: "unknown/openai",
|
|
wantErr: true,
|
|
},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
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)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// --- end-to-end: native passthrough through Gateway.ServeHTTP ---
|
|
|
|
// TestHandlePassthrough_ClientKeySwappedForProviderKey proves the client's
|
|
// own gateway credential never reaches the upstream provider, and the
|
|
// provider's own configured key is injected in its place.
|
|
func TestHandlePassthrough_ClientKeySwappedForProviderKey(t *testing.T) {
|
|
var gotAuth, gotAPIKeyHeader string
|
|
var gotBody []byte
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotAuth = r.Header.Get("Authorization")
|
|
gotAPIKeyHeader = r.Header.Get("x-api-key")
|
|
gotBody, _ = io.ReadAll(r.Body)
|
|
w.Header().Set("Content-Type", "text/plain")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte("upstream-ok"))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := CreateConfig()
|
|
cfg.Providers = map[string]*ProviderConfig{
|
|
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-provider-real"},
|
|
}
|
|
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.MethodPost, "/openai/v1/audio/speech", strings.NewReader(`{"foo":"bar"}`))
|
|
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 rec.Body.String() != "upstream-ok" {
|
|
t.Errorf("body = %q, want verbatim upstream body", rec.Body.String())
|
|
}
|
|
if strings.Contains(gotAuth, "sk-alice") {
|
|
t.Errorf("upstream saw the client's gateway key in Authorization: %q", gotAuth)
|
|
}
|
|
if gotAuth != "Bearer sk-provider-real" {
|
|
t.Errorf("Authorization = %q, want the provider's own key injected", gotAuth)
|
|
}
|
|
if gotAPIKeyHeader != "" {
|
|
t.Errorf("x-api-key = %q, want empty (openai adapter never sets it)", gotAPIKeyHeader)
|
|
}
|
|
if string(gotBody) != `{"foo":"bar"}` {
|
|
t.Errorf("upstream body = %q, want the client body preserved", gotBody)
|
|
}
|
|
}
|
|
|
|
// TestHandlePassthrough_OpenAI_ClientAPIKeyHeaderStripped proves the
|
|
// client's gateway credential, presented via x-api-key rather than
|
|
// Authorization, never reaches an OpenAI passthrough upstream. openai's
|
|
// injectAuth only ever sets Authorization, so if gatewayCredentialHeaders'
|
|
// X-Api-Key entry were ever dropped from the strip set, this exact header
|
|
// would pass straight through copyHeadersExcept to the upstream (I1: the
|
|
// sibling test above authenticates via Authorization and never sends
|
|
// x-api-key at all, so it never exercised this strip entry; the anthropic
|
|
// coverage sends x-api-key but targets a provider whose own injectAuth
|
|
// always overwrites it regardless of whether the strip ran).
|
|
func TestHandlePassthrough_OpenAI_ClientAPIKeyHeaderStripped(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{
|
|
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-provider-real"},
|
|
}
|
|
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
|
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "gateway-key-must-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, "/openai/v1/audio/speech", strings.NewReader(`{"foo":"bar"}`))
|
|
// Gateway auth presented via x-api-key only (no Authorization header) —
|
|
// presentedKey falls back to x-api-key when Authorization is absent.
|
|
req.Header.Set("x-api-key", "gateway-key-must-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 != "" {
|
|
t.Errorf("upstream saw x-api-key = %q, want empty (openai adapter only ever sets Authorization; the client's gateway credential must be stripped)", gotAPIKeyHeader)
|
|
}
|
|
}
|
|
|
|
// syncFlushWriter is an http.ResponseWriter/http.Flusher test double that
|
|
// signals writeCh (non-blocking) after every Write, so a test can prove
|
|
// two upstream chunks reached the client as two distinct Write calls
|
|
// instead of one coalesced Read — a plain httptest.Recorder or the
|
|
// unified routes' recordingWriter can't rule out coalescing, since
|
|
// nothing stops both upstream flushes landing in a single client-side
|
|
// Read before either type's Write is ever inspected. Guarded by mu since
|
|
// the handler goroutine writes concurrently with the test goroutine's
|
|
// reads of buf/writes/flushes.
|
|
type syncFlushWriter struct {
|
|
hdr http.Header
|
|
writeCh chan struct{}
|
|
buf bytes.Buffer
|
|
mu sync.Mutex
|
|
status int
|
|
writes int
|
|
flushes int
|
|
}
|
|
|
|
func newSyncFlushWriter() *syncFlushWriter {
|
|
return &syncFlushWriter{hdr: http.Header{}, writeCh: make(chan struct{}, 8)}
|
|
}
|
|
|
|
func (s *syncFlushWriter) Header() http.Header { return s.hdr }
|
|
|
|
func (s *syncFlushWriter) WriteHeader(status int) {
|
|
s.mu.Lock()
|
|
s.status = status
|
|
s.mu.Unlock()
|
|
}
|
|
|
|
func (s *syncFlushWriter) Write(b []byte) (int, error) {
|
|
s.mu.Lock()
|
|
if s.status == 0 {
|
|
s.status = http.StatusOK
|
|
}
|
|
n, _ := s.buf.Write(b)
|
|
s.writes++
|
|
s.mu.Unlock()
|
|
select {
|
|
case s.writeCh <- struct{}{}:
|
|
default:
|
|
}
|
|
return n, nil
|
|
}
|
|
|
|
func (s *syncFlushWriter) Flush() {
|
|
s.mu.Lock()
|
|
s.flushes++
|
|
s.mu.Unlock()
|
|
}
|
|
|
|
func (s *syncFlushWriter) snapshot() (status, writes, flushes int, body string) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
return s.status, s.writes, s.flushes, s.buf.String()
|
|
}
|
|
|
|
// TestHandlePassthrough_StreamingSSE_FlushesIncrementally proves a raw
|
|
// SSE passthrough is forwarded to the client chunk by chunk, flushing
|
|
// after each one, instead of buffered until the response completes. The
|
|
// fake upstream blocks between its two chunks on continueCh, which the
|
|
// test only closes after observing (via writeCh) that the first chunk
|
|
// already reached the client as its own Write — this rules out the two
|
|
// chunks getting coalesced into a single client-side Read/Write, which a
|
|
// plain "write and flush twice with no synchronization" fake upstream
|
|
// cannot rule out.
|
|
func TestHandlePassthrough_StreamingSSE_FlushesIncrementally(t *testing.T) {
|
|
continueCh := make(chan struct{})
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
fl, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Fatal("fake upstream ResponseWriter does not support Flush")
|
|
}
|
|
_, _ = w.Write([]byte("data: chunk-1\n\n"))
|
|
fl.Flush()
|
|
<-continueCh
|
|
_, _ = w.Write([]byte("data: chunk-2\n\n"))
|
|
fl.Flush()
|
|
}))
|
|
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/native-stream", nil)
|
|
req.Header.Set("Authorization", "Bearer sk-alice")
|
|
sw := newSyncFlushWriter()
|
|
done := make(chan struct{})
|
|
go func() {
|
|
h.ServeHTTP(sw, req)
|
|
close(done)
|
|
}()
|
|
|
|
select {
|
|
case <-sw.writeCh:
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("timed out waiting for the first chunk to reach the client")
|
|
}
|
|
_, writesAfterFirst, _, bodyAfterFirst := sw.snapshot()
|
|
if writesAfterFirst != 1 {
|
|
t.Fatalf("writes after first chunk = %d, want exactly 1 (chunks must not coalesce)", writesAfterFirst)
|
|
}
|
|
if !strings.Contains(bodyAfterFirst, "chunk-1") || strings.Contains(bodyAfterFirst, "chunk-2") {
|
|
t.Fatalf("body after first chunk = %q, want chunk-1 only", bodyAfterFirst)
|
|
}
|
|
|
|
close(continueCh)
|
|
select {
|
|
case <-done:
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("timed out waiting for ServeHTTP to finish")
|
|
}
|
|
|
|
status, writes, flushes, body := sw.snapshot()
|
|
if status != http.StatusOK {
|
|
t.Fatalf("status = %d, want 200", status)
|
|
}
|
|
if writes < 2 {
|
|
t.Fatalf("want >=2 writes (one per upstream chunk), got %d", writes)
|
|
}
|
|
if flushes < writes {
|
|
t.Fatalf("want a flush per write (incremental delivery), got %d writes and %d flushes", writes, flushes)
|
|
}
|
|
if !strings.Contains(body, "chunk-1") || !strings.Contains(body, "chunk-2") {
|
|
t.Errorf("body missing streamed chunks, got %q", body)
|
|
}
|
|
}
|
|
|
|
// TestHandlePassthrough_HopByHopHeadersStripped proves hop-by-hop headers
|
|
// are stripped in both directions, while ordinary headers pass through
|
|
// untouched.
|
|
func TestHandlePassthrough_HopByHopHeadersStripped(t *testing.T) {
|
|
var gotHeaders http.Header
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotHeaders = r.Header.Clone()
|
|
w.Header().Set("Content-Type", "text/plain")
|
|
w.Header().Set("Connection", "close") // hop-by-hop; must not reach the client
|
|
w.Header().Set("X-Custom", "keep-me") // ordinary header; must reach the client
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte("ok"))
|
|
}))
|
|
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", nil)
|
|
req.Header.Set("Authorization", "Bearer sk-alice")
|
|
req.Header.Set("Connection", "keep-alive")
|
|
req.Header.Set("Te", "trailers")
|
|
req.Header.Set("X-Client-Custom", "keep-me-too")
|
|
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 gotHeaders.Get("Te") != "" {
|
|
t.Errorf("upstream saw hop-by-hop Te header: %q", gotHeaders.Get("Te"))
|
|
}
|
|
if gotHeaders.Get("X-Client-Custom") != "keep-me-too" {
|
|
t.Errorf("upstream did not see ordinary header X-Client-Custom")
|
|
}
|
|
if rec.Header().Get("Connection") != "" {
|
|
t.Errorf("client saw hop-by-hop Connection header from upstream response: %q", rec.Header().Get("Connection"))
|
|
}
|
|
if rec.Header().Get("X-Custom") != "keep-me" {
|
|
t.Errorf("client did not see ordinary response header X-Custom")
|
|
}
|
|
}
|
|
|
|
// TestHandlePassthrough_NonStreamJSON_AccountsUsage proves a non-streaming
|
|
// application/json response's usage is extracted and accounted against
|
|
// the caller's own limiter counters.
|
|
func TestHandlePassthrough_NonStreamJSON_AccountsUsage(t *testing.T) {
|
|
const respBody = `{"model":"gpt-native-x","choices":[],"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")
|
|
}
|
|
// SHOULD-5 (v0.22 review round): production spawns recordProviderAttempt's
|
|
// store write in its own goroutine, off this route's TTFB path — this
|
|
// test asserts on the resulting counter immediately after ServeHTTP
|
|
// returns, so it overrides spawn to run synchronously instead (the
|
|
// deterministic-test half of that dependency-injection field; see
|
|
// limiter.spawn's own doc comment, limits.go).
|
|
gw.limiter.spawn = func(f func()) { f() }
|
|
|
|
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, body=%s", rec.Code, rec.Body.String())
|
|
}
|
|
if rec.Body.String() != respBody {
|
|
t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody)
|
|
}
|
|
|
|
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
|
if !ok || tokIn != 7 {
|
|
t.Errorf("user tokin/day counter = %d (ok=%v), want 7", tokIn, ok)
|
|
}
|
|
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
|
|
if !ok || tokOut != 3 {
|
|
t.Errorf("user tokout/day counter = %d (ok=%v), want 3", tokOut, ok)
|
|
}
|
|
|
|
// withTotalScope (handlePassthrough, routes_passthrough.go) must have
|
|
// appended the synthetic total scope too (v0.2 data-layer task).
|
|
totalTokIn, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricTokIn, windowDay, time.Now())
|
|
if !ok || totalTokIn != 7 {
|
|
t.Errorf("total tokin/day counter = %d (ok=%v), want 7", totalTokIn, ok)
|
|
}
|
|
totalTokOut, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricTokOut, windowDay, time.Now())
|
|
if !ok || totalTokOut != 3 {
|
|
t.Errorf("total tokout/day counter = %d (ok=%v), want 3", totalTokOut, ok)
|
|
}
|
|
|
|
// Feature A (v0.22): handlePassthrough wraps r's context with an
|
|
// attemptRecorder before calling g.proxyUpstream — provider-level
|
|
// only (model="": the upstream model here lives in the response body,
|
|
// read only after this attempt already resolved, per
|
|
// recordProviderAttempt's own doc comment, limits.go).
|
|
attempts, ok := gw.limiter.getCounter(kindProvider, "openai", metricProvAttempt, windowDay, time.Now())
|
|
if !ok || attempts != 1 {
|
|
t.Errorf("provider attempts/day = %d (ok=%v), want 1", attempts, ok)
|
|
}
|
|
fails, _ := gw.limiter.getCounter(kindProvider, "openai", metricProvFail, windowDay, time.Now())
|
|
if fails != 0 {
|
|
t.Errorf("provider fails/day = %d, want 0", fails)
|
|
}
|
|
}
|
|
|
|
// TestHandlePassthrough_DeadUpstream_RecordsProviderFailure proves a
|
|
// connection-level failure (proxyUpstream's client.Do returning a network
|
|
// error) still reports one provider-level attempt AND one failure —
|
|
// Feature A (v0.22): passthrough makes no retry.go attempt at all, so
|
|
// this exercises proxyUpstream's own attemptRecorderFromContext call
|
|
// directly, not retryPolicy.do's.
|
|
func TestHandlePassthrough_DeadUpstream_RecordsProviderFailure(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
|
|
deadURL := srv.URL
|
|
srv.Close() // closed before any request: every dial fails
|
|
|
|
cfg := CreateConfig()
|
|
cfg.Providers = map[string]*ProviderConfig{
|
|
"openai": {Type: "openai", BaseURL: deadURL, 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")
|
|
}
|
|
// SHOULD-5 (v0.22 review round): see the identical override in
|
|
// TestHandlePassthrough_NonStreamJSON_AccountsUsage above.
|
|
gw.limiter.spawn = func(f func()) { f() }
|
|
|
|
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.StatusBadGateway {
|
|
t.Fatalf("status = %d, want 502, body=%s", rec.Code, rec.Body.String())
|
|
}
|
|
|
|
attempts, ok := gw.limiter.getCounter(kindProvider, "openai", metricProvAttempt, windowDay, time.Now())
|
|
if !ok || attempts != 1 {
|
|
t.Errorf("provider attempts/day = %d (ok=%v), want 1", attempts, ok)
|
|
}
|
|
fails, ok := gw.limiter.getCounter(kindProvider, "openai", metricProvFail, windowDay, time.Now())
|
|
if !ok || fails != 1 {
|
|
t.Errorf("provider fails/day = %d (ok=%v), want 1", fails, ok)
|
|
}
|
|
}
|
|
|
|
// TestHandlePassthrough_GroupDeniesProvider_Returns403 covers a provider
|
|
// the group's providers allowlist does not include.
|
|
func TestHandlePassthrough_GroupDeniesProvider_Returns403(t *testing.T) {
|
|
cfg := CreateConfig()
|
|
cfg.Providers = map[string]*ProviderConfig{
|
|
"openai": {Type: "openai", APIKey: "k"},
|
|
"anthropic": {Type: "anthropic", APIKey: "k"},
|
|
}
|
|
cfg.Groups = map[string]*GroupConfig{"default": {Providers: []string{"anthropic"}}}
|
|
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", nil)
|
|
req.Header.Set("Authorization", "Bearer sk-alice")
|
|
rec := httptest.NewRecorder()
|
|
h.ServeHTTP(rec, req)
|
|
|
|
if rec.Code != http.StatusForbidden {
|
|
t.Fatalf("status = %d, want 403, body=%s", rec.Code, rec.Body.String())
|
|
}
|
|
}
|
|
|
|
// TestServeHTTP_UnknownProviderPrefix_Returns404 covers a first path
|
|
// segment that names no configured provider: the request must fall
|
|
// through to the ordinary unknown-route 404, never reaching auth.
|
|
func TestServeHTTP_UnknownProviderPrefix_Returns404(t *testing.T) {
|
|
cfg := CreateConfig()
|
|
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
|
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, "/notaprovider/v1/thing", nil)
|
|
rec := httptest.NewRecorder()
|
|
h.ServeHTTP(rec, req)
|
|
|
|
if rec.Code != http.StatusNotFound {
|
|
t.Fatalf("status = %d, want 404, body=%s", rec.Code, rec.Body.String())
|
|
}
|
|
}
|
|
|
|
// TestHandlePassthrough_UpgradeHeader_Returns501 covers a websocket
|
|
// upgrade request against a passthrough route, which this gateway cannot
|
|
// proxy.
|
|
func TestHandlePassthrough_UpgradeHeader_Returns501(t *testing.T) {
|
|
cfg := CreateConfig()
|
|
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
|
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/realtime", nil)
|
|
req.Header.Set("Authorization", "Bearer sk-alice")
|
|
req.Header.Set("Upgrade", "websocket")
|
|
req.Header.Set("Connection", "Upgrade")
|
|
rec := httptest.NewRecorder()
|
|
h.ServeHTTP(rec, req)
|
|
|
|
if rec.Code != http.StatusNotImplemented {
|
|
t.Fatalf("status = %d, want 501, body=%s", rec.Code, rec.Body.String())
|
|
}
|
|
}
|
|
|
|
// TestHandlePassthrough_DeadUpstream_Returns502 covers a configured
|
|
// provider whose upstream is unreachable.
|
|
func TestHandlePassthrough_DeadUpstream_Returns502(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
|
|
deadURL := srv.URL
|
|
srv.Close() // upstream is now unreachable
|
|
|
|
cfg := CreateConfig()
|
|
cfg.Providers = map[string]*ProviderConfig{
|
|
"openai": {Type: "openai", BaseURL: deadURL, 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", nil)
|
|
req.Header.Set("Authorization", "Bearer sk-alice")
|
|
rec := httptest.NewRecorder()
|
|
h.ServeHTTP(rec, req)
|
|
|
|
if rec.Code != http.StatusBadGateway {
|
|
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)
|
|
}
|
|
|
|
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
|
if !ok || tokIn != 7 {
|
|
t.Errorf("user tokin/day counter = %d (ok=%v), want 7 — usage must still be accounted through gzip", tokIn, ok)
|
|
}
|
|
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
|
|
if !ok || tokOut != 3 {
|
|
t.Errorf("user tokout/day counter = %d (ok=%v), want 3 — usage must still be accounted through gzip", tokOut, 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))
|
|
}
|
|
|
|
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
|
if !ok || tokIn != 0 {
|
|
t.Errorf("user tokin/day counter = %d (ok=%v), want 0 — usage accounting must be skipped over the tee cap", tokIn, ok)
|
|
}
|
|
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
|
|
if !ok || tokOut != 0 {
|
|
t.Errorf("user tokout/day counter = %d (ok=%v), want 0 — usage accounting must be skipped over the tee cap", tokOut, 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)
|
|
}
|
|
}
|
|
|
|
// TestHandlePassthrough_ContextDeadlineExceeded_RecordsProviderFailure is
|
|
// the passthrough half of the route-level regression the review's
|
|
// blocker demanded (v0.22, round 2) — a real context deadline against a
|
|
// real sleeping httptest upstream. Passthrough was NEVER actually broken
|
|
// by the multi-%w bug (proxyUpstream's client.Do error reaches
|
|
// attemptRecorderFromContext bare, never re-wrapped by upstreamBytes —
|
|
// see recordProviderAttempt's own call site here, routes_passthrough.go),
|
|
// but the review round proved that by observation, not a route-level
|
|
// assertion; this pins it so a future refactor cannot silently regress
|
|
// it the same way the unified route regressed.
|
|
func TestHandlePassthrough_ContextDeadlineExceeded_RecordsProviderFailure(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
time.Sleep(600 * time.Millisecond) // well past the request's own 60ms deadline below
|
|
}))
|
|
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")
|
|
}
|
|
gw.limiter.spawn = func(f func()) { f() }
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Millisecond)
|
|
defer cancel()
|
|
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(`{}`)).WithContext(ctx)
|
|
req.Header.Set("Authorization", "Bearer sk-alice")
|
|
rec := httptest.NewRecorder()
|
|
h.ServeHTTP(rec, req)
|
|
|
|
now := time.Now()
|
|
attempts, _ := gw.limiter.getCounter(kindProvider, "openai", metricProvAttempt, windowDay, now)
|
|
if attempts != 1 {
|
|
t.Errorf("provider attempts/day = %d, want 1", attempts)
|
|
}
|
|
fails, _ := gw.limiter.getCounter(kindProvider, "openai", metricProvFail, windowDay, now)
|
|
if fails != 1 {
|
|
t.Errorf("provider fails/day = %d, want 1 — a real context-deadline timeout must count as a provider-health failure (SHOULD-1)", fails)
|
|
}
|
|
}
|