diff --git a/llmgateway.go b/llmgateway.go index aab5c9d..d766fd2 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -268,6 +268,18 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } + if providerName, rest, ok := passthroughRoute(r.URL.Path); ok { + if _, known := g.adapters[providerName]; known { + u, grp, ok := g.auth.identify(r) + if !ok { + writeOAIError(sw, http.StatusUnauthorized, "authentication_error", "invalid or missing API key") + return + } + g.handlePassthrough(sw, r, u, grp, providerName, rest) + return + } + } + if g.cfg.PassthroughUnknown { g.next.ServeHTTP(sw, r) return diff --git a/provider_anthropic.go b/provider_anthropic.go index 3d24b50..50f6184 100644 --- a/provider_anthropic.go +++ b/provider_anthropic.go @@ -52,6 +52,9 @@ func (a *anthropicAdapter) typeName() string { return providerTypeAnthropic } // base implements providerAdapter. func (a *anthropicAdapter) base() string { return a.baseURL } +// httpClient implements providerAdapter. +func (a *anthropicAdapter) httpClient() *http.Client { return a.client } + // injectAuth implements providerAdapter: Anthropic's x-api-key scheme, // plus the anthropic-version header every request must carry. Unlike // openai-type's injectAuth, this is never a no-op — the constructor diff --git a/provider_gemini.go b/provider_gemini.go index dfe414a..21b4a71 100644 --- a/provider_gemini.go +++ b/provider_gemini.go @@ -45,6 +45,9 @@ func (a *geminiAdapter) typeName() string { return providerTypeGemini } // base implements providerAdapter. func (a *geminiAdapter) base() string { return a.baseURL } +// httpClient implements providerAdapter. +func (a *geminiAdapter) httpClient() *http.Client { return a.client } + // injectAuth implements providerAdapter: Gemini's x-goog-api-key header // scheme. Never a no-op — the constructor already rejected an empty apiKey. func (a *geminiAdapter) injectAuth(r *http.Request) { diff --git a/provider_openai.go b/provider_openai.go index b326968..dd9bce7 100644 --- a/provider_openai.go +++ b/provider_openai.go @@ -41,6 +41,9 @@ func (a *openaiAdapter) typeName() string { return providerTypeOpenAI } // base implements providerAdapter. func (a *openaiAdapter) base() string { return a.baseURL } +// httpClient implements providerAdapter. +func (a *openaiAdapter) httpClient() *http.Client { return a.client } + // injectAuth implements providerAdapter: OpenAI's bearer-token scheme. A // no-op when a.apiKey is empty, per the keyless-upstream ruling — a request // against a keyless upstream must carry no Authorization header at all, diff --git a/providers.go b/providers.go index 5612b9b..4172870 100644 --- a/providers.go +++ b/providers.go @@ -109,6 +109,13 @@ type providerAdapter interface { embeddings(ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) // listModels returns the provider's available model ids, for discovery. listModels(ctx context.Context) ([]string, error) + // httpClient returns the adapter's shared *http.Client, so a caller + // outside this file (the native passthrough route, routes_passthrough.go) + // can issue upstream requests through the same connection-pooled client + // every other adapter call uses, instead of building a fresh one per + // request. Named httpClient rather than client to avoid colliding with + // every implementation's own "client" struct field of the same name. + httpClient() *http.Client } // newAdapterHTTPClient returns an *http.Client for a provider adapter to diff --git a/registry_test.go b/registry_test.go index 0750cdb..5f42f67 100644 --- a/registry_test.go +++ b/registry_test.go @@ -40,6 +40,7 @@ func (f *fakeAdapter) listModels(ctx context.Context) ([]string, error) { } return f.listModelsFn(ctx) } +func (f *fakeAdapter) httpClient() *http.Client { return nil } // recordingLog is a thread-safe log func for tests, capturing formatted // lines instead of writing to stderr. diff --git a/routes_passthrough.go b/routes_passthrough.go new file mode 100644 index 0000000..eb13a22 --- /dev/null +++ b/routes_passthrough.go @@ -0,0 +1,237 @@ +package traefikllmgateway + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "strings" +) + +// maxPassthroughBytes caps a client passthrough request body: 32MiB, +// generous enough for a native multimodal payload (audio, image, or a +// large document) that the unified routes' 10MiB maxRequestBytes does not +// need to accommodate, while still bounding memory against an oversized +// or malicious body. +const maxPassthroughBytes = 32 << 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 +// message, and must never be relayed by an intermediary. +var hopByHopHeaders = map[string]bool{ + "Connection": true, + "Keep-Alive": true, + "Transfer-Encoding": true, + "Upgrade": true, + "Proxy-Authorization": true, + "Proxy-Connection": true, + "Te": true, + "Trailer": true, +} + +// gatewayCredentialHeaders are the client's own gateway-authentication +// headers — auth.go's presentedKey reads exactly these two. Stripped from +// the outgoing upstream request so a client's gateway API key never +// reaches the upstream provider; adapter.injectAuth sets the provider's +// own credential afterward. +var gatewayCredentialHeaders = map[string]bool{ + "Authorization": true, + "X-Api-Key": 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 +// Transport are already canonicalized by net/http, so a plain map lookup +// against the canonical names above is exact. +func copyHeadersExcept(dst, src http.Header, excepts ...map[string]bool) { + for k, vs := range src { + skip := false + for _, ex := range excepts { + if ex[k] { + skip = true + break + } + } + if skip { + continue + } + dst[k] = append([]string(nil), vs...) + } +} + +// 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. +func passthroughRoute(path string) (providerName, rest string, ok bool) { + trimmed := strings.TrimPrefix(path, "/") + if trimmed == "" { + return "", "", false + } + providerName, rest, _ = strings.Cut(trimmed, "/") + if providerName == "" { + return "", "", false + } + return providerName, rest, true +} + +// 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 +// carry (Gemini's native response carries neither "model" nor any of the +// URL-supplied model id). +type passthroughUsagePayload struct { + Model string `json:"model"` + Usage struct { + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + } `json:"usage"` + UsageMetadata struct { + PromptTokenCount int64 `json:"promptTokenCount"` + CandidatesTokenCount int64 `json:"candidatesTokenCount"` + } `json:"usageMetadata"` +} + +// 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) { + var payload passthroughUsagePayload + _ = json.Unmarshal(body, &payload) // best-effort; zero usage/empty model on decode failure + + model := payload.Model + if model == "" { + model = "unknown/" + providerName + } + + switch typeName { + case providerTypeOpenAI: + return usage{prompt: payload.Usage.PromptTokens, completion: payload.Usage.CompletionTokens}, model + case providerTypeAnthropic: + return usage{prompt: payload.Usage.InputTokens, completion: payload.Usage.OutputTokens}, model + case providerTypeGemini: + return usage{prompt: payload.UsageMetadata.PromptTokenCount, completion: payload.UsageMetadata.CandidatesTokenCount}, model + default: + return usage{}, model + } +} + +// handlePassthrough implements the native provider passthrough route: +// "/{providerName}/{rest...}" reverse-proxies rest, verbatim, to +// providerName's configured upstream, with the client's gateway +// credential swapped for the provider's own. It is a raw proxy, not a +// translation layer like the unified routes — a non-2xx upstream status +// is forwarded to the client exactly as the upstream sent it, not wrapped +// in the gateway's own error envelope the way the unified routes wrap a +// providerHTTPError. +// +// 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 +// structure, including its "/v1beta/models/{model}:generateContent" +// 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 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") + return + } + + scopes := buildLimitScopes(u, grp) + if violation := g.limiter.checkAndCount(scopes); violation != nil { + writeLimitViolation(w, violation) + return + } + + // providerName is only ever passed here by ServeHTTP's dispatch, which + // already confirmed it is a key of g.adapters before calling in. + adapter := g.adapters[providerName] + + upstreamURL := adapter.base() + "/" + rest + if r.URL.RawQuery != "" { + upstreamURL += "?" + r.URL.RawQuery + } + + bodyReader := io.LimitReader(r.Body, maxPassthroughBytes) + // gosec G704 (SSRF via taint analysis) flags upstreamURL as + // request-derived: it is, by design — this is a reverse proxy, and its + // whole job is to forward a client-supplied path/query onto the + // upstream. The host component is never request-derived: adapter.base() + // is a fixed, operator-configured URL from Config.Providers, and + // 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 + if err != nil { + g.errorf("passthrough: build upstream request (provider %q): %v", providerName, err) + writeOAIError(w, http.StatusBadGateway, "server_error", "upstream connection error") + return + } + upstreamReq.ContentLength = r.ContentLength + if r.ContentLength < 0 || r.ContentLength > maxPassthroughBytes { + // Unknown or over-cap length: let the transport negotiate chunked + // transfer instead of advertising a Content-Length the capped + // bodyReader above may not actually deliver. + upstreamReq.ContentLength = -1 + } + copyHeadersExcept(upstreamReq.Header, r.Header, hopByHopHeaders, gatewayCredentialHeaders) + adapter.injectAuth(upstreamReq) + + resp, err := adapter.httpClient().Do(upstreamReq) //nolint:gosec // same upstreamReq built above; host is operator-fixed, 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) + return + } + g.errorf("passthrough: upstream connection error (provider %q): %v", providerName, err) + writeOAIError(w, http.StatusBadGateway, "server_error", "upstream connection error") + return + } + defer resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure + + copyHeadersExcept(w.Header(), resp.Header, hopByHopHeaders) + 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) + 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) + } +} diff --git a/routes_passthrough_test.go b/routes_passthrough_test.go new file mode 100644 index 0000000..4769578 --- /dev/null +++ b/routes_passthrough_test.go @@ -0,0 +1,494 @@ +package traefikllmgateway + +import ( + "bytes" + "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 + }{ + { + 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", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotUsage, gotModel := 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) + } + }) + } +} + +// --- 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) + } +} + +// 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") + } + + 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) + } + + 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", tok, 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()) + } +}