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:
parent
0d3036f785
commit
3632967404
7 changed files with 423 additions and 18 deletions
58
mcp/event.go
58
mcp/event.go
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
14
mcp/sse.go
14
mcp/sse.go
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue