mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-20 05:56:56 +00:00
fix(security): bound fan-out and method surface and scrub relayed text
- Add a process-wide gate on MCP tools/list backend connections. The per-request limit alone let N callers open N times as many. - Reject CONNECT and TRACE on native passthrough and on MCP and A2A targets with 405. - Scrub URL query strings and the target URL from a backend JSON-RPC error message before it reaches the caller. - Replace the value of redis.address in GET /admin/api/config. The users file path stays visible by design. - Escape line breaks in every log line so an upstream string cannot forge log entries. - Add stripClientAnthropicBeta. When true, POST /v1/messages does not forward the client anthropic-beta header. The default stays false.
This commit is contained in:
@@ -381,6 +381,7 @@ config accepts them as YAML, which decodes to the same JSON shape.
|
|||||||
| `metrics` | `*MetricsConfig` | `nil` (disabled) | Prometheus text-exposition endpoint — see [Metrics](#metrics). `nil` or `enabled: false` means the metrics route is not registered at all. |
|
| `metrics` | `*MetricsConfig` | `nil` (disabled) | Prometheus text-exposition endpoint — see [Metrics](#metrics). `nil` or `enabled: false` means the metrics route is not registered at all. |
|
||||||
| `modelAliases` | `map[string]string` | `{}` | Operator-defined alias id → target model id — see [Model aliases](#model-aliases). Omitted or empty means no aliases, byte-identical to a gateway built before this field existed. |
|
| `modelAliases` | `map[string]string` | `{}` | Operator-defined alias id → target model id — see [Model aliases](#model-aliases). Omitted or empty means no aliases, byte-identical to a gateway built before this field existed. |
|
||||||
| `requireModelPrefix` | `bool` | `false` | When `true`, a model resolves only as `provider/model` or as a `modelAliases` id. A bare id returns 404 `unknown model` and `GET /v1/models` lists no bare ids. `false` keeps the sorted-first bare-id collision winner. |
|
| `requireModelPrefix` | `bool` | `false` | When `true`, a model resolves only as `provider/model` or as a `modelAliases` id. A bare id returns 404 `unknown model` and `GET /v1/models` lists no bare ids. `false` keeps the sorted-first bare-id collision winner. |
|
||||||
|
| `stripClientAnthropicBeta` | `bool` | `false` | When `true`, `POST /v1/messages` does not forward the client's `anthropic-beta` header to the provider, as native passthrough already does. `false` forwards it, so a client can enable beta features on the operator's own Anthropic account. |
|
||||||
| `modelMeta` | `map[string]*ModelMetaConfig` | `{}` | Per-model context-window and cost overrides, keyed by exact `provider/model` or a bare model id — see [Model metadata](#model-metadata). Omitted or empty means no overrides; every model's metadata still resolves through discovery and the built-in table. |
|
| `modelMeta` | `map[string]*ModelMetaConfig` | `{}` | Per-model context-window and cost overrides, keyed by exact `provider/model` or a bare model id — see [Model metadata](#model-metadata). Omitted or empty means no overrides; every model's metadata still resolves through discovery and the built-in table. |
|
||||||
| `passthroughUnknown` | `bool` | `false` | `false`: a request matching none of the plugin's routes gets a 404 JSON envelope. `true`: it falls through to the router's own backing service. **Does not cover the three media routes** — see the upgrade note below. |
|
| `passthroughUnknown` | `bool` | `false` | `false`: a request matching none of the plugin's routes gets a 404 JSON envelope. `true`: it falls through to the router's own backing service. **Does not cover the three media routes** — see the upgrade note below. |
|
||||||
| `maxInFlightBodyRequests` | `int` | `0` (self-tuned) | Caps how many unified/media requests may concurrently be inside their own request-body read+decode step — see [Body-admission limit](#body-admission-limit). `0` self-tunes from the host's own `GOMAXPROCS`; a positive value pins an exact cap instead, clamped to a maximum of 10,000. |
|
| `maxInFlightBodyRequests` | `int` | `0` (self-tuned) | Caps how many unified/media requests may concurrently be inside their own request-body read+decode step — see [Body-admission limit](#body-admission-limit). `0` self-tunes from the host's own `GOMAXPROCS`; a positive value pins an exact cap instead, clamped to a maximum of 10,000. |
|
||||||
@@ -540,6 +541,7 @@ byte-identically to before multi-group support existed.
|
|||||||
| `enabled` | `bool` | `false` | `false`: no caching, `/v1/chat/completions` and `/v1/embeddings` behave byte-identically to a gateway built before this field existed. `true`: caching per [Caching](#caching) below — requires `redis` to be configured too, or it silently stays off (one warning logged). |
|
| `enabled` | `bool` | `false` | `false`: no caching, `/v1/chat/completions` and `/v1/embeddings` behave byte-identically to a gateway built before this field existed. `true`: caching per [Caching](#caching) below — requires `redis` to be configured too, or it silently stays off (one warning logged). |
|
||||||
| `ttl` | `string` (Go duration) | `5m` | How long a cached response stays valid. A malformed duration string is a construction error. Zero or negative is accepted, not refused — construction logs one warning, since the store floors it to a 1-second TTL rather than actually disabling the cache; set `enabled: false` to turn caching off. |
|
| `ttl` | `string` (Go duration) | `5m` | How long a cached response stays valid. A malformed duration string is a construction error. Zero or negative is accepted, not refused — construction logs one warning, since the store floors it to a 1-second TTL rather than actually disabling the cache; set `enabled: false` to turn caching off. |
|
||||||
| `maxBodyBytes` | `int` | `1048576` (1MiB) | Largest response body still eligible for caching; a larger one is skipped (never stored, never an error). Maximum accepted value is `8388608` (8MiB) — above that is a construction error. |
|
| `maxBodyBytes` | `int` | `1048576` (1MiB) | Largest response body still eligible for caching; a larger one is skipped (never stored, never an error). Maximum accepted value is `8388608` (8MiB) — above that is a construction error. |
|
||||||
|
| `accountHits` | `bool` | `false` | When `true`, a cache hit adds the usage stored with the cached response to the requesting user's and group's token and cost counters, so budgets also bind repeated identical prompts. `false` counts a hit as a request only. |
|
||||||
|
|
||||||
### `BreakerConfig`
|
### `BreakerConfig`
|
||||||
|
|
||||||
|
|||||||
@@ -44,6 +44,14 @@ var redactedConfigKeys = map[string]bool{
|
|||||||
"password": true,
|
"password": true,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// hiddenConfigKeys names JSON keys whose non-empty value is replaced
|
||||||
|
// outright: RedisConfig.Address ("address"), which can embed credentials in
|
||||||
|
// a redis:// URL. UsersConfig.File stays visible on purpose, as a path
|
||||||
|
// reveals where user keys live without holding one.
|
||||||
|
var hiddenConfigKeys = map[string]bool{
|
||||||
|
"address": true,
|
||||||
|
}
|
||||||
|
|
||||||
// baseURLLikeConfigKeys names every JSON object key redactConfig passes
|
// baseURLLikeConfigKeys names every JSON object key redactConfig passes
|
||||||
// through sanitizeBaseURL (admin.go) instead of the literal/env/file
|
// through sanitizeBaseURL (admin.go) instead of the literal/env/file
|
||||||
// redaction rule above: ProviderConfig.BaseURL ("baseUrl") and
|
// redaction rule above: ProviderConfig.BaseURL ("baseUrl") and
|
||||||
@@ -102,6 +110,10 @@ func redactValue(v any) {
|
|||||||
switch {
|
switch {
|
||||||
case redactedConfigKeys[k] && isString:
|
case redactedConfigKeys[k] && isString:
|
||||||
t[k] = redactSecretString(s)
|
t[k] = redactSecretString(s)
|
||||||
|
case hiddenConfigKeys[k] && isString:
|
||||||
|
if s != "" {
|
||||||
|
t[k] = redactedSecretPlaceholder
|
||||||
|
}
|
||||||
case baseURLLikeConfigKeys[k] && isString:
|
case baseURLLikeConfigKeys[k] && isString:
|
||||||
t[k] = sanitizeBaseURL(s)
|
t[k] = sanitizeBaseURL(s)
|
||||||
default:
|
default:
|
||||||
|
|||||||
@@ -107,7 +107,7 @@ func TestRedactConfig_LiteralAndIndirectSecrets(t *testing.T) {
|
|||||||
}
|
}
|
||||||
body := string(raw)
|
body := string(raw)
|
||||||
|
|
||||||
for _, leaked := range []string{"literal-key", "redis-literal-pw", "alice-literal-key", "s3cr3t", "user:pw", "abc123"} {
|
for _, leaked := range []string{"literal-key", "redis-literal-pw", "alice-literal-key", "s3cr3t", "user:pw", "abc123", "redis:6379"} {
|
||||||
if strings.Contains(body, leaked) {
|
if strings.Contains(body, leaked) {
|
||||||
t.Errorf("redacted config still contains %q: %s", leaked, body)
|
t.Errorf("redacted config still contains %q: %s", leaked, body)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -172,6 +172,11 @@ type Config struct {
|
|||||||
// callers that set one — and nothing for a model that has a price.
|
// callers that set one — and nothing for a model that has a price.
|
||||||
// Prefer adding the missing "pricing" entry over setting this.
|
// Prefer adding the missing "pricing" entry over setting this.
|
||||||
AllowUnpricedWithCostBudget bool `json:"allowUnpricedWithCostBudget,omitempty"`
|
AllowUnpricedWithCostBudget bool `json:"allowUnpricedWithCostBudget,omitempty"`
|
||||||
|
// StripClientAnthropicBeta stops POST /v1/messages from forwarding the
|
||||||
|
// client's anthropic-beta header to the provider, as native passthrough
|
||||||
|
// already does. Without it a tenant can enable beta features on the
|
||||||
|
// operator's own Anthropic account. Defaults to false.
|
||||||
|
StripClientAnthropicBeta bool `json:"stripClientAnthropicBeta,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// ProviderConfig describes one upstream LLM provider.
|
// ProviderConfig describes one upstream LLM provider.
|
||||||
@@ -767,6 +772,7 @@ type Gateway struct {
|
|||||||
events *eventLog
|
events *eventLog
|
||||||
bodyAdmission chan struct{}
|
bodyAdmission chan struct{}
|
||||||
teeBudget *teeBudget
|
teeBudget *teeBudget
|
||||||
|
fanoutGate chan struct{}
|
||||||
failoverHealth *requestHealthTracker
|
failoverHealth *requestHealthTracker
|
||||||
targetHealth *targetHealthTracker
|
targetHealth *targetHealthTracker
|
||||||
limiter *limiter
|
limiter *limiter
|
||||||
@@ -1056,6 +1062,7 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str
|
|||||||
}
|
}
|
||||||
g.bodyAdmission = make(chan struct{}, bodyAdmissionCap)
|
g.bodyAdmission = make(chan struct{}, bodyAdmissionCap)
|
||||||
g.teeBudget = newTeeBudget(maxAccountingTeeTotalBytes)
|
g.teeBudget = newTeeBudget(maxAccountingTeeTotalBytes)
|
||||||
|
g.fanoutGate = make(chan struct{}, mcpFederatedFanoutGlobalConcurrency)
|
||||||
|
|
||||||
if config.Users != nil && config.Users.File != "" {
|
if config.Users != nil && config.Users.File != "" {
|
||||||
if err = attachUsersFile(auth, config.Users.File, g); err != nil {
|
if err = attachUsersFile(auth, config.Users.File, g); err != nil {
|
||||||
|
|||||||
@@ -4,12 +4,22 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
|
"strings"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// logSafe escapes the line breaks in msg. A value an upstream controls, such
|
||||||
|
// as a panic or error string, must not forge extra log lines.
|
||||||
|
func logSafe(msg string) string {
|
||||||
|
if !strings.ContainsAny(msg, "\r\n") {
|
||||||
|
return msg
|
||||||
|
}
|
||||||
|
return strings.NewReplacer("\r", "\\r", "\n", "\\n").Replace(msg)
|
||||||
|
}
|
||||||
|
|
||||||
// logf writes an info-level log line to stderr. No timestamps — Traefik
|
// logf writes an info-level log line to stderr. No timestamps — Traefik
|
||||||
// adds its own. Never log key material.
|
// adds its own. Never log key material.
|
||||||
func (g *Gateway) logf(format string, args ...any) {
|
func (g *Gateway) logf(format string, args ...any) {
|
||||||
fmt.Fprintf(os.Stderr, "llmgw[%s] INFO %s\n", g.name, fmt.Sprintf(format, args...))
|
fmt.Fprintf(os.Stderr, "llmgw[%s] INFO %s\n", g.name, logSafe(fmt.Sprintf(format, args...)))
|
||||||
}
|
}
|
||||||
|
|
||||||
// warnf writes a warning-level log line to stderr. No timestamps —
|
// warnf writes a warning-level log line to stderr. No timestamps —
|
||||||
@@ -30,7 +40,7 @@ func (g *Gateway) logf(format string, args ...any) {
|
|||||||
// see noteConfigWarning's own doc comment for why.
|
// see noteConfigWarning's own doc comment for why.
|
||||||
func (g *Gateway) warnf(format string, args ...any) {
|
func (g *Gateway) warnf(format string, args ...any) {
|
||||||
msg := fmt.Sprintf(format, args...)
|
msg := fmt.Sprintf(format, args...)
|
||||||
fmt.Fprintf(os.Stderr, "llmgw[%s] WARN %s\n", g.name, msg)
|
fmt.Fprintf(os.Stderr, "llmgw[%s] WARN %s\n", g.name, logSafe(msg))
|
||||||
g.noteConfigWarning(msg)
|
g.noteConfigWarning(msg)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -80,7 +90,7 @@ func (g *Gateway) configWarningsSnapshot() (warnings []string, dropped int) {
|
|||||||
// errorf writes an error-level log line to stderr. No timestamps — Traefik
|
// errorf writes an error-level log line to stderr. No timestamps — Traefik
|
||||||
// adds its own. Never log key material.
|
// adds its own. Never log key material.
|
||||||
func (g *Gateway) errorf(format string, args ...any) {
|
func (g *Gateway) errorf(format string, args ...any) {
|
||||||
fmt.Fprintf(os.Stderr, "llmgw[%s] ERROR %s\n", g.name, fmt.Sprintf(format, args...))
|
fmt.Fprintf(os.Stderr, "llmgw[%s] ERROR %s\n", g.name, logSafe(fmt.Sprintf(format, args...)))
|
||||||
}
|
}
|
||||||
|
|
||||||
// pricingWarn is g's own per-instance pricing-warning sink, passed to
|
// pricingWarn is g's own per-instance pricing-warning sink, passed to
|
||||||
|
|||||||
@@ -196,6 +196,9 @@ func (g *Gateway) handleTargetProxy(w http.ResponseWriter, r *http.Request, u *u
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if !proxyMethodAllowed(w, r) {
|
||||||
|
return
|
||||||
|
}
|
||||||
if r.Header.Get("Upgrade") != "" {
|
if r.Header.Get("Upgrade") != "" {
|
||||||
writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported")
|
writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported")
|
||||||
return
|
return
|
||||||
|
|||||||
+40
-1
@@ -9,6 +9,7 @@ import (
|
|||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/url"
|
"net/url"
|
||||||
|
"regexp"
|
||||||
"sort"
|
"sort"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
@@ -1137,6 +1138,10 @@ func lastSSEDataLine(body []byte) []byte {
|
|||||||
// concurrency or backend count — see its own doc comment.
|
// concurrency or backend count — see its own doc comment.
|
||||||
const mcpFederatedFanoutConcurrency = 8
|
const mcpFederatedFanoutConcurrency = 8
|
||||||
|
|
||||||
|
// mcpFederatedFanoutGlobalConcurrency caps the backend connections every
|
||||||
|
// concurrent federated tools/list call may hold together.
|
||||||
|
const mcpFederatedFanoutGlobalConcurrency = 32
|
||||||
|
|
||||||
// mcpFederatedToolsList implements the federated "tools/list": it fans
|
// mcpFederatedToolsList implements the federated "tools/list": it fans
|
||||||
// out one per-server fetch (below) per server in names (handleMCPFederated's
|
// out one per-server fetch (below) per server in names (handleMCPFederated's
|
||||||
// own allowedMCPServerNames(grp) call, made once there so the exact set
|
// own allowedMCPServerNames(grp) call, made once there so the exact set
|
||||||
@@ -1281,6 +1286,23 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, format mcpRespons
|
|||||||
}
|
}
|
||||||
defer func() { <-sem }()
|
defer func() { <-sem }()
|
||||||
|
|
||||||
|
// The per-request semaphore above bounds one tools/list call;
|
||||||
|
// fanoutGate bounds every concurrent call together, so N
|
||||||
|
// callers cannot open N x mcpFederatedFanoutConcurrency
|
||||||
|
// connections to the backends. A nil gate (a Gateway built
|
||||||
|
// without New) is unbounded, as before.
|
||||||
|
if g.fanoutGate != nil {
|
||||||
|
select {
|
||||||
|
case g.fanoutGate <- struct{}{}:
|
||||||
|
defer func() { <-g.fanoutGate }()
|
||||||
|
case <-fanoutCtx.Done():
|
||||||
|
mu.Lock()
|
||||||
|
failed = append(failed, name)
|
||||||
|
mu.Unlock()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// F10: this goroutine won the semaphore, so it is about to
|
// F10: this goroutine won the semaphore, so it is about to
|
||||||
// actually dial the server — recorded now, unconditionally,
|
// actually dial the server — recorded now, unconditionally,
|
||||||
// regardless of what the paginated fetch below returns, so
|
// regardless of what the paginated fetch below returns, so
|
||||||
@@ -1629,5 +1651,22 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, format mcpRespons
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
writeJSONRPCEnvelope(w, format, jsonrpcResponse{ID: req.ID, Result: resp.Result, Error: resp.Error})
|
relayed := resp.Error
|
||||||
|
if relayed != nil {
|
||||||
|
scrubbed := *relayed
|
||||||
|
scrubbed.Message = scrubRelayedBackendMessage(relayed.Message, g.cfg.MCPServers[serverName].URL)
|
||||||
|
relayed = &scrubbed
|
||||||
|
}
|
||||||
|
writeJSONRPCEnvelope(w, format, jsonrpcResponse{ID: req.ID, Result: resp.Result, Error: relayed})
|
||||||
|
}
|
||||||
|
|
||||||
|
// urlQueryPattern matches the query string of any URL in free text.
|
||||||
|
var urlQueryPattern = regexp.MustCompile(`(https?://[^\s?#"']+)\?[^\s#"']*`)
|
||||||
|
|
||||||
|
// scrubRelayedBackendMessage removes credentials a backend's own error text
|
||||||
|
// can carry before the text reaches a caller: the configured target URL in
|
||||||
|
// any form, and the query string of every URL, where API keys travel.
|
||||||
|
func scrubRelayedBackendMessage(msg, targetURL string) string {
|
||||||
|
msg = sanitizeProviderErr(msg, targetURL)
|
||||||
|
return urlQueryPattern.ReplaceAllString(msg, "$1?[redacted]")
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ import (
|
|||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
"strconv"
|
"strconv"
|
||||||
|
"strings"
|
||||||
)
|
)
|
||||||
|
|
||||||
// messagesPath is the inbound Anthropic Messages API route this gateway
|
// messagesPath is the inbound Anthropic Messages API route this gateway
|
||||||
@@ -212,6 +213,9 @@ func (g *Gateway) callAnthropicMessagesPassthrough(ctx context.Context, w http.R
|
|||||||
hdr := http.Header{}
|
hdr := http.Header{}
|
||||||
hdr.Set("Content-Type", "application/json")
|
hdr.Set("Content-Type", "application/json")
|
||||||
for _, name := range anthropicPassthroughForwardedHeaders {
|
for _, name := range anthropicPassthroughForwardedHeaders {
|
||||||
|
if g.cfg.StripClientAnthropicBeta && strings.EqualFold(name, "anthropic-beta") {
|
||||||
|
continue
|
||||||
|
}
|
||||||
for _, v := range clientHeaders.Values(name) {
|
for _, v := range clientHeaders.Values(name) {
|
||||||
hdr.Add(name, v)
|
hdr.Add(name, v)
|
||||||
}
|
}
|
||||||
|
|||||||
+17
-1
@@ -423,8 +423,8 @@ func hasTraversalSegment(rest string) bool {
|
|||||||
// parse afterward. truncated records whether the cap was hit, so the
|
// 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.
|
// caller knows a parse would only fail on cut-off JSON and skips it.
|
||||||
type cappedAccountingBuffer struct {
|
type cappedAccountingBuffer struct {
|
||||||
buf bytes.Buffer
|
|
||||||
budget *teeBudget
|
budget *teeBudget
|
||||||
|
buf bytes.Buffer
|
||||||
reserved int64
|
reserved int64
|
||||||
truncated bool
|
truncated bool
|
||||||
}
|
}
|
||||||
@@ -1225,6 +1225,9 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if !proxyMethodAllowed(w, r) {
|
||||||
|
return
|
||||||
|
}
|
||||||
if r.Header.Get("Upgrade") != "" {
|
if r.Header.Get("Upgrade") != "" {
|
||||||
writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported")
|
writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported")
|
||||||
return
|
return
|
||||||
@@ -1798,3 +1801,16 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream
|
|||||||
}
|
}
|
||||||
return proxyResult{isJSON: isJSON, isSSE: isSSE, tee: tee, sseTee: sseTee, status: resp.StatusCode, requestBytes: requestBytes}, true
|
return proxyResult{isJSON: isJSON, isSSE: isSSE, tee: tee, sseTee: sseTee, status: resp.StatusCode, requestBytes: requestBytes}, true
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// proxyMethodAllowed rejects the methods a reverse proxy to a configured
|
||||||
|
// backend must never relay. CONNECT opens a tunnel and TRACE echoes the
|
||||||
|
// request back, and no supported provider or target needs either.
|
||||||
|
func proxyMethodAllowed(w http.ResponseWriter, r *http.Request) bool {
|
||||||
|
switch r.Method {
|
||||||
|
case http.MethodConnect, http.MethodTrace:
|
||||||
|
w.Header().Set("Allow", "GET, HEAD, POST, PUT, PATCH, DELETE, OPTIONS")
|
||||||
|
writeOAIError(w, http.StatusMethodNotAllowed, "invalid_request_error", "method not allowed")
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|||||||
@@ -2,9 +2,12 @@ package traefikllmgateway
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/json"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
@@ -178,3 +181,173 @@ func TestCacheHit_AccountHits_ChargesServingUser(t *testing.T) {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestMessages_StripClientAnthropicBeta(t *testing.T) {
|
||||||
|
for _, tt := range []struct {
|
||||||
|
name string
|
||||||
|
wantBeta string
|
||||||
|
strip bool
|
||||||
|
}{
|
||||||
|
{"default forwards the header", "context-1m", false},
|
||||||
|
{"flag strips the header", "", true},
|
||||||
|
} {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
var gotBeta string
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
gotBeta = r.Header.Get("anthropic-beta")
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = w.Write([]byte(`{"id":"m","type":"message","role":"assistant","content":[],"model":"claude-real","usage":{"input_tokens":1,"output_tokens":1}}`))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
gw := newMessagesTestGateway(t, srv.URL, "")
|
||||||
|
gw.cfg.StripClientAnthropicBeta = tt.strip
|
||||||
|
|
||||||
|
body := map[string]any{"model": "claude-test", "max_tokens": 10, "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
|
||||||
|
req := newMessagesRequest(t, body, "x-api-key", "sk-alice")
|
||||||
|
req.Header.Set("anthropic-beta", "context-1m")
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
gw.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||||
|
assert.Equal(t, tt.wantBeta, gotBeta)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProxyMethodAllowed(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
tests := []struct {
|
||||||
|
method string
|
||||||
|
allowed bool
|
||||||
|
}{
|
||||||
|
{http.MethodGet, true},
|
||||||
|
{http.MethodPost, true},
|
||||||
|
{http.MethodPut, true},
|
||||||
|
{http.MethodPatch, true},
|
||||||
|
{http.MethodDelete, true},
|
||||||
|
{http.MethodOptions, true},
|
||||||
|
{http.MethodConnect, false},
|
||||||
|
{http.MethodTrace, false},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.method, func(t *testing.T) {
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
got := proxyMethodAllowed(rec, httptest.NewRequest(tt.method, "/x", nil))
|
||||||
|
assert.Equal(t, tt.allowed, got)
|
||||||
|
if !tt.allowed {
|
||||||
|
assert.Equal(t, http.StatusMethodNotAllowed, rec.Code)
|
||||||
|
assert.NotEmpty(t, rec.Header().Get("Allow"))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandlePassthrough_TraceMethod_Returns405_UpstreamUntouched(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
called := false
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { called = true }))
|
||||||
|
defer srv.Close()
|
||||||
|
gw := newPassthroughUsageGateway(t, `{}`, nil, &LimitsConfig{}, nil)
|
||||||
|
gw.cfg.Providers["openai"].BaseURL = srv.URL
|
||||||
|
|
||||||
|
req := httptest.NewRequest(http.MethodTrace, "/openai/v1/x", nil)
|
||||||
|
req.Header.Set("Authorization", "Bearer sk-alice")
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
gw.ServeHTTP(rec, req)
|
||||||
|
|
||||||
|
assert.Equal(t, http.StatusMethodNotAllowed, rec.Code)
|
||||||
|
assert.False(t, called)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestScrubRelayedBackendMessage(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
tests := []struct {
|
||||||
|
name, in, target, want string
|
||||||
|
}{
|
||||||
|
{"plain text untouched", "tool execution failed", "http://b", "tool execution failed"},
|
||||||
|
{"query string redacted", "dial http://backend.internal?api-key=SECRET rejected", "http://other", "dial http://backend.internal?[redacted] rejected"},
|
||||||
|
{"configured url with query redacted", "GET http://b/mcp?token=abc failed", "http://b/mcp?token=abc", "GET http://b/mcp failed"},
|
||||||
|
{"empty", "", "http://b", ""},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
got := scrubRelayedBackendMessage(tt.in, tt.target)
|
||||||
|
assert.Equal(t, tt.want, got)
|
||||||
|
assert.NotContains(t, got, "SECRET")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandleMCPFederated_ToolsCall_BackendErrorSecretScrubbed(t *testing.T) {
|
||||||
|
alpha := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
var req jsonrpcRequest
|
||||||
|
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "dial http://backend.internal?api-key=SECRET rejected"}})
|
||||||
|
}))
|
||||||
|
defer alpha.Close()
|
||||||
|
cfg := newFederationTestConfig(alpha.URL, "http://beta.invalid", false)
|
||||||
|
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_lookup"})
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("5"), Params: callParams}))
|
||||||
|
|
||||||
|
require.Equal(t, http.StatusOK, rec.Code)
|
||||||
|
assert.NotContains(t, rec.Body.String(), "SECRET")
|
||||||
|
assert.Contains(t, rec.Body.String(), "[redacted]")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandleMCPFederated_ToolsList_GlobalFanoutGateBoundsConnections(t *testing.T) {
|
||||||
|
var inFlight, peak int32
|
||||||
|
release := make(chan struct{})
|
||||||
|
backend := func() *httptest.Server {
|
||||||
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
n := atomic.AddInt32(&inFlight, 1)
|
||||||
|
for {
|
||||||
|
p := atomic.LoadInt32(&peak)
|
||||||
|
if n <= p || atomic.CompareAndSwapInt32(&peak, p, n) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
<-release
|
||||||
|
atomic.AddInt32(&inFlight, -1)
|
||||||
|
var req jsonrpcRequest
|
||||||
|
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: json.RawMessage(`{"tools":[]}`)})
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
alpha, beta := backend(), backend()
|
||||||
|
defer alpha.Close()
|
||||||
|
defer beta.Close()
|
||||||
|
|
||||||
|
cfg := newFederationTestConfig(alpha.URL, beta.URL, false)
|
||||||
|
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||||
|
require.NoError(t, err)
|
||||||
|
gw, ok := h.(*Gateway)
|
||||||
|
require.True(t, ok)
|
||||||
|
gw.fanoutGate = make(chan struct{}, 3)
|
||||||
|
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
for i := 0; i < 5; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}))
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
time.Sleep(300 * time.Millisecond)
|
||||||
|
assert.LessOrEqual(t, atomic.LoadInt32(&peak), int32(3), "5 callers x 2 backends must not exceed the global gate")
|
||||||
|
close(release)
|
||||||
|
wg.Wait()
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLogSafe(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
assert.Equal(t, "plain", logSafe("plain"))
|
||||||
|
assert.Equal(t, `a\nllmgw[x] ERROR forged`, logSafe("a\nllmgw[x] ERROR forged"))
|
||||||
|
assert.Equal(t, `a\r\nb`, logSafe("a\r\nb"))
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user