diff --git a/README.md b/README.md index c4408f3..d8f2ee1 100644 --- a/README.md +++ b/README.md @@ -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. | | `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. | +| `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. | | `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. | @@ -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). | | `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. | +| `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` diff --git a/admin_config.go b/admin_config.go index bf1374b..69e47a8 100644 --- a/admin_config.go +++ b/admin_config.go @@ -44,6 +44,14 @@ var redactedConfigKeys = map[string]bool{ "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 // through sanitizeBaseURL (admin.go) instead of the literal/env/file // redaction rule above: ProviderConfig.BaseURL ("baseUrl") and @@ -102,6 +110,10 @@ func redactValue(v any) { switch { case redactedConfigKeys[k] && isString: t[k] = redactSecretString(s) + case hiddenConfigKeys[k] && isString: + if s != "" { + t[k] = redactedSecretPlaceholder + } case baseURLLikeConfigKeys[k] && isString: t[k] = sanitizeBaseURL(s) default: diff --git a/admin_config_test.go b/admin_config_test.go index ec2f301..ab54cdd 100644 --- a/admin_config_test.go +++ b/admin_config_test.go @@ -107,7 +107,7 @@ func TestRedactConfig_LiteralAndIndirectSecrets(t *testing.T) { } 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) { t.Errorf("redacted config still contains %q: %s", leaked, body) } diff --git a/llmgateway.go b/llmgateway.go index 9677f9d..1a24f89 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -172,6 +172,11 @@ type Config struct { // callers that set one — and nothing for a model that has a price. // Prefer adding the missing "pricing" entry over setting this. 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. @@ -767,6 +772,7 @@ type Gateway struct { events *eventLog bodyAdmission chan struct{} teeBudget *teeBudget + fanoutGate chan struct{} failoverHealth *requestHealthTracker targetHealth *targetHealthTracker 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.teeBudget = newTeeBudget(maxAccountingTeeTotalBytes) + g.fanoutGate = make(chan struct{}, mcpFederatedFanoutGlobalConcurrency) if config.Users != nil && config.Users.File != "" { if err = attachUsersFile(auth, config.Users.File, g); err != nil { diff --git a/logger.go b/logger.go index 7acf3d3..b254f7c 100644 --- a/logger.go +++ b/logger.go @@ -4,12 +4,22 @@ import ( "fmt" "net/http" "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 // adds its own. Never log key material. 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 — @@ -30,7 +40,7 @@ func (g *Gateway) logf(format string, args ...any) { // see noteConfigWarning's own doc comment for why. func (g *Gateway) warnf(format string, args ...any) { 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) } @@ -80,7 +90,7 @@ func (g *Gateway) configWarningsSnapshot() (warnings []string, dropped int) { // errorf writes an error-level log line to stderr. No timestamps — Traefik // adds its own. Never log key material. 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 diff --git a/mcp_a2a.go b/mcp_a2a.go index c1a545f..b18af2d 100644 --- a/mcp_a2a.go +++ b/mcp_a2a.go @@ -196,6 +196,9 @@ func (g *Gateway) handleTargetProxy(w http.ResponseWriter, r *http.Request, u *u return } + if !proxyMethodAllowed(w, r) { + return + } if r.Header.Get("Upgrade") != "" { writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported") return diff --git a/mcp_federation.go b/mcp_federation.go index a3735a4..2ee423d 100644 --- a/mcp_federation.go +++ b/mcp_federation.go @@ -9,6 +9,7 @@ import ( "io" "net/http" "net/url" + "regexp" "sort" "strings" "sync" @@ -1137,6 +1138,10 @@ func lastSSEDataLine(body []byte) []byte { // concurrency or backend count — see its own doc comment. 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 // out one per-server fetch (below) per server in names (handleMCPFederated's // 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 }() + // 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 // actually dial the server — recorded now, unconditionally, // regardless of what the paginated fetch below returns, so @@ -1629,5 +1651,22 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, format mcpRespons 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]") } diff --git a/routes_messages.go b/routes_messages.go index 416c578..f6a82d4 100644 --- a/routes_messages.go +++ b/routes_messages.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "strconv" + "strings" ) // 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.Set("Content-Type", "application/json") for _, name := range anthropicPassthroughForwardedHeaders { + if g.cfg.StripClientAnthropicBeta && strings.EqualFold(name, "anthropic-beta") { + continue + } for _, v := range clientHeaders.Values(name) { hdr.Add(name, v) } diff --git a/routes_passthrough.go b/routes_passthrough.go index 2522b90..cdbfb05 100644 --- a/routes_passthrough.go +++ b/routes_passthrough.go @@ -423,8 +423,8 @@ func hasTraversalSegment(rest string) bool { // parse afterward. truncated records whether the cap was hit, so the // caller knows a parse would only fail on cut-off JSON and skips it. type cappedAccountingBuffer struct { - buf bytes.Buffer budget *teeBudget + buf bytes.Buffer reserved int64 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") != "" { writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported") 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 } + +// 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 +} diff --git a/security_fixes_test.go b/security_fixes_test.go index b180372..86b9c4c 100644 --- a/security_fixes_test.go +++ b/security_fixes_test.go @@ -2,9 +2,12 @@ package traefikllmgateway import ( "context" + "encoding/json" "net/http" "net/http/httptest" "strings" + "sync" + "sync/atomic" "testing" "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")) +}