diff --git a/CHANGELOG.md b/CHANGELOG.md index afafcb1a..11668f0b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,24 @@ Runtime behaviour does change (rendered SQL, error codes and text, parse acceptance, CLI refusals), so read **Behavior changes** before upgrading. ### Behavior changes +- **`gosqlx-mcp` HTTP hardening**: rate limiting now keys on the TCP peer + address and no longer trusts `X-Forwarded-For` (which let any client evade + limits and grow limiter memory without bound). Behind a proxy, set + `GOSQLX_MCP_TRUSTED_PROXY_HEADER` (e.g. `CF-Connecting-IP`, as `render.yaml` + does) or `GOSQLX_MCP_TRUSTED_PROXY_HOPS`. Limiter memory is capped with + least-recently-seen eviction; GET (SSE) streams are capped at 4 per IP and + no longer tighten limits for everyone; JSON-RPC batches are charged per tool + call. Denials are HTTP 429 with `Retry-After` and the request id (a denied + `initialize` keeps HTTP 200 for mcp-go clients). Oversized bodies (over + 64 KB) get 413, bodies not received within 30s get 408 or a closed + connection on any route, idle connections close after 120s and headers are + capped at 32 KB. Bearer tokens are compared in constant time and the + `Bearer` scheme is case-insensitive. Tool input problems (missing, empty or + non-string `sql`, parse/format failures, wrongly typed options) are now + `isError` tool results naming the argument instead of JSON-RPC -32603 + internal errors, and `indent_size` must be an integer from 0 to 16. + New `pkg/mcp` API: `NewRateLimitMiddleware(cfg, next)`, + `Config.TrustedProxyHeader`, `Config.TrustedProxyHops`. - **Comments now count toward `MaxTokens`**: tokens plus collected comments may not exceed `MaxTokens` (1,000,000). Past that the tokenizer returns the existing E1007 token-limit error at the comment or token that crosses the diff --git a/cmd/gosqlx-mcp/main.go b/cmd/gosqlx-mcp/main.go index 0756875c..ad9c16cf 100644 --- a/cmd/gosqlx-mcp/main.go +++ b/cmd/gosqlx-mcp/main.go @@ -23,6 +23,8 @@ // GOSQLX_MCP_PORT bind port (default: 8080) // GOSQLX_MCP_AUTH_TOKEN bearer token; empty disables auth // 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 // // Pass --stdio to serve MCP over stdin/stdout instead of HTTP. // @@ -41,6 +43,7 @@ import ( "net/http" "os" "os/signal" + "strings" "syscall" "time" @@ -49,6 +52,18 @@ import ( mcpserver "github.com/mark3labs/mcp-go/server" ) +// HTTP server limits. See runHTTP for why there is no Read/WriteTimeout. +const ( + readHeaderTimeout = 10 * time.Second + idleTimeout = 120 * time.Second + maxHeaderBytes = 32 << 10 // 32KB + // otherBodyTimeout bounds request-body reads on routes other than /mcp, + // which never read a body. net/http drains an unread body (up to 256KB) + // before replying, so without a deadline a trickled body holds the + // connection open forever. /mcp sets its own deadline in the rate limiter. + otherBodyTimeout = 30 * time.Second +) + func main() { if err := run(); err != nil { fmt.Fprintf(os.Stderr, "gosqlx-mcp: %v\n", err) @@ -92,11 +107,7 @@ func runHTTP() error { srv := gosqlxmcp.New(cfg) - httpSrv := &http.Server{ - Addr: cfg.Addr(), - Handler: newMux(cfg, srv), - ReadHeaderTimeout: 10 * time.Second, - } + httpSrv := newHTTPServer(cfg, srv) go func() { <-ctx.Done() @@ -105,13 +116,54 @@ func runHTTP() error { _ = httpSrv.Shutdown(shutdownCtx) }() - log.Printf("gosqlx-mcp: listening on %s (auth=%v)\n", cfg.Addr(), cfg.AuthEnabled()) + log.Printf("gosqlx-mcp: listening on %s (auth=%v, client-ip=%s)\n", cfg.Addr(), cfg.AuthEnabled(), clientIPSource(cfg)) if err := httpSrv.ListenAndServe(); err != nil && err != http.ErrServerClosed { return fmt.Errorf("server error: %w", err) } return nil } +// newHTTPServer builds the http.Server for runHTTP. +// +// No ReadTimeout/WriteTimeout: they would cut off long-lived GET (SSE) +// streams. POST bodies get their own read deadline in the rate-limit +// middleware; IdleTimeout closes idle keep-alive connections. +func newHTTPServer(cfg *gosqlxmcp.Config, srv *gosqlxmcp.Server) *http.Server { + return &http.Server{ + Addr: cfg.Addr(), + Handler: withBodyDeadline(newMux(cfg, srv), otherBodyTimeout), + ReadHeaderTimeout: readHeaderTimeout, + IdleTimeout: idleTimeout, + MaxHeaderBytes: maxHeaderBytes, + } +} + +// withBodyDeadline sets a read deadline of timeout on every request outside +// /mcp, so slow or stalled bodies sent to /health, the challenge route or an +// unknown path cannot keep connections open. /mcp is left alone: its POST +// bodies get a deadline in the rate limiter and its GET (SSE) streams must +// not be cut off. +func withBodyDeadline(next http.Handler, timeout time.Duration) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/mcp" && !strings.HasPrefix(r.URL.Path, "/mcp/") { + _ = http.NewResponseController(w).SetReadDeadline(time.Now().Add(timeout)) + } + next.ServeHTTP(w, r) + }) +} + +// clientIPSource describes where rate limiting takes the client IP from. +func clientIPSource(cfg *gosqlxmcp.Config) string { + switch { + case cfg.TrustedProxyHeader != "": + return cfg.TrustedProxyHeader + case cfg.TrustedProxyHops > 0: + return fmt.Sprintf("X-Forwarded-For[-%d]", cfg.TrustedProxyHops) + default: + return "RemoteAddr" + } +} + // newMux builds the HTTP routes: // // /mcp, /mcp/ MCP streamable HTTP (rate limit → auth → MCP) @@ -121,7 +173,7 @@ func runHTTP() error { // Every other path (including OAuth discovery probes from connector wizards) // returns 404, which tells clients that no OAuth is configured. func newMux(cfg *gosqlxmcp.Config, srv *gosqlxmcp.Server) *http.ServeMux { - handler := gosqlxmcp.RateLimitMiddleware(srv.Handler()) + handler := gosqlxmcp.NewRateLimitMiddleware(cfg, srv.Handler()) mux := http.NewServeMux() mux.Handle("/mcp", handler) diff --git a/cmd/gosqlx-mcp/main_test.go b/cmd/gosqlx-mcp/main_test.go index e73b8807..1b29357d 100644 --- a/cmd/gosqlx-mcp/main_test.go +++ b/cmd/gosqlx-mcp/main_test.go @@ -17,12 +17,15 @@ package main import ( "context" "encoding/json" + "errors" "fmt" "io" + "net" "net/http" "net/http/httptest" "strings" "testing" + "time" "github.com/ajitpratap0/GoSQLX/pkg/gosqlx" gosqlxmcp "github.com/ajitpratap0/GoSQLX/pkg/mcp" @@ -190,6 +193,8 @@ func TestNewMux_HTTPChain(t *testing.T) { } t.Run("rate limit", func(t *testing.T) { + // X-Forwarded-For is ignored by default, so the limit keys on the + // loopback peer; the header only proves spoofing does not help. const xff = "198.51.100.7" resp, _ := mcpPost(t, ts.URL+"/mcp", token, "", xff, initializeRequest) sid := resp.Header.Get("Mcp-Session-Id") @@ -199,11 +204,17 @@ func TestNewMux_HTTPChain(t *testing.T) { resp, body := mcpPost(t, ts.URL+"/mcp", token, sid, xff, call) if strings.Contains(body, "Rate limit exceeded") { limited = true - if !strings.Contains(body, `"code":-32000`) { - t.Errorf("rate-limit response is not a JSON-RPC -32000 error: %s", body) + if resp.StatusCode != http.StatusTooManyRequests { + t.Errorf("rate-limit status = %d, want 429", resp.StatusCode) } - if resp.Header.Get("X-RateLimit-Remaining") != "0" { - t.Errorf("X-RateLimit-Remaining = %q, want 0", resp.Header.Get("X-RateLimit-Remaining")) + if !strings.Contains(body, `"code":-32000`) || !strings.Contains(body, `"id":2`) { + t.Errorf("rate-limit response is not a JSON-RPC -32000 error for id 2: %s", body) + } + if resp.Header.Get("Retry-After") == "" { + t.Error("rate-limit response has no Retry-After header") + } + if resp.Header.Get("X-RateLimit-Remaining") == "" { + t.Error("rate-limit response has no X-RateLimit-Remaining header") } } } @@ -225,3 +236,125 @@ func TestNewMux_HTTPChain(t *testing.T) { } }) } + +func TestNewHTTPServer_Limits(t *testing.T) { + cfg := gosqlxmcp.DefaultConfig() + s := newHTTPServer(cfg, gosqlxmcp.New(cfg)) + tests := []struct { + name string + got any + want any + }{ + {"ReadHeaderTimeout", s.ReadHeaderTimeout, readHeaderTimeout}, + {"IdleTimeout", s.IdleTimeout, idleTimeout}, + {"MaxHeaderBytes", s.MaxHeaderBytes, maxHeaderBytes}, + // A global ReadTimeout/WriteTimeout would cut off GET SSE streams. + {"ReadTimeout unset", s.ReadTimeout, time.Duration(0)}, + {"WriteTimeout unset", s.WriteTimeout, time.Duration(0)}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if tt.got != tt.want { + t.Errorf("%s = %v, want %v", tt.name, tt.got, tt.want) + } + }) + } +} + +func TestNewMux_RateLimitIgnoresSpoofedXFF(t *testing.T) { + cfg := gosqlxmcp.DefaultConfig() + ts := httptest.NewServer(newMux(cfg, gosqlxmcp.New(cfg))) + t.Cleanup(ts.Close) + + body := `{"jsonrpc":"2.0","id":1,"method":"tools/list"}` + var last int + for i := 0; i < 15; i++ { + req, err := http.NewRequest(http.MethodPost, ts.URL+"/mcp", strings.NewReader(body)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set("X-Forwarded-For", fmt.Sprintf("11.0.%d.1", i)) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("POST: %v", err) + } + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() + last = resp.StatusCode + } + if last != http.StatusTooManyRequests { + t.Fatalf("15 rapid requests with rotating X-Forwarded-For: last status %d, want 429", last) + } +} + +func TestClientIPSource(t *testing.T) { + tests := []struct { + name string + cfg gosqlxmcp.Config + want string + }{ + {"default", gosqlxmcp.Config{}, "RemoteAddr"}, + {"header", gosqlxmcp.Config{TrustedProxyHeader: "Cf-Connecting-Ip"}, "Cf-Connecting-Ip"}, + {"hops", gosqlxmcp.Config{TrustedProxyHops: 2}, "X-Forwarded-For[-2]"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := clientIPSource(&tt.cfg); got != tt.want { + t.Errorf("clientIPSource() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestWithBodyDeadline_SlowBodyOutsideMCP(t *testing.T) { + tests := []struct { + name string + wrap bool + wantClose bool + }{ + {"with deadline", true, true}, + // Control: without the wrapper the server keeps waiting for the body, + // proving the test detects the slowloris case. + {"without deadline", false, false}, + } + for _, tt := range tests { + for _, path := range []string{"/health", "/nope"} { + t.Run(tt.name+path, func(t *testing.T) { + cfg := gosqlxmcp.DefaultConfig() + var h http.Handler = newMux(cfg, gosqlxmcp.New(cfg)) + if tt.wrap { + h = withBodyDeadline(h, 200*time.Millisecond) + } + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + srv := &http.Server{Handler: h, ReadHeaderTimeout: time.Second} + go func() { _ = srv.Serve(ln) }() + t.Cleanup(func() { _ = srv.Close() }) + + conn, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + // Promise a body that never arrives. + req := "POST " + path + " HTTP/1.1\r\nHost: x\r\nContent-Length: 1000\r\n\r\nabc" + if _, err := conn.Write([]byte(req)); err != nil { + t.Fatalf("write: %v", err) + } + _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) + // The server replies, then must close the connection once the + // body-drain deadline passes instead of waiting for the body. + _, err = io.ReadAll(conn) + var ne net.Error + timedOut := errors.As(err, &ne) && ne.Timeout() + if closed := err == nil; closed != tt.wantClose || (!tt.wantClose && !timedOut) { + t.Errorf("connection closed by server = %v (err %v), want %v", closed, err, tt.wantClose) + } + }) + } + } +} diff --git a/docs/MCP_GUIDE.md b/docs/MCP_GUIDE.md index 5f312188..9a448014 100644 --- a/docs/MCP_GUIDE.md +++ b/docs/MCP_GUIDE.md @@ -84,6 +84,12 @@ 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 +- **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`. + +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. ### Health Check @@ -210,6 +216,9 @@ The `gosqlx-mcp` server is configured exclusively via environment variables. No | `GOSQLX_MCP_HOST` | `127.0.0.1` | string | Any valid bind address | | `GOSQLX_MCP_PORT` | `8080` | integer | 1–65535; non-integer or out-of-range → startup error | | `GOSQLX_MCP_AUTH_TOKEN` | *(empty)* | string | Empty = auth disabled; value is whitespace-trimmed | +| `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` | **Notes:** @@ -217,6 +226,21 @@ The `gosqlx-mcp` server is configured exclusively via environment variables. No - `GOSQLX_MCP_AUTH_TOKEN` enables bearer token auth for all requests when set to a non-empty string. - MCP server configuration is independent of `.gosqlx.yml` - the YAML config is not read by `gosqlx-mcp`. +### Client IP and reverse proxies + +Rate limits are tracked per client IP (IPv6 clients per `/64`). By default the server uses the TCP peer address (`RemoteAddr`) and ignores `X-Forwarded-For`, `X-Real-IP` and similar headers, because any client can send them: trusting the left-most `X-Forwarded-For` entry would let a client pick a fresh rate-limit bucket on every request. + +Behind a reverse proxy every request comes from the proxy, so all users would share one bucket. Opt in to a forwarded address: + +- `GOSQLX_MCP_TRUSTED_PROXY_HEADER=CF-Connecting-IP`: use a single-IP header that your proxy overwrites on every request. Cloudflare sets `CF-Connecting-IP` to the visitor address. The hosted server runs on Render, whose edge is Cloudflare, and `render.yaml` sets this. +- `GOSQLX_MCP_TRUSTED_PROXY_HOPS=N`: use the entry `N` positions from the right of `X-Forwarded-For` (the address appended by the outermost of `N` trusted proxies). Entries further left are client-supplied and are ignored. + +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. + +### 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. + --- ## Authentication @@ -237,13 +261,19 @@ curl -s -X POST http://127.0.0.1:8080/mcp \ -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"validate_sql","arguments":{"sql":"SELECT 1"}}}' ``` -Requests missing or carrying an incorrect token receive HTTP `401 Unauthorized`. The `BearerAuthMiddleware` wraps the streamable HTTP handler and is a no-op when auth is disabled. +Requests missing or carrying an incorrect token receive HTTP `401 Unauthorized`. The `Bearer` scheme is matched case-insensitively and surrounding whitespace is ignored. Tokens are compared in constant time (SHA-256 digests with `crypto/subtle`). The `BearerAuthMiddleware` wraps the streamable HTTP handler and is a no-op when auth is disabled. --- ## Tools Reference -All tools accept a required `sql` string parameter. The server returns tool-semantic failures (e.g., invalid SQL) as a valid JSON result with `valid: false` rather than as a protocol error. Protocol errors (missing required parameter, server fault) return an MCP error response. +All tools accept a required `sql` string parameter. Errors are reported at three levels: + +- **Result data**: `validate_sql` reports invalid SQL as a normal result with `valid: false`; `analyze_sql` reports sub-tool failures under its `errors` key. +- **Tool errors** (`isError: true`, message in the text content): input problems the model can fix, such as a missing, empty or non-string `sql`, SQL that `parse_sql`, `extract_metadata` or `format_sql` cannot parse, or an optional argument of the wrong type or out of range (the message names the argument). +- **JSON-RPC errors**: unknown tools, malformed requests and genuine server faults. + +Optional arguments must match their schema type. Absent or `null` means the default. For clients that stringify scalars, `"true"`/`"false"` and decimal-integer strings are also accepted. --- @@ -301,7 +331,7 @@ Invalid SQL: | Parameter | Type | Required | Default | Description | |-----------|------|----------|---------|-------------| | `sql` | string | yes | - | The SQL string to format | -| `indent_size` | integer | no | `2` | Spaces per indent level | +| `indent_size` | integer | no | `2` | Spaces per indent level (0–16) | | `uppercase_keywords` | boolean | no | `false` | Uppercase SQL keywords | | `add_semicolon` | boolean | no | `false` | Append a trailing semicolon | @@ -698,6 +728,8 @@ func main() { | `New` | `(cfg *Config) *Server` | Create server with all 7 tools registered | | `(*Server).Start` | `(ctx context.Context) error` | Bind, serve, block until ctx cancelled | | `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 | | `(*Config).Addr` | `() string` | Returns `"host:port"` | | `(*Config).AuthEnabled` | `() bool` | Reports whether auth token is set | @@ -708,6 +740,10 @@ type Config struct { Host string // GOSQLX_MCP_HOST (default "127.0.0.1") Port int // GOSQLX_MCP_PORT (default 8080, range 1–65535) AuthToken string // GOSQLX_MCP_AUTH_TOKEN (default "" = auth disabled) + + 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) } ``` diff --git a/pkg/mcp/config.go b/pkg/mcp/config.go index 0b56ec57..ee998956 100644 --- a/pkg/mcp/config.go +++ b/pkg/mcp/config.go @@ -32,10 +32,15 @@ // GOSQLX_MCP_AUTH_TOKEN bearer token; empty disables auth // GOSQLX_MCP_OPENAI_APPS_CHALLENGE ChatGPT app domain-verification token, // served at /.well-known/openai-apps-challenge +// GOSQLX_MCP_TRUSTED_PROXY_HEADER header set by a trusted proxy that holds +// 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) package mcp import ( "fmt" + "net/http" "os" "strconv" "strings" @@ -62,8 +67,26 @@ type Config struct { // served as plain text at /.well-known/openai-apps-challenge. // Source: GOSQLX_MCP_OPENAI_APPS_CHALLENGE (default: "" - not served) OpenAIAppsChallenge string + + // TrustedProxyHeader names a header, written by a trusted reverse proxy, + // that carries the client IP used as the rate-limit key (for example + // "CF-Connecting-IP" behind Cloudflare). Only set it when every request + // reaches the server through that proxy, since clients can send any + // header themselves. Invalid values fall back to RemoteAddr. + // Source: GOSQLX_MCP_TRUSTED_PROXY_HEADER (default: "" - not used) + TrustedProxyHeader string + + // TrustedProxyHops is the number of trusted proxies in front of the + // server that append to X-Forwarded-For. When > 0 the client IP is the + // entry that many positions from the right; entries further left are + // client-supplied and ignored. + // Source: GOSQLX_MCP_TRUSTED_PROXY_HOPS (default: 0 - X-Forwarded-For ignored) + TrustedProxyHops int } +// maxTrustedProxyHops bounds GOSQLX_MCP_TRUSTED_PROXY_HOPS. +const maxTrustedProxyHops = 16 + // LoadConfig reads configuration from environment variables, applying defaults // for any variables that are unset or empty. func LoadConfig() (*Config, error) { @@ -92,9 +115,46 @@ func LoadConfig() (*Config, error) { cfg.OpenAIAppsChallenge = v } + if v := strings.TrimSpace(os.Getenv("GOSQLX_MCP_TRUSTED_PROXY_HEADER")); v != "" { + if !validHeaderName(v) { + return nil, fmt.Errorf("GOSQLX_MCP_TRUSTED_PROXY_HEADER: %q is not a valid HTTP header name", v) + } + canonical := http.CanonicalHeaderKey(v) + if canonical == "X-Forwarded-For" { + return nil, fmt.Errorf("GOSQLX_MCP_TRUSTED_PROXY_HEADER: use GOSQLX_MCP_TRUSTED_PROXY_HOPS for X-Forwarded-For") + } + cfg.TrustedProxyHeader = canonical + } + + if v := strings.TrimSpace(os.Getenv("GOSQLX_MCP_TRUSTED_PROXY_HOPS")); v != "" { + hops, err := strconv.Atoi(v) + if err != nil { + return nil, fmt.Errorf("GOSQLX_MCP_TRUSTED_PROXY_HOPS: expected integer, got %q", v) + } + if hops < 0 || hops > maxTrustedProxyHops { + return nil, fmt.Errorf("GOSQLX_MCP_TRUSTED_PROXY_HOPS: %d is out of range (0–%d)", hops, maxTrustedProxyHops) + } + cfg.TrustedProxyHops = hops + } + return cfg, nil } +// validHeaderName reports whether s is a valid HTTP header field name +// (an RFC 9110 token). +func validHeaderName(s string) bool { + for i := 0; i < len(s); i++ { + c := s[i] + switch { + case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z', c >= '0' && c <= '9': + case strings.IndexByte("!#$%&'*+-.^_`|~", c) >= 0: + default: + return false + } + } + return s != "" +} + // DefaultConfig returns a Config with all defaults applied (auth disabled). func DefaultConfig() *Config { return &Config{ diff --git a/pkg/mcp/config_test.go b/pkg/mcp/config_test.go index 96d8e286..b63abc23 100644 --- a/pkg/mcp/config_test.go +++ b/pkg/mcp/config_test.go @@ -142,3 +142,44 @@ func TestLoadConfig_OpenAIAppsChallenge(t *testing.T) { }) } } + +func TestLoadConfig_TrustedProxy(t *testing.T) { + tests := []struct { + name string + header string + hops string + wantHeader string + wantHops int + wantErr bool + }{ + {name: "unset", wantHeader: "", wantHops: 0}, + {name: "cloudflare header canonicalised", header: " cf-connecting-ip ", wantHeader: "Cf-Connecting-Ip"}, + {name: "hops", hops: "2", wantHops: 2}, + {name: "both", header: "True-Client-IP", hops: "1", wantHeader: "True-Client-Ip", wantHops: 1}, + {name: "invalid header name", header: "CF Connecting IP", wantErr: true}, + {name: "XFF as header rejected", header: "x-forwarded-for", wantErr: true}, + {name: "hops non-numeric", hops: "one", wantErr: true}, + {name: "hops negative", hops: "-1", wantErr: true}, + {name: "hops too large", hops: "17", wantErr: true}, + } + 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) + cfg, err := LoadConfig() + if tt.wantErr { + if err == nil { + t.Fatalf("expected error, got config %+v", cfg) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cfg.TrustedProxyHeader != tt.wantHeader || cfg.TrustedProxyHops != tt.wantHops { + t.Errorf("got header=%q hops=%d, want header=%q hops=%d", + cfg.TrustedProxyHeader, cfg.TrustedProxyHops, tt.wantHeader, tt.wantHops) + } + }) + } +} diff --git a/pkg/mcp/middleware.go b/pkg/mcp/middleware.go index 8076bfa4..5819048f 100644 --- a/pkg/mcp/middleware.go +++ b/pkg/mcp/middleware.go @@ -15,6 +15,8 @@ package mcp import ( + "crypto/sha256" + "crypto/subtle" "net/http" "strings" ) @@ -23,12 +25,17 @@ import ( // authentication when cfg.AuthEnabled() is true. When auth is disabled it // passes all requests through unchanged. On failure it responds 401 with a // WWW-Authenticate header. +// +// Tokens are compared as SHA-256 digests with crypto/subtle, so the +// comparison time reveals neither the token's content nor its length. func BearerAuthMiddleware(cfg *Config, next http.Handler) http.Handler { if !cfg.AuthEnabled() { return next } + want := sha256.Sum256([]byte(cfg.AuthToken)) return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if extractBearerToken(r) != cfg.AuthToken { + got := sha256.Sum256([]byte(extractBearerToken(r))) + if subtle.ConstantTimeCompare(got[:], want[:]) != 1 { w.Header().Set("WWW-Authenticate", `Bearer realm="gosqlx-mcp"`) http.Error(w, "Unauthorized", http.StatusUnauthorized) return @@ -37,13 +44,14 @@ func BearerAuthMiddleware(cfg *Config, next http.Handler) http.Handler { }) } -// extractBearerToken parses the "Authorization: Bearer " header. -// Returns an empty string if the header is absent or malformed. +// extractBearerToken parses the "Authorization: Bearer " header. The +// scheme is matched case-insensitively (RFC 9110 §11.1) and surrounding +// whitespace is trimmed. Returns an empty string if the header is absent or +// malformed. func extractBearerToken(r *http.Request) string { - const prefix = "Bearer " - auth := r.Header.Get("Authorization") - if !strings.HasPrefix(auth, prefix) { + scheme, token, ok := strings.Cut(strings.TrimSpace(r.Header.Get("Authorization")), " ") + if !ok || !strings.EqualFold(scheme, "Bearer") { return "" } - return auth[len(prefix):] + return strings.TrimSpace(token) } diff --git a/pkg/mcp/middleware_test.go b/pkg/mcp/middleware_test.go index f3ab7173..3a439cb5 100644 --- a/pkg/mcp/middleware_test.go +++ b/pkg/mcp/middleware_test.go @@ -140,6 +140,9 @@ func TestExtractBearerToken(t *testing.T) { {"no bearer", "Token mytoken", ""}, {"basic", "Basic dXNlcjpwYXNz", ""}, {"bearer no token", "Bearer ", ""}, + {"lowercase scheme", "bearer mytoken", "mytoken"}, + {"trimmed", " Bearer mytoken ", "mytoken"}, + {"no space", "Bearermytoken", ""}, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { @@ -154,3 +157,36 @@ func TestExtractBearerToken(t *testing.T) { }) } } + +func TestBearerAuthMiddleware_HeaderVariants(t *testing.T) { + cfg := &Config{Host: "127.0.0.1", Port: 8080, AuthToken: "secret"} + tests := []struct { + name string + header string + wantStatus int + }{ + {"canonical", "Bearer secret", http.StatusOK}, + {"lowercase scheme", "bearer secret", http.StatusOK}, + {"uppercase scheme", "BEARER secret", http.StatusOK}, + {"surrounding whitespace", " Bearer secret ", http.StatusOK}, + {"prefix of token", "Bearer secre", http.StatusUnauthorized}, + {"token with suffix", "Bearer secretx", http.StatusUnauthorized}, + {"empty token", "Bearer ", http.StatusUnauthorized}, + {"scheme only", "Bearer", http.StatusUnauthorized}, + {"wrong case token", "Bearer SECRET", http.StatusUnauthorized}, + } + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/mcp", nil) + req.Header.Set("Authorization", tc.header) + rr := httptest.NewRecorder() + BearerAuthMiddleware(cfg, next).ServeHTTP(rr, req) + if rr.Code != tc.wantStatus { + t.Errorf("status = %d, want %d", rr.Code, tc.wantStatus) + } + }) + } +} diff --git a/pkg/mcp/protocol_test.go b/pkg/mcp/protocol_test.go index 04d150f0..c19555a7 100644 --- a/pkg/mcp/protocol_test.go +++ b/pkg/mcp/protocol_test.go @@ -263,9 +263,14 @@ func TestProtocol_CallTool_MissingSQLReturnsError(t *testing.T) { req.Params.Name = "validate_sql" req.Params.Arguments = map[string]any{} - _, err := c.CallTool(ctx, req) - if err == nil { - t.Fatal("expected error for missing sql param via protocol") + // Input problems are reported as an isError tool result so the model + // can correct its call, not as a JSON-RPC internal error. + res, err := c.CallTool(ctx, req) + if err != nil { + t.Fatalf("expected isError tool result, got protocol error: %v", err) + } + if res == nil || !res.IsError { + t.Fatalf("expected isError tool result for missing sql param, got %+v", res) } } diff --git a/pkg/mcp/ratelimit.go b/pkg/mcp/ratelimit.go index 9fc90a68..e53d49d9 100644 --- a/pkg/mcp/ratelimit.go +++ b/pkg/mcp/ratelimit.go @@ -16,12 +16,16 @@ package mcp import ( "bytes" + "container/list" "encoding/json" + "errors" "fmt" "hash/fnv" "io" + "math" "net" "net/http" + "strconv" "strings" "sync" "sync/atomic" @@ -50,55 +54,116 @@ const ( numShards = 16 // number of shards for concurrent map access cleanupInterval = 5 * time.Minute // how often to run stale entry cleanup staleTimeout = 10 * time.Minute // idle time before an entry is removed - maxBodySize = 64 * 1024 // 64KB max request body size + + // maxBucketsPerShard caps how many clients each shard tracks. When a shard + // is full the least-recently-seen client is evicted, so memory stays bounded + // (numShards*maxBucketsPerShard buckets) no matter how many distinct + // addresses send traffic. + maxBucketsPerShard = 4096 + + // maxStreamsPerIP caps concurrent GET (SSE) streams from one client. + maxStreamsPerIP = 4 + // streamWeight is the sustained-budget cost of opening a GET stream. + streamWeight = 1 + + // bodyReadTimeout bounds how long a client may take to send a request + // body, so slow-body (slowloris) connections cannot be held open forever. + bodyReadTimeout = 30 * time.Second + + // 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. + maxBodySize = 64 * 1024 ) -// bucket tracks rate limit state for a single IP address. +// bucket tracks rate limit state for a single client key. type bucket struct { + key string + + // lastSeen is guarded by the owning shard's mutex. + lastSeen time.Time + mu sync.Mutex burstTokens float64 lastBurst time.Time weightedCount int windowStart time.Time - lastSeen time.Time } -// shard is one segment of the sharded map, holding its own lock and bucket map. +// shard is one segment of the sharded map, holding its own lock, bucket map +// and an LRU list (front = most recently seen) used for eviction and cleanup. type shard struct { - sync.RWMutex - buckets map[string]*bucket + mu sync.Mutex + buckets map[string]*list.Element + lru *list.List +} + +// streamCounter tracks open GET streams per client key. Entries are removed +// when their count drops to zero, so its size is bounded by open connections. +type streamCounter struct { + mu sync.Mutex + open map[string]int +} + +// acquire reserves a stream slot for key, reporting false when key already +// holds max streams. +func (s *streamCounter) acquire(key string, max int) bool { + s.mu.Lock() + defer s.mu.Unlock() + if s.open[key] >= max { + return false + } + s.open[key]++ + return true +} + +// release frees a slot reserved by acquire. +func (s *streamCounter) release(key string) { + s.mu.Lock() + defer s.mu.Unlock() + if s.open[key] <= 1 { + delete(s.open, key) + return + } + s.open[key]-- } // rateLimiter implements a three-layer rate limiting strategy: // per-IP burst limiting, tool-weighted sustained limiting, and adaptive load scaling. type rateLimiter struct { shards [numShards]shard + maxPerShard int activeReqs atomic.Int64 + streams streamCounter stopCleanup chan struct{} + closeOnce sync.Once } // NewRateLimiter creates a new rateLimiter, initializes all shards, // and starts the background cleanup goroutine. func NewRateLimiter() *rateLimiter { rl := &rateLimiter{ + maxPerShard: maxBucketsPerShard, + streams: streamCounter{open: make(map[string]int)}, stopCleanup: make(chan struct{}), } for i := range rl.shards { - rl.shards[i].buckets = make(map[string]*bucket) + rl.shards[i].buckets = make(map[string]*list.Element) + rl.shards[i].lru = list.New() } go rl.cleanupLoop() return rl } -// Close stops the background cleanup goroutine. +// Close stops the background cleanup goroutine. It is safe to call more than once. func (rl *rateLimiter) Close() { - close(rl.stopCleanup) + rl.closeOnce.Do(func() { close(rl.stopCleanup) }) } -// shardFor returns the shard index for the given IP using FNV hashing. -func shardFor(ip string) int { +// shardFor returns the shard index for the given key using FNV hashing. +func shardFor(key string) int { h := fnv.New32a() - _, _ = h.Write([]byte(ip)) + _, _ = h.Write([]byte(key)) return int(h.Sum32() % numShards) } @@ -116,36 +181,77 @@ func currentLimit(rl *rateLimiter) int { } } -// allow checks whether a request from the given IP with the given weight -// should be permitted. It returns: -// - allowed: whether the request is permitted -// - limit: the current sustained limit -// - remaining: how many weighted units remain in the window -// - resetTime: when the current window resets -func (rl *rateLimiter) allow(ip string, weight int) (bool, int, int, time.Time) { - idx := shardFor(ip) - sh := &rl.shards[idx] - - now := time.Now() - limit := currentLimit(rl) +// bucketFor returns the bucket for key, creating it (and evicting the +// least-recently-seen bucket when the shard is full) if needed. +func (rl *rateLimiter) bucketFor(key string, now time.Time) *bucket { + sh := &rl.shards[shardFor(key)] + sh.mu.Lock() + defer sh.mu.Unlock() + + if el, ok := sh.buckets[key]; ok { + b, _ := el.Value.(*bucket) + b.lastSeen = now + sh.lru.MoveToFront(el) + return b + } - sh.Lock() - b, ok := sh.buckets[ip] - if !ok { - b = &bucket{ - burstTokens: float64(burstLimit), - lastBurst: now, - windowStart: now, - lastSeen: now, + for sh.lru.Len() >= rl.maxPerShard { + oldest := sh.lru.Back() + if oldest == nil { + break + } + if old, ok := oldest.Value.(*bucket); ok { + delete(sh.buckets, old.key) } - sh.buckets[ip] = b + sh.lru.Remove(oldest) + } + + b := &bucket{ + key: key, + lastSeen: now, + burstTokens: float64(burstLimit), + lastBurst: now, + windowStart: now, } - sh.Unlock() + sh.buckets[key] = sh.lru.PushFront(b) + return b +} + +// decision is the outcome of a rate-limit check. +type decision struct { + allowed bool + // limit is the current sustained limit (weighted requests per minute). + limit int + // remaining is the unused sustained budget in the current window. + remaining int + // reset is when the current sustained window ends. + reset time.Time + // retryAfter is how long a denied client should wait; zero when allowed. + retryAfter time.Duration +} + +// allow checks whether a request from the given client key with the given +// weight should be permitted, consuming budget when it is. +func (rl *rateLimiter) allow(key string, weight int) decision { + now := time.Now() + limit := currentLimit(rl) + b := rl.bucketFor(key, now) b.mu.Lock() defer b.mu.Unlock() - b.lastSeen = now + // Reset the sustained window if a minute has passed. + if now.Sub(b.windowStart) >= time.Minute { + b.weightedCount = 0 + b.windowStart = now + } + d := decision{limit: limit, reset: b.windowStart.Add(time.Minute)} + remaining := func() int { + if r := limit - b.weightedCount; r > 0 { + return r + } + return 0 + } // --- Layer 1: Burst check (token bucket) --- elapsed := now.Sub(b.lastBurst).Seconds() @@ -156,47 +262,25 @@ func (rl *rateLimiter) allow(ip string, weight int) (bool, int, int, time.Time) b.lastBurst = now if b.burstTokens < 1.0 { - resetTime := now.Add(time.Duration((1.0-b.burstTokens)/float64(burstLimit)*1000) * time.Millisecond) - return false, limit, 0, resetTime - } - - // Weight-0 requests (initialize, list_tools) bypass sustained limiting. - if weight == 0 { - b.burstTokens -= 1.0 - remaining := limit - b.weightedCount - if remaining < 0 { - remaining = 0 - } - resetTime := b.windowStart.Add(time.Minute) - return true, limit, remaining, resetTime + d.remaining = remaining() + d.retryAfter = time.Duration((1.0 - b.burstTokens) / float64(burstLimit) * float64(time.Second)) + return d } // --- Layer 2: Sustained weighted check --- - // Reset window if a minute has passed. - if now.Sub(b.windowStart) >= time.Minute { - b.weightedCount = 0 - b.windowStart = now - } - - resetTime := b.windowStart.Add(time.Minute) - - if b.weightedCount+weight > limit { - remaining := limit - b.weightedCount - if remaining < 0 { - remaining = 0 - } - return false, limit, remaining, resetTime + // Weight-0 requests (initialize, tools/list) bypass sustained limiting. + if weight > 0 && b.weightedCount+weight > limit { + d.remaining = remaining() + d.retryAfter = d.reset.Sub(now) + return d } // Allowed: consume tokens. b.burstTokens -= 1.0 b.weightedCount += weight - - remaining := limit - b.weightedCount - if remaining < 0 { - remaining = 0 - } - return true, limit, remaining, resetTime + d.allowed = true + d.remaining = remaining() + return d } // cleanupLoop periodically removes entries that have been idle for longer @@ -209,56 +293,128 @@ func (rl *rateLimiter) cleanupLoop() { case <-rl.stopCleanup: return case <-ticker.C: - rl.cleanup() + rl.cleanup(time.Now()) } } } -// cleanup removes stale bucket entries from all shards. -// Collects stale IPs first, then deletes - avoids holding bucket mutex -// while modifying the shard map. -func (rl *rateLimiter) cleanup() { - now := time.Now() +// cleanup removes buckets not seen since now-staleTimeout. Each shard's LRU +// list is ordered by lastSeen, so only the stale tail is visited. +func (rl *rateLimiter) cleanup(now time.Time) { for i := range rl.shards { sh := &rl.shards[i] - sh.Lock() - var staleIPs []string - for ip, b := range sh.buckets { - b.mu.Lock() - stale := now.Sub(b.lastSeen) > staleTimeout - b.mu.Unlock() - if stale { - staleIPs = append(staleIPs, ip) + sh.mu.Lock() + for el := sh.lru.Back(); el != nil; { + b, ok := el.Value.(*bucket) + if ok && now.Sub(b.lastSeen) <= staleTimeout { + break } + prev := el.Prev() + if ok { + delete(sh.buckets, b.key) + } + sh.lru.Remove(el) + el = prev } - for _, ip := range staleIPs { - delete(sh.buckets, ip) - } - sh.Unlock() + sh.mu.Unlock() } } -// extractIP extracts the client IP address from the request. -// It checks X-Forwarded-For first (taking the first IP in the list), -// then falls back to RemoteAddr. -func extractIP(r *http.Request) string { - if xff := r.Header.Get("X-Forwarded-For"); xff != "" { - parts := strings.Split(xff, ",") - ip := strings.TrimSpace(parts[0]) - if ip != "" { - return ip +// clientIPResolver derives the rate-limit key for a request. By default it +// uses the TCP peer address only; forwarding headers are client-controlled +// and are honoured only when the operator opts in via Config. +type clientIPResolver struct { + // header, when set, is a single-IP header written by a trusted proxy + // (e.g. CF-Connecting-IP). + header string + // 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 +} + +// newClientIPResolver builds a resolver from cfg (nil means RemoteAddr only). +func newClientIPResolver(cfg *Config) clientIPResolver { + if cfg == nil { + return clientIPResolver{} + } + return clientIPResolver{header: cfg.TrustedProxyHeader, hops: cfg.TrustedProxyHops} +} + +// clientKey returns the normalised client address for r. Invalid or missing +// forwarded values fall back to RemoteAddr. +func (c clientIPResolver) clientKey(r *http.Request) string { + if c.header != "" { + if vals := r.Header.Values(c.header); len(vals) > 0 { + // The trusted proxy's value is the last one it wrote. + entries := splitList(vals[len(vals)-1]) + if len(entries) > 0 { + if key, ok := normalizeIP(entries[len(entries)-1]); ok { + return key + } + } + } + } + if c.hops > 0 { + var entries []string + for _, v := range r.Header.Values("X-Forwarded-For") { + entries = append(entries, splitList(v)...) + } + if len(entries) >= c.hops { + if key, ok := normalizeIP(entries[len(entries)-c.hops]); ok { + return key + } } } host, _, err := net.SplitHostPort(r.RemoteAddr) if err != nil { - return r.RemoteAddr + host = r.RemoteAddr + } + if key, ok := normalizeIP(host); ok { + return key } return host } +// splitList splits a comma-separated header value into trimmed, non-empty entries. +func splitList(v string) []string { + parts := strings.Split(v, ",") + out := parts[:0] + for _, p := range parts { + if p = strings.TrimSpace(p); p != "" { + out = append(out, p) + } + } + return out +} + +// ipv6KeyMask groups IPv6 clients by /64, the smallest prefix normally +// assigned to one subscriber, so rotating addresses inside it does not +// yield fresh budgets. +var ipv6KeyMask = net.CIDRMask(64, 128) + +// normalizeIP parses s (optionally with a port or brackets) and returns its +// canonical rate-limit key: the dotted IPv4 address, or the IPv6 /64 prefix. +func normalizeIP(s string) (string, bool) { + ip := net.ParseIP(s) + if ip == nil { + host, _, err := net.SplitHostPort(s) + if err != nil { + return "", false + } + if ip = net.ParseIP(host); ip == nil { + return "", false + } + } + if v4 := ip.To4(); v4 != nil { + return v4.String(), true + } + return ip.Mask(ipv6KeyMask).String() + "/64", true +} + // jsonRPCRequest is used to partially decode a JSON-RPC request to extract -// the method and tool name. +// the id, method and tool name. type jsonRPCRequest struct { + ID json.RawMessage `json:"id"` Method string `json:"method"` Params json.RawMessage `json:"params"` } @@ -268,83 +424,206 @@ type toolCallParams struct { Name string `json:"name"` } -// extractToolWeight reads the request body to determine the tool weight. -// It re-buffers the body so downstream handlers can read it again. -// Returns the weight (0 for non-tool requests) and the modified request. -func extractToolWeight(r *http.Request) (int, *http.Request) { - if r.Body == nil { - return 0, r +// requestInfo is what the rate limiter needs to know about a JSON-RPC body. +type requestInfo struct { + // weight is the sustained-budget cost (0 for non-tool requests). + weight int + // id is the JSON-RPC id to echo in an error, or nil when unknown. + id json.RawMessage + // method is the JSON-RPC method of a single (non-batch) request. + method string +} + +// weightOf returns the cost of one decoded JSON-RPC message. +func weightOf(req jsonRPCRequest) int { + if req.Method != "tools/call" { + return 0 + } + var params toolCallParams + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + return 1 } + if w, ok := toolWeights[params.Name]; ok { + return w + } + return 1 +} - limited := http.MaxBytesReader(nil, r.Body, maxBodySize) - body, err := io.ReadAll(limited) - if err != nil { - // On read error (e.g., body too large), re-buffer what we got - // and treat as weight 1 (default cost). - r.Body = io.NopCloser(bytes.NewReader(body)) - return 1, r +// 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. +func readRequestInfo(w http.ResponseWriter, r *http.Request) (requestInfo, *http.Request, error) { + if r.Body == nil || r.Body == http.NoBody { + return requestInfo{}, r, nil } - // Re-buffer the body for downstream handlers. + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, maxBodySize)) + if err != nil { + return requestInfo{}, r, fmt.Errorf("read request body: %w", err) + } r.Body = io.NopCloser(bytes.NewReader(body)) + trimmed := bytes.TrimSpace(body) + 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 + } + info := requestInfo{} + for _, m := range batch { + info.weight += weightOf(m) + } + return info, r, nil + } + var req jsonRPCRequest if err := json.Unmarshal(body, &req); err != nil { - return 0, r + return requestInfo{}, r, nil + } + info := requestInfo{weight: weightOf(req), method: req.Method} + if len(req.ID) > 0 && !bytes.Equal(req.ID, []byte("null")) { + info.id = req.ID } + return info, r, nil +} - if req.Method != "tools/call" { - return 0, r +// writeJSONRPCError writes a JSON-RPC error response with the given HTTP status. +func writeJSONRPCError(w http.ResponseWriter, status int, id json.RawMessage, code int, msg string) { + if len(id) == 0 { + id = json.RawMessage("null") + } + resp, err := json.Marshal(struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Error struct { + Code int `json:"code"` + Message string `json:"message"` + } `json:"error"` + }{ + JSONRPC: "2.0", + ID: id, + Error: struct { + Code int `json:"code"` + Message string `json:"message"` + }{Code: code, Message: msg}, + }) + if err != nil { + http.Error(w, msg, status) + return } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _, _ = w.Write(resp) +} - var params toolCallParams - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - return 1, r +// writeBodyError rejects a request whose body could not be read. +func writeBodyError(w http.ResponseWriter, err error, readTimeout time.Duration) { + var tooLarge *http.MaxBytesError + var netErr net.Error + switch { + case errors.As(err, &tooLarge): + writeJSONRPCError(w, http.StatusRequestEntityTooLarge, nil, -32600, + fmt.Sprintf("Request body exceeds the %d-byte (%d KB) limit.", maxBodySize, maxBodySize/1024)) + case errors.As(err, &netErr) && netErr.Timeout(): + writeJSONRPCError(w, http.StatusRequestTimeout, nil, -32600, + fmt.Sprintf("Request body was not received within %s.", readTimeout)) + default: + writeJSONRPCError(w, http.StatusBadRequest, nil, -32600, "Request body could not be read.") } +} - if w, ok := toolWeights[params.Name]; ok { - return w, r +// retryAfterSeconds rounds d up to whole seconds, with a minimum of 1. +func retryAfterSeconds(d time.Duration) int { + s := int(math.Ceil(d.Seconds())) + if s < 1 { + return 1 } - return 1, r + return s } -// RateLimitMiddleware creates a rate limiter and returns an http.Handler -// that enforces the three-layer rate limiting strategy. Requests that exceed -// limits receive an HTTP 200 response with a JSON-RPC error (MCP clients -// expect JSON-RPC responses, not HTTP 429). +// RateLimitMiddleware is NewRateLimitMiddleware with the default +// configuration: clients are keyed by their TCP peer address only. func RateLimitMiddleware(next http.Handler) http.Handler { + return NewRateLimitMiddleware(nil, next) +} + +// NewRateLimitMiddleware creates a rate limiter and returns an http.Handler +// that enforces the three-layer rate limiting strategy. Clients are keyed by +// 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. +func NewRateLimitMiddleware(cfg *Config, next http.Handler) http.Handler { + h, _ := newRateLimitHandler(cfg, next, bodyReadTimeout) + return h +} + +// newRateLimitHandler is NewRateLimitMiddleware with a configurable body +// read timeout; it also returns the limiter so tests can inspect it. +func newRateLimitHandler(cfg *Config, next http.Handler, readTimeout time.Duration) (http.Handler, *rateLimiter) { rl := NewRateLimiter() + resolver := newClientIPResolver(cfg) return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - ip := extractIP(r) - weight, r := extractToolWeight(r) - - rl.activeReqs.Add(1) - defer rl.activeReqs.Add(-1) + key := resolver.clientKey(r) + isStream := r.Method == http.MethodGet + + var info requestInfo + if isStream { + 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)) + var err error + info, r, err = readRequestInfo(w, r) + if err != nil { + writeBodyError(w, err, readTimeout) + return + } + _ = rc.SetReadDeadline(time.Time{}) + rl.activeReqs.Add(1) + defer rl.activeReqs.Add(-1) + } - allowed, limit, remaining, resetTime := rl.allow(ip, weight) + d := rl.allow(key, info.weight) - // Always set rate limit headers. - w.Header().Set("X-RateLimit-Limit", fmt.Sprintf("%d", limit)) - w.Header().Set("X-RateLimit-Remaining", fmt.Sprintf("%d", remaining)) - w.Header().Set("X-RateLimit-Reset", fmt.Sprintf("%d", resetTime.Unix())) + h := w.Header() + h.Set("X-RateLimit-Limit", strconv.Itoa(d.limit)) + h.Set("X-RateLimit-Remaining", strconv.Itoa(d.remaining)) + h.Set("X-RateLimit-Reset", strconv.FormatInt(d.reset.Unix(), 10)) - if !allowed { - retryAfter := time.Until(resetTime).Seconds() - if retryAfter < 1 { - retryAfter = 1 + if !d.allowed { + secs := retryAfterSeconds(d.retryAfter) + h.Set("Retry-After", strconv.Itoa(secs)) + status := http.StatusTooManyRequests + if info.method == "initialize" { + status = http.StatusOK } - - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusOK) - resp := fmt.Sprintf( - `{"jsonrpc":"2.0","id":null,"error":{"code":-32000,"message":"Rate limit exceeded. Try again in %.0f seconds."}}`, - retryAfter, - ) - _, _ = w.Write([]byte(resp)) + writeJSONRPCError(w, status, info.id, -32000, + fmt.Sprintf("Rate limit exceeded. Try again in %d seconds.", secs)) 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 cdab1565..4ed4ba13 100644 --- a/pkg/mcp/ratelimit_test.go +++ b/pkg/mcp/ratelimit_test.go @@ -15,20 +15,30 @@ package mcp import ( + "bufio" "bytes" + "context" "encoding/json" + "errors" + "fmt" "io" + "net" "net/http" "net/http/httptest" + "os" + "strconv" "strings" "testing" + "testing/iotest" + "time" ) func TestRateLimitAllowsNormalTraffic(t *testing.T) { rl := NewRateLimiter() defer rl.Close() - allowed, limit, remaining, _ := rl.allow("192.168.1.1", 1) + d := rl.allow("192.168.1.1", 1) + allowed, limit, remaining := d.allowed, d.limit, d.remaining if !allowed { t.Fatal("expected first request to be allowed") } @@ -47,14 +57,14 @@ func TestRateLimitBlocksBurst(t *testing.T) { ip := "10.0.0.1" // Send burstLimit requests - all should be allowed. for i := 0; i < burstLimit; i++ { - allowed, _, _, _ := rl.allow(ip, 1) + allowed := rl.allow(ip, 1).allowed if !allowed { t.Fatalf("request %d should have been allowed", i+1) } } // The 11th request should be burst-limited. - allowed, _, _, _ := rl.allow(ip, 1) + allowed := rl.allow(ip, 1).allowed if allowed { t.Fatal("expected 11th request to be blocked by burst limit") } @@ -68,7 +78,7 @@ func TestRateLimitToolWeighting(t *testing.T) { // analyze_sql has weight 5. After 10 requests at burst limit, // we should have consumed 50 of 120 weighted budget. for i := 0; i < burstLimit; i++ { - allowed, _, _, _ := rl.allow(ip, 5) // analyze_sql weight + allowed := rl.allow(ip, 5).allowed // analyze_sql weight if !allowed { t.Fatalf("request %d should have been allowed", i+1) } @@ -81,12 +91,10 @@ func TestRateLimitToolWeighting(t *testing.T) { ip2 := "10.0.0.3" count := 0 for { - allowed, _, remaining, _ := rl.allow(ip2, 5) - if !allowed { + if !rl.allow(ip2, 5).allowed { break } count++ - _ = remaining // Safety: break if we somehow loop too many times if count > 200 { t.Fatal("too many requests allowed, expected sustained limit to kick in") @@ -107,7 +115,7 @@ func TestRateLimitNonToolUnlimited(t *testing.T) { // Weight 0 requests bypass sustained limiting but still hit burst limit. // Send burstLimit weight-0 requests. for i := 0; i < burstLimit; i++ { - allowed, _, _, _ := rl.allow(ip, 0) + allowed := rl.allow(ip, 0).allowed if !allowed { t.Fatalf("weight-0 request %d should have been allowed", i+1) } @@ -117,9 +125,12 @@ func TestRateLimitNonToolUnlimited(t *testing.T) { // But verify the sustained budget (weightedCount) was NOT consumed. idx := shardFor(ip) sh := &rl.shards[idx] - sh.RLock() - b := sh.buckets[ip] - sh.RUnlock() + sh.mu.Lock() + b, ok := sh.buckets[ip].Value.(*bucket) + sh.mu.Unlock() + if !ok { + t.Fatal("bucket not found") + } b.mu.Lock() wc := b.weightedCount @@ -181,14 +192,17 @@ func TestRateLimitResponse(t *testing.T) { } } - // Next request should be rate-limited but still return 200 with JSON-RPC error. + // Next request should be rate-limited: HTTP 429 with a JSON-RPC error. req := httptest.NewRequest("POST", "/mcp", makeBody()) req.RemoteAddr = ip + ":12345" rec := httptest.NewRecorder() handler.ServeHTTP(rec, req) - if rec.Code != http.StatusOK { - t.Fatalf("rate-limited response should be HTTP 200, got %d", rec.Code) + if rec.Code != http.StatusTooManyRequests { + t.Fatalf("rate-limited response should be HTTP 429, got %d", rec.Code) + } + if got := rec.Header().Get("Retry-After"); got != "1" { + t.Fatalf("Retry-After = %q, want 1", got) } var resp struct { @@ -205,6 +219,9 @@ func TestRateLimitResponse(t *testing.T) { if resp.JSONRPC != "2.0" { t.Fatalf("expected jsonrpc 2.0, got %s", resp.JSONRPC) } + if resp.ID == nil || *resp.ID != 1 { + t.Fatalf("expected the request id 1 to be echoed, got %v", resp.ID) + } if resp.Error.Code != -32000 { t.Fatalf("expected error code -32000, got %d", resp.Error.Code) } @@ -224,126 +241,435 @@ func TestRateLimitResponse(t *testing.T) { } } -func TestExtractIP(t *testing.T) { +func TestClientKey(t *testing.T) { tests := []struct { name string - xff string + header string + hops int + xff []string + cfIP string remoteAddr string want string }{ - { - name: "X-Forwarded-For single IP", - xff: "203.0.113.50", - remoteAddr: "10.0.0.1:1234", - want: "203.0.113.50", - }, - { - name: "X-Forwarded-For multiple IPs", - xff: "203.0.113.50, 70.41.3.18, 150.172.238.178", - remoteAddr: "10.0.0.1:1234", - want: "203.0.113.50", - }, - { - name: "no XFF, use RemoteAddr with port", - xff: "", - remoteAddr: "192.168.1.100:5678", - want: "192.168.1.100", - }, - { - name: "no XFF, RemoteAddr without port", - xff: "", - remoteAddr: "192.168.1.100", - want: "192.168.1.100", - }, - { - name: "XFF with spaces", - xff: " 10.0.0.5 , 10.0.0.6", - remoteAddr: "127.0.0.1:9999", - want: "10.0.0.5", - }, + {name: "default ignores spoofed XFF", xff: []string{"203.0.113.50"}, remoteAddr: "10.0.0.1:1234", want: "10.0.0.1"}, + {name: "default ignores trusted-looking header", cfIP: "203.0.113.9", remoteAddr: "10.0.0.1:1234", want: "10.0.0.1"}, + {name: "RemoteAddr without port", remoteAddr: "192.168.1.100", want: "192.168.1.100"}, + {name: "trusted header used", header: "Cf-Connecting-Ip", cfIP: "203.0.113.9", remoteAddr: "10.0.0.1:1234", want: "203.0.113.9"}, + {name: "trusted header invalid falls back", header: "Cf-Connecting-Ip", cfIP: "not-an-ip", remoteAddr: "10.0.0.1:1234", want: "10.0.0.1"}, + {name: "trusted header oversized falls back", header: "Cf-Connecting-Ip", cfIP: strings.Repeat("a", 2048), remoteAddr: "10.0.0.1:1234", want: "10.0.0.1"}, + {name: "trusted header missing falls back", header: "Cf-Connecting-Ip", remoteAddr: "10.0.0.1:1234", want: "10.0.0.1"}, + {name: "trusted header list takes rightmost", header: "Cf-Connecting-Ip", cfIP: "1.1.1.1, 203.0.113.9", remoteAddr: "10.0.0.1:1234", want: "203.0.113.9"}, + {name: "hops=1 takes rightmost XFF", hops: 1, xff: []string{"6.6.6.6, 203.0.113.50"}, remoteAddr: "10.0.0.1:1234", want: "203.0.113.50"}, + {name: "hops=2 skips one proxy", hops: 2, xff: []string{"6.6.6.6, 203.0.113.50, 70.41.3.18"}, remoteAddr: "10.0.0.1:1234", want: "203.0.113.50"}, + {name: "hops spans header lines", hops: 2, xff: []string{"6.6.6.6", "203.0.113.50", "70.41.3.18"}, remoteAddr: "10.0.0.1:1234", want: "203.0.113.50"}, + {name: "hops beyond entries falls back", hops: 3, xff: []string{"203.0.113.50"}, remoteAddr: "10.0.0.1:1234", want: "10.0.0.1"}, + {name: "hops invalid entry falls back", hops: 1, xff: []string{"garbage"}, remoteAddr: "10.0.0.1:1234", want: "10.0.0.1"}, + {name: "XFF entry with port", hops: 1, xff: []string{"203.0.113.50:4711"}, remoteAddr: "10.0.0.1:1234", want: "203.0.113.50"}, + {name: "IPv6 grouped by /64", remoteAddr: "[2001:db8:1:2:aaaa::1]:443", want: "2001:db8:1:2::/64"}, + {name: "IPv4-mapped IPv6 normalised", header: "Cf-Connecting-Ip", cfIP: "::ffff:203.0.113.9", remoteAddr: "10.0.0.1:1234", want: "203.0.113.9"}, } - for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - req := httptest.NewRequest("GET", "/", nil) + req := httptest.NewRequest(http.MethodPost, "/mcp", nil) req.RemoteAddr = tt.remoteAddr - if tt.xff != "" { - req.Header.Set("X-Forwarded-For", tt.xff) + for _, v := range tt.xff { + req.Header.Add("X-Forwarded-For", v) } - got := extractIP(req) - if got != tt.want { - t.Errorf("extractIP() = %q, want %q", got, tt.want) + if tt.cfIP != "" { + req.Header.Set("CF-Connecting-IP", tt.cfIP) + } + r := newClientIPResolver(&Config{TrustedProxyHeader: tt.header, TrustedProxyHops: tt.hops}) + if got := r.clientKey(req); got != tt.want { + t.Errorf("clientKey() = %q, want %q", got, tt.want) } }) } } -func TestExtractToolWeight(t *testing.T) { +// okHandler is a downstream handler that always succeeds. +var okHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + w.WriteHeader(http.StatusOK) +}) + +func toolCallBody(id, tool string) string { + return `{"jsonrpc":"2.0","id":` + id + `,"method":"tools/call","params":{"name":"` + tool + `","arguments":{"sql":"SELECT 1"}}}` +} + +func TestRateLimitMiddleware_SpoofedXFFStillLimited(t *testing.T) { + handler := RateLimitMiddleware(okHandler) + var last int + for i := 0; i <= burstLimit; i++ { + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(toolCallBody("1", "validate_sql"))) + req.RemoteAddr = "198.51.100.7:5555" + req.Header.Set("X-Forwarded-For", fmt.Sprintf("11.0.%d.1", i)) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + last = rec.Code + } + if last != http.StatusTooManyRequests { + t.Fatalf("request %d with rotating X-Forwarded-For: status %d, want 429", burstLimit+1, last) + } +} + +func TestRateLimiter_BucketsBounded(t *testing.T) { + tests := []struct { + name string + maxPerShard int + keys int + }{ + {name: "small cap flood", maxPerShard: 8, keys: 10000}, + {name: "default cap flood", maxPerShard: maxBucketsPerShard, keys: 200000}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rl := NewRateLimiter() + defer rl.Close() + rl.maxPerShard = tt.maxPerShard + + for i := 0; i < tt.keys; i++ { + rl.allow(fmt.Sprintf("10.%d.%d.%d", i>>16&255, i>>8&255, i&255), 1) + } + if got, max := rl.bucketCount(), numShards*tt.maxPerShard; got > max { + t.Fatalf("bucketCount = %d after %d distinct keys, want <= %d", got, tt.keys, max) + } + }) + } +} + +func TestRateLimiter_EvictsLeastRecentlySeen(t *testing.T) { + rl := NewRateLimiter() + defer rl.Close() + rl.maxPerShard = 2 + + // Find three keys that land in the same shard. + var keys []string + for i := 0; len(keys) < 3; i++ { + k := fmt.Sprintf("10.0.0.%d", i) + if len(keys) == 0 || shardFor(k) == shardFor(keys[0]) { + keys = append(keys, k) + } + } + rl.allow(keys[0], 1) + rl.allow(keys[1], 1) + rl.allow(keys[0], 1) // keys[0] is now most recently seen + rl.allow(keys[2], 1) // evicts keys[1] + + sh := &rl.shards[shardFor(keys[0])] + sh.mu.Lock() + _, has0 := sh.buckets[keys[0]] + _, has1 := sh.buckets[keys[1]] + _, has2 := sh.buckets[keys[2]] + sh.mu.Unlock() + if !has0 || has1 || !has2 { + t.Fatalf("after eviction: has[0]=%v has[1]=%v has[2]=%v, want true false true", has0, has1, has2) + } +} + +func TestRateLimiter_CleanupRemovesStale(t *testing.T) { + rl := NewRateLimiter() + defer rl.Close() + for i := 0; i < 100; i++ { + rl.allow(fmt.Sprintf("10.1.0.%d", i), 1) + } + rl.cleanup(time.Now()) + if got := rl.bucketCount(); got != 100 { + t.Fatalf("fresh buckets removed: count = %d, want 100", got) + } + rl.cleanup(time.Now().Add(staleTimeout + time.Minute)) + if got := rl.bucketCount(); got != 0 { + t.Fatalf("stale buckets kept: count = %d, want 0", got) + } +} + +func TestRateLimitMiddleware_Denials(t *testing.T) { tests := []struct { name string body string - wantWeight int + wantStatus int + wantID string + }{ + {name: "tools/call echoes numeric id", body: toolCallBody("7", "validate_sql"), wantStatus: http.StatusTooManyRequests, wantID: "7"}, + {name: "tools/call echoes string id", body: toolCallBody(`"req-abc"`, "validate_sql"), wantStatus: http.StatusTooManyRequests, wantID: `"req-abc"`}, + {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"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + handler := RateLimitMiddleware(okHandler) + var rec *httptest.ResponseRecorder + for i := 0; i <= burstLimit; i++ { + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(tt.body)) + req.RemoteAddr = "10.9.9.9:1" + rec = httptest.NewRecorder() + handler.ServeHTTP(rec, req) + } + if rec.Code != tt.wantStatus { + t.Fatalf("status = %d, want %d", rec.Code, tt.wantStatus) + } + if rec.Header().Get("Retry-After") == "" { + t.Fatal("missing Retry-After header") + } + var resp struct { + ID json.RawMessage `json:"id"` + Error struct { + Code int `json:"code"` + } `json:"error"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode %q: %v", rec.Body.String(), err) + } + if string(resp.ID) != tt.wantID { + t.Errorf("id = %s, want %s", resp.ID, tt.wantID) + } + if resp.Error.Code != -32000 { + t.Errorf("error code = %d, want -32000", resp.Error.Code) + } + }) + } +} + +func TestRateLimitMiddleware_HeadersOnDenial(t *testing.T) { + t.Run("burst denial reports sustained remaining", func(t *testing.T) { + handler := RateLimitMiddleware(okHandler) + var rec *httptest.ResponseRecorder + for i := 0; i <= burstLimit; i++ { + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(toolCallBody("1", "validate_sql"))) + req.RemoteAddr = "10.8.0.1:1" + rec = httptest.NewRecorder() + handler.ServeHTTP(rec, req) + } + if rec.Code != http.StatusTooManyRequests { + t.Fatalf("status = %d, want 429", rec.Code) + } + want := strconv.Itoa(sustainedLimit - burstLimit) + if got := rec.Header().Get("X-RateLimit-Remaining"); got != want { + t.Errorf("X-RateLimit-Remaining = %s, want %s", got, want) + } + if got := rec.Header().Get("Retry-After"); got != "1" { + t.Errorf("Retry-After = %s, want 1", got) + } + }) + t.Run("sustained denial waits for window reset", func(t *testing.T) { + rl := NewRateLimiter() + defer rl.Close() + key := "10.8.0.2" + // Exhaust the sustained budget without tripping the burst limiter. + sh := &rl.shards[shardFor(key)] + rl.allow(key, 1) + sh.mu.Lock() + b, _ := sh.buckets[key].Value.(*bucket) + sh.mu.Unlock() + b.mu.Lock() + b.weightedCount = sustainedLimit + b.mu.Unlock() + + d := rl.allow(key, 1) + if d.allowed { + t.Fatal("expected sustained denial") + } + if d.remaining != 0 { + t.Errorf("remaining = %d, want 0", d.remaining) + } + if secs := retryAfterSeconds(d.retryAfter); secs < 50 || secs > 60 { + t.Errorf("Retry-After = %ds, want ~60s", secs) + } + }) +} + +func TestRateLimitMiddleware_BodyLimits(t *testing.T) { + tests := []struct { + name string + size int + wantStatus int }{ - { - name: "validate_sql", - body: `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"validate_sql","arguments":{"sql":"SELECT 1"}}}`, - wantWeight: 1, - }, - { - name: "analyze_sql", - body: `{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"analyze_sql","arguments":{"sql":"SELECT 1"}}}`, - wantWeight: 5, - }, - { - name: "security_scan", - body: `{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"security_scan","arguments":{"sql":"SELECT 1"}}}`, - wantWeight: 2, - }, - { - name: "initialize method (not tools/call)", - body: `{"jsonrpc":"2.0","id":4,"method":"initialize","params":{}}`, - wantWeight: 0, - }, - { - name: "list_tools method", - body: `{"jsonrpc":"2.0","id":5,"method":"tools/list","params":{}}`, - wantWeight: 0, - }, - { - name: "unknown tool", - body: `{"jsonrpc":"2.0","id":6,"method":"tools/call","params":{"name":"unknown_tool","arguments":{}}}`, - wantWeight: 1, - }, - { - name: "empty body", - body: "", - wantWeight: 0, - }, - { - name: "invalid JSON", - body: `not json at all`, - wantWeight: 0, - }, + {name: "at limit accepted", size: maxBodySize, wantStatus: http.StatusOK}, + {name: "over limit rejected with 413", size: maxBodySize + 1, wantStatus: http.StatusRequestEntityTooLarge}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + prefix := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"validate_sql","arguments":{"sql":"` + suffix := `"}}}` + body := prefix + strings.Repeat("x", tt.size-len(prefix)-len(suffix)) + suffix + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + rec := httptest.NewRecorder() + RateLimitMiddleware(okHandler).ServeHTTP(rec, req) + if rec.Code != tt.wantStatus { + t.Fatalf("status = %d, want %d", rec.Code, tt.wantStatus) + } + if tt.wantStatus == http.StatusRequestEntityTooLarge && + !strings.Contains(rec.Body.String(), "exceeds the 65536-byte") { + t.Errorf("413 body does not explain the limit: %s", rec.Body.String()) + } + }) + } +} + +func TestRateLimitMiddleware_SlowBodyTimesOut(t *testing.T) { + h, rl := newRateLimitHandler(nil, okHandler, 200*time.Millisecond) + defer rl.Close() + ts := httptest.NewServer(h) + defer ts.Close() + + conn, err := net.Dial("tcp", ts.Listener.Addr().String()) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + _, _ = fmt.Fprintf(conn, "POST /mcp HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\nContent-Length: 1000\r\n\r\n{") + + _ = conn.SetReadDeadline(time.Now().Add(5 * time.Second)) + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + if err != nil { + t.Fatalf("no response to a stalled body within 5s: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusRequestTimeout { + t.Fatalf("status = %d, want 408", resp.StatusCode) + } +} + +func TestRateLimitMiddleware_ReadDeadlineClearedAfterBody(t *testing.T) { + // A POST whose response streams for longer than the body read timeout + // must not be cancelled by that deadline. + slow := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + select { + case <-time.After(500 * time.Millisecond): + case <-r.Context().Done(): + http.Error(w, "cancelled", http.StatusInternalServerError) + return + } + _, _ = w.Write([]byte("done")) + }) + h, rl := newRateLimitHandler(nil, slow, 100*time.Millisecond) + 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) + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK || string(b) != "done" { + t.Fatalf("status %d body %q, want 200 done", resp.StatusCode, b) + } +} + +func TestRateLimitMiddleware_GETStreams(t *testing.T) { + started := make(chan struct{}, 16) + stream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _ = http.NewResponseController(w).Flush() + started <- struct{}{} + <-r.Context().Done() + }) + 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() + open := 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 + } + + var held []*http.Response + for i := 0; i < maxStreamsPerIP; i++ { + resp := open() + if resp.StatusCode != http.StatusOK { + t.Fatalf("stream %d: status %d, want 200", i+1, resp.StatusCode) + } + <-started + held = append(held, resp) + } + if got := rl.activeReqs.Load(); got != 0 { + t.Errorf("open GET streams counted as active requests: %d", got) + } + if got := currentLimit(rl); got != sustainedLimit { + t.Errorf("open GET streams changed the sustained limit to %d", got) } + extra := open() + extra.Body.Close() + if extra.StatusCode != http.StatusTooManyRequests { + t.Fatalf("stream %d: status %d, want 429", maxStreamsPerIP+1, extra.StatusCode) + } + + // Closing a stream frees its slot. + held[0].Body.Close() + deadline := time.Now().Add(2 * time.Second) + for { + rl.streams.mu.Lock() + n := rl.streams.open["127.0.0.1"] + rl.streams.mu.Unlock() + if n < maxStreamsPerIP || time.Now().After(deadline) { + break + } + time.Sleep(10 * time.Millisecond) + } + resp := open() + if resp.StatusCode != http.StatusOK { + t.Fatalf("stream after close: status %d, want 200", resp.StatusCode) + } + <-started + resp.Body.Close() + for _, r := range held[1:] { + r.Body.Close() + } +} + +func TestReadRequestInfo(t *testing.T) { + tests := []struct { + name string + body string + wantWeight int + wantID string + wantMethod string + }{ + {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: "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: "empty body", body: "", wantWeight: 0}, + {name: "invalid JSON", body: `not json at all`, wantWeight: 0}, + } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { var body io.Reader if tt.body != "" { body = strings.NewReader(tt.body) } - req := httptest.NewRequest("POST", "/mcp", body) - - weight, newReq := extractToolWeight(req) - if weight != tt.wantWeight { - t.Errorf("extractToolWeight() weight = %d, want %d", weight, tt.wantWeight) + req := httptest.NewRequest(http.MethodPost, "/mcp", body) + info, newReq, err := readRequestInfo(httptest.NewRecorder(), req) + 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) } - - // Verify the body was re-buffered (downstream can read it). if tt.body != "" { rebuffered, err := io.ReadAll(newReq.Body) if err != nil { - t.Fatalf("failed to read re-buffered body: %v", err) + t.Fatalf("read re-buffered body: %v", err) } if !bytes.Equal(rebuffered, []byte(tt.body)) { t.Errorf("re-buffered body mismatch: got %q, want %q", rebuffered, tt.body) @@ -352,3 +678,36 @@ func TestExtractToolWeight(t *testing.T) { }) } } + +func TestRateLimitMiddleware_RejectsUnreadableBody(t *testing.T) { + tests := []struct { + name string + err error + wantStatus int + }{ + {name: "transport error", err: errors.New("boom"), wantStatus: http.StatusBadRequest}, + {name: "timeout", err: os.ErrDeadlineExceeded, wantStatus: http.StatusRequestTimeout}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/mcp", iotest.ErrReader(tt.err)) + rec := httptest.NewRecorder() + RateLimitMiddleware(okHandler).ServeHTTP(rec, req) + if rec.Code != tt.wantStatus { + t.Fatalf("status = %d, want %d", rec.Code, tt.wantStatus) + } + }) + } +} + +// bucketCount reports how many client buckets are currently tracked. +func (rl *rateLimiter) bucketCount() int { + n := 0 + for i := range rl.shards { + sh := &rl.shards[i] + sh.mu.Lock() + n += len(sh.buckets) + sh.mu.Unlock() + } + return n +} diff --git a/pkg/mcp/tools.go b/pkg/mcp/tools.go index 32d52a67..43f99569 100644 --- a/pkg/mcp/tools.go +++ b/pkg/mcp/tools.go @@ -18,6 +18,8 @@ import ( "context" "encoding/json" "fmt" + "math" + "strconv" "strings" "sync" @@ -380,94 +382,219 @@ func toolResult(data map[string]any) (*mcpmcp.CallToolResult, error) { // MCP handler functions // --------------------------------------------------------------------------- -// handleValidateSQL is the MCP tool handler for "validate_sql". -func handleValidateSQL(ctx context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { - sql := req.GetString("sql", "") +// maxIndentSize bounds format_sql's indent_size so a request cannot ask the +// formatter for multi-megabyte indentation strings. +const maxIndentSize = 16 + +// errSQLRequired is the tool error for a missing, null or empty sql argument. +const errSQLRequired = "parameter 'sql' is required and must not be empty" + +// jsonTypeName names the JSON type of a decoded argument value for error messages. +func jsonTypeName(v any) string { + switch v.(type) { + case nil: + return "null" + case string: + return "string" + case bool: + return "boolean" + case float64, float32, int, int64, int32, json.Number: + return "number" + case []any: + return "array" + case map[string]any: + return "object" + default: + return fmt.Sprintf("%T", v) + } +} + +// sqlArg returns the required "sql" argument. When it is missing, null, +// empty or not a string, it returns an isError tool result instead, so the +// model sees an actionable message rather than a JSON-RPC internal error. +func sqlArg(req mcpmcp.CallToolRequest) (string, *mcpmcp.CallToolResult) { + v, ok := req.GetArguments()["sql"] + if !ok || v == nil { + return "", mcpmcp.NewToolResultError(errSQLRequired) + } + sql, ok := v.(string) + if !ok { + return "", mcpmcp.NewToolResultError(fmt.Sprintf("parameter 'sql' must be a string, got %s", jsonTypeName(v))) + } if sql == "" { - return nil, fmt.Errorf("parameter 'sql' is required and must not be empty") + return "", mcpmcp.NewToolResultError(errSQLRequired) + } + return sql, nil +} + +// optionalString returns the string argument name, or def when it is absent or null. +func optionalString(req mcpmcp.CallToolRequest, name, def string) (string, error) { + v, ok := req.GetArguments()[name] + if !ok || v == nil { + return def, nil + } + s, ok := v.(string) + if !ok { + return "", fmt.Errorf("parameter '%s' must be a string, got %s", name, jsonTypeName(v)) + } + return s, nil +} + +// optionalBool returns the boolean argument name, or def when it is absent or +// null. The strings "true" and "false" are accepted because some clients +// stringify scalar arguments; anything else is an error. +func optionalBool(req mcpmcp.CallToolRequest, name string, def bool) (bool, error) { + v, ok := req.GetArguments()[name] + if !ok || v == nil { + return def, nil + } + switch b := v.(type) { + case bool: + return b, nil + case string: + switch strings.ToLower(strings.TrimSpace(b)) { + case "true": + return true, nil + case "false": + return false, nil + } } - dialect := req.GetString("dialect", "") - result, err := validateSQLInternal(sql, dialect) + return false, fmt.Errorf("parameter '%s' must be a boolean, got %s %s", name, jsonTypeName(v), shortJSON(v)) +} + +// optionalInt returns the integer argument name, or def when it is absent or +// null. Whole-number JSON numbers and decimal strings are accepted; the value +// must lie in [min, max]. +func optionalInt(req mcpmcp.CallToolRequest, name string, def, min, max int) (int, error) { + v, ok := req.GetArguments()[name] + if !ok || v == nil { + return def, nil + } + var f float64 + valid := true + switch n := v.(type) { + case float64: + f = n + case int: + f = float64(n) + case int64: + f = float64(n) + case json.Number: + parsed, err := n.Float64() + f, valid = parsed, err == nil + case string: + parsed, err := strconv.Atoi(strings.TrimSpace(n)) + f, valid = float64(parsed), err == nil + default: + valid = false + } + if !valid || f != math.Trunc(f) || f < float64(min) || f > float64(max) { + return 0, fmt.Errorf("parameter '%s' must be an integer between %d and %d, got %s %s", + name, min, max, jsonTypeName(v), shortJSON(v)) + } + return int(f), nil +} + +// shortJSON renders v for an error message, truncated to keep messages small. +func shortJSON(v any) string { + b, err := json.Marshal(v) if err != nil { - return nil, err + return "?" + } + const maxLen = 40 + if len(b) > maxLen { + return string(b[:maxLen]) + "…" + } + return string(b) +} + +// toolOutcome converts an internal result into a tool result. Errors from +// the *Internal functions are input problems (unparseable or unformattable +// SQL), so they become isError results; only marshalling failures in +// toolResult remain Go errors (JSON-RPC internal errors). +func toolOutcome(result map[string]any, err error) (*mcpmcp.CallToolResult, error) { + if err != nil { + return mcpmcp.NewToolResultError(err.Error()), nil } return toolResult(result) } +// handleValidateSQL is the MCP tool handler for "validate_sql". +func handleValidateSQL(ctx context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + sql, errRes := sqlArg(req) + if errRes != nil { + return errRes, nil + } + dialect, err := optionalString(req, "dialect", "") + if err != nil { + return mcpmcp.NewToolResultError(err.Error()), nil + } + return toolOutcome(validateSQLInternal(sql, dialect)) +} + // handleFormatSQL is the MCP tool handler for "format_sql". func handleFormatSQL(ctx context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { - sql := req.GetString("sql", "") - if sql == "" { - return nil, fmt.Errorf("parameter 'sql' is required and must not be empty") + sql, errRes := sqlArg(req) + if errRes != nil { + return errRes, nil } - indentSize := req.GetInt("indent_size", 2) - uppercaseKeywords := req.GetBool("uppercase_keywords", false) - addSemicolon := req.GetBool("add_semicolon", false) - result, err := formatSQLInternal(sql, indentSize, uppercaseKeywords, addSemicolon) + indentSize, err := optionalInt(req, "indent_size", 2, 0, maxIndentSize) if err != nil { - return nil, err + return mcpmcp.NewToolResultError(err.Error()), nil } - return toolResult(result) + uppercaseKeywords, err := optionalBool(req, "uppercase_keywords", false) + if err != nil { + return mcpmcp.NewToolResultError(err.Error()), nil + } + addSemicolon, err := optionalBool(req, "add_semicolon", false) + if err != nil { + return mcpmcp.NewToolResultError(err.Error()), nil + } + return toolOutcome(formatSQLInternal(sql, indentSize, uppercaseKeywords, addSemicolon)) } // handleParseSQL is the MCP tool handler for "parse_sql". func handleParseSQL(ctx context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { - sql := req.GetString("sql", "") - if sql == "" { - return nil, fmt.Errorf("parameter 'sql' is required and must not be empty") - } - result, err := parseSQLInternal(sql) - if err != nil { - return nil, err + sql, errRes := sqlArg(req) + if errRes != nil { + return errRes, nil } - return toolResult(result) + return toolOutcome(parseSQLInternal(sql)) } // handleExtractMetadata is the MCP tool handler for "extract_metadata". func handleExtractMetadata(ctx context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { - sql := req.GetString("sql", "") - if sql == "" { - return nil, fmt.Errorf("parameter 'sql' is required and must not be empty") + sql, errRes := sqlArg(req) + if errRes != nil { + return errRes, nil } - result, err := extractMetadataInternal(sql) - if err != nil { - return nil, err - } - return toolResult(result) + return toolOutcome(extractMetadataInternal(sql)) } // handleSecurityScan is the MCP tool handler for "security_scan". func handleSecurityScan(ctx context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { - sql := req.GetString("sql", "") - if sql == "" { - return nil, fmt.Errorf("parameter 'sql' is required and must not be empty") - } - result, err := securityScanInternal(sql) - if err != nil { - return nil, err + sql, errRes := sqlArg(req) + if errRes != nil { + return errRes, nil } - return toolResult(result) + return toolOutcome(securityScanInternal(sql)) } // handleLintSQL is the MCP tool handler for "lint_sql". func handleLintSQL(ctx context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { - sql := req.GetString("sql", "") - if sql == "" { - return nil, fmt.Errorf("parameter 'sql' is required and must not be empty") - } - result, err := lintSQLInternal(sql) - if err != nil { - return nil, err + sql, errRes := sqlArg(req) + if errRes != nil { + return errRes, nil } - return toolResult(result) + return toolOutcome(lintSQLInternal(sql)) } // handleAnalyzeSQL is the MCP tool handler for "analyze_sql". // It fans out all six analysis tools concurrently and merges the results. func handleAnalyzeSQL(ctx context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { - sql := req.GetString("sql", "") - if sql == "" { - return nil, fmt.Errorf("parameter 'sql' is required and must not be empty") + sql, errRes := sqlArg(req) + if errRes != nil { + return errRes, nil } type namedResult struct { diff --git a/pkg/mcp/tools_test.go b/pkg/mcp/tools_test.go index 93460c62..8a5c546d 100644 --- a/pkg/mcp/tools_test.go +++ b/pkg/mcp/tools_test.go @@ -17,6 +17,7 @@ package mcp import ( "context" "encoding/json" + "strings" "testing" mcpmcp "github.com/mark3labs/mcp-go/mcp" @@ -54,6 +55,29 @@ func unmarshalResult(t *testing.T, res *mcpmcp.CallToolResult) map[string]any { return out } +// requireToolError asserts that a handler reported an input problem as an +// isError tool result (not a Go error, which becomes JSON-RPC -32603) whose +// text contains want. +func requireToolError(t *testing.T, res *mcpmcp.CallToolResult, err error, want string) { + t.Helper() + if err != nil { + t.Fatalf("expected isError tool result, got Go error: %v", err) + } + if res == nil || !res.IsError { + t.Fatalf("expected isError tool result, got %+v", res) + } + if len(res.Content) == 0 { + t.Fatal("tool error has no content") + } + text, ok := res.Content[0].(mcpmcp.TextContent) + if !ok { + t.Fatalf("expected TextContent, got %T", res.Content[0]) + } + if !strings.Contains(text.Text, want) { + t.Errorf("tool error %q does not contain %q", text.Text, want) + } +} + // --------------------------------------------------------------------------- // handleValidateSQL // --------------------------------------------------------------------------- @@ -90,20 +114,16 @@ func TestHandleValidateSQL(t *testing.T) { } }) - t.Run("empty sql returns protocol error", func(t *testing.T) { + t.Run("empty sql returns tool error", func(t *testing.T) { req := makeReq(map[string]any{"sql": ""}) - _, err := handleValidateSQL(ctx, req) - if err == nil { - t.Fatal("expected error for empty sql, got nil") - } + res, err := handleValidateSQL(ctx, req) + requireToolError(t, res, err, "required") }) - t.Run("missing sql param returns protocol error", func(t *testing.T) { + t.Run("missing sql param returns tool error", func(t *testing.T) { req := makeReq(map[string]any{}) - _, err := handleValidateSQL(ctx, req) - if err == nil { - t.Fatal("expected error for missing sql param, got nil") - } + res, err := handleValidateSQL(ctx, req) + requireToolError(t, res, err, "required") }) t.Run("valid SQL with dialect", func(t *testing.T) { @@ -175,20 +195,16 @@ func TestHandleFormatSQL(t *testing.T) { } }) - t.Run("empty sql returns protocol error", func(t *testing.T) { + t.Run("empty sql returns tool error", func(t *testing.T) { req := makeReq(map[string]any{"sql": ""}) - _, err := handleFormatSQL(ctx, req) - if err == nil { - t.Fatal("expected error for empty sql, got nil") - } + res, err := handleFormatSQL(ctx, req) + requireToolError(t, res, err, "required") }) - t.Run("missing sql param returns protocol error", func(t *testing.T) { + t.Run("missing sql param returns tool error", func(t *testing.T) { req := makeReq(map[string]any{}) - _, err := handleFormatSQL(ctx, req) - if err == nil { - t.Fatal("expected error for missing sql param, got nil") - } + res, err := handleFormatSQL(ctx, req) + requireToolError(t, res, err, "required") }) } @@ -218,20 +234,16 @@ func TestHandleParseSQL(t *testing.T) { } }) - t.Run("empty sql returns protocol error", func(t *testing.T) { + t.Run("empty sql returns tool error", func(t *testing.T) { req := makeReq(map[string]any{"sql": ""}) - _, err := handleParseSQL(ctx, req) - if err == nil { - t.Fatal("expected error for empty sql, got nil") - } + res, err := handleParseSQL(ctx, req) + requireToolError(t, res, err, "required") }) - t.Run("missing sql param returns protocol error", func(t *testing.T) { + t.Run("missing sql param returns tool error", func(t *testing.T) { req := makeReq(map[string]any{}) - _, err := handleParseSQL(ctx, req) - if err == nil { - t.Fatal("expected error for missing sql param, got nil") - } + res, err := handleParseSQL(ctx, req) + requireToolError(t, res, err, "required") }) } @@ -256,20 +268,16 @@ func TestHandleExtractMetadata(t *testing.T) { } }) - t.Run("empty sql returns protocol error", func(t *testing.T) { + t.Run("empty sql returns tool error", func(t *testing.T) { req := makeReq(map[string]any{"sql": ""}) - _, err := handleExtractMetadata(ctx, req) - if err == nil { - t.Fatal("expected error for empty sql, got nil") - } + res, err := handleExtractMetadata(ctx, req) + requireToolError(t, res, err, "required") }) - t.Run("missing sql param returns protocol error", func(t *testing.T) { + t.Run("missing sql param returns tool error", func(t *testing.T) { req := makeReq(map[string]any{}) - _, err := handleExtractMetadata(ctx, req) - if err == nil { - t.Fatal("expected error for missing sql param, got nil") - } + res, err := handleExtractMetadata(ctx, req) + requireToolError(t, res, err, "required") }) } @@ -295,20 +303,16 @@ func TestHandleSecurityScan(t *testing.T) { } }) - t.Run("empty sql returns protocol error", func(t *testing.T) { + t.Run("empty sql returns tool error", func(t *testing.T) { req := makeReq(map[string]any{"sql": ""}) - _, err := handleSecurityScan(ctx, req) - if err == nil { - t.Fatal("expected error for empty sql, got nil") - } + res, err := handleSecurityScan(ctx, req) + requireToolError(t, res, err, "required") }) - t.Run("missing sql param returns protocol error", func(t *testing.T) { + t.Run("missing sql param returns tool error", func(t *testing.T) { req := makeReq(map[string]any{}) - _, err := handleSecurityScan(ctx, req) - if err == nil { - t.Fatal("expected error for missing sql param, got nil") - } + res, err := handleSecurityScan(ctx, req) + requireToolError(t, res, err, "required") }) } @@ -334,20 +338,16 @@ func TestHandleLintSQL(t *testing.T) { } }) - t.Run("empty sql returns protocol error", func(t *testing.T) { + t.Run("empty sql returns tool error", func(t *testing.T) { req := makeReq(map[string]any{"sql": ""}) - _, err := handleLintSQL(ctx, req) - if err == nil { - t.Fatal("expected error for empty sql, got nil") - } + res, err := handleLintSQL(ctx, req) + requireToolError(t, res, err, "required") }) - t.Run("missing sql param returns protocol error", func(t *testing.T) { + t.Run("missing sql param returns tool error", func(t *testing.T) { req := makeReq(map[string]any{}) - _, err := handleLintSQL(ctx, req) - if err == nil { - t.Fatal("expected error for missing sql param, got nil") - } + res, err := handleLintSQL(ctx, req) + requireToolError(t, res, err, "required") }) } @@ -406,20 +406,16 @@ func TestHandleAnalyzeSQL(t *testing.T) { } }) - t.Run("empty sql returns protocol error", func(t *testing.T) { + t.Run("empty sql returns tool error", func(t *testing.T) { req := makeReq(map[string]any{"sql": ""}) - _, err := handleAnalyzeSQL(ctx, req) - if err == nil { - t.Fatal("expected error for empty sql, got nil") - } + res, err := handleAnalyzeSQL(ctx, req) + requireToolError(t, res, err, "required") }) - t.Run("missing sql param returns protocol error", func(t *testing.T) { + t.Run("missing sql param returns tool error", func(t *testing.T) { req := makeReq(map[string]any{}) - _, err := handleAnalyzeSQL(ctx, req) - if err == nil { - t.Fatal("expected error for missing sql param, got nil") - } + res, err := handleAnalyzeSQL(ctx, req) + requireToolError(t, res, err, "required") }) } @@ -686,10 +682,8 @@ func TestHandleFormatSQL_InvalidSQL(t *testing.T) { req := makeReq(map[string]any{ "sql": ")))((( @@@ !!!", }) - _, err := handleFormatSQL(ctx, req) - if err == nil { - t.Fatal("expected error for unparseable SQL, got nil") - } + res, err := handleFormatSQL(ctx, req) + requireToolError(t, res, err, "format failed") } func TestHandleAnalyzeSQL_InvalidSQLPartialResults(t *testing.T) { @@ -740,3 +734,95 @@ func TestHandleAnalyzeSQL_InvalidSQLPartialResults(t *testing.T) { } } } + +func TestHandlers_InputProblemsAreToolErrors(t *testing.T) { + ctx := context.Background() + type handler func(context.Context, mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) + handlers := map[string]handler{ + "validate_sql": handleValidateSQL, + "format_sql": handleFormatSQL, + "parse_sql": handleParseSQL, + "extract_metadata": handleExtractMetadata, + "security_scan": handleSecurityScan, + "lint_sql": handleLintSQL, + "analyze_sql": handleAnalyzeSQL, + } + tests := []struct { + name string + tool string + args map[string]any + want string + }{ + {name: "missing sql", tool: "validate_sql", args: map[string]any{}, want: "parameter 'sql' is required"}, + {name: "null sql", tool: "lint_sql", args: map[string]any{"sql": nil}, want: "parameter 'sql' is required"}, + {name: "empty sql", tool: "analyze_sql", args: map[string]any{"sql": ""}, want: "parameter 'sql' is required"}, + {name: "numeric sql", tool: "validate_sql", args: map[string]any{"sql": 123.0}, want: "parameter 'sql' must be a string, got number"}, + {name: "object sql", tool: "security_scan", args: map[string]any{"sql": map[string]any{}}, want: "must be a string, got object"}, + {name: "parse failure", tool: "parse_sql", args: map[string]any{"sql": "SELEC"}, want: "parse failed"}, + {name: "metadata parse failure", tool: "extract_metadata", args: map[string]any{"sql": "SELEC"}, want: "parse failed"}, + {name: "format failure", tool: "format_sql", args: map[string]any{"sql": "SELEC"}, want: "format failed"}, + {name: "indent_size word", tool: "format_sql", args: map[string]any{"sql": "SELECT 1", "indent_size": "four"}, want: "parameter 'indent_size' must be an integer between 0 and 16, got string \"four\""}, + {name: "indent_size negative", tool: "format_sql", args: map[string]any{"sql": "SELECT 1", "indent_size": -3.0}, want: "'indent_size' must be an integer between 0 and 16"}, + {name: "indent_size absurd", tool: "format_sql", args: map[string]any{"sql": "SELECT 1", "indent_size": 1e9}, want: "'indent_size' must be an integer"}, + {name: "indent_size fractional", tool: "format_sql", args: map[string]any{"sql": "SELECT 1", "indent_size": 2.5}, want: "'indent_size' must be an integer"}, + {name: "indent_size bool", tool: "format_sql", args: map[string]any{"sql": "SELECT 1", "indent_size": true}, want: "got boolean"}, + {name: "uppercase_keywords yes", tool: "format_sql", args: map[string]any{"sql": "SELECT 1", "uppercase_keywords": "yes"}, want: "parameter 'uppercase_keywords' must be a boolean"}, + {name: "add_semicolon number", tool: "format_sql", args: map[string]any{"sql": "SELECT 1", "add_semicolon": 1.0}, want: "parameter 'add_semicolon' must be a boolean, got number"}, + {name: "dialect number", tool: "validate_sql", args: map[string]any{"sql": "SELECT 1", "dialect": 5.0}, want: "parameter 'dialect' must be a string, got number"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h, ok := handlers[tt.tool] + if !ok { + t.Fatalf("no handler for %s", tt.tool) + } + res, err := h(ctx, makeReq(tt.args)) + requireToolError(t, res, err, tt.want) + }) + } +} + +func TestHandleFormatSQL_OptionalArgsAccepted(t *testing.T) { + ctx := context.Background() + tests := []struct { + name string + args map[string]any + wantOpts map[string]any + wantInSQL string + }{ + {name: "defaults", args: map[string]any{}, wantOpts: map[string]any{"indent_size": 2.0, "uppercase_keywords": false, "add_semicolon": false}}, + {name: "null treated as absent", args: map[string]any{"indent_size": nil, "add_semicolon": nil}, wantOpts: map[string]any{"indent_size": 2.0, "add_semicolon": false}}, + {name: "native types", args: map[string]any{"indent_size": 4.0, "uppercase_keywords": true, "add_semicolon": true}, wantOpts: map[string]any{"indent_size": 4.0, "uppercase_keywords": true, "add_semicolon": true}, wantInSQL: ";"}, + {name: "stringified scalars", args: map[string]any{"indent_size": "4", "uppercase_keywords": "TRUE", "add_semicolon": "false"}, wantOpts: map[string]any{"indent_size": 4.0, "uppercase_keywords": true, "add_semicolon": false}}, + {name: "bounds inclusive", args: map[string]any{"indent_size": 0.0}, wantOpts: map[string]any{"indent_size": 0.0}}, + {name: "upper bound", args: map[string]any{"indent_size": float64(maxIndentSize)}, wantOpts: map[string]any{"indent_size": float64(maxIndentSize)}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + args := map[string]any{"sql": "select a from t"} + for k, v := range tt.args { + args[k] = v + } + res, err := handleFormatSQL(ctx, makeReq(args)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if res.IsError { + t.Fatalf("unexpected tool error: %+v", res.Content) + } + data := unmarshalResult(t, res) + opts, ok := data["options"].(map[string]any) + if !ok { + t.Fatalf("missing options in %v", data) + } + for k, want := range tt.wantOpts { + if opts[k] != want { + t.Errorf("options[%s] = %v, want %v", k, opts[k], want) + } + } + if sql, _ := data["formatted_sql"].(string); tt.wantInSQL != "" && !strings.Contains(sql, tt.wantInSQL) { + t.Errorf("formatted_sql %q does not contain %q", sql, tt.wantInSQL) + } + }) + } +} diff --git a/plugins/gosqlx/skills/sql-formatting/SKILL.md b/plugins/gosqlx/skills/sql-formatting/SKILL.md index 1137204e..420bcce5 100644 --- a/plugins/gosqlx/skills/sql-formatting/SKILL.md +++ b/plugins/gosqlx/skills/sql-formatting/SKILL.md @@ -14,7 +14,7 @@ Input: `$ARGUMENTS` - a SQL string or a path to a `.sql` file. 1. Get the SQL (read the file if a path was given). 2. Call `format_sql` (`mcp__plugin_gosqlx_gosqlx__format_sql`) with: - `sql` (required) - - `indent_size` - default 2. `0` puts each statement on one line. Clauses always start at column 0; larger values only indent nested blocks such as CREATE TABLE column lists + - `indent_size` - integer 0–16, default 2. `0` puts each statement on one line. Clauses always start at column 0; larger values only indent nested blocks such as CREATE TABLE column lists - `uppercase_keywords` - `true` upper-cases all keywords (`AND`, `IN`, `CASE`, ...) but not data type names such as `int`. The default `false` still upper-cases clause keywords (`SELECT`, `FROM`, `WHERE`, ...) and leaves others such as `and`/`or` as written; the MCP tool cannot lowercase keywords - `add_semicolon` - default `false` 3. If formatting fails, the SQL probably does not parse: call `validate_sql` and report the error instead of guessing a fix. diff --git a/render.yaml b/render.yaml index 87e195ec..fb689654 100644 --- a/render.yaml +++ b/render.yaml @@ -10,6 +10,15 @@ services: value: "0.0.0.0" - key: GOSQLX_MCP_PORT value: "8080" + # Rate-limit key. Render's edge (and the mcp.gosqlx.dev zone) is + # Cloudflare, which overwrites CF-Connecting-IP with the visitor IP on + # every request, so clients cannot spoof it. X-Forwarded-For is NOT used: + # Render appends to a client-supplied value instead of resetting it, so + # its left-most entry is attacker-controlled. RemoteAddr alone would be + # Render's proxy, putting every user in one bucket. Only set this when + # all traffic reaches the server through Cloudflare. + - key: GOSQLX_MCP_TRUSTED_PROXY_HEADER + value: "CF-Connecting-IP" # Domain-verification token for the ChatGPT app directory, issued by the # OpenAI Platform during submission. Set the value in the Render # dashboard; it is served at /.well-known/openai-apps-challenge.