internal/json: json nesting depth check (#1200)

Added a lightweight JSON depth-checking pass before the actual parsing.

internal/json/json.go

@@ -8,11 +8,21 @@ package json
 
 import (
 	"bytes"
+	"fmt"
 	"io"
 
 	"github.com/segmentio/encoding/json"
 )
 
+// defaultMaxDepth is the maximum JSON nesting depth accepted by [Unmarshal].
+// Realistic MCP payloads nest only a handful of levels, so 1000 is far more
+// than any legitimate message requires while keeping worst-case parse cost
+// bounded.
+const defaultMaxDepth = 1000
+
+// errMaxDepthExceeded is returned by [Unmarshal] when the input nests deeper than [defaultMaxDepth].
+var errMaxDepthExceeded = fmt.Errorf("json: exceeded maximum nesting depth of %d", defaultMaxDepth)
+
 type Decoder struct {
 	dec *json.Decoder
 }
@@ -28,5 +38,47 @@ func (d *Decoder) Decode(v any) error {
 }
 
 func Unmarshal(data []byte, v any) error {
+	if err := checkMaxDepth(data, defaultMaxDepth); err != nil {
+		return err
+	}
 	return NewDecoder(bytes.NewReader(data)).Decode(v)
 }
+
+// checkMaxDepth scans data once and reports [errMaxDepthExceeded] if the
+// nesting of JSON objects and arrays exceeds maxDepth. It is a lightweight pass which
+// tracks '{' and '[' against '}' and ']' while skipping over the contents of strings.
+func checkMaxDepth(data []byte, maxDepth int) error {
+	var (
+		depth    int
+		inString bool
+		escaped  bool
+	)
+	for _, b := range data {
+		if inString {
+			switch {
+			case escaped:
+				escaped = false
+			case b == '\\':
+				escaped = true
+			case b == '"':
+				inString = false
+			}
+			continue
+		}
+
+		switch b {
+		case '"':
+			inString = true
+		case '{', '[':
+			depth++
+			if depth > maxDepth {
+				return errMaxDepthExceeded
+			}
+		case '}', ']':
+			if depth > 0 {
+				depth--
+			}
+		}
+	}
+	return nil
+}

internal/json/json_test.go

@@ -7,9 +7,11 @@ package json
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 )
@@ -140,6 +142,51 @@ func TestUnmarshalNullCharacter(t *testing.T) {
 	}
 }
 
+func TestUnmarshalChecksNestingDepth(t *testing.T) {
+	createNesting := func(depth int) []byte {
+		return []byte(strings.Repeat("[", depth) + strings.Repeat("]", depth))
+	}
+
+	tests := []struct {
+		name    string
+		data    []byte
+		wantErr bool
+	}{
+		{"flat object", []byte(`{"a":1}`), false},
+		{"at limit", createNesting(defaultMaxDepth), false},
+		{"one over limit", createNesting(defaultMaxDepth + 1), true},
+		{"brackets in strings do not count", []byte(`{"a":"[[[[[[[[[["}`), false},
+		{"escaped quote in string", []byte(`{"a":"\"[[[[[["}`), false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var v any
+			err := Unmarshal(tt.data, &v)
+			gotDepthErr := errors.Is(err, errMaxDepthExceeded)
+			if gotDepthErr != tt.wantErr {
+				t.Fatalf("Unmarshal() error = %v (err=%v), want depth error = %v", gotDepthErr, err, tt.wantErr)
+			}
+		})
+	}
+}
+
+func TestUnmarshalMaxDepthExitsFast(t *testing.T) {
+	data := []byte(strings.Repeat("[", 500_000) + strings.Repeat("]", 500_000))
+	done := make(chan error, 1)
+	go func() {
+		var v any
+		done <- Unmarshal(data, &v)
+	}()
+	select {
+	case err := <-done:
+		if !errors.Is(err, errMaxDepthExceeded) {
+			t.Fatalf("Unmarshal error = %v, want %v", err, errMaxDepthExceeded)
+		}
+	case <-time.After(500 * time.Millisecond):
+		t.Fatal("Unmarshal did not reject deeply nested payload promptly")
+	}
+}
+
 func FuzzUnmarshalKeyComparison(f *testing.F) {
 	// Add seed corpus for expected valid key and the known problematic cases.
 	f.Add("field")