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.