From c433f6cf90ca29498a6f7bcf96710104bb217a3f Mon Sep 17 00:00:00 2001
From: meraklbz <lbz2770828522@gmail.com>
Date: Sun, 16 Aug 2026 22:10:24 +0800
Subject: [PATCH] mcp: reject deeply nested tool arguments before validation

Tool input is unmarshaled with a JSON decoder that imposes no nesting
limit of its own, and schema validation can become quadratic in the
nesting depth of the input. A small adversarial payload of a few tens of
kilobytes can therefore burn minutes of CPU during validation.

Scan the raw input for nesting depth before unmarshaling or validating,
and reject input deeper than 1000 levels with an error. The scan is
linear in the length of the input and ignores string literal contents.
---
 mcp/tool.go      | 45 +++++++++++++++++++++++++++++++++++++++++++++
 mcp/tool_test.go | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+)

diff --git a/mcp/tool.go b/mcp/tool.go
index 04cf551a9..5789be0c9 100644
--- a/mcp/tool.go
+++ b/mcp/tool.go
@@ -62,6 +62,44 @@ type serverTool struct {
 	handler ToolHandler
 }
 
+// maxJSONDepth is the maximum nesting depth accepted for tool arguments and
+// output. Deeply nested JSON is rejected up front to avoid pathological
+// validation costs on adversarial input.
+const maxJSONDepth = 1000
+
+// checkJSONDepth scans data and reports an error if its nesting depth exceeds
+// maxJSONDepth. It is linear in the length of data and ignores the contents of
+// string literals.
+func checkJSONDepth(data json.RawMessage) error {
+	depth := 0
+	inString := false
+	escaped := false
+	for _, b := range data {
+		if inString {
+			if escaped {
+				escaped = false
+			} else if b == '\\' {
+				escaped = true
+			} else if b == '"' {
+				inString = false
+			}
+			continue
+		}
+		switch b {
+		case '"':
+			inString = true
+		case '{', '[':
+			depth++
+			if depth > maxJSONDepth {
+				return fmt.Errorf("JSON nesting exceeds maximum allowed depth of %d", maxJSONDepth)
+			}
+		case '}', ']':
+			depth--
+		}
+	}
+	return nil
+}
+
 // applySchema validates whether data is valid JSON according to the provided
 // schema, after applying schema defaults.
 //
@@ -76,6 +114,13 @@ func applySchema(data json.RawMessage, resolved *jsonschema.Resolved, forOutput
 	// TODO: use reflection to create the struct type to unmarshal into.
 	// Separate validation from assignment.
 
+	// Reject deeply nested input before unmarshaling or validating. Deep
+	// nesting can make schema validation quadratic in the input depth, and the
+	// JSON decoder used here does not impose a nesting limit of its own.
+	if err := checkJSONDepth(data); err != nil {
+		return nil, err
+	}
+
 	// Use default JSON marshalling for validation.
 	//
 	// This avoids inconsistent representation due to custom marshallers, such as
diff --git a/mcp/tool_test.go b/mcp/tool_test.go
index 7a82c96dd..81a68e263 100644
--- a/mcp/tool_test.go
+++ b/mcp/tool_test.go
@@ -59,6 +59,49 @@ func TestApplySchema(t *testing.T) {
 	}
 }
 
+func TestApplySchemaRejectsDeeplyNestedInput(t *testing.T) {
+	schema := &jsonschema.Schema{Type: "object"}
+	resolved, err := schema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	nested := func(depth int) string {
+		var b strings.Builder
+		b.WriteString("{\"a\":")
+		for i := 0; i < depth; i++ {
+			b.WriteString("[")
+		}
+		b.WriteString("1")
+		for i := 0; i < depth; i++ {
+			b.WriteString("]")
+		}
+		b.WriteString("}")
+		return b.String()
+	}
+
+	// Deeply nested input must be rejected quickly, regardless of the schema.
+	// nested(depth) builds one object wrapper plus `depth` nested arrays, so the
+	// total nesting is depth+1.
+	for _, depth := range []int{maxJSONDepth, 2000, 20000} {
+		raw := json.RawMessage(nested(depth))
+		if _, err := applySchema(raw, resolved, false); err == nil {
+			t.Errorf("applySchema with depth %d: got nil error, want depth error", depth)
+		}
+	}
+
+	// Nesting at or below the limit must still validate.
+	if _, err := applySchema(json.RawMessage(nested(maxJSONDepth-1)), resolved, false); err != nil {
+		t.Errorf("applySchema with depth %d: unexpected error: %v", maxJSONDepth, err)
+	}
+
+	// Braces inside string literals must not count toward nesting depth.
+	data := `{"a":"[[[[[[[[[[[[[[[[[[[[","b":[1,2,3]}`
+	if _, err := applySchema(json.RawMessage(data), resolved, false); err != nil {
+		t.Errorf("applySchema with deep string literal: unexpected error: %v", err)
+	}
+}
+
 func TestApplySchemaOutput(t *testing.T) {
 	// SEP-2106: when forOutput is true, the schema may have a non-object root.
 	for _, tc := range []struct {