From bb9087d2526b216b902095abbf64a5ab728ef8a4 Mon Sep 17 00:00:00 2001
From: Rishabh Nrupnarayan <aryannr97@gmail.com>
Date: Wed, 6 Aug 2025 21:10:39 +0530
Subject: [PATCH] mcp: handle bad trailing stdio input graciously (#179) (#192)

You can make MCP server log error, send response and then abruptly exit, with a json input malformed at the end.

While using json.RawMessage, once a valid first json is decoded, trailing bad input is silently ignored.
Without graciously handling this input, mcp server is currently sending response as well as encountering error.
It should just report error without further processing of request.

Fixes #179.
---
 mcp/transport.go      | 13 +++++++++++++
 mcp/transport_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/mcp/transport.go b/mcp/transport.go
index dccc920..f2d5c72 100644
--- a/mcp/transport.go
+++ b/mcp/transport.go
@@ -308,6 +308,19 @@ func newIOConn(rwc io.ReadWriteCloser) *ioConn {
 		for {
 			var raw json.RawMessage
 			err := dec.Decode(&raw)
+			// If decoding was successful, check for trailing data at the end of the stream.
+			if err == nil {
+				// Read the next byte to check if there is trailing data.
+				var tr [1]byte
+				if n, readErr := dec.Buffered().Read(tr[:]); n > 0 {
+					// If read byte is not a newline, it is an error.
+					if tr[0] != '\n' {
+						err = fmt.Errorf("invalid trailing data at the end of stream")
+					}
+				} else if readErr != nil && readErr != io.EOF {
+					err = readErr
+				}
+			}
 			select {
 			case incoming <- msgOrErr{msg: raw, err: err}:
 			case <-closed:
diff --git a/mcp/transport_test.go b/mcp/transport_test.go
index c63b84e..18a326e 100644
--- a/mcp/transport_test.go
+++ b/mcp/transport_test.go
@@ -7,6 +7,7 @@ package mcp
 import (
 	"context"
 	"io"
+	"strings"
 	"testing"
 
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
@@ -51,3 +52,41 @@ func TestBatchFraming(t *testing.T) {
 		}
 	}
 }
+
+func TestIOConnRead(t *testing.T) {
+	tests := []struct {
+		name  string
+		input string
+		want  string
+	}{
+
+		{
+			name:  "valid json input",
+			input: `{"jsonrpc":"2.0","id":1,"method":"test","params":{}}`,
+			want:  "",
+		},
+
+		{
+			name: "newline at the end of first valid json input",
+			input: `{"jsonrpc":"2.0","id":1,"method":"test","params":{}}
+			`,
+			want: "",
+		},
+		{
+			name:  "bad data at the end of first valid json input",
+			input: `{"jsonrpc":"2.0","id":1,"method":"test","params":{}},`,
+			want:  "invalid trailing data at the end of stream",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tr := newIOConn(rwc{
+				rc: io.NopCloser(strings.NewReader(tt.input)),
+			})
+			_, err := tr.Read(context.Background())
+			if err != nil && err.Error() != tt.want {
+				t.Errorf("ioConn.Read() = %v, want %v", err.Error(), tt.want)
+			}
+		})
+	}
+}