diff --git a/AGENTS.md b/AGENTS.md index e924a9f..4bd9069 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -44,6 +44,7 @@ Server.runWithIO() ← dispatch loop - **Tests use pipes.** Integration tests simulate stdio with `io.Pipe()`. E2E tests spawn a real subprocess via `os/exec`. - **Error codes follow JSON-RPC 2.0.** `-32700` = parse error (broken JSON), `-32600` = invalid request (well-formed JSON that is not a Request object, id `null`), `-32601` = method not found, `-32602` = invalid params, `-32000` = application error. - **Bad input never kills the loop.** A malformed line is answered in-band and the server keeps serving; `RunWithIO` returns only on EOF or a read/write failure. +- **Inbound messages are size-capped.** One line may carry at most `Server.MaxRequestBytes` bytes (default `DefaultMaxRequestBytes`, 10 MiB; negative disables — not recommended). An oversized line is answered with `-32600` (id null), its remainder discarded, and the loop keeps serving — a single line can never exhaust memory. - **Go naming.** Exported types are PascalCase. Unexported internals are camelCase. Test functions are `TestXxx`. - **Protocol version pinned.** `2024-11-05` hardcoded — update manually when MCP spec revs. diff --git a/gomcp/server.go b/gomcp/server.go index 40fc9e4..a374cae 100644 --- a/gomcp/server.go +++ b/gomcp/server.go @@ -23,6 +23,17 @@ import ( // DefaultProtocolVersion is the MCP protocol version this server speaks. const DefaultProtocolVersion = "2025-03-26" +// DefaultMaxRequestBytes caps a single inbound JSON-RPC message when +// Server.MaxRequestBytes is unset. 10 MiB matches what a typical MCP client +// accepts for a response — requests larger than that are almost certainly +// hostile or malformed, and reading them unboundedly lets one line OOM the +// server process. +const DefaultMaxRequestBytes int64 = 10 << 20 + +// errMessageTooLarge reports that an inbound line exceeded the configured +// size cap. The rest of the line has already been discarded by readMessage. +var errMessageTooLarge = errors.New("message exceeds maximum size") + // Server is an MCP server that communicates over stdio using JSON-RPC 2.0. // It handles the MCP handshake and dispatches tools, resources, and prompts // to registered handlers. @@ -35,6 +46,13 @@ type Server struct { prompts map[string]Prompt initialized bool mu sync.Mutex + + // MaxRequestBytes caps one inbound JSON-RPC message (one newline- + // delimited line). Zero selects DefaultMaxRequestBytes; a negative value + // disables the cap (not recommended — a single unbounded line can + // exhaust memory). An oversized message is answered in-band with a + // -32600 error (id null) and the dispatch loop keeps serving. + MaxRequestBytes int64 } // NewServer creates a new MCP server with the given name and version. @@ -97,12 +115,29 @@ func (s *Server) RunWithIO(r io.Reader, w io.Writer) error { // the next read, so nothing retains a reference to it across iterations. var lineBuf []byte + s.mu.Lock() + maxReq := s.MaxRequestBytes + s.mu.Unlock() + if maxReq == 0 { + maxReq = DefaultMaxRequestBytes + } + for { - line, err := readMessage(br, lineBuf) + line, err := readMessage(br, lineBuf, maxReq) if err != nil { if err == io.EOF { return nil } + if errors.Is(err, errMessageTooLarge) { + // Bad input never kills the loop: answer in-band and keep + // serving. The oversized line is unparseable by definition, + // so the response carries a null id. + if werr := encoder.Encode(NewJSONRPCError(nil, -32600, + fmt.Sprintf("Invalid Request: message exceeds maximum size of %d bytes", maxReq))); werr != nil { + return fmt.Errorf("write error response: %w", werr) + } + continue + } return fmt.Errorf("read error: %w", err) } lineBuf = line[:0] // reclaim capacity for the next iteration @@ -167,13 +202,23 @@ func (s *Server) RunWithIO(r io.Reader, w io.Writer) error { // readMessage reads one newline-terminated message from br into buf, // returning the bytes without the trailing newline. A final message not -// terminated by EOF is still returned; a subsequent call then reports -// io.EOF. Lines longer than the reader's buffer are accumulated, so there -// is no message-size limit (matching the previous json.Decoder behavior). -func readMessage(br *bufio.Reader, buf []byte) ([]byte, error) { +// terminated by newline is still returned at EOF; a subsequent call then +// reports io.EOF. Lines longer than the reader's buffer are accumulated up +// to max bytes; beyond that the remainder of the line is discarded +// (constant memory) and errMessageTooLarge is returned, so a hostile +// client cannot exhaust memory with a single oversized line. +func readMessage(br *bufio.Reader, buf []byte, max int64) ([]byte, error) { buf = buf[:0] for { chunk, err := br.ReadSlice('\n') + if max > 0 && int64(len(buf))+int64(len(chunk)) > max { + // Discard through the end of this line so the stream stays + // framed and the next message parses normally. + for err == bufio.ErrBufferFull { + _, err = br.ReadSlice('\n') + } + return nil, errMessageTooLarge + } buf = append(buf, chunk...) switch err { case nil: diff --git a/gomcp/server_test.go b/gomcp/server_test.go index 356513d..572b617 100644 --- a/gomcp/server_test.go +++ b/gomcp/server_test.go @@ -1,6 +1,7 @@ package gomcp import ( + "bytes" "context" "encoding/json" "fmt" @@ -1063,3 +1064,90 @@ func TestPromptsListEncoderError(t *testing.T) { t.Fatal("expected encoder write error, got nil") } } + +// ----- Request size cap (2026-08 audit) ----- + +// An oversized line (no newline within MaxRequestBytes) must be rejected +// in-band with -32600 and the remainder of the line discarded, so the +// dispatch loop survives and the next well-formed message is still served. +// Before the fix, readMessage accumulated the line unboundedly — a client +// could OOM the server with a single multi-gigabyte line. +func TestRunWithIO_RequestSizeCap(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + + srv := NewServer("test-server", "1.0.0") + srv.MaxRequestBytes = 1024 + + go func() { + srv.RunWithIO(inReader, outWriter) + }() + + go func() { + // Oversized line: 8 KiB of junk followed by a valid request on the + // next line. + inWriter.Write(bytes.Repeat([]byte("a"), 8*1024)) + inWriter.Write([]byte("\n")) + inWriter.Write([]byte(`{"jsonrpc":"2.0","id":7,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{}}}` + "\n")) + inWriter.Close() + }() + + dec := json.NewDecoder(outReader) + + var errResp map[string]any + if err := dec.Decode(&errResp); err != nil { + t.Fatalf("read first response: %v", err) + } + if errObj, ok := errResp["error"].(map[string]any); !ok { + t.Fatalf("expected an error object, got: %v", errResp) + } else { + if code, _ := errObj["code"].(float64); code != -32600 { + t.Errorf("error code = %v, want -32600", errObj["code"]) + } + if msg, _ := errObj["message"].(string); !strings.Contains(msg, "maximum size") { + t.Errorf("error message = %q, want it to name the maximum size", msg) + } + } + if id, present := errResp["id"]; !present || id != nil { + t.Errorf("error response id = %v (present=%v), want null", id, present) + } + + // The loop must keep serving after the oversized line. + var initResp map[string]any + if err := dec.Decode(&initResp); err != nil { + t.Fatalf("server did not survive the oversized line: %v", err) + } + if result, ok := initResp["result"].(map[string]any); !ok || result["serverInfo"] == nil { + t.Fatalf("expected initialize result after oversized line, got: %v", initResp) + } +} + +// A line under the cap is served normally, and a zero MaxRequestBytes +// selects the default cap rather than disabling the limit. +func TestRunWithIO_UnderCapStillServed(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + + srv := NewServer("test-server", "1.0.0") + srv.MaxRequestBytes = 0 // default (10 MiB) — a 4 KiB line must pass + + go func() { + srv.RunWithIO(inReader, outWriter) + }() + + go func() { + payload := bytes.Repeat([]byte(" "), 4*1024) + inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"ping","params":{`)) + inWriter.Write(payload) + inWriter.Write([]byte("}}\n")) + inWriter.Close() + }() + + var resp map[string]any + if err := json.NewDecoder(outReader).Decode(&resp); err != nil { + t.Fatalf("read ping response: %v", err) + } + if _, ok := resp["result"]; !ok { + t.Fatalf("expected ping result, got: %v", resp) + } +}