mcp: enforce body size limit in Stremable connection (#1055)

This PR enforces a body size limit on Streamable connection to prevent
OOM attacks

mcp/streamable.go

@@ -194,8 +194,25 @@ type StreamableHTTPOptions struct {
 	//   protection := http.NewCrossOriginProtection()
 	//   protectedHandler := protection.Handler(handler)
 	CrossOriginProtection *http.CrossOriginProtection
+
+	// MaxRequestBodyBytes limits the number of bytes read from any incoming
+	// HTTP request body. Requests that exceed this limit are rejected with
+	// 413 Request Entity Too Large.
+	//
+	// The limit is enforced during the read, so it applies uniformly to
+	// requests using Content-Length, Transfer-Encoding: chunked, or HTTP/2
+	// (which has no Content-Length).
+	//
+	// If zero, [DefaultMaxRequestBodyBytes] is used.
+	// A negative value disables the limit entirely; do not use this on
+	// servers exposed to untrusted clients.
+	MaxRequestBodyBytes int64
 }
 
+// DefaultMaxRequestBodyBytes is the default value used for
+// [StreamableHTTPOptions.MaxRequestBodyBytes] when it is left at zero.
+const DefaultMaxRequestBodyBytes = 4 << 20 // 4 MiB
+
 // NewStreamableHTTPHandler returns a new [StreamableHTTPHandler].
 //
 // The getServer function is used to create or look up servers for new
@@ -216,6 +233,10 @@ func NewStreamableHTTPHandler(getServer func(*http.Request) *Server, opts *Strea
 		h.opts.CrossOriginProtection = &http.CrossOriginProtection{}
 	}
 
+	if h.opts.MaxRequestBodyBytes == 0 {
+		h.opts.MaxRequestBodyBytes = DefaultMaxRequestBodyBytes
+	}
+
 	return h
 }
 
@@ -308,6 +329,11 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 		}
 	}
 
+	// Bound the request body to protect against OOM attacks.
+	if req.Body != nil && h.opts.MaxRequestBodyBytes > 0 {
+		req.Body = http.MaxBytesReader(w, req.Body, h.opts.MaxRequestBodyBytes)
+	}
+
 	// [§2.7] of the spec (2025-06-18): validate the MCP-Protocol-Version
 	// header. If provided, it must be a supported version. If absent, the
 	// version is unknown (the request may be an initialize for any version).
