diff --git a/cmd/gosqlx-mcp/main.go b/cmd/gosqlx-mcp/main.go index ad9c16cf..25c05c6d 100644 --- a/cmd/gosqlx-mcp/main.go +++ b/cmd/gosqlx-mcp/main.go @@ -25,6 +25,9 @@ // GOSQLX_MCP_OPENAI_APPS_CHALLENGE ChatGPT app domain-verification token // GOSQLX_MCP_TRUSTED_PROXY_HEADER trusted client-IP header (e.g. CF-Connecting-IP) // GOSQLX_MCP_TRUSTED_PROXY_HOPS trusted proxies appending to X-Forwarded-For +// GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER header carrying the proxy shared secret +// GOSQLX_MCP_TRUSTED_PROXY_SECRET shared secret; forwarded IPs are trusted only +// on requests that carry it // // Pass --stdio to serve MCP over stdin/stdout instead of HTTP. // @@ -117,6 +120,9 @@ func runHTTP() error { }() log.Printf("gosqlx-mcp: listening on %s (auth=%v, client-ip=%s)\n", cfg.Addr(), cfg.AuthEnabled(), clientIPSource(cfg)) + if msg := sharedClientIPWarning(cfg); msg != "" { + log.Print(msg) + } if err := httpSrv.ListenAndServe(); err != nil && err != http.ErrServerClosed { return fmt.Errorf("server error: %w", err) } @@ -152,16 +158,35 @@ func withBodyDeadline(next http.Handler, timeout time.Duration) http.Handler { }) } -// clientIPSource describes where rate limiting takes the client IP from. +// clientIPSource describes where rate limiting takes the client IP from. It +// names the secret header when one gates the forwarded address, but never +// the secret itself. func clientIPSource(cfg *gosqlxmcp.Config) string { + var src string switch { case cfg.TrustedProxyHeader != "": - return cfg.TrustedProxyHeader + src = cfg.TrustedProxyHeader case cfg.TrustedProxyHops > 0: - return fmt.Sprintf("X-Forwarded-For[-%d]", cfg.TrustedProxyHops) + src = fmt.Sprintf("X-Forwarded-For[-%d]", cfg.TrustedProxyHops) default: return "RemoteAddr" } + if cfg.TrustedProxySecretHeader != "" { + src += " when " + cfg.TrustedProxySecretHeader + " matches, else RemoteAddr" + } + return src +} + +// sharedClientIPWarning returns the startup warning logged when rate limiting +// is keyed on RemoteAddr only, or "" when a trusted proxy is configured. +func sharedClientIPWarning(cfg *gosqlxmcp.Config) string { + if cfg.TrustedProxyHeader != "" || cfg.TrustedProxyHops > 0 { + return "" + } + return "gosqlx-mcp: warning: no trusted proxy configured, so rate limits are keyed on the TCP peer address. " + + "If this server runs behind a reverse proxy or load balancer, every client shares the proxy's " + + "rate-limit bucket and per-client GET (SSE) stream cap. Set GOSQLX_MCP_TRUSTED_PROXY_HEADER " + + "or GOSQLX_MCP_TRUSTED_PROXY_HOPS (see docs/MCP_GUIDE.md, \"Client IP and reverse proxies\")." } // newMux builds the HTTP routes: diff --git a/cmd/gosqlx-mcp/main_test.go b/cmd/gosqlx-mcp/main_test.go index 1b29357d..833cba12 100644 --- a/cmd/gosqlx-mcp/main_test.go +++ b/cmd/gosqlx-mcp/main_test.go @@ -298,6 +298,8 @@ func TestClientIPSource(t *testing.T) { {"default", gosqlxmcp.Config{}, "RemoteAddr"}, {"header", gosqlxmcp.Config{TrustedProxyHeader: "Cf-Connecting-Ip"}, "Cf-Connecting-Ip"}, {"hops", gosqlxmcp.Config{TrustedProxyHops: 2}, "X-Forwarded-For[-2]"}, + {"secret-gated header", gosqlxmcp.Config{TrustedProxyHeader: "Cf-Connecting-Ip", TrustedProxySecretHeader: "X-Origin-Secret", TrustedProxySecret: "s3cret-0123456789abcdef"}, + "Cf-Connecting-Ip when X-Origin-Secret matches, else RemoteAddr"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -358,3 +360,43 @@ func TestWithBodyDeadline_SlowBodyOutsideMCP(t *testing.T) { } } } + +func TestClientIPSource_NeverLogsSecret(t *testing.T) { + cfg := gosqlxmcp.Config{TrustedProxyHops: 1, TrustedProxySecretHeader: "X-Origin-Secret", TrustedProxySecret: "s3cret-0123456789abcdef"} + if got := clientIPSource(&cfg); strings.Contains(got, cfg.TrustedProxySecret) { + t.Fatalf("clientIPSource leaks the secret: %q", got) + } +} + +func TestSharedClientIPWarning(t *testing.T) { + tests := []struct { + name string + cfg gosqlxmcp.Config + wantWarn bool + }{ + {"no trusted proxy warns", gosqlxmcp.Config{Host: "0.0.0.0", Port: 8080}, true}, + {"loopback still warns (local reverse proxy)", gosqlxmcp.Config{Host: "127.0.0.1", Port: 8080}, true}, + {"trusted header silences", gosqlxmcp.Config{TrustedProxyHeader: "Cf-Connecting-Ip"}, false}, + {"trusted hops silences", gosqlxmcp.Config{TrustedProxyHops: 1}, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := sharedClientIPWarning(&tt.cfg) + if (got != "") != tt.wantWarn { + t.Fatalf("sharedClientIPWarning() = %q, want warning=%v", got, tt.wantWarn) + } + if !tt.wantWarn { + return + } + for _, want := range []string{"warning", "share", "rate-limit bucket", "stream cap", + "GOSQLX_MCP_TRUSTED_PROXY_HEADER", "GOSQLX_MCP_TRUSTED_PROXY_HOPS"} { + if !strings.Contains(got, want) { + t.Errorf("warning %q does not mention %q", got, want) + } + } + if strings.Contains(got, "\n") { + t.Errorf("warning should be a single log line: %q", got) + } + }) + } +} diff --git a/docs/CHATGPT_APP.md b/docs/CHATGPT_APP.md index 267f55bb..fa4bdf4a 100644 --- a/docs/CHATGPT_APP.md +++ b/docs/CHATGPT_APP.md @@ -185,6 +185,19 @@ HTML gains remote `src`/`href` URLs, `@import`, `fetch(`, `XMLHttpRequest` or | `/health` | `{"status":"ok","version":,"tools":}` | | `/.well-known/openai-apps-challenge` | Plain-text domain-verification token. Served only when `GOSQLX_MCP_OPENAI_APPS_CHALLENGE` is set; returns 404 otherwise | +### Rate limits are shared + +Rate limits are per client IP (see [MCP_GUIDE.md, Rate Limits](MCP_GUIDE.md#rate-limits)). +ChatGPT calls the MCP server from OpenAI's servers, not from each user's browser, so +the "client" the limiter sees is an OpenAI egress address. Many ChatGPT users can +therefore share one budget of 10 requests/second and 120 weighted requests/minute, +and heavy use by some can make tool calls fail for others with a rate-limit error +(HTTP 429, or HTTP 200 with a JSON-RPC error for `initialize`) until the window +resets. The same applies to any proxy in front of a self-hosted server that is not +configured as trusted: without `GOSQLX_MCP_TRUSTED_PROXY_HEADER` or +`GOSQLX_MCP_TRUSTED_PROXY_HOPS`, every client behind it shares one bucket and one +4-stream `GET` cap, and `gosqlx-mcp` logs a warning at startup. + ### Sessions The streamable HTTP transport runs in mcp-go's default stateless mode. Session IDs @@ -197,9 +210,10 @@ invalidates an ID: a later request with the same ID simply starts new state. mcp-go's informational transport logs, which include session IDs, are discarded; only transport errors are logged, so nothing identifies a session in the logs. -The HTTP server sets `ReadHeaderTimeout`. `Server.Start` (library use) also sets -`IdleTimeout` but applies only bearer auth, not the rate limiter or the 64 KB body -cap; `cmd/gosqlx-mcp` wraps the handler in `RateLimitMiddleware`. Call +The HTTP server sets `ReadHeaderTimeout`. `Server.Start` and `Server.Handler` +(library use) apply only bearer auth, not `RateLimitMiddleware`, so they have no rate +limit, stream cap, body read deadline or 64 KB body cap; `cmd/gosqlx-mcp` wraps the +handler in `NewRateLimitMiddleware`. Call `Server.Shutdown` when you discard a `Server` whose `Handler` you used, to stop the session sweeper goroutine. diff --git a/docs/MCP_GUIDE.md b/docs/MCP_GUIDE.md index 9a448014..3cd07b02 100644 --- a/docs/MCP_GUIDE.md +++ b/docs/MCP_GUIDE.md @@ -84,10 +84,11 @@ The public server has smart rate limiting to ensure fair usage: - **Sustained**: 120 weighted-requests/minute per IP - Heavier tools (`analyze_sql`) consume more budget than simple ones (`validate_sql`) - Limits tighten automatically under high server load -- **Streams**: at most 4 concurrent `GET /mcp` (SSE) streams per IP; each costs 1 unit of sustained budget +- **Streams**: at most 4 concurrent `GET /mcp` (SSE) streams per IP; each opened stream costs 1 unit of sustained budget. A `GET` refused because 4 streams are already open costs nothing +- **Malformed bodies**: a non-empty POST body that is not a JSON-RPC message or a non-empty batch costs 1 unit, so malformed requests are not free - **Request body**: at most 64 KB, and it must arrive within 30 seconds -A denied request gets HTTP `429 Too Many Requests` with a `Retry-After` header (seconds) and a JSON-RPC error (code `-32000`) carrying the request's `id`. A denied `initialize` is the exception: it returns HTTP 200 with the same JSON-RPC error, because mcp-go clients treat any 4xx on `initialize` as a legacy-SSE server. Every response carries `X-RateLimit-Limit`, `X-RateLimit-Remaining` (sustained budget left in the current minute) and `X-RateLimit-Reset` (Unix time the minute ends). An oversized body gets HTTP `413` with a JSON-RPC error that names the limit; a body that stalls gets HTTP `408`. +A denied request gets HTTP `429 Too Many Requests` with a `Retry-After` header (seconds) and a JSON-RPC error (code `-32000`) carrying the request's `id`. A denied `initialize`, alone or inside a JSON-RPC batch, is the exception: it returns HTTP 200 with the same JSON-RPC error, because mcp-go clients treat any 4xx on `initialize` as a legacy-SSE server. A `GET` refused by the stream cap gets HTTP 429 with `Retry-After: 60`. Every response to a request the limiter evaluates carries `X-RateLimit-Limit`, `X-RateLimit-Remaining` (sustained budget left in the current minute) and `X-RateLimit-Reset` (Unix time the minute ends); stream-cap, body-size and body-timeout rejections, which happen before that check, do not. An oversized body gets HTTP `413` with a JSON-RPC error that names the limit; a body that stalls gets HTTP `408`. The 64 KB limit is deliberate: `analyze_sql` allocates roughly 500× the SQL size (about 33 MB for 64 KB of SQL, 530 MB for 1 MB), so larger bodies would let a handful of requests exhaust a small instance. @@ -219,6 +220,8 @@ The `gosqlx-mcp` server is configured exclusively via environment variables. No | `GOSQLX_MCP_OPENAI_APPS_CHALLENGE` | *(empty)* | string | ChatGPT app domain-verification token served at `/.well-known/openai-apps-challenge` | | `GOSQLX_MCP_TRUSTED_PROXY_HEADER` | *(empty)* | string | Header a trusted proxy sets to the client IP (e.g. `CF-Connecting-IP`); must be a valid header name and not `X-Forwarded-For` | | `GOSQLX_MCP_TRUSTED_PROXY_HOPS` | `0` | integer | 0–16; number of trusted proxies that append to `X-Forwarded-For` | +| `GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER` | *(empty)* | string | Optional. Header the proxy sets to the shared secret; set together with `GOSQLX_MCP_TRUSTED_PROXY_SECRET`. Must be a valid header name and not `X-Forwarded-For`, `Authorization` or the client-IP header. Requires the trusted proxy header or hops | +| `GOSQLX_MCP_TRUSTED_PROXY_SECRET` | *(empty)* | string | Optional. Shared secret, at least 16 characters; whitespace-trimmed. Never logged | **Notes:** @@ -237,6 +240,30 @@ Behind a reverse proxy every request comes from the proxy, so all users would sh If the header is set, it takes precedence. A missing or unparseable value falls back to `RemoteAddr`. Only enable either option when every request reaches the server through that proxy; otherwise clients can set the header directly. +**Startup warning.** When neither option is set, `gosqlx-mcp` logs one warning at startup: behind a reverse proxy or load balancer every client then shares the proxy's address, so all of them share one rate-limit bucket and one 4-stream `GET` cap. The warning is expected for a server that clients reach directly; behind a proxy, set one of the options above. + +#### Requiring a proxy secret (optional) + +A trusted client-IP header is only as trustworthy as the path to the server. If the origin can also be reached without the proxy (for example through Render's `*.onrender.com` hostname rather than the Cloudflare zone), a client can send its own `CF-Connecting-IP` and pick its rate-limit key. To close that gap, have the proxy add a secret header and tell the server to trust forwarded addresses only on requests that carry it: + +```bash +GOSQLX_MCP_TRUSTED_PROXY_HEADER=CF-Connecting-IP +GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER=X-Origin-Secret +GOSQLX_MCP_TRUSTED_PROXY_SECRET= +``` + +With both secret variables set, `GOSQLX_MCP_TRUSTED_PROXY_HEADER` and `GOSQLX_MCP_TRUSTED_PROXY_HOPS` apply only when the last value of the secret header equals the secret (compared in constant time). Requests with a missing or wrong secret are not rejected; they are keyed by `RemoteAddr`, so traffic that bypasses the proxy shares one bucket instead of choosing its own. The secret header is removed before the request reaches the MCP handler, and the secret never appears in logs or configuration errors. With neither variable set, behaviour is unchanged. + +To set the header with a Cloudflare Transform Rule: + +1. Generate a secret, for example `openssl rand -hex 32`. +2. In the Cloudflare dashboard for the zone that serves the MCP hostname, open **Rules → Transform Rules → Modify Request Header** and create a rule. +3. Match the MCP hostname, for example *Hostname equals `mcp.example.com`*. +4. Add the action **Set static**, header name `X-Origin-Secret`, value the secret. *Set* overwrites any value a client sent, so clients cannot supply their own. +5. Deploy the rule, then set `GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER` and `GOSQLX_MCP_TRUSTED_PROXY_SECRET` on the server as secret environment variables (on Render, in the dashboard rather than in `render.yaml`) and restart it. + +Deploy the rule before the server variables: until both sides agree, requests fall back to the shared `RemoteAddr` bucket rather than failing. Rotate the secret the same way. The rule must cover every hostname that should get per-client limits. + ### HTTP server limits `gosqlx-mcp` sets `ReadHeaderTimeout` 10 s, `IdleTimeout` 120 s and `MaxHeaderBytes` 32 KB. It sets no global `ReadTimeout`/`WriteTimeout`, because those would cut long-lived `GET /mcp` SSE streams. POST bodies instead get a 30 s read deadline in the rate-limit middleware, so slow-body (slowloris) connections are closed with HTTP 408. @@ -726,7 +753,8 @@ func main() { | `LoadConfig` | `() (*Config, error)` | Load from env vars | | `DefaultConfig` | `() *Config` | Defaults: `127.0.0.1:8080`, auth disabled | | `New` | `(cfg *Config) *Server` | Create server with all 7 tools registered | -| `(*Server).Start` | `(ctx context.Context) error` | Bind, serve, block until ctx cancelled | +| `(*Server).Start` | `(ctx context.Context) error` | Bind, serve, block until ctx cancelled. Applies bearer auth only, not `RateLimitMiddleware` | +| `(*Server).Handler` | `() http.Handler` | MCP transport behind bearer auth, without `RateLimitMiddleware`; wrap it with `NewRateLimitMiddleware` before exposing it to untrusted clients | | `BearerAuthMiddleware` | `(cfg *Config, next http.Handler) http.Handler` | Auth wrapper; no-op when auth is disabled | | `NewRateLimitMiddleware` | `(cfg *Config, next http.Handler) http.Handler` | Rate limiter keyed by client IP, honouring `cfg`'s trusted-proxy settings | | `RateLimitMiddleware` | `(next http.Handler) http.Handler` | `NewRateLimitMiddleware` keyed on `RemoteAddr` only | @@ -744,6 +772,10 @@ type Config struct { OpenAIAppsChallenge string // GOSQLX_MCP_OPENAI_APPS_CHALLENGE (default "" = not served) TrustedProxyHeader string // GOSQLX_MCP_TRUSTED_PROXY_HEADER (default "" = not used) TrustedProxyHops int // GOSQLX_MCP_TRUSTED_PROXY_HOPS (default 0 = X-Forwarded-For ignored) + + // Optional: trust the two fields above only when this header carries the secret. + TrustedProxySecretHeader string // GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER (default "" = not required) + TrustedProxySecret string // GOSQLX_MCP_TRUSTED_PROXY_SECRET } ``` diff --git a/pkg/mcp/config.go b/pkg/mcp/config.go index ee998956..6c78e244 100644 --- a/pkg/mcp/config.go +++ b/pkg/mcp/config.go @@ -36,6 +36,12 @@ // the client IP (e.g. CF-Connecting-IP); empty = unused // GOSQLX_MCP_TRUSTED_PROXY_HOPS number of trusted proxies appending to // X-Forwarded-For (0-16, default 0 = ignore XFF) +// GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER optional header the trusted proxy +// sets to GOSQLX_MCP_TRUSTED_PROXY_SECRET; when both are +// set, forwarded client IPs are trusted only on +// requests that carry the secret +// GOSQLX_MCP_TRUSTED_PROXY_SECRET shared secret for the header above +// (at least 16 characters) package mcp import ( @@ -82,11 +88,29 @@ type Config struct { // client-supplied and ignored. // Source: GOSQLX_MCP_TRUSTED_PROXY_HOPS (default: 0 - X-Forwarded-For ignored) TrustedProxyHops int + + // TrustedProxySecretHeader and TrustedProxySecret, when both set, make + // TrustedProxyHeader and TrustedProxyHops apply only to requests whose + // TrustedProxySecretHeader value equals TrustedProxySecret (compared in + // constant time). Other requests are keyed by RemoteAddr. This stops a + // client that reaches the origin directly, bypassing the proxy, from + // choosing its own rate-limit key. The proxy must overwrite the header on + // every request (for example with a Cloudflare Transform Rule). The + // header is removed before the request reaches the MCP handler. If only + // one of the two is set, forwarded addresses are never trusted. + // Source: GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER and + // GOSQLX_MCP_TRUSTED_PROXY_SECRET (default: "" - not required) + TrustedProxySecretHeader string + TrustedProxySecret string } // maxTrustedProxyHops bounds GOSQLX_MCP_TRUSTED_PROXY_HOPS. const maxTrustedProxyHops = 16 +// minTrustedProxySecretLen is the shortest accepted +// GOSQLX_MCP_TRUSTED_PROXY_SECRET. +const minTrustedProxySecretLen = 16 + // LoadConfig reads configuration from environment variables, applying defaults // for any variables that are unset or empty. func LoadConfig() (*Config, error) { @@ -137,9 +161,44 @@ func LoadConfig() (*Config, error) { cfg.TrustedProxyHops = hops } + if err := loadTrustedProxySecret(cfg); err != nil { + return nil, err + } + return cfg, nil } +// loadTrustedProxySecret reads GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER and +// GOSQLX_MCP_TRUSTED_PROXY_SECRET into cfg. It must run after the trusted +// proxy header and hops are loaded. Errors never include the secret. +func loadTrustedProxySecret(cfg *Config) error { + header := strings.TrimSpace(os.Getenv("GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER")) + secret := strings.TrimSpace(os.Getenv("GOSQLX_MCP_TRUSTED_PROXY_SECRET")) + if header == "" && secret == "" { + return nil + } + if header == "" || secret == "" { + return fmt.Errorf("GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER and GOSQLX_MCP_TRUSTED_PROXY_SECRET must be set together") + } + if !validHeaderName(header) { + return fmt.Errorf("GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER: %q is not a valid HTTP header name", header) + } + canonical := http.CanonicalHeaderKey(header) + switch canonical { + case "X-Forwarded-For", "Authorization", cfg.TrustedProxyHeader: + return fmt.Errorf("GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER: %q cannot be used as the secret header", canonical) + } + if len(secret) < minTrustedProxySecretLen { + return fmt.Errorf("GOSQLX_MCP_TRUSTED_PROXY_SECRET: must be at least %d characters", minTrustedProxySecretLen) + } + if cfg.TrustedProxyHeader == "" && cfg.TrustedProxyHops == 0 { + return fmt.Errorf("GOSQLX_MCP_TRUSTED_PROXY_SECRET requires GOSQLX_MCP_TRUSTED_PROXY_HEADER or GOSQLX_MCP_TRUSTED_PROXY_HOPS") + } + cfg.TrustedProxySecretHeader = canonical + cfg.TrustedProxySecret = secret + return nil +} + // validHeaderName reports whether s is a valid HTTP header field name // (an RFC 9110 token). func validHeaderName(s string) bool { diff --git a/pkg/mcp/config_test.go b/pkg/mcp/config_test.go index b63abc23..b0ec27dc 100644 --- a/pkg/mcp/config_test.go +++ b/pkg/mcp/config_test.go @@ -16,6 +16,7 @@ package mcp import ( "os" + "strings" "testing" ) @@ -183,3 +184,62 @@ func TestLoadConfig_TrustedProxy(t *testing.T) { }) } } + +func TestLoadConfig_TrustedProxySecret(t *testing.T) { + const secret = "s3cret-0123456789abcdef" + tests := []struct { + name string + header string + hops string + secretHdr string + secret string + wantHeader string + wantErr string + }{ + {name: "unset keeps default", header: "CF-Connecting-IP"}, + {name: "header and secret", header: "CF-Connecting-IP", secretHdr: " x-origin-secret ", secret: secret, wantHeader: "X-Origin-Secret"}, + {name: "with hops", hops: "1", secretHdr: "X-Origin-Secret", secret: secret, wantHeader: "X-Origin-Secret"}, + {name: "secret without header", header: "CF-Connecting-IP", secret: secret, wantErr: "must be set together"}, + {name: "header without secret", header: "CF-Connecting-IP", secretHdr: "X-Origin-Secret", wantErr: "must be set together"}, + {name: "invalid header name", header: "CF-Connecting-IP", secretHdr: "X Origin", secret: secret, wantErr: "not a valid HTTP header name"}, + {name: "XFF rejected", hops: "1", secretHdr: "X-Forwarded-For", secret: secret, wantErr: "cannot be used"}, + {name: "Authorization rejected", header: "CF-Connecting-IP", secretHdr: "authorization", secret: secret, wantErr: "cannot be used"}, + {name: "same as client IP header rejected", header: "CF-Connecting-IP", secretHdr: "cf-connecting-ip", secret: secret, wantErr: "cannot be used"}, + {name: "short secret rejected", header: "CF-Connecting-IP", secretHdr: "X-Origin-Secret", secret: "short", wantErr: "at least 16"}, + {name: "no trusted proxy rejected", secretHdr: "X-Origin-Secret", secret: secret, wantErr: "requires GOSQLX_MCP_TRUSTED_PROXY_HEADER"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("GOSQLX_MCP_TRUSTED_PROXY_HEADER", tt.header) + t.Setenv("GOSQLX_MCP_TRUSTED_PROXY_HOPS", tt.hops) + t.Setenv("GOSQLX_MCP_TRUSTED_PROXY_SECRET_HEADER", tt.secretHdr) + t.Setenv("GOSQLX_MCP_TRUSTED_PROXY_SECRET", tt.secret) + cfg, err := LoadConfig() + if tt.wantErr != "" { + if err == nil { + t.Fatalf("expected error containing %q, got config %+v", tt.wantErr, cfg) + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Errorf("error = %q, want it to contain %q", err, tt.wantErr) + } + if tt.secret != "" && strings.Contains(err.Error(), tt.secret) { + t.Errorf("error leaks the secret: %q", err) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cfg.TrustedProxySecretHeader != tt.wantHeader { + t.Errorf("TrustedProxySecretHeader = %q, want %q", cfg.TrustedProxySecretHeader, tt.wantHeader) + } + wantSecret := "" + if tt.wantHeader != "" { + wantSecret = tt.secret + } + if cfg.TrustedProxySecret != wantSecret { + t.Errorf("TrustedProxySecret not loaded as expected (set=%v)", cfg.TrustedProxySecret != "") + } + }) + } +} diff --git a/pkg/mcp/ratelimit.go b/pkg/mcp/ratelimit.go index e53d49d9..3cfe4c68 100644 --- a/pkg/mcp/ratelimit.go +++ b/pkg/mcp/ratelimit.go @@ -17,11 +17,14 @@ package mcp import ( "bytes" "container/list" + "crypto/sha256" + "crypto/subtle" "encoding/json" "errors" "fmt" "hash/fnv" "io" + "log" "math" "net" "net/http" @@ -70,6 +73,11 @@ const ( // body, so slow-body (slowloris) connections cannot be held open forever. bodyReadTimeout = 30 * time.Second + // minUnparsedWeight is the sustained-budget cost of a non-empty POST body + // that is not a decodable JSON-RPC message or batch, so malformed + // requests are not free. + minUnparsedWeight = 1 + // maxBodySize is the largest accepted request body. analyze_sql on 64KB of // SQL allocates ~33MB; 1MB of SQL allocates ~530MB, more than the 512MB // hosted instance has, so the limit is kept at 64KB. @@ -137,6 +145,14 @@ type rateLimiter struct { streams streamCounter stopCleanup chan struct{} closeOnce sync.Once + + // deadlineErrs counts failures to set or clear the body read deadline. + // Only the first is logged; see noteDeadlineErr. + deadlineErrs atomic.Int64 + deadlineLogged atomic.Bool + // logf logs operational warnings. It never receives session IDs, request + // bodies or client addresses. + logf func(format string, v ...any) } // NewRateLimiter creates a new rateLimiter, initializes all shards, @@ -146,6 +162,7 @@ func NewRateLimiter() *rateLimiter { maxPerShard: maxBucketsPerShard, streams: streamCounter{open: make(map[string]int)}, stopCleanup: make(chan struct{}), + logf: log.Printf, } for i := range rl.shards { rl.shards[i].buckets = make(map[string]*list.Element) @@ -283,6 +300,27 @@ func (rl *rateLimiter) allow(key string, weight int) decision { return d } +// noteDeadlineErr records a failure to set or clear the request body read +// deadline. The first failure is logged with a fixed message (the error text +// can contain connection addresses, so only its kind is reported); later +// ones are only counted, so a misbehaving connection cannot flood the log. +func (rl *rateLimiter) noteDeadlineErr(err error) { + if err == nil { + return + } + rl.deadlineErrs.Add(1) + if !rl.deadlineLogged.CompareAndSwap(false, true) { + return + } + if errors.Is(err, http.ErrNotSupported) { + rl.logf("gosqlx-mcp: warning: the HTTP response writer does not support read deadlines; " + + "slow request bodies are not timed out. Further deadline errors are counted, not logged.") + return + } + rl.logf("gosqlx-mcp: warning: could not set a request body read deadline (%T); "+ + "further deadline errors are counted, not logged.", err) +} + // cleanupLoop periodically removes entries that have been idle for longer // than staleTimeout. func (rl *rateLimiter) cleanupLoop() { @@ -330,6 +368,17 @@ type clientIPResolver struct { // hops, when > 0, is the number of trusted proxies that append to // X-Forwarded-For; the client is the hops-th entry from the right. hops int + + // secretHeader, when set, gates header and hops: they are used only for + // requests whose secretHeader value hashes to secretHash. + secretHeader string + // secretGated is true when either secret setting is present. With only + // one present, secretOK is false and forwarded values are never trusted. + secretGated bool + secretOK bool + // secretHash is the SHA-256 of the configured secret. Comparing digests + // keeps the comparison constant-time regardless of the presented length. + secretHash [sha256.Size]byte } // newClientIPResolver builds a resolver from cfg (nil means RemoteAddr only). @@ -337,12 +386,49 @@ func newClientIPResolver(cfg *Config) clientIPResolver { if cfg == nil { return clientIPResolver{} } - return clientIPResolver{header: cfg.TrustedProxyHeader, hops: cfg.TrustedProxyHops} + c := clientIPResolver{header: cfg.TrustedProxyHeader, hops: cfg.TrustedProxyHops} + if cfg.TrustedProxySecretHeader != "" || cfg.TrustedProxySecret != "" { + c.secretGated = true + c.secretHeader = http.CanonicalHeaderKey(cfg.TrustedProxySecretHeader) + c.secretOK = cfg.TrustedProxySecretHeader != "" && cfg.TrustedProxySecret != "" + c.secretHash = sha256.Sum256([]byte(cfg.TrustedProxySecret)) + } + return c +} + +// fromTrustedProxy reports whether forwarded client addresses on r may be +// used: always when no secret is configured, otherwise only when the last +// value of the secret header matches the configured secret. +func (c clientIPResolver) fromTrustedProxy(r *http.Request) bool { + if !c.secretGated { + return true + } + if !c.secretOK { + return false + } + vals := r.Header.Values(c.secretHeader) + if len(vals) == 0 { + return false + } + got := sha256.Sum256([]byte(vals[len(vals)-1])) + return subtle.ConstantTimeCompare(got[:], c.secretHash[:]) == 1 +} + +// stripSecret removes the proxy secret header from r so it never reaches +// downstream handlers. +func (c clientIPResolver) stripSecret(r *http.Request) { + if c.secretHeader != "" { + r.Header.Del(c.secretHeader) + } } // clientKey returns the normalised client address for r. Invalid or missing -// forwarded values fall back to RemoteAddr. +// forwarded values, and requests without the configured proxy secret, fall +// back to RemoteAddr. func (c clientIPResolver) clientKey(r *http.Request) string { + if (c.header != "" || c.hops > 0) && !c.fromTrustedProxy(r) { + return remoteAddrKey(r) + } if c.header != "" { if vals := r.Header.Values(c.header); len(vals) > 0 { // The trusted proxy's value is the last one it wrote. @@ -365,6 +451,11 @@ func (c clientIPResolver) clientKey(r *http.Request) string { } } } + return remoteAddrKey(r) +} + +// remoteAddrKey returns the normalised TCP peer address of r. +func remoteAddrKey(r *http.Request) string { host, _, err := net.SplitHostPort(r.RemoteAddr) if err != nil { host = r.RemoteAddr @@ -432,6 +523,8 @@ type requestInfo struct { id json.RawMessage // method is the JSON-RPC method of a single (non-batch) request. method string + // initialize is true for an initialize request, alone or in a batch. + initialize bool } // weightOf returns the cost of one decoded JSON-RPC message. @@ -451,9 +544,10 @@ func weightOf(req jsonRPCRequest) int { // readRequestInfo reads the request body (at most maxBodySize bytes) to // determine its weight and JSON-RPC id, then re-buffers it so downstream -// handlers can read it again. A non-nil error means the body could not be -// read in full (too large, too slow, or a transport error) and the request -// must be rejected. +// handlers can read it again. A non-empty body that is not a JSON-RPC +// message or a non-empty batch costs minUnparsedWeight. A non-nil error +// means the body could not be read in full (too large, too slow, or a +// transport error) and the request must be rejected. func readRequestInfo(w http.ResponseWriter, r *http.Request) (requestInfo, *http.Request, error) { if r.Body == nil || r.Body == http.NoBody { return requestInfo{}, r, nil @@ -469,21 +563,24 @@ func readRequestInfo(w http.ResponseWriter, r *http.Request) (requestInfo, *http if len(trimmed) > 0 && trimmed[0] == '[' { // JSON-RPC batch: charge for every message in it. var batch []jsonRPCRequest - if err := json.Unmarshal(trimmed, &batch); err != nil { - return requestInfo{}, r, nil + if err := json.Unmarshal(trimmed, &batch); err != nil || len(batch) == 0 { + return requestInfo{weight: minUnparsedWeight}, r, nil } info := requestInfo{} for _, m := range batch { info.weight += weightOf(m) + if m.Method == "initialize" { + info.initialize = true + } } return info, r, nil } var req jsonRPCRequest if err := json.Unmarshal(body, &req); err != nil { - return requestInfo{}, r, nil + return requestInfo{weight: minUnparsedWeight}, r, nil } - info := requestInfo{weight: weightOf(req), method: req.Method} + info := requestInfo{weight: weightOf(req), method: req.Method, initialize: req.Method == "initialize"} if len(req.ID) > 0 && !bytes.Equal(req.ID, []byte("null")) { info.id = req.ID } @@ -545,7 +642,9 @@ func retryAfterSeconds(d time.Duration) int { } // RateLimitMiddleware is NewRateLimitMiddleware with the default -// configuration: clients are keyed by their TCP peer address only. +// configuration: clients are keyed by their TCP peer address only. Behind a +// reverse proxy that means every client shares the proxy's budget and stream +// cap; use NewRateLimitMiddleware with a trusted-proxy Config there. func RateLimitMiddleware(next http.Handler) http.Handler { return NewRateLimitMiddleware(nil, next) } @@ -555,11 +654,12 @@ func RateLimitMiddleware(next http.Handler) http.Handler { // RemoteAddr unless cfg names a trusted proxy header or hop count. // // Denied requests get HTTP 429 with Retry-After and a JSON-RPC error that -// echoes the request id. A denied initialize gets HTTP 200 instead, because -// mcp-go clients treat any 4xx on initialize as "legacy SSE server". -// POST bodies must arrive within bodyReadTimeout and fit in maxBodySize. -// GET (SSE) streams are capped per client and excluded from the load metric, -// since they stay open for the life of the session. +// echoes the request id. A denied initialize (alone or in a batch) gets +// HTTP 200 instead, because mcp-go clients treat any 4xx on initialize as +// "legacy SSE server". POST bodies must arrive within bodyReadTimeout and fit +// in maxBodySize. GET (SSE) streams are capped per client and excluded from +// the load metric, since they stay open for the life of the session; a GET +// refused by the stream cap is not charged against the rate limit. func NewRateLimitMiddleware(cfg *Config, next http.Handler) http.Handler { h, _ := newRateLimitHandler(cfg, next, bodyReadTimeout) return h @@ -573,24 +673,35 @@ func newRateLimitHandler(cfg *Config, next http.Handler, readTimeout time.Durati return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { key := resolver.clientKey(r) + resolver.stripSecret(r) isStream := r.Method == http.MethodGet var info requestInfo if isStream { + // Reserve a stream slot before charging the rate limit, so a + // GET refused by the cap costs no budget. A GET later denied by + // the rate limit releases its slot through the deferred release. + if !rl.streams.acquire(key, maxStreamsPerIP) { + w.Header().Set("Retry-After", "60") + writeJSONRPCError(w, http.StatusTooManyRequests, nil, -32000, + fmt.Sprintf("Too many open streams: at most %d concurrent GET streams per client.", maxStreamsPerIP)) + return + } + defer rl.streams.release(key) info.weight = streamWeight } else { // Bound the body read. On success clear the deadline so a // streamed (SSE) POST response is not cut off by it; on failure // keep it, so net/http's drain of the unread body cannot block. rc := http.NewResponseController(w) - _ = rc.SetReadDeadline(time.Now().Add(readTimeout)) + rl.noteDeadlineErr(rc.SetReadDeadline(time.Now().Add(readTimeout))) var err error info, r, err = readRequestInfo(w, r) if err != nil { writeBodyError(w, err, readTimeout) return } - _ = rc.SetReadDeadline(time.Time{}) + rl.noteDeadlineErr(rc.SetReadDeadline(time.Time{})) rl.activeReqs.Add(1) defer rl.activeReqs.Add(-1) } @@ -606,7 +717,7 @@ func newRateLimitHandler(cfg *Config, next http.Handler, readTimeout time.Durati secs := retryAfterSeconds(d.retryAfter) h.Set("Retry-After", strconv.Itoa(secs)) status := http.StatusTooManyRequests - if info.method == "initialize" { + if info.initialize { status = http.StatusOK } writeJSONRPCError(w, status, info.id, -32000, @@ -614,16 +725,6 @@ func newRateLimitHandler(cfg *Config, next http.Handler, readTimeout time.Durati return } - if isStream { - if !rl.streams.acquire(key, maxStreamsPerIP) { - h.Set("Retry-After", "60") - writeJSONRPCError(w, http.StatusTooManyRequests, nil, -32000, - fmt.Sprintf("Too many open streams: at most %d concurrent GET streams per client.", maxStreamsPerIP)) - return - } - defer rl.streams.release(key) - } - next.ServeHTTP(w, r) }), rl } diff --git a/pkg/mcp/ratelimit_test.go b/pkg/mcp/ratelimit_test.go index 4ed4ba13..3ee24e71 100644 --- a/pkg/mcp/ratelimit_test.go +++ b/pkg/mcp/ratelimit_test.go @@ -28,6 +28,7 @@ import ( "os" "strconv" "strings" + "sync" "testing" "testing/iotest" "time" @@ -394,6 +395,8 @@ func TestRateLimitMiddleware_Denials(t *testing.T) { {name: "notification gets null id", body: `{"jsonrpc":"2.0","method":"notifications/initialized"}`, wantStatus: http.StatusTooManyRequests, wantID: "null"}, {name: "unparseable body gets null id", body: `{not json`, wantStatus: http.StatusTooManyRequests, wantID: "null"}, {name: "initialize stays HTTP 200", body: `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}`, wantStatus: http.StatusOK, wantID: "1"}, + {name: "batched initialize stays HTTP 200", body: `[{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}]`, wantStatus: http.StatusOK, wantID: "null"}, + {name: "batch without initialize gets 429", body: "[" + toolCallBody("1", "validate_sql") + "]", wantStatus: http.StatusTooManyRequests, wantID: "null"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -639,17 +642,24 @@ func TestReadRequestInfo(t *testing.T) { wantWeight int wantID string wantMethod string + wantInit bool }{ {name: "validate_sql", body: toolCallBody("1", "validate_sql"), wantWeight: 1, wantID: "1", wantMethod: "tools/call"}, {name: "analyze_sql", body: toolCallBody("2", "analyze_sql"), wantWeight: 5, wantID: "2", wantMethod: "tools/call"}, {name: "security_scan", body: toolCallBody(`"s"`, "security_scan"), wantWeight: 2, wantID: `"s"`, wantMethod: "tools/call"}, - {name: "initialize", body: `{"jsonrpc":"2.0","id":4,"method":"initialize","params":{}}`, wantWeight: 0, wantID: "4", wantMethod: "initialize"}, + {name: "initialize", body: `{"jsonrpc":"2.0","id":4,"method":"initialize","params":{}}`, wantWeight: 0, wantID: "4", wantMethod: "initialize", wantInit: true}, {name: "tools/list", body: `{"jsonrpc":"2.0","id":5,"method":"tools/list","params":{}}`, wantWeight: 0, wantID: "5", wantMethod: "tools/list"}, {name: "unknown tool", body: toolCallBody("6", "unknown_tool"), wantWeight: 1, wantID: "6", wantMethod: "tools/call"}, {name: "null id", body: `{"jsonrpc":"2.0","id":null,"method":"ping"}`, wantWeight: 0, wantMethod: "ping"}, {name: "batch sums weights", body: "[" + toolCallBody("1", "analyze_sql") + "," + toolCallBody("2", "lint_sql") + "]", wantWeight: 7}, + {name: "batched initialize", body: `[{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}},{"jsonrpc":"2.0","method":"notifications/initialized"}]`, wantWeight: 0, wantInit: true}, {name: "empty body", body: "", wantWeight: 0}, - {name: "invalid JSON", body: `not json at all`, wantWeight: 0}, + {name: "invalid JSON costs minimum weight", body: `not json at all`, wantWeight: minUnparsedWeight}, + {name: "truncated object costs minimum weight", body: `{"jsonrpc":"2.0","method":"tools/call"`, wantWeight: minUnparsedWeight}, + {name: "non-object JSON costs minimum weight", body: `42`, wantWeight: minUnparsedWeight}, + {name: "whitespace-only body costs minimum weight", body: " \n\t", wantWeight: minUnparsedWeight}, + {name: "invalid batch costs minimum weight", body: `[{"method":`, wantWeight: minUnparsedWeight}, + {name: "empty batch costs minimum weight", body: `[]`, wantWeight: minUnparsedWeight}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -662,9 +672,9 @@ func TestReadRequestInfo(t *testing.T) { if err != nil { t.Fatalf("readRequestInfo: %v", err) } - if info.weight != tt.wantWeight || string(info.id) != tt.wantID || info.method != tt.wantMethod { - t.Errorf("got weight=%d id=%s method=%q, want weight=%d id=%s method=%q", - info.weight, info.id, info.method, tt.wantWeight, tt.wantID, tt.wantMethod) + if info.weight != tt.wantWeight || string(info.id) != tt.wantID || info.method != tt.wantMethod || info.initialize != tt.wantInit { + t.Errorf("got weight=%d id=%s method=%q initialize=%v, want weight=%d id=%s method=%q initialize=%v", + info.weight, info.id, info.method, info.initialize, tt.wantWeight, tt.wantID, tt.wantMethod, tt.wantInit) } if tt.body != "" { rebuffered, err := io.ReadAll(newReq.Body) @@ -711,3 +721,322 @@ func (rl *rateLimiter) bucketCount() int { } return n } + +// weightedCount reports the sustained budget key has used in its window. +func (rl *rateLimiter) weightedCount(t *testing.T, key string) int { + t.Helper() + sh := &rl.shards[shardFor(key)] + sh.mu.Lock() + el, ok := sh.buckets[key] + sh.mu.Unlock() + if !ok { + return 0 + } + b, ok := el.Value.(*bucket) + if !ok { + t.Fatalf("bucket for %s has type %T", key, el.Value) + } + b.mu.Lock() + defer b.mu.Unlock() + return b.weightedCount +} + +// openStreams reports how many GET stream slots key holds. +func (rl *rateLimiter) openStreams(key string) int { + rl.streams.mu.Lock() + defer rl.streams.mu.Unlock() + return rl.streams.open[key] +} + +// discardLogs stops rl from writing warnings to the test output. +func discardLogs(rl *rateLimiter) { rl.logf = func(string, ...any) {} } + +func TestRateLimitMiddleware_StreamCapRejectionIsFree(t *testing.T) { + started := make(chan struct{}, 16) + stream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _ = http.NewResponseController(w).Flush() + started <- struct{}{} + <-r.Context().Done() + return + } + _, _ = io.Copy(io.Discard, r.Body) + w.WriteHeader(http.StatusOK) + }) + h, rl := newRateLimitHandler(nil, stream, time.Second) + defer rl.Close() + ts := httptest.NewServer(h) + defer ts.Close() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + get := func() *http.Response { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, ts.URL+"/mcp", nil) + if err != nil { + t.Fatalf("new request: %v", err) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("GET: %v", err) + } + return resp + } + + for i := 0; i < maxStreamsPerIP; i++ { + resp := get() + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("stream %d: status %d, want 200", i+1, resp.StatusCode) + } + <-started + } + const key = "127.0.0.1" + if got := rl.weightedCount(t, key); got != maxStreamsPerIP*streamWeight { + t.Fatalf("after %d streams weightedCount = %d, want %d", maxStreamsPerIP, got, maxStreamsPerIP*streamWeight) + } + + // More refused GETs than the burst limit allows: none may cost budget. + for i := 0; i < 2*burstLimit; i++ { + resp := get() + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusTooManyRequests || !strings.Contains(string(body), "Too many open streams") { + t.Fatalf("refused GET %d: status %d body %q, want 429 stream-cap error", i+1, resp.StatusCode, body) + } + } + if got := rl.weightedCount(t, key); got != maxStreamsPerIP*streamWeight { + t.Errorf("stream-cap rejections were charged: weightedCount = %d, want %d", got, maxStreamsPerIP*streamWeight) + } + + // The client can still make tool calls with its full remaining budget. + resp, err := http.Post(ts.URL+"/mcp", "application/json", strings.NewReader(toolCallBody("1", "validate_sql"))) + if err != nil { + t.Fatalf("post: %v", err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("POST after refused GETs: status %d, want 200", resp.StatusCode) + } + want := strconv.Itoa(sustainedLimit - maxStreamsPerIP*streamWeight - 1) + if got := resp.Header.Get("X-RateLimit-Remaining"); got != want { + t.Errorf("X-RateLimit-Remaining = %s, want %s", got, want) + } +} + +func TestRateLimitMiddleware_RateLimitedGETReleasesStreamSlot(t *testing.T) { + h, rl := newRateLimitHandler(nil, okHandler, time.Second) + defer rl.Close() + discardLogs(rl) + const key = "10.7.0.1" + rl.allow(key, 1) + sh := &rl.shards[shardFor(key)] + sh.mu.Lock() + b, _ := sh.buckets[key].Value.(*bucket) + sh.mu.Unlock() + b.mu.Lock() + b.weightedCount = sustainedLimit + b.mu.Unlock() + + for i := 0; i < maxStreamsPerIP+1; i++ { + req := httptest.NewRequest(http.MethodGet, "/mcp", nil) + req.RemoteAddr = key + ":1" + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusTooManyRequests || !strings.Contains(rec.Body.String(), "Rate limit exceeded") { + t.Fatalf("GET %d: status %d body %q, want 429 rate-limit error", i+1, rec.Code, rec.Body.String()) + } + } + if got := rl.openStreams(key); got != 0 { + t.Fatalf("rate-limited GETs kept %d stream slots, want 0", got) + } +} + +func TestRateLimitMiddleware_ReadDeadlineErrors(t *testing.T) { + t.Run("unsupported writer is logged once and counted", func(t *testing.T) { + h, rl := newRateLimitHandler(nil, okHandler, time.Second) + defer rl.Close() + var logs []string + var mu sync.Mutex + rl.logf = func(format string, v ...any) { + mu.Lock() + defer mu.Unlock() + logs = append(logs, fmt.Sprintf(format, v...)) + } + + const sessionID = "mcp-session-0b5e7c1a-1111-2222-3333-444455556666" + const sqlMarker = "SELECT secret_marker_column" + for i := 0; i < 3; i++ { + body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"validate_sql","arguments":{"sql":"` + sqlMarker + `"}}}` + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + req.RemoteAddr = "10.6.0.1:1" + req.Header.Set("Mcp-Session-Id", sessionID) + rec := httptest.NewRecorder() // does not support read deadlines + h.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("request %d: status %d, want 200", i+1, rec.Code) + } + } + // Each request fails to set and to clear the deadline. + if got := rl.deadlineErrs.Load(); got != 6 { + t.Errorf("deadlineErrs = %d, want 6", got) + } + mu.Lock() + defer mu.Unlock() + if len(logs) != 1 { + t.Fatalf("logged %d warnings, want exactly 1: %q", len(logs), logs) + } + if !strings.Contains(logs[0], "does not support read deadlines") { + t.Errorf("warning does not explain the problem: %q", logs[0]) + } + for _, secret := range []string{sessionID, "mcp-session-", sqlMarker, "10.6.0.1"} { + if strings.Contains(logs[0], secret) { + t.Errorf("warning leaks %q: %q", secret, logs[0]) + } + } + }) + t.Run("other errors log only the error type", func(t *testing.T) { + rl := NewRateLimiter() + defer rl.Close() + var logs []string + rl.logf = func(format string, v ...any) { logs = append(logs, fmt.Sprintf(format, v...)) } + rl.noteDeadlineErr(nil) + if rl.deadlineErrs.Load() != 0 || len(logs) != 0 { + t.Fatalf("nil error was recorded: count=%d logs=%q", rl.deadlineErrs.Load(), logs) + } + rl.noteDeadlineErr(&net.OpError{Op: "set", Net: "tcp", Addr: &net.TCPAddr{IP: net.ParseIP("192.0.2.44"), Port: 443}, Err: errors.New("use of closed network connection")}) + rl.noteDeadlineErr(errors.New("again")) + if got := rl.deadlineErrs.Load(); got != 2 { + t.Errorf("deadlineErrs = %d, want 2", got) + } + if len(logs) != 1 { + t.Fatalf("logged %d warnings, want 1: %q", len(logs), logs) + } + if strings.Contains(logs[0], "192.0.2.44") || !strings.Contains(logs[0], "*net.OpError") { + t.Errorf("warning = %q, want the error type without its address", logs[0]) + } + }) + t.Run("real connection has no deadline errors", func(t *testing.T) { + h, rl := newRateLimitHandler(nil, okHandler, time.Second) + defer rl.Close() + ts := httptest.NewServer(h) + defer ts.Close() + resp, err := http.Post(ts.URL+"/mcp", "application/json", strings.NewReader(toolCallBody("1", "validate_sql"))) + if err != nil { + t.Fatalf("post: %v", err) + } + resp.Body.Close() + if got := rl.deadlineErrs.Load(); got != 0 { + t.Errorf("deadlineErrs = %d, want 0", got) + } + }) +} + +func TestRateLimitMiddleware_UnparseableBodyCostsBudget(t *testing.T) { + h, rl := newRateLimitHandler(nil, okHandler, time.Second) + defer rl.Close() + discardLogs(rl) + const key = "10.5.0.1" + for _, body := range []string{`{not json`, `[]`, `[{"method":`} { + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + req.RemoteAddr = key + ":1" + h.ServeHTTP(httptest.NewRecorder(), req) + } + if got := rl.weightedCount(t, key); got != 3*minUnparsedWeight { + t.Fatalf("weightedCount = %d, want %d", got, 3*minUnparsedWeight) + } + // A body-less request (e.g. DELETE to end a session) stays free. + req := httptest.NewRequest(http.MethodDelete, "/mcp", nil) + req.RemoteAddr = key + ":1" + h.ServeHTTP(httptest.NewRecorder(), req) + if got := rl.weightedCount(t, key); got != 3*minUnparsedWeight { + t.Fatalf("DELETE was charged: weightedCount = %d, want %d", got, 3*minUnparsedWeight) + } +} + +const testProxySecret = "s3cret-0123456789abcdef" + +func TestClientKey_ProxySecret(t *testing.T) { + gated := &Config{TrustedProxyHeader: "Cf-Connecting-Ip", TrustedProxySecretHeader: "X-Origin-Secret", TrustedProxySecret: testProxySecret} + gatedHops := &Config{TrustedProxyHops: 1, TrustedProxySecretHeader: "X-Origin-Secret", TrustedProxySecret: testProxySecret} + tests := []struct { + name string + cfg *Config + secrets []string + want string + }{ + {name: "matching secret trusts header", cfg: gated, secrets: []string{testProxySecret}, want: "203.0.113.9"}, + {name: "missing secret falls back", cfg: gated, want: "10.0.0.1"}, + {name: "wrong secret falls back", cfg: gated, secrets: []string{"s3cret-0123456789abcdeX"}, want: "10.0.0.1"}, + {name: "secret prefix falls back", cfg: gated, secrets: []string{testProxySecret[:8]}, want: "10.0.0.1"}, + {name: "secret with suffix falls back", cfg: gated, secrets: []string{testProxySecret + "x"}, want: "10.0.0.1"}, + {name: "empty secret falls back", cfg: gated, secrets: []string{""}, want: "10.0.0.1"}, + {name: "last value is the proxy's", cfg: gated, secrets: []string{"forged", testProxySecret}, want: "203.0.113.9"}, + {name: "forged last value falls back", cfg: gated, secrets: []string{testProxySecret, "forged"}, want: "10.0.0.1"}, + {name: "matching secret trusts XFF hops", cfg: gatedHops, secrets: []string{testProxySecret}, want: "198.51.100.4"}, + {name: "missing secret ignores XFF hops", cfg: gatedHops, want: "10.0.0.1"}, + {name: "secret header without secret never trusts", cfg: &Config{TrustedProxyHeader: "Cf-Connecting-Ip", TrustedProxySecretHeader: "X-Origin-Secret"}, secrets: []string{""}, want: "10.0.0.1"}, + {name: "secret without header never trusts", cfg: &Config{TrustedProxyHeader: "Cf-Connecting-Ip", TrustedProxySecret: testProxySecret}, want: "10.0.0.1"}, + {name: "no secret configured keeps old behaviour", cfg: &Config{TrustedProxyHeader: "Cf-Connecting-Ip"}, want: "203.0.113.9"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/mcp", nil) + req.RemoteAddr = "10.0.0.1:1234" + req.Header.Set("CF-Connecting-IP", "203.0.113.9") + req.Header.Set("X-Forwarded-For", "6.6.6.6, 198.51.100.4") + for _, v := range tt.secrets { + req.Header.Add("X-Origin-Secret", v) + } + if got := newClientIPResolver(tt.cfg).clientKey(req); got != tt.want { + t.Errorf("clientKey() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestRateLimitMiddleware_ProxySecret(t *testing.T) { + cfg := &Config{TrustedProxyHeader: "Cf-Connecting-Ip", TrustedProxySecretHeader: "X-Origin-Secret", TrustedProxySecret: testProxySecret} + var seen []string + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seen = append(seen, r.Header.Get("X-Origin-Secret")) + w.WriteHeader(http.StatusOK) + }) + h, rl := newRateLimitHandler(cfg, next, time.Second) + defer rl.Close() + discardLogs(rl) + + send := func(cfIP, secret string) int { + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(toolCallBody("1", "validate_sql"))) + req.RemoteAddr = "10.4.0.1:1" + req.Header.Set("CF-Connecting-IP", cfIP) + if secret != "" { + req.Header.Set("X-Origin-Secret", secret) + } + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec.Code + } + + // Without the secret, rotating CF-Connecting-IP does not escape the + // RemoteAddr bucket. + last := 0 + for i := 0; i <= burstLimit; i++ { + last = send(fmt.Sprintf("11.0.%d.1", i), "") + } + if last != http.StatusTooManyRequests { + t.Fatalf("rotating CF-Connecting-IP without the secret: status %d, want 429", last) + } + // With the secret, each forwarded client gets its own bucket. + for i := 0; i <= burstLimit; i++ { + if code := send(fmt.Sprintf("12.0.%d.1", i), testProxySecret); code != http.StatusOK { + t.Fatalf("client %d with the secret: status %d, want 200", i, code) + } + } + for _, v := range seen { + if v != "" { + t.Fatalf("secret header reached the MCP handler: %q", v) + } + } +} diff --git a/pkg/mcp/server.go b/pkg/mcp/server.go index 8471922b..50fd3aaa 100644 --- a/pkg/mcp/server.go +++ b/pkg/mcp/server.go @@ -94,6 +94,12 @@ func (s *Server) ToolCount() int { // without starting an HTTP server. Use this when you need to wrap // the handler with additional middleware (e.g., rate limiting). // +// The returned handler does NOT apply RateLimitMiddleware or +// NewRateLimitMiddleware, so it has no per-client rate limit, GET stream +// cap, body read deadline or 64KB body cap. Wrap it with +// NewRateLimitMiddleware (as cmd/gosqlx-mcp does) before exposing it to +// untrusted clients. +// // Every call shares one streamable HTTP transport, whose idle-session sweeper // runs in a background goroutine until Shutdown is called. func (s *Server) Handler() http.Handler { @@ -232,9 +238,11 @@ func (s *Server) MCPServer() *mcpserver.MCPServer { // It blocks until ctx is cancelled or a fatal error occurs; either way the // HTTP server and the transport's session sweeper are shut down. // -// Start applies bearer auth only. It does NOT apply rate limiting or the -// request body size cap: wrap Handler with RateLimitMiddleware (as -// cmd/gosqlx-mcp does) when the server is exposed to untrusted clients. +// Start applies bearer auth only. It does NOT apply RateLimitMiddleware or +// NewRateLimitMiddleware, so there is no rate limit, GET stream cap, body +// read deadline or request body size cap. When the server is exposed to +// untrusted clients, build your own http.Server around Handler wrapped with +// NewRateLimitMiddleware (as cmd/gosqlx-mcp does) instead of calling Start. func (s *Server) Start(ctx context.Context) error { httpSrv := &http.Server{ Addr: s.cfg.Addr(),