From be9f20edd7d9b34ebe145bd9bd3d9e15165eebf2 Mon Sep 17 00:00:00 2001 From: tangerg Date: Fri, 6 Nov 2026 17:12:50 +0800 Subject: [PATCH] mcp: make the ioConn frame cap independent of read boundaries frameLimitReader reset its count after each message, but json.Decoder reads ahead: when it finishes a message it may already hold bytes of the next one, and those bytes were counted toward the finished message. The next message could then exceed MaxLineLength by as many bytes as had been read ahead, so whether it was rejected depended on how the reader split the stream. The reader now counts the bytes it has read, and each frame's budget starts at the decoder's InputOffset, so whether a message fits depends only on the message. That is what the MaxLineLength documentation already describes, so the documentation is unchanged. Line terminators and blank lines before a message count toward it, since they follow the previous message. Without the change, four cases of the new table test fail, and the fuzz test finds a counterexample within a second. Fixes #1509 AI assistance: prepared with Claude; I reviewed the change and can explain it. --- mcp/transport.go | 25 +++++---- mcp/transport_test.go | 115 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 131 insertions(+), 9 deletions(-) diff --git a/mcp/transport.go b/mcp/transport.go index d475152b..adb21906 100644 --- a/mcp/transport.go +++ b/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 "" } diff --git a/mcp/transport_test.go b/mcp/transport_test.go index b6551cfa..6c4d239b 100644 --- a/mcp/transport_test.go +++ b/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.