Skip to content

Commit 93873ad

Browse files
committed
fix: reply -32700 on stdio parse errors instead of closing
1 parent 88df9af commit 93873ad

1 file changed

Lines changed: 123 additions & 24 deletions

File tree

crates/rmcp/src/transport/async_rw.rs

Lines changed: 123 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,15 @@
11
use std::{marker::PhantomData, sync::Arc};
22

3-
// use crate::schema::*;
4-
use futures::{SinkExt, StreamExt};
3+
use futures::SinkExt;
54
use serde::{Serialize, de::DeserializeOwned};
65
use thiserror::Error;
76
use tokio::{
8-
io::{AsyncRead, AsyncWrite},
7+
io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader},
98
sync::Mutex,
109
};
1110
use tokio_util::{
1211
bytes::{Buf, BufMut, BytesMut},
13-
codec::{Decoder, Encoder, FramedRead, FramedWrite},
12+
codec::{Decoder, Encoder, FramedWrite},
1413
};
1514

1615
use super::{IntoTransport, Transport};
@@ -47,8 +46,10 @@ where
4746
pub type TransportWriter<Role, W> = FramedWrite<W, JsonRpcMessageCodec<TxJsonRpcMessage<Role>>>;
4847

4948
pub struct AsyncRwTransport<Role: ServiceRole, R: AsyncRead, W: AsyncWrite> {
50-
read: FramedRead<R, JsonRpcMessageCodec<RxJsonRpcMessage<Role>>>,
49+
read: BufReader<R>,
50+
line_buf: Vec<u8>,
5151
write: Arc<Mutex<Option<TransportWriter<Role, W>>>>,
52+
_role: PhantomData<fn() -> Role>,
5253
}
5354

5455
impl<Role: ServiceRole, R, W> AsyncRwTransport<Role, R, W>
@@ -57,15 +58,17 @@ where
5758
W: Send + AsyncWrite + Unpin + 'static,
5859
{
5960
pub fn new(read: R, write: W) -> Self {
60-
let read = FramedRead::new(
61-
read,
62-
JsonRpcMessageCodec::<RxJsonRpcMessage<Role>>::default(),
63-
);
61+
let read = BufReader::new(read);
6462
let write = Arc::new(Mutex::new(Some(FramedWrite::new(
6563
write,
6664
JsonRpcMessageCodec::<TxJsonRpcMessage<Role>>::default(),
6765
))));
68-
Self { read, write }
66+
Self {
67+
read,
68+
line_buf: Vec::new(),
69+
write,
70+
_role: PhantomData,
71+
}
6972
}
7073
}
7174

@@ -116,15 +119,42 @@ where
116119
}
117120
}
118121

