mcp: add an sse event size cap (#1205)

update sse I/O code to limit the amount of memory a single
non-terminated event can take

---------

Co-authored-by: Guglielmo Colombo <guglielmoc@google.com>
This commit is contained in:
Yaroslav 2026-09-02 11:35:23 +02:00 • committed by GitHub
commit 3632967404
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 423 additions and 18 deletions

View file

@ -65,13 +65,21 @@ func writeEvent(w http.ResponseWriter, evt Event) (int, error) {
return n, err
}
// scanEvents iterates SSE events in the given scanner. The iterated error is
// terminal: if encountered, the stream is corrupt or broken and should no
// longer be used.
// DefaultMaxEventSize is the default maximum number of bytes buffered while
// reading a single server-sent event before the input is rejected.
const DefaultMaxEventSize = 16 << 20 // 16 MiB
// scanEventsLimited iterates SSE events in the given reader, buffering at most
// maxEventSize bytes per event. When maxEventSize > 0, an event is reported as
// [errMalformedEvent] once its size exceeds it; a non-positive maxEventSize
// disables the cap.
//
// The iterated error is terminal: if encountered, the stream is corrupt or
// broken and should no longer be used.
//
// TODO(rfindley): consider a different API here that makes failure modes more
// apparent.
func scanEvents(r io.Reader) iter.Seq2[Event, error] {
func scanEventsLimited(r io.Reader, maxEventSize int) iter.Seq2[Event, error] {
reader := bufio.NewReader(r)
// TODO: investigate proper behavior when events are out of order, or have
@ -94,8 +102,9 @@ func scanEvents(r io.Reader) iter.Seq2[Event, error] {
// - Lines starting with ":" are ignored.
// - Records are terminated with two consecutive newlines.
var (
evt Event
dataBuf *bytes.Buffer // if non-nil, preceding field was also data
evt Event
dataBuf *bytes.Buffer // if non-nil, preceding field was also data
eventBytes int // bytes read for the current, in-progress event
)
yieldEvent := func() bool {
if dataBuf != nil {
@ -112,11 +121,20 @@ func scanEvents(r io.Reader) iter.Seq2[Event, error] {
return true
}
for {
line, err := reader.ReadBytes('\n')
budget := -1
if maxEventSize > 0 {
budget = maxEventSize - eventBytes
}
line, err := readEventLine(reader, budget)
if errors.Is(err, errEventTooLarge) {
yield(Event{}, fmt.Errorf("%w: SSE event exceeded %d bytes without terminating", errMalformedEvent, maxEventSize))
return
}
if err != nil && !errors.Is(err, io.EOF) {
yield(Event{}, fmt.Errorf("error reading event: %v", err))
return
}
eventBytes += len(line)
line = bytes.TrimRight(line, "\r\n")
isEOF := errors.Is(err, io.EOF)
@ -124,6 +142,7 @@ func scanEvents(r io.Reader) iter.Seq2[Event, error] {
if !yieldEvent() {
return
}
eventBytes = 0 // reset the budget between events
if isEOF {
return
}
@ -322,6 +341,31 @@ var ErrEventsPurged = errors.New("data purged")
// transient I/O errors which may be retryable.
var errMalformedEvent = errors.New("malformed event")
// errEventTooLarge is returned by [readEventLine] when a single line would
// exceed the remaining per-event byte budget.
var errEventTooLarge = errors.New("SSE event exceeded maximum size")
// readEventLine reads a single '\n'-terminated line from r. When budget >= 0 it
// reads at most budget bytes, returning [errEventTooLarge] once that budget is
// exceeded before a newline arrives.
func readEventLine(r *bufio.Reader, budget int) ([]byte, error) {
if budget < 0 {
return r.ReadBytes('\n')
}
var line []byte
for {
frag, err := r.ReadSlice('\n')
if len(line)+len(frag) > budget {
return nil, errEventTooLarge
}
line = append(line, frag...)
if errors.Is(err, bufio.ErrBufferFull) {
continue // looking for delim
}
return line, err
}
}
// After implements [EventStore.After].
func (s *MemoryEventStore) After(_ context.Context, sessionID, streamID string, index int) iter.Seq2[[]byte, error] {
// Return the data items to yield.

View file

@ -5,9 +5,12 @@
package mcp
import (
"bytes"
"context"
"crypto/rand"
"errors"
"fmt"
"io"
"slices"
"strings"
"testing"
@ -134,7 +137,7 @@ func TestScanEvents(t *testing.T) {
r := strings.NewReader(tt.input)
var got []Event
var err error
for e, err2 := range scanEvents(r) {
for e, err2 := range scanEventsLimited(r, DefaultMaxEventSize) {
if err2 != nil {
err = err2
break
@ -175,6 +178,136 @@ func TestScanEvents(t *testing.T) {
}
}
// endlessReader streams a fixed prefix once, then repeats fill forever.
type endlessReader struct {
prefix string
sent bool
repeat string
repIndex int
}
func (r *endlessReader) Read(p []byte) (int, error) {
if !r.sent {
n := copy(p, r.prefix)
r.sent = true
return n, nil
}
for i := range p {
p[i] = r.repeat[r.repIndex%len(r.repeat)]
r.repIndex++
}
return len(p), nil
}
func (r *endlessReader) Close() error { return nil }
// TestScanEventsMaxEventSize verifies that scanEvents bounds the bytes
// buffered for a single event.
func TestScanEventsMaxEventSize(t *testing.T) {
wantByte := byte('A')
eventOfSize := func(size int) []byte {
var buf bytes.Buffer
buf.WriteString("data: ")
buf.Write(bytes.Repeat([]byte{wantByte}, size))
buf.WriteString("\n\n")
return buf.Bytes()
}
tests := []struct {
name string
reader io.Reader
maxEventSize int
wantErr bool
wantDataLengths []int
}{
{
name: "negative budget disables the cap",
reader: bytes.NewReader(eventOfSize(512)),
maxEventSize: -1,
wantDataLengths: []int{512},
},
{
name: "long line under cap", // bufio.Reader default buf size is 1 << 12
reader: bytes.NewReader(eventOfSize(1 << 14)),
maxEventSize: 1 << 15,
wantDataLengths: []int{1 << 14},
},
{
name: "unbounded rejected",
reader: &endlessReader{prefix: "data: ", repeat: string(wantByte)},
maxEventSize: 1024,
wantErr: true,
},
{
name: "consecutive data fields longer than limit",
reader: strings.NewReader(func() string {
var b strings.Builder
for range 2048 {
b.WriteString("data: \n")
}
return b.String()
}()),
maxEventSize: 1024,
wantErr: true,
},
{
name: "budget resets between events",
reader: bytes.NewReader(append(eventOfSize(3*1024), eventOfSize(3*1024)...)),
maxEventSize: 4 * 1024,
wantDataLengths: []int{3 * 1024, 3 * 1024},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
type result struct {
events []Event
err error
}
done := make(chan result, 1)
go func() {
var events []Event
for e, err := range scanEventsLimited(tt.reader, tt.maxEventSize) {
if err != nil {
done <- result{events, err}
return
}
events = append(events, e)
}
done <- result{events, nil}
}()
var res result
select {
case res = <-done:
case <-time.After(5 * time.Second):
t.Fatal("scanEvents did not return")
}
if tt.wantErr {
if !errors.Is(res.err, errMalformedEvent) {
t.Fatalf("got error %v, want errMalformedEvent", res.err)
}
return
}
if res.err != nil {
t.Fatalf("unexpected error: %v", res.err)
}
if len(res.events) != len(tt.wantDataLengths) {
t.Fatalf("got %d events, want %d", len(res.events), len(tt.wantDataLengths))
}
for i, wantLen := range tt.wantDataLengths {
wantRepeated := string(wantByte)
if got := res.events[i].Data; string(got) != strings.Repeat(wantRepeated, wantLen) {
t.Errorf("event %d: got %d data bytes, want %d bytes of %s", i, len(got), wantLen, wantRepeated)
}
}
})
}
}
func TestMemoryEventStoreState(t *testing.T) {
ctx := context.Background()

View file

@ -397,6 +397,11 @@ type SSEClientTransport struct {
// HTTPClient is the client to use for making HTTP requests. If nil,
// http.DefaultClient is used.
HTTPClient *http.Client
// MaxEventSize bounds the number of bytes buffered while reading a single
// server-sent event. A value of 0 selects [DefaultMaxEventSize], a negative
// value disables the cap.
MaxEventSize int
}
// Connect connects through the client endpoint.
@ -427,9 +432,14 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) {
return nil, fmt.Errorf("failed to connect: %s", http.StatusText(resp.StatusCode))
}
maxEventSize := c.MaxEventSize
if maxEventSize == 0 {
maxEventSize = DefaultMaxEventSize
}
msgEndpoint, err := func() (*url.URL, error) {
var evt Event
for evt, err = range scanEvents(resp.Body) {
for evt, err = range scanEventsLimited(resp.Body, maxEventSize) {
break
}
if err != nil {
@ -458,7 +468,7 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) {
go func() {
defer s.Close() // close the transport when the GET exits
for evt, err := range scanEvents(resp.Body) {
for evt, err := range scanEventsLimited(resp.Body, maxEventSize) {
if err != nil {
return
}

View file

@ -2018,6 +2018,11 @@ type StreamableClientTransport struct {
// OAuthHandler is an optional field that, if provided, will be used to authorize the requests.
OAuthHandler auth.OAuthHandler
// MaxEventSize bounds the number of bytes buffered while reading a single
// server-sent event. A value of 0 selects [DefaultMaxEventSize], a negative
// value disables the cap.
MaxEventSize int
// TODO(rfindley): propose exporting these.
// If strict is set, the transport is in 'strict mode', where any violation
// of the MCP spec causes a failure.
@ -2105,6 +2110,7 @@ func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, er
failed: make(chan struct{}),
disableStandaloneSSE: t.DisableStandaloneSSE,
oauthHandler: t.OAuthHandler,
maxEventSize: t.MaxEventSize,
}
return conn, nil
}
@ -2126,6 +2132,10 @@ type streamableClientConn struct {
// oauthHandler is the OAuth handler for the connection.
oauthHandler auth.OAuthHandler // from [StreamableClientTransport.OAuthHandler]
// maxEventSize bounds the number of bytes buffered while reading a single
// server-sent event. Resolved from [StreamableClientTransport.MaxEventSize].
maxEventSize int
// Guard calls to Close, as it may be called multiple times.
closeOnce sync.Once
closeErr error
@ -2635,7 +2645,7 @@ func (c *streamableClientConn) processStream(ctx context.Context, requestSummary
io.Copy(io.Discard, resp.Body)
resp.Body.Close()
}()
for evt, err := range scanEvents(resp.Body) {
for evt, err := range scanEventsLimited(resp.Body, c.maxEventSize) {
if err != nil {
if ctx.Err() != nil {
return "", 0, true // don't reconnect: client cancelled

View file

@ -1544,7 +1544,7 @@ func (s streamableRequest) do(ctx context.Context, serverURL, sessionID string,
var respBody []byte
if contentType == "text/event-stream" {
r := readerInto{resp.Body, new(bytes.Buffer)}
for evt, err := range scanEvents(r) {
for evt, err := range scanEventsLimited(r, DefaultMaxEventSize) {
if err != nil {
return newSessionID, resp.StatusCode, nil, fmt.Errorf("reading events: %v", err)
}
@ -2856,6 +2856,69 @@ data: {"jsonrpc":"2.0","id":1,"result":{}}
}
}
// TestProcessStreamMaxEventSize verifies that a streamableClientConn honors its
// configured maxEventSize.
func TestProcessStreamMaxEventSize(t *testing.T) {
jsonrpcEventOfSize := func(padSize int) string {
msg := fmt.Sprintf(`{"jsonrpc":"2.0","id":1,"result":{"pad":%q}}`, strings.Repeat("A", padSize))
return "data: " + msg + "\n\n"
}
tests := []struct {
name string
maxEventSize int
event string
wantFail bool
}{
{
name: "event over cap is rejected",
event: jsonrpcEventOfSize(4096),
maxEventSize: 1024,
wantFail: true,
},
{
name: "event under cap is accepted",
event: jsonrpcEventOfSize(1024),
maxEventSize: 4096,
wantFail: false,
},
{
name: "negative cap disables the limit",
event: jsonrpcEventOfSize(4096),
maxEventSize: -1,
wantFail: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx := t.Context()
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(tt.event)),
}
conn := &streamableClientConn{
ctx: ctx,
done: make(chan struct{}),
incoming: make(chan jsonrpc.Message, 10),
failed: make(chan struct{}),
logger: ensureLogger(nil),
maxEventSize: tt.maxEventSize,
}
conn.processStream(ctx, "test", resp, nil)
err := conn.failure()
if tt.wantFail && err == nil {
t.Fatal("failure() = nil, want a non-nil error for an oversized event")
}
if !tt.wantFail && err != nil {
t.Fatalf("failure() = %v, want nil", err)
}
})
}
}
// TestScanEventsPingFiltering is a unit test for the low-level event scanning
// with ping events to verify scanEvents properly parses all event types.
func TestScanEventsPingFiltering(t *testing.T) {
@ -2878,7 +2941,7 @@ data: {"jsonrpc":"2.0","method":"test2","params":{}}
var events []Event
// Scan all events
for evt, err := range scanEvents(reader) {
for evt, err := range scanEventsLimited(reader, DefaultMaxEventSize) {
if err != nil {
if err != io.EOF {
t.Fatalf("scanEvents error: %v", err)

View file

@ -119,13 +119,22 @@ type serverConnection interface {
sessionUpdated(ServerSessionState)
}
// DefaultMaxLineLength is the default maximum number of bytes buffered while
// decoding a single inbound JSON-RPC frame.
const DefaultMaxLineLength = 16 * 1024 * 1024
// A StdioTransport is a [Transport] that communicates over stdin/stdout using
// newline-delimited JSON.
type StdioTransport struct{}
type StdioTransport struct {
// MaxLineLength bounds the number of bytes that may be buffered while
// decoding a single inbound JSON-RPC frame. A value of 0 selects [DefaultMaxLineLength],
// a negative value disables the cap.
MaxLineLength int
}
// Connect implements the [Transport] interface.
func (*StdioTransport) Connect(context.Context) (Connection, error) {
return newIOConn(rwc{os.Stdin, nopCloserWriter{os.Stdout}}), nil
func (t *StdioTransport) Connect(context.Context) (Connection, error) {
return newIOConnLimited(rwc{os.Stdin, nopCloserWriter{os.Stdout}}, t.MaxLineLength), nil
}
// nopCloserWriter is an io.WriteCloser with a trivial Close method.
@ -140,11 +149,15 @@ func (nopCloserWriter) Close() error { return nil }
type IOTransport struct {
Reader io.ReadCloser
Writer io.WriteCloser
// MaxLineLength bounds the number of bytes that may be buffered while
// decoding a single inbound JSON-RPC frame. A value of 0 selects [DefaultMaxLineLength],
// a negative value disables the cap.
MaxLineLength int
}
// Connect implements the [Transport] interface.
func (t *IOTransport) Connect(context.Context) (Connection, error) {
return newIOConn(rwc{t.Reader, t.Writer}), nil
return newIOConnLimited(rwc{t.Reader, t.Writer}, t.MaxLineLength), nil
}
// An InMemoryTransport is a [Transport] that communicates over an in-memory
@ -478,6 +491,17 @@ type msgOrErr struct {
}
func newIOConn(rwc io.ReadWriteCloser) *ioConn {
return newIOConnLimited(rwc, DefaultMaxLineLength)
}
// newIOConnLimited builds an [ioConn] over rwc that bounds the number of bytes
// buffered while decoding a single inbound JSON-RPC frame to maxLineLength.
// maxLineLength == 0 selects [DefaultMaxLineLength], a negative value means no cap.
func newIOConnLimited(rwc io.ReadWriteCloser, maxLineLength int) *ioConn {
limit := maxLineLength
if limit == 0 {
limit = DefaultMaxLineLength
}
var (
incoming = make(chan msgOrErr)
closed = make(chan struct{})
@ -489,7 +513,15 @@ func newIOConn(rwc io.ReadWriteCloser) *ioConn {
// but that is unavoidable since AFAIK there is no (easy and portable) way to
// guarantee that reads of stdin are unblocked when closed.
go func() {
dec := json.NewDecoder(rwc)
var (
reader io.Reader = rwc
limiter *frameLimitReader
)
if limit > 0 {
limiter = &frameLimitReader{r: rwc, limit: limit}
reader = limiter
}
dec := json.NewDecoder(reader)
for {
var raw json.RawMessage
err := dec.Decode(&raw)
@ -515,6 +547,9 @@ func newIOConn(rwc io.ReadWriteCloser) *ioConn {
if err != nil {
return
}
if limiter != nil {
limiter.resetFrame()
}
}
}()
return &ioConn{
@ -524,6 +559,32 @@ func newIOConn(rwc io.ReadWriteCloser) *ioConn {
}
}
// errFrameTooLarge means that a single inbound JSON-RPC frame exceeded the configured byte
// limit before the value was completely received.
var errFrameTooLarge = errors.New("inbound JSON-RPC frame exceeded the configured maximum line length")
// frameLimitReader bounds the number of bytes [json.Decoder] may buffer while
// decoding a single JSON value. Read returns [errFrameTooLarge] once the budget is exhausted.
type frameLimitReader struct {
r io.Reader
limit int
count int
}
func (r *frameLimitReader) Read(p []byte) (int, error) {
if r.count >= r.limit {
return 0, errFrameTooLarge
}
if len(p) > r.limit-r.count {
p = p[:r.limit-r.count]
}
n, err := r.r.Read(p)
r.count += n
return n, err
}
func (r *frameLimitReader) resetFrame() { r.count = 0 }
func (c *ioConn) SessionID() string { return "" }
func (c *ioConn) sessionUpdated(state ServerSessionState) {

View file

@ -6,9 +6,12 @@ package mcp
import (
"context"
"errors"
"fmt"
"io"
"strings"
"testing"
"time"
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
@ -167,3 +170,84 @@ func TestIOConnRead_EmptyMethod(t *testing.T) {
t.Errorf("ID = %v, want 5", req.ID.Raw())
}
}
// TestIOConnFrameCap verifies single inbound frame cap enforcement.
func TestIOConnFrameCap(t *testing.T) {
tests := []struct {
name string
r io.ReadCloser
limit int
wantMessageCount int
wantErr bool
}{
{
name: "infinite string",
r: &endlessReader{prefix: `"`, repeat: "A"},
limit: 1024,
wantErr: true,
},
{
name: "cap does not reset per new line",
r: &endlessReader{prefix: "[", repeat: "0,\n"},
limit: 1024,
wantErr: true,
},
{
name: "cap resets per message",
limit: 1024, // < 600 * 3
r: io.NopCloser(strings.NewReader(func() string {
var b strings.Builder
for i := range 3 {
fmt.Fprintf(&b, `{"jsonrpc":"2.0","id":%d,"method":"test","params":{"pad":"`, i)
b.WriteString(strings.Repeat("A", 600))
b.WriteString(`"}}`)
b.WriteByte('\n')
}
return b.String()
}())),
wantMessageCount: 3,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tr := newIOConnLimited(rwc{rc: tt.r}, tt.limit)
t.Cleanup(func() { tr.Close() })
read := func() (jsonrpc.Message, error) {
type result struct {
msg jsonrpc.Message
err error
}
done := make(chan result, 1)
go func() {
msg, err := tr.Read(context.Background())
done <- result{msg, err}
}()
select {
case r := <-done:
return r.msg, r.err
case <-time.After(10 * time.Second):
t.Fatal("Read() did not return: frame cap failed to trip")
return nil, nil
}
}
for i := range tt.wantMessageCount {
msg, err := read()
if err != nil {
t.Fatalf("Read() #%d error = %v, want nil", i, err)
}
if msg == nil {
t.Fatalf("Read() #%d returned nil message", i)
}
}
if tt.wantErr {
if _, err := read(); !errors.Is(err, errFrameTooLarge) {
t.Fatalf("Read() error = %v, want errFrameTooLarge", err)
}
}
})
}
}