fix(mcp): harden gosqlx-mcp HTTP server (rate limiting, timeouts, auth, tool errors)

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

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)

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)
+				}
+			})
+		}
+	}
+}