119-
fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<Role>>> {
120-
let next = self.read.next();
121-
async {
122-
next.await.and_then(|e| {
123-
e.inspect_err(|e| {
122+
async fn receive(&mut self) -> Option<RxJsonRpcMessage<Role>> {
123+
loop {
124+
self.line_buf.clear();
125+
match self.read.read_until(b'\n', &mut self.line_buf).await {
126+
Ok(0) => return None,
127+
Ok(_) => {}
128+
Err(e) => {
124129
tracing::error!("Error reading from stream: {}", e);
125-
})
126-
.ok()
127-
})
130+
return None;
131+
}
132+
}
133+
let line = without_carriage_return(
134+
self.line_buf.strip_suffix(b"\n").unwrap_or(&self.line_buf),
135+
);
136+
if line.is_empty() {
137+
continue;
138+
}
139+
match try_parse_with_compatibility::<RxJsonRpcMessage<Role>>(line, "receive") {
140+
Ok(Some(msg)) => return Some(msg),
141+
Ok(None) => continue,
142+
Err(JsonRpcMessageCodecError::Serde(e)) => {
143+
tracing::debug!("Parse error on incoming message: {e}");
144+
let mut write = self.write.lock().await;
145+
let framed = write.as_mut()?;
146+
let inner = framed.get_mut();
147+
if inner.write_all(PARSE_ERROR_RESPONSE).await.is_err()
148+
|| inner.flush().await.is_err()
149+
{
150+
return None;
151+
}
152+
}
153+
Err(e) => {
154+
tracing::error!("Error reading from stream: {}", e);
155+
return None;
156+
}
157+
}
128158
}
129159
}
130160

@@ -172,13 +202,18 @@ impl<T> JsonRpcMessageCodec<T> {
172202
}
173203

174204
fn without_carriage_return(s: &[u8]) -> &[u8] {
175-
if let Some(&b'\r') = s.last() {
176-
&s[..s.len() - 1]
177-
} else {
178-
s
179-
}
205+
s.strip_suffix(b"\r").unwrap_or(s)
180206
}
181207

208+
/// UTF-8 byte order mark. RFC 8259 §8.1 allows JSON parsers to ignore a leading BOM.
209+
const UTF8_BOM: &[u8; 3] = b"\xEF\xBB\xBF";
210+
211+
// JSON-RPC 2.0 §5.1: https://www.jsonrpc.org/specification#error_object
212+
// Hardcoded bytes because `RequestId` has no `Null` variant — we can't
213+
// build an `id: null` JsonRpcError through the typed codec.
214+
const PARSE_ERROR_RESPONSE: &[u8] =
215+
b"{\"jsonrpc\":\"2.0\",\"error\":{\"code\":-32700,\"message\":\"Parse error\"},\"id\":null}\n";
216+
182217
/// Check if a method is a standard MCP method (request, response, or notification).
183218
/// This includes both requests and notifications defined in the MCP specification.
184219
///
@@ -247,6 +282,7 @@ fn try_parse_with_compatibility<T: serde::de::DeserializeOwned>(
247282
line: &[u8],
248283
context: &str,
249284
) -> Result<Option<T>, JsonRpcMessageCodecError> {
285+
let line = line.strip_prefix(UTF8_BOM.as_slice()).unwrap_or(line);
250286
if let Ok(line_str) = std::str::from_utf8(line) {
251287
match serde_json::from_slice(line) {
252288
Ok(item) => Ok(Some(item)),
@@ -406,7 +442,8 @@ impl<T: Serialize> Encoder<T> for JsonRpcMessageCodec<T> {
406442

407443
#[cfg(test)]
408444
mod test {
409-
use futures::{Sink, Stream};
445+
use futures::{Sink, Stream, StreamExt};
446+
use tokio_util::codec::FramedRead;
410447

411448
use super::*;
412449
fn from_async_read<T: DeserializeOwned, R: AsyncRead>(reader: R) -> impl Stream<Item = T> {
@@ -555,4 +592,66 @@ mod test {
555592

556593
println!("Standard notifications are preserved, non-standard are handled gracefully");
557594
}
595+
596+
#[tokio::test]
597+
async fn test_decode_strips_utf8_bom() {
598+
use futures::StreamExt;
599+
use tokio::io::BufReader;
600+
601+
// Valid JSON-RPC message preceded by a UTF-8 BOM (EF BB BF). Some Windows
602+
// tooling and editors prepend this; the codec should ignore it per RFC 8259 §8.1.
603+
let mut data = Vec::new();
604+
data.extend_from_slice(UTF8_BOM);
605+
data.extend_from_slice(br#"{"jsonrpc":"2.0","method":"ping","id":1}"#);
606+
data.push(b'\n');
607+
608+
let mut cursor = BufReader::new(&data[..]);
609+
let mut stream = from_async_read::<serde_json::Value, _>(&mut cursor);
610+
611+
let item = stream
612+
.next()
613+
.await
614+
.expect("should decode BOM-prefixed line");
615+
assert_eq!(
616+
item,
617+
serde_json::json!({"jsonrpc": "2.0", "method": "ping", "id": 1})
618+
);
619+
}
620+
621+
#[cfg(feature = "server")]
622+
#[tokio::test]
623+
async fn receive_recovers_from_parse_error() {
624+
use tokio::io::AsyncReadExt;
625+
626+
use crate::{RoleServer, transport::Transport};
627+
628+
// Two paired streams: `server_io` is wrapped by the transport; the test
629+
// drives `client_io` to act as the peer.
630+
let (server_io, client_io) = tokio::io::duplex(4096);
631+
let (server_r, server_w) = tokio::io::split(server_io);
632+
let (mut client_r, mut client_w) = tokio::io::split(client_io);
633+
634+
let mut transport = AsyncRwTransport::<RoleServer, _, _>::new(server_r, server_w);
635+
636+
client_w
637+
.write_all(
638+
b"not json\n{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}\n",
639+
)
640+
.await
641+
.unwrap();
642+
643+
let received = transport
644+
.receive()
645+
.await
646+
.expect("transport should recover and yield the next valid message");
647+
648+
let mut reply = vec![0u8; PARSE_ERROR_RESPONSE.len()];
649+
client_r.read_exact(&mut reply).await.unwrap();
650+
651+
assert_eq!(reply, PARSE_ERROR_RESPONSE);
652+
assert_eq!(
653+
serde_json::to_value(&received).unwrap()["method"],
654+
"notifications/initialized",
655+
);
656+
}
558657
}

0 commit comments

Comments
 (0)