mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-16 05:37:28 +00:00
Merge pull request #6 from lukaszraczylo/security/fixes-oct-2026
fix(security): harden accounting, fan-out and relayed text
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. |
|
||||
| `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`
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -1824,6 +1824,23 @@ func (l *limiter) accountWith(scopes []limitScope, u usage, costMicros int64, x
|
||||
// misbehaving or hostile upstream is visible rather than silently
|
||||
// rewriting a budget. Clamped per-component, so one bad field does not
|
||||
// discard the other two, which may be perfectly valid.
|
||||
// The positive direction needs a ceiling too: one absurd figure from an
|
||||
// upstream (or a saturated cost) would otherwise put a counter past
|
||||
// every budget until its window, up to 400 days on Redis, expires. No
|
||||
// real request reports figures anywhere near these bounds, so an
|
||||
// out-of-range figure is dropped like a negative one.
|
||||
if u.prompt > maxUsageTokensPerRequest || u.completion > maxUsageTokensPerRequest || costMicros > maxCostMicrosPerRequest {
|
||||
l.logf("%s", fmt.Sprintf("limits: dropping implausible usage reported by an upstream (prompt=%d completion=%d costMicros=%d)", u.prompt, u.completion, costMicros))
|
||||
if u.prompt > maxUsageTokensPerRequest {
|
||||
u.prompt = 0
|
||||
}
|
||||
if u.completion > maxUsageTokensPerRequest {
|
||||
u.completion = 0
|
||||
}
|
||||
if costMicros > maxCostMicrosPerRequest {
|
||||
costMicros = 0
|
||||
}
|
||||
}
|
||||
if u.prompt < 0 || u.completion < 0 || costMicros < 0 {
|
||||
l.logf("%s", fmt.Sprintf("limits: dropping negative usage reported by an upstream (prompt=%d completion=%d costMicros=%d); these counters are increment-only", u.prompt, u.completion, costMicros))
|
||||
if u.prompt < 0 {
|
||||
@@ -2250,6 +2267,15 @@ const (
|
||||
// scope can never collide with a user, group, provider or target scope.
|
||||
const kindModel = "model"
|
||||
|
||||
// maxUsageTokensPerRequest and maxCostMicrosPerRequest bound what one
|
||||
// upstream response may add to a token or cost counter. 2^30 tokens is
|
||||
// far above any model context window, and 2^40 micro-USD is about
|
||||
// 1.1 million USD.
|
||||
const (
|
||||
maxUsageTokensPerRequest = 1 << 30
|
||||
maxCostMicrosPerRequest = 1 << 40
|
||||
)
|
||||
|
||||
// maxModelScopeIDLen bounds the "provider/model" id withModelScope will
|
||||
// accept as a counter-scope id. Generous next to any real model id —
|
||||
// the longest in the built-in table is well under 40 bytes — and far
|
||||
|
||||
+19
-5
@@ -55,11 +55,6 @@ type Config struct {
|
||||
// (validateModelAliases, registry.go).
|
||||
ModelAliases map[string]string `json:"modelAliases,omitempty"`
|
||||
Providers map[string]*ProviderConfig `json:"providers,omitempty"`
|
||||
// RequireModelPrefix, when true, makes every model addressable only as
|
||||
// "provider/model" or as a ModelAliases id: a bare discovered id is not
|
||||
// resolved (404 unknown model) and is not listed by GET /v1/models. It
|
||||
// removes the sorted-first bare-id collision winner. Defaults to false.
|
||||
RequireModelPrefix bool `json:"requireModelPrefix,omitempty"`
|
||||
// ModelMeta declares per-model metadata overrides (feature v0.23):
|
||||
// context window size and per-token cost, keyed by an exact
|
||||
// "provider/model" id or a bare model id (applying wherever that bare
|
||||
@@ -172,6 +167,16 @@ 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"`
|
||||
// RequireModelPrefix, when true, makes every model addressable only as
|
||||
// "provider/model" or as a ModelAliases id: a bare discovered id is not
|
||||
// resolved (404 unknown model) and is not listed by GET /v1/models. It
|
||||
// removes the sorted-first bare-id collision winner. Defaults to false.
|
||||
RequireModelPrefix bool `json:"requireModelPrefix,omitempty"`
|
||||
}
|
||||
|
||||
// ProviderConfig describes one upstream LLM provider.
|
||||
@@ -447,6 +452,11 @@ type CacheConfig struct {
|
||||
// value above maxCacheMaxBodyBytes is a construction error.
|
||||
MaxBodyBytes int `json:"maxBodyBytes,omitempty"`
|
||||
Enabled bool `json:"enabled,omitempty"`
|
||||
// AccountHits charges the usage a cached response originally reported
|
||||
// to the user and group that receive the hit, so token and cost
|
||||
// budgets also bind repeated identical prompts. Defaults to false: a
|
||||
// hit then adds only to the request counters.
|
||||
AccountHits bool `json:"accountHits,omitempty"`
|
||||
}
|
||||
|
||||
// BreakerConfig configures the per-provider discovery circuit breaker
|
||||
@@ -761,6 +771,8 @@ type Gateway struct {
|
||||
next http.Handler
|
||||
events *eventLog
|
||||
bodyAdmission chan struct{}
|
||||
teeBudget *teeBudget
|
||||
fanoutGate chan struct{}
|
||||
failoverHealth *requestHealthTracker
|
||||
targetHealth *targetHealthTracker
|
||||
limiter *limiter
|
||||
@@ -1049,6 +1061,8 @@ 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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+40
-1
@@ -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]")
|
||||
}
|
||||
|
||||
+10
-3
@@ -664,9 +664,6 @@ type modelRegistry struct {
|
||||
// cheap nil-map lookup (always false/zero) that adds no branching
|
||||
// cost to the v0.1 no-aliases case.
|
||||
aliases map[string]string
|
||||
// requirePrefix is Config.RequireModelPrefix: resolve and listFor treat
|
||||
// a bare id as unaddressable, leaving "provider/model" and aliases.
|
||||
requirePrefix bool
|
||||
// modelMeta is the validated Config.ModelMeta config-override map
|
||||
// (feature v0.23), built once by newModelRegistry via
|
||||
// validateModelMeta and never mutated afterward. nil (not merely
|
||||
@@ -722,6 +719,9 @@ type modelRegistry struct {
|
||||
warnedMu sync.Mutex
|
||||
// modelsCacheMu guards modelsCache.
|
||||
modelsCacheMu sync.Mutex
|
||||
// requirePrefix is Config.RequireModelPrefix: resolve and listFor treat
|
||||
// a bare id as unaddressable, leaving "provider/model" and aliases.
|
||||
requirePrefix bool
|
||||
}
|
||||
|
||||
// modelsCacheEntry is one cached, encoded GET /v1/models response body
|
||||
@@ -1555,6 +1555,13 @@ func (m *modelRegistry) bareWinner(id string, grp *group, extraIDs ...string) (s
|
||||
return visible[0], true
|
||||
}
|
||||
|
||||
// providerHasModel reports whether provider's explicit or discovered model
|
||||
// set contains id.
|
||||
func (m *modelRegistry) providerHasModel(provider, id string) bool {
|
||||
st, ok := m.states[provider]
|
||||
return ok && st.hasModel(id)
|
||||
}
|
||||
|
||||
// resolve maps a client-requested model id to the adapter that serves it.
|
||||
// It returns the adapter, the upstream model id to send that adapter (the
|
||||
// id form the provider itself knows, with any gateway-side provider prefix
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
+80
-4
@@ -7,7 +7,6 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
@@ -29,6 +28,11 @@ const maxPassthroughBytes = 32 << 20
|
||||
// of any size — without ever waiting for the full body to buffer first.
|
||||
const maxAccountingTeeBytes = 4 << 20
|
||||
|
||||
// maxAccountingTeeTotalBytes caps the bytes all concurrent accounting tees
|
||||
// may retain together. A response that cannot reserve room is accounted
|
||||
// from the truncated-body estimate instead of its parsed usage.
|
||||
const maxAccountingTeeTotalBytes = 256 << 20
|
||||
|
||||
// sseAccountingHeadBytes/sseAccountingTailBytes cap sseAccountingBuffer
|
||||
// (verify-core fix, round 4): an SSE stream used to be teed through the
|
||||
// same 4MiB-capped cappedAccountingBuffer as a JSON body, which held
|
||||
@@ -419,10 +423,44 @@ 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 {
|
||||
budget *teeBudget
|
||||
buf bytes.Buffer
|
||||
reserved int64
|
||||
truncated bool
|
||||
}
|
||||
|
||||
// teeBudget bounds the bytes every in-flight accounting tee may hold at
|
||||
// once, so N concurrent large responses cannot add up to N x
|
||||
// maxAccountingTeeBytes. Counters use the sync/atomic function forms, the
|
||||
// style this package already uses under Yaegi.
|
||||
type teeBudget struct {
|
||||
used int64
|
||||
limit int64
|
||||
}
|
||||
|
||||
func newTeeBudget(limit int64) *teeBudget { return &teeBudget{limit: limit} }
|
||||
|
||||
// reserve claims n bytes and reports whether the budget had room.
|
||||
func (b *teeBudget) reserve(n int64) bool {
|
||||
if atomic.AddInt64(&b.used, n) > b.limit {
|
||||
atomic.AddInt64(&b.used, -n)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (b *teeBudget) free(n int64) { atomic.AddInt64(&b.used, -n) }
|
||||
|
||||
// release returns the tee's reserved bytes to the budget. It is safe on a
|
||||
// nil receiver and idempotent.
|
||||
func (c *cappedAccountingBuffer) release() {
|
||||
if c == nil || c.budget == nil || c.reserved == 0 {
|
||||
return
|
||||
}
|
||||
c.budget.free(c.reserved)
|
||||
c.reserved = 0
|
||||
}
|
||||
|
||||
// Write implements io.Writer. It always reports success for the full
|
||||
// input length — even once truncated is set and further bytes are being
|
||||
// discarded — because this type is driven by io.TeeReader, whose Read
|
||||
@@ -444,6 +482,13 @@ func (c *cappedAccountingBuffer) Write(p []byte) (int, error) {
|
||||
p = p[:remaining]
|
||||
c.truncated = true
|
||||
}
|
||||
if c.budget != nil {
|
||||
if !c.budget.reserve(int64(len(p))) {
|
||||
c.truncated = true
|
||||
return n, nil
|
||||
}
|
||||
c.reserved += int64(len(p))
|
||||
}
|
||||
c.buf.Write(p) //nolint:errcheck // bytes.Buffer.Write never errors
|
||||
return n, nil
|
||||
}
|
||||
@@ -1180,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
|
||||
@@ -1224,9 +1272,11 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
|
||||
// was already restored via io.MultiReader after that first call) so
|
||||
// this check runs independently of whether the group restricts
|
||||
// models at all.
|
||||
var gatedModel string
|
||||
if !g.cfg.AllowUnpricedWithCostBudget && scopesHaveCostBudget(scopes) {
|
||||
if model, hasModel, _, peekErr := peekPassthroughModel(r); peekErr == nil && hasModel {
|
||||
canonical := providerName + "/" + model
|
||||
gatedModel = model
|
||||
if !priceKnown(canonical, model, g.cfg.Pricing, g.cfg.ModelMeta) {
|
||||
g.logf("passthrough: refusing model %q: it has no configured price, and a cost budget applies to this caller that cannot be enforced without one", model)
|
||||
g.recordUnpricedRefusalEvent(scopes, canonical, providerName, routePassthrough) // F3 hook 5 (item 5, this round)
|
||||
@@ -1271,6 +1321,7 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
|
||||
}
|
||||
|
||||
result, complete := g.proxyUpstream(w, r, upstreamURL, adapter.httpClient(), adapter.injectAuth, providerCredentialRetargetHeaders, true, "passthrough (provider "+providerName+")", adapter.requestTimeout())
|
||||
defer result.tee.release()
|
||||
g.recordProxyEvent(scopes, providerName, routePassthrough, result) // F3 hook 3 (v0.3 dashboard task)
|
||||
|
||||
// Finding 2 fix (review-routes.md): a build/dial failure that never
|
||||
@@ -1335,7 +1386,7 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
|
||||
if !canEstimate || result.requestBytes <= 0 {
|
||||
return
|
||||
}
|
||||
respUsage = usage{prompt: int64(math.Ceil(float64(result.requestBytes) / 4)), estimated: true}
|
||||
respUsage = usage{prompt: estimatePromptTokens(int64(result.requestBytes)), estimated: true}
|
||||
if respModel == "" {
|
||||
respModel = "unknown/" + providerName
|
||||
}
|
||||
@@ -1343,7 +1394,19 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
|
||||
}
|
||||
|
||||
canonical := providerName + "/" + respModel
|
||||
cost := unifiedCostMicros(canonical, respModel, respUsage, g.cfg.Pricing, g.cfg.ModelMeta, g.pricingWarn)
|
||||
billModel := respModel
|
||||
// Gate and bill on the same price: an upstream that echoes a model id
|
||||
// the gate never saw must not turn a priced request into a free one.
|
||||
if gatedModel != "" && respModel != gatedModel && !priceKnown(canonical, respModel, g.cfg.Pricing, g.cfg.ModelMeta) {
|
||||
billModel = gatedModel
|
||||
}
|
||||
cost := unifiedCostMicros(providerName+"/"+billModel, billModel, respUsage, g.cfg.Pricing, g.cfg.ModelMeta, g.pricingWarn)
|
||||
// Only ids the provider is known to serve get a model scope: the id is
|
||||
// read from the upstream response body, so any other value would mint
|
||||
// counter keys without bound.
|
||||
if !g.registry.providerHasModel(providerName, respModel) {
|
||||
canonical = ""
|
||||
}
|
||||
// withModelScope drops the model scope when respModel is empty (the
|
||||
// upstream reported no model id), so a passthrough reply the gateway
|
||||
// cannot attribute still accounts to user/group/total as before.
|
||||
@@ -1714,7 +1777,7 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream
|
||||
var reader io.Reader = resp.Body
|
||||
switch {
|
||||
case isJSON:
|
||||
tee = &cappedAccountingBuffer{}
|
||||
tee = &cappedAccountingBuffer{budget: g.teeBudget}
|
||||
reader = io.TeeReader(resp.Body, tee)
|
||||
case isSSE:
|
||||
sseTee = &sseAccountingBuffer{}
|
||||
@@ -1738,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
|
||||
}
|
||||
|
||||
@@ -3355,7 +3355,7 @@ func TestHandlePassthrough_ModelLatencyAndUModelSuppressed_ProviderHistogramsOnl
|
||||
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{
|
||||
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"},
|
||||
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-native-x"}},
|
||||
}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{
|
||||
|
||||
+21
-3
@@ -6,7 +6,6 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -501,6 +500,10 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop
|
||||
// requires for a cache-enabled, admin-ENABLED-but-stats-off
|
||||
// case. Gating here, not just inside recordCacheHit, is what
|
||||
// actually saves the parse.
|
||||
if g.cfg.Cache.AccountHits {
|
||||
hitUsage := cachedUsage(cached.Body)
|
||||
g.limiter.accountWith(scopes, hitUsage, unifiedCostMicros(cand.canonical, cand.upstreamModel, hitUsage, g.cfg.Pricing, g.cfg.ModelMeta, g.pricingWarn), accountExtras{provider: cand.providerName})
|
||||
}
|
||||
if g.limiter.statsAdmin {
|
||||
saved := unifiedCostMicros(cand.canonical, cand.upstreamModel, cachedUsage(cached.Body), g.cfg.Pricing, g.cfg.ModelMeta, func(string) {})
|
||||
g.limiter.recordCacheHit(cand.canonical, saved)
|
||||
@@ -691,13 +694,13 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop
|
||||
// never streamed at all, or a failover pair is double-
|
||||
// billed: the failed candidate's phantom estimate plus
|
||||
// the next candidate's real usage.
|
||||
result.prompt = int64(math.Ceil(float64(len(body)) / 4))
|
||||
result.prompt = estimatePromptTokens(int64(len(body)))
|
||||
result.estimated = true
|
||||
g.logf("%s: streaming response from model %q failed or was canceled after starting to write to the client with zero usage reported; accounting an estimate derived from request body size", logPrefix, cand.canonical)
|
||||
case streaming:
|
||||
g.logf("%s: zero usage reported for a streaming response from model %q; accounting the request only", logPrefix, cand.canonical)
|
||||
case callErr == nil:
|
||||
result.prompt = int64(math.Ceil(float64(len(body)) / 4))
|
||||
result.prompt = estimatePromptTokens(int64(len(body)))
|
||||
result.estimated = true
|
||||
}
|
||||
}
|
||||
@@ -1435,3 +1438,18 @@ func (w *cacheCaptureWriter) Write(b []byte) (int, error) {
|
||||
}
|
||||
return w.statusTrackingWriter.Write(b)
|
||||
}
|
||||
|
||||
// maxEstimatedPromptTokens caps a body-size estimate, so a padded request
|
||||
// cannot bill, or overshoot a budget by, more than a very large context
|
||||
// window would legitimately hold.
|
||||
const maxEstimatedPromptTokens = 1 << 20
|
||||
|
||||
// estimatePromptTokens estimates prompt tokens as one per four body bytes,
|
||||
// rounded up and capped at maxEstimatedPromptTokens.
|
||||
func estimatePromptTokens(bodyBytes int64) int64 {
|
||||
est := (bodyBytes + 3) / 4
|
||||
if est > maxEstimatedPromptTokens {
|
||||
return maxEstimatedPromptTokens
|
||||
}
|
||||
return est
|
||||
}
|
||||
|
||||
@@ -0,0 +1,353 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// newPassthroughUsageGateway builds a Gateway with one openai provider
|
||||
// pointed at an upstream that always answers upstreamBody, and one user
|
||||
// with the given limits and pricing.
|
||||
func newPassthroughUsageGateway(t *testing.T, upstreamBody string, models []string, limits *LimitsConfig, pricing map[string]*ModelPricing) *Gateway {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(upstreamBody))
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{
|
||||
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: models},
|
||||
}
|
||||
cfg.Pricing = pricing
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{
|
||||
{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: limits},
|
||||
}}
|
||||
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.limiter.spawn = func(f func()) { f() }
|
||||
return gw
|
||||
}
|
||||
|
||||
func doPassthrough(t *testing.T, gw *Gateway) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(`{"model":"gpt-1"}`))
|
||||
req.Header.Set("Authorization", "Bearer sk-alice")
|
||||
rec := httptest.NewRecorder()
|
||||
gw.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
func TestPassthrough_EchoedUnknownModel_MintsNoModelScope(t *testing.T) {
|
||||
t.Parallel()
|
||||
gw := newPassthroughUsageGateway(t, `{"model":"echo-unknown","usage":{"prompt_tokens":4,"completion_tokens":2}}`, []string{"gpt-known"}, &LimitsConfig{}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, doPassthrough(t, gw).Code)
|
||||
|
||||
reqs, _ := gw.limiter.getCounter(kindModel, "openai/echo-unknown", metricReq, windowDay, time.Now())
|
||||
assert.Zero(t, reqs, "an id the provider does not serve must not get a model scope")
|
||||
tokIn, _ := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
||||
assert.EqualValues(t, 4, tokIn, "user accounting must be unchanged")
|
||||
}
|
||||
|
||||
func TestPassthrough_EchoedKnownModel_KeepsModelScope(t *testing.T) {
|
||||
t.Parallel()
|
||||
gw := newPassthroughUsageGateway(t, `{"model":"gpt-known","usage":{"prompt_tokens":4,"completion_tokens":2}}`, []string{"gpt-known"}, &LimitsConfig{}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, doPassthrough(t, gw).Code)
|
||||
|
||||
req, ok := gw.limiter.getCounter(kindModel, "openai/gpt-known", metricReq, windowDay, time.Now())
|
||||
assert.True(t, ok)
|
||||
assert.EqualValues(t, 1, req)
|
||||
}
|
||||
|
||||
func TestPassthrough_ImplausibleUsage_DoesNotLockOutBudget(t *testing.T) {
|
||||
t.Parallel()
|
||||
gw := newPassthroughUsageGateway(t, `{"model":"gpt-known","usage":{"prompt_tokens":9000000000000000000,"completion_tokens":1}}`, []string{"gpt-known"}, &LimitsConfig{TokensPerDay: 1000}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, doPassthrough(t, gw).Code)
|
||||
assert.Equal(t, http.StatusOK, doPassthrough(t, gw).Code, "an absurd reported figure must not exhaust the budget")
|
||||
}
|
||||
|
||||
func TestPassthrough_EchoedUnpricedModel_BilledAtRequestModelPrice(t *testing.T) {
|
||||
t.Parallel()
|
||||
pricing := map[string]*ModelPricing{"openai/gpt-1": {InputPerM: 10, OutputPerM: 10}}
|
||||
gw := newPassthroughUsageGateway(t, `{"model":"unpriced-model","usage":{"prompt_tokens":1000,"completion_tokens":500}}`, nil, &LimitsConfig{CostPerDayUSD: 1}, pricing)
|
||||
|
||||
require.Equal(t, http.StatusOK, doPassthrough(t, gw).Code)
|
||||
|
||||
cost, ok := gw.limiter.getCounter("user", "alice", metricCost, windowDay, time.Now())
|
||||
assert.True(t, ok, "cost must accrue even though the upstream echoed an unpriced model")
|
||||
assert.EqualValues(t, 15000, cost, "1500 tokens at 10 USD per million = 15000 micro-USD")
|
||||
}
|
||||
|
||||
func TestLimiterAccountWith_ImplausibleUsageDropped(t *testing.T) {
|
||||
t.Parallel()
|
||||
l := newLimiter(nil, true)
|
||||
scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{TokensPerDay: 1000, CostPerDayUSD: 1}}}
|
||||
|
||||
l.account(scopes, usage{prompt: maxUsageTokensPerRequest + 1, completion: 7}, maxCostMicrosPerRequest+1)
|
||||
|
||||
assert.Nil(t, l.checkAndCount(scopes), "implausible figures must not push the budgets over")
|
||||
out, _ := l.getCounter("user", "u", metricTokOut, windowDay, time.Now())
|
||||
assert.EqualValues(t, 7, out, "a plausible component of the same request is kept")
|
||||
}
|
||||
|
||||
func TestEstimatePromptTokens(t *testing.T) {
|
||||
t.Parallel()
|
||||
tests := []struct {
|
||||
name string
|
||||
bytes int64
|
||||
want int64
|
||||
}{
|
||||
{"empty", 0, 0},
|
||||
{"rounds up", 5, 2},
|
||||
{"exact", 400, 100},
|
||||
{"capped", 32 << 20, maxEstimatedPromptTokens},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
assert.Equal(t, tt.want, estimatePromptTokens(tt.bytes))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCappedAccountingBuffer_SharedBudgetBoundsConcurrentTees(t *testing.T) {
|
||||
t.Parallel()
|
||||
budget := newTeeBudget(10)
|
||||
a := &cappedAccountingBuffer{budget: budget}
|
||||
b := &cappedAccountingBuffer{budget: budget}
|
||||
|
||||
_, _ = a.Write([]byte("0123456789"))
|
||||
_, err := b.Write([]byte("abc"))
|
||||
require.NoError(t, err)
|
||||
assert.False(t, a.truncated)
|
||||
assert.True(t, b.truncated, "a tee that cannot reserve budget must stop retaining bytes")
|
||||
assert.Zero(t, b.buf.Len())
|
||||
|
||||
a.release()
|
||||
c := &cappedAccountingBuffer{budget: budget}
|
||||
_, _ = c.Write([]byte("abc"))
|
||||
assert.False(t, c.truncated, "released bytes return to the budget")
|
||||
a.release()
|
||||
var nilTee *cappedAccountingBuffer
|
||||
nilTee.release()
|
||||
}
|
||||
|
||||
func TestCacheHit_AccountHits_ChargesServingUser(t *testing.T) {
|
||||
const respBody = `{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
accountHits bool
|
||||
wantTokIn int64
|
||||
}{
|
||||
{"off: a hit is free", false, 10},
|
||||
{"on: a hit is charged", true, 20},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(respBody))
|
||||
}))
|
||||
defer srv.Close()
|
||||
gw := newCacheTestGateway(t, srv, nil, defaultCacheMaxBodyBytes)
|
||||
gw.cfg.Cache.AccountHits = tt.accountHits
|
||||
body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
|
||||
|
||||
for _, want := range []string{"miss", "hit"} {
|
||||
rec := httptest.NewRecorder()
|
||||
gw.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body))
|
||||
require.Equal(t, want, rec.Header().Get("X-Llmgw-Cache"))
|
||||
}
|
||||
|
||||
tokIn, _ := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
|
||||
assert.Equal(t, tt.wantTokIn, tokIn)
|
||||
reqs, _ := gw.limiter.getCounter("user", "alice", metricReq, windowDay, time.Now())
|
||||
assert.EqualValues(t, 2, reqs, "the request counter must not double count a hit")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
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