diff --git a/mcp/tool.go b/mcp/tool.go index 04cf551a..5789be0c 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 7a82c96d..81a68e26 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 {