mcp: make the ioConn frame cap independent of read boundaries

mcp/transport.go

@@ -607,7 +607,7 @@ func newIOConnLimited(rwc io.ReadWriteCloser, maxLineLength int) *ioConn {
 			limiter *frameLimitReader
 		)
 		if limit > 0 {
-			limiter = &frameLimitReader{r: rwc, limit: limit}
+			limiter = &frameLimitReader{r: rwc, limit: int64(limit)}
 			reader = limiter
 		}
 		dec := json.NewDecoder(reader)
@@ -637,7 +637,7 @@ func newIOConnLimited(rwc io.ReadWriteCloser, maxLineLength int) *ioConn {
 				return
 			}
 			if limiter != nil {
-				limiter.resetFrame()
+				limiter.beginFrame(dec.InputOffset())
 			}
 		}
 	}()
@@ -655,25 +655,32 @@ var errFrameTooLarge = errors.New("inbound JSON-RPC frame exceeded the configure
 
 // frameLimitReader bounds the number of bytes [json.Decoder] may buffer while
 // decoding a single JSON value. Read returns [errFrameTooLarge] once the budget is exhausted.
+//
+// The decoder reads ahead, so when it finishes a value it may already hold
+// bytes of the next one. Those bytes belong to the next frame's budget, which
+// therefore starts at the decoder's input offset rather than at the reader's.
 type frameLimitReader struct {
 	r     io.Reader
-	limit int
-	count int
+	limit int64
+	read  int64 // bytes read from r
+	start int64 // offset in r of the current frame
 }
 
 func (r *frameLimitReader) Read(p []byte) (int, error) {
-	if r.count >= r.limit {
+	remaining := r.limit - (r.read - r.start)
+	if remaining <= 0 {
 		return 0, errFrameTooLarge
 	}
-	if len(p) > r.limit-r.count {
-		p = p[:r.limit-r.count]
+	if int64(len(p)) > remaining {
+		p = p[:remaining]
 	}
 	n, err := r.r.Read(p)
-	r.count += n
+	r.read += int64(n)
 	return n, err
 }
 
-func (r *frameLimitReader) resetFrame() { r.count = 0 }
+// beginFrame starts the budget of a frame that begins at offset in r.
+func (r *frameLimitReader) beginFrame(offset int64) { r.start = offset }
 
 func (c *ioConn) SessionID() string { return "" }
 

mcp/transport_test.go

@@ -593,6 +593,121 @@ func TestIOConnFrameCap(t *testing.T) {
 	}
 }
 
+// TestIOConnFrameCapIgnoresReadBoundaries checks that whether a message fits
+// within the frame cap depends only on the message, not on how the reader
+// splits the stream or on what the decoder read ahead while decoding the
+// message before it.
+func TestIOConnFrameCapIgnoresReadBoundaries(t *testing.T) {
+	const limit = 128
+	chunks := []int{1, 2, 3, 7, 64, limit - 1, limit, limit + 1, 1 << 16}
+	tests := []struct {
+		name   string
+		frames []capFrame
+	}{
+		{"first message at the cap", []capFrame{{"", limit}}},
+		{"first message over the cap", []capFrame{{"", limit + 1}}},
+		{"message at the cap after a short one", []capFrame{{"", 50}, {"\n", limit}}},
+		{"message over the cap after a short one", []capFrame{{"", 50}, {"\n", limit + 1}}},
+		{"CRLF counts toward the next message", []capFrame{{"", 50}, {"\r\n", limit}, {"\r\n", limit + 1}}},
+		{"blank lines count toward the next message", []capFrame{{"", 50}, {"\n\n\n", limit}, {"\n\n", limit + 1}}},
+		{"each message has the whole cap", []capFrame{{"", limit}, {"\n", limit}, {"\n", limit}, {"\n", limit}}},
+		{"short messages then one over the cap", []capFrame{{"", 50}, {"\n", 50}, {"\n", 50}, {"\n", limit + 1}}},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			checkFrameCap(t, limit, tt.frames, chunks)
+		})
+	}
+}
+
+func FuzzIOConnFrameCap(f *testing.F) {
+	f.Add(uint8(64), uint8(0), []byte{0, 1, 2, 3})
+	f.Add(uint8(0), uint8(63), []byte{4, 255, 17, 200, 9})
+	f.Add(uint8(255), uint8(255), []byte{252, 253, 254, 255})
+	f.Fuzz(func(t *testing.T, limitByte, chunkByte uint8, layout []byte) {
+		const maxFrames = 16
+		if len(layout) > maxFrames {
+			layout = layout[:maxFrames]
+		}
+		separators := []string{"\n", "\r\n", "\n\n", "\r\n\r\n"}
+		frames := make([]capFrame, len(layout))
+		for i, b := range layout {
+			sep := separators[b%4]
+			if i == 0 {
+				sep = ""
+			}
+			frames[i] = capFrame{sep, len(sep) + len(capMessage("")) + 5*int(b/4)}
+		}
+		checkFrameCap(t, 64+int(limitByte), frames, []int{1, 1 + int(chunkByte), 1 << 16})
+	})
+}
+
+// A capFrame describes a message in a stream given to the frame cap.
+type capFrame struct {
+	sep  string // line terminators before the message
+	size int    // bytes the cap counts: len(sep) plus the message
+}
+
+// capMessage returns a JSON-RPC notification carrying pad.
+func capMessage(pad string) string {
+	return `{"jsonrpc":"2.0","method":"m","params":{"p":"` + pad + `"}}`
+}
+
+// checkFrameCap reads the stream described by frames, splitting it into
+// reads of each of the chunk sizes in turn, and checks that the connection
+// delivers exactly the leading messages whose size is at most limit.
+func checkFrameCap(t *testing.T, limit int, frames []capFrame, chunks []int) {
+	t.Helper()
+	var b strings.Builder
+	for _, f := range frames {
+		pad := f.size - len(f.sep) - len(capMessage(""))
+		if pad < 0 {
+			t.Fatalf("frame size %d is too small for a message after %q", f.size, f.sep)
+		}
+		b.WriteString(f.sep)
+		b.WriteString(capMessage(strings.Repeat("x", pad)))
+	}
+	b.WriteString("\n")
+	stream := b.String()
+
+	want, wantErr := 0, io.EOF
+	for want < len(frames) && frames[want].size <= limit {
+		want++
+	}
+	if want < len(frames) {
+		wantErr = errFrameTooLarge
+	}
+	for _, chunk := range chunks {
+		r := &chunkReader{r: strings.NewReader(stream), n: chunk}
+		tr := newIOConnLimited(rwc{rc: io.NopCloser(r)}, limit)
+		got := 0
+		var err error
+		for {
+			if _, err = tr.Read(context.Background()); err != nil {
+				break
+			}
+			got++
+		}
+		tr.Close()
+		if got != want || !errors.Is(err, wantErr) {
+			t.Errorf("reading %d bytes at a time: got %d messages, then %v; want %d, then %v", chunk, got, err, want, wantErr)
+		}
+	}
+}
+
+// chunkReader returns at most n bytes from each Read.
+type chunkReader struct {
+	r io.Reader
+	n int
+}
+
+func (r *chunkReader) Read(p []byte) (int, error) {
+	if len(p) > r.n {
+		p = p[:r.n]
+	}
+	return r.r.Read(p)
+}
+
 // TestIOConnDropResponse checks that ioConn discards the response to a call
 // the peer cancelled on a session using protocol version 2026-07-28 or later,
 // and still sends it on earlier versions.