@@ -368,6 +394,11 @@ func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http.
 
 	info, err := h.ephemeralConnectOpts(req)
 	if err != nil {
+		var mbe *http.MaxBytesError
+		if errors.As(err, &mbe) {
+			http.Error(w, fmt.Sprintf("request body exceeds %d bytes", mbe.Limit), http.StatusRequestEntityTooLarge)
+			return
+		}
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
@@ -435,7 +466,8 @@ func (h *StreamableHTTPHandler) ephemeralConnectOpts(req *http.Request) (*epheme
 	var hasInitialize, hasInitialized, usesNewProtocol, isSubscriptionsListen bool
 	body, err := io.ReadAll(req.Body)
 	if err != nil {
-		return nil, fmt.Errorf("failed to read body")
+		// Preserve *http.MaxBytesError so serveStateless can respond with 413.
+		return nil, fmt.Errorf("failed to read body: %w", err)
 	}
 	req.Body.Close()
 	req.Body = io.NopCloser(bytes.NewBuffer(body))
@@ -1374,6 +1406,11 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 	// Read incoming messages.
 	body, err := io.ReadAll(req.Body)
 	if err != nil {
+		var mbe *http.MaxBytesError
+		if errors.As(err, &mbe) {
+			http.Error(w, fmt.Sprintf("request body exceeds %d bytes", mbe.Limit), http.StatusRequestEntityTooLarge)
+			return
+		}
 		http.Error(w, "failed to read body", http.StatusBadRequest)
 		return
 	}

mcp/streamable_test.go

@@ -3802,3 +3802,105 @@ func TestStreamableServerRejectsDuplicateInFlightRequestID(t *testing.T) {
 	default:
 	}
 }
+
+// TestStreamableMaxRequestBodyBytes verifies that the streamable HTTP handler
+// enforces StreamableHTTPOptions.MaxRequestBodyBytes on incoming request
+// bodies. The limit must apply uniformly regardless of transfer encoding:
+// Content-Length, Transfer-Encoding: chunked, or HTTP/2 (which has no
+// Content-Length).
+func TestStreamableMaxRequestBodyBytes(t *testing.T) {
+	// A minimally-valid JSON-RPC initialize request. The body-size check runs
+	// before parsing, so exact contents don't matter for the reject cases,
+	// but we use a real one for the "under limit" success case.
+	validInit := []byte(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"t","version":"v1"}}}`)
+
+	// Build a payload of the requested size by padding a JSON object. The
+	// server doesn't need to parse it; the MaxBytesReader trips first.
+	oversizedBody := func(n int) []byte {
+		b := make([]byte, n)
+		b[0] = '{'
+		for i := 1; i < n-1; i++ {
+			b[i] = ' '
+		}
+		b[n-1] = '}'
+		return b
+	}
+
+	tests := []struct {
+		name           string
+		limit          int64 // value assigned to MaxRequestBodyBytes
+		body           []byte
+		wantStatus     int
+		wantStatusDesc string
+	}{
+		{
+			name:           "default limit rejects oversized",
+			limit:          0, // triggers DefaultMaxRequestBodyBytes (4 MiB)
+			body:           oversizedBody(DefaultMaxRequestBodyBytes + 1),
+			wantStatus:     http.StatusRequestEntityTooLarge,
+			wantStatusDesc: "413 for body > 4MiB default",
+		},
+		{
+			name:           "custom limit rejects oversized",
+			limit:          1024,
+			body:           oversizedBody(2048),
+			wantStatus:     http.StatusRequestEntityTooLarge,
+			wantStatusDesc: "413 for body > custom 1KiB limit",
+		},
+		{
+			name:           "chunked encoding does not bypass limit",
+			limit:          1024,
+			body:           oversizedBody(4096),
+			wantStatus:     http.StatusRequestEntityTooLarge,
+			wantStatusDesc: "413 for chunked body > limit (regression: Content-Length-only checks are bypassable)",
+		},
+		{
+			name:           "negative limit disables the check",
+			limit:          -1,
+			body:           oversizedBody(2 * DefaultMaxRequestBodyBytes),
+			wantStatus:     http.StatusBadRequest, // body is not valid JSON-RPC; but the size check must not fire
+			wantStatusDesc: "not 413 when limit is disabled",
+		},
+		{
+			name:           "under limit succeeds",
+			limit:          64 * 1024,
+			body:           validInit,
+			wantStatus:     http.StatusOK,
+			wantStatusDesc: "200 for a valid initialize under the limit",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			server := NewServer(testImpl, nil)
+			handler := NewStreamableHTTPHandler(
+				func(*http.Request) *Server { return server },
+				&StreamableHTTPOptions{
+					Stateless:           true,
+					MaxRequestBodyBytes: tc.limit,
+				},
+			)
+			ts := httptest.NewServer(handler)
+			defer ts.Close()
+
+			req, err := http.NewRequest(http.MethodPost, ts.URL, bytes.NewReader(tc.body))
+			if err != nil {
+				t.Fatalf("NewRequest: %v", err)
+			}
+			req.Header.Set("Content-Type", "application/json")
+			req.Header.Set("Accept", "application/json, text/event-stream")
+
+			resp, err := ts.Client().Do(req)
+			if err != nil {
+				t.Fatalf("Do: %v", err)
+			}
+			defer resp.Body.Close()
+			body, _ := io.ReadAll(resp.Body)
+
+			if resp.StatusCode != tc.wantStatus {
+				t.Fatalf("%s: status = %d, want %d (body=%q)",
+					tc.wantStatusDesc, resp.StatusCode, tc.wantStatus, string(body))
+			}
+		})
+	}
+}