Skip to main content

mcpmesh_codec/
lib.rs

1//! mcpmesh-codec: the family's ONE wire codec — compact JSON, UTF-8,
2//! one frame per `\n`, 16 MiB cap.
3//!
4//! Both ends of every wire share THIS implementation: the daemon side re-exports it as
5//! `mcpmesh_net::framing`, and the no-iroh client side (kb, loc, the host shell) as
6//! `mcpmesh_local_api::codec` — so the two ends can never drift apart. Deliberately tiny:
7//! serde_json + tokio io traits only, NO iroh. Session POLICY (violation strikes,
8//! synthesized error frames) stays with the session owners (`mcpmesh-net`); only the pure
9//! codec lives here.
10use serde_json::Value;
11use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
12
13/// 16 MiB — the family frame cap. App payloads are ≤1 MiB by policy.
14pub const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
15
16/// One inbound frame, or a framing violation (kept distinct so callers can log/strike).
17///
18/// Deliberately EXHAUSTIVE (no `#[non_exhaustive]`): a reader yields exactly these
19/// two outcomes, and every caller must decide both — a silently-ignorable third
20/// outcome would be a codec bug, not an extension point.
21#[derive(Debug)]
22pub enum Inbound {
23    Frame(Value),
24    Violation(Violation),
25}
26
27/// How an inbound frame violated the codec. `#[non_exhaustive]` so a future
28/// violation kind is not a breaking change — match with a wildcard arm.
29#[derive(Debug, PartialEq, Eq)]
30#[non_exhaustive]
31pub enum Violation {
32    TooLarge,
33    InvalidJson,
34}
35
36/// NDJSON frame reader over any `AsyncRead`. Reads byte-at-a-time from an INTERNAL
37/// `BufReader` (misuse-proof: callers cannot forget the buffering and pay a
38/// poll-per-byte penalty).
39///
40/// `next()` is cancellation-safe: all parse state lives in `self`; dropping the
41/// future loses no bytes.
42#[derive(Debug)]
43pub struct FrameReader<R> {
44    reader: BufReader<R>,
45    max_frame: usize,
46    buf: Vec<u8>,
47    /// true while skipping the remainder of an oversized line
48    discarding: bool,
49}
50
51impl<R: AsyncRead + Unpin> FrameReader<R> {
52    pub fn new(reader: R, max_frame: usize) -> Self {
53        debug_assert!(max_frame > 0);
54        Self {
55            reader: BufReader::new(reader),
56            max_frame,
57            buf: Vec::new(),
58            discarding: false,
59        }
60    }
61
62    /// Unwrap to the BUFFERED reader half. Returns the internal `BufReader` — NOT the
63    /// raw stream — so bytes already read ahead (e.g. a frame the peer pipelined behind
64    /// the one just parsed) travel WITH the reader instead of being silently dropped.
65    /// Call this only BETWEEN frames: any partially-accumulated frame bytes in the
66    /// codec's own line buffer are discarded (immediately after a successful `next()`
67    /// that buffer is empty, so the hand-off is lossless).
68    pub fn into_inner(self) -> BufReader<R> {
69        self.reader
70    }
71
72    /// `Ok(None)` = clean EOF. Violations do not consume the following frames.
73    ///
74    /// EOF while discarding an oversized line still reports `TooLarge` (the next
75    /// call returns `Ok(None)`). A truncated ORDINARY frame at EOF is silently
76    /// dropped (matches MCP stdio convention — deliberate).
77    pub async fn next(&mut self) -> std::io::Result<Option<Inbound>> {
78        let mut byte = [0u8; 1];
79        loop {
80            let n = self.reader.read(&mut byte).await?;
81            if n == 0 {
82                if std::mem::take(&mut self.discarding) {
83                    return Ok(Some(Inbound::Violation(Violation::TooLarge)));
84                }
85                return Ok(None);
86            }
87            if byte[0] == b'\n' {
88                if std::mem::take(&mut self.discarding) {
89                    return Ok(Some(Inbound::Violation(Violation::TooLarge)));
90                }
91                // Re-grows from zero each frame by design: no 16MiB pinned per idle
92                // session; do not "optimize" into buffer retention.
93                let line = std::mem::take(&mut self.buf);
94                // An empty line is InvalidJson (-32700) and strikes: recorded decision.
95                return Ok(Some(match serde_json::from_slice::<Value>(&line) {
96                    Ok(v) => Inbound::Frame(v),
97                    Err(_) => Inbound::Violation(Violation::InvalidJson),
98                }));
99            }
100            if self.discarding {
101                continue;
102            }
103            if self.buf.len() >= self.max_frame {
104                self.discarding = true;
105                self.buf.clear();
106                continue;
107            }
108            self.buf.push(byte[0]);
109        }
110    }
111}
112
113/// Serialize `frame` and write it terminated by `\n`, flushing (a BufWriter-safe no-op
114/// for unbuffered writers — makes BufWriter wrapping misuse-proof).
115pub async fn write_frame<W: AsyncWrite + Unpin>(w: &mut W, frame: &Value) -> std::io::Result<()> {
116    let mut line = serde_json::to_vec(frame)?;
117    line.push(b'\n');
118    w.write_all(&line).await?;
119    w.flush().await
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125    use serde_json::json;
126
127    fn reader_over(bytes: &[u8], cap: usize) -> FrameReader<&[u8]> {
128        FrameReader::new(bytes, cap)
129    }
130
131    #[tokio::test]
132    async fn valid_frames_roundtrip() {
133        let mut buf = Vec::new();
134        write_frame(&mut buf, &json!({"jsonrpc":"2.0","id":1,"method":"ping"}))
135            .await
136            .unwrap();
137        write_frame(&mut buf, &json!({"jsonrpc":"2.0","id":2,"method":"pong"}))
138            .await
139            .unwrap();
140        let mut r = reader_over(&buf, 1024);
141        assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 1));
142        assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 2));
143        assert!(r.next().await.unwrap().is_none()); // clean EOF
144    }
145
146    #[tokio::test]
147    async fn oversized_frame_is_discarded_and_reported_and_stream_continues() {
148        let big = format!("{{\"pad\":\"{}\"}}\n", "x".repeat(100));
149        let mut bytes = big.into_bytes();
150        bytes.extend_from_slice(b"{\"jsonrpc\":\"2.0\",\"id\":7,\"method\":\"ok\"}\n");
151        let mut r = reader_over(&bytes, 64);
152        assert!(matches!(
153            r.next().await.unwrap().unwrap(),
154            Inbound::Violation(Violation::TooLarge)
155        ));
156        assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 7));
157    }
158
159    #[tokio::test]
160    async fn invalid_json_is_a_violation() {
161        let mut r = reader_over(b"not json at all\n", 1024);
162        assert!(matches!(
163            r.next().await.unwrap().unwrap(),
164            Inbound::Violation(Violation::InvalidJson)
165        ));
166    }
167
168    #[tokio::test]
169    async fn oversized_frame_truncated_by_eof_is_still_reported() {
170        let bytes = format!("{{\"pad\":\"{}\"}}", "x".repeat(100)).into_bytes(); // no newline
171        let mut r = reader_over(&bytes, 64);
172        assert!(matches!(
173            r.next().await.unwrap().unwrap(),
174            Inbound::Violation(Violation::TooLarge)
175        ));
176        assert!(r.next().await.unwrap().is_none());
177    }
178
179    #[tokio::test]
180    async fn back_to_back_oversized_frames_then_valid_frame() {
181        let big = format!("{{\"pad\":\"{}\"}}\n", "x".repeat(100));
182        let mut bytes = big.clone().into_bytes();
183        bytes.extend_from_slice(big.as_bytes());
184        bytes.extend_from_slice(b"{\"jsonrpc\":\"2.0\",\"id\":9,\"method\":\"ok\"}\n");
185        let mut r = reader_over(&bytes, 64);
186        assert!(matches!(
187            r.next().await.unwrap().unwrap(),
188            Inbound::Violation(Violation::TooLarge)
189        ));
190        assert!(matches!(
191            r.next().await.unwrap().unwrap(),
192            Inbound::Violation(Violation::TooLarge)
193        ));
194        assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 9));
195    }
196
197    #[tokio::test]
198    async fn frame_of_exactly_max_frame_bytes_passes_one_more_is_too_large() {
199        // A bare JSON string is a valid frame; 62 x's + 2 quotes = exactly 64 bytes
200        // before the newline terminator (the newline is not counted against the cap).
201        let payload = "x".repeat(62);
202        let exact = format!("\"{payload}\"\n");
203        let mut r = reader_over(exact.as_bytes(), 64);
204        assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v == payload));
205
206        let over = format!("\"{}\"\n", "x".repeat(63)); // 65 bytes before the newline
207        let mut r = reader_over(over.as_bytes(), 64);
208        assert!(matches!(
209            r.next().await.unwrap().unwrap(),
210            Inbound::Violation(Violation::TooLarge)
211        ));
212    }
213
214    /// `into_inner` hands back the internal `BufReader`, so read-ahead the codec already
215    /// pulled off the stream (a pipelined second frame) survives the hand-off and is
216    /// readable through a NEW FrameReader over the returned reader. This is the
217    /// lossless-rebox contract kb's mesh opener relies on (mcpmesh-local/1 OpenSession).
218    #[tokio::test]
219    async fn into_inner_preserves_pipelined_read_ahead() {
220        let bytes = b"{\"id\":1}\n{\"id\":2}\n";
221        let mut r = reader_over(bytes, 1024);
222        // Frame 1 consumed; the internal BufReader has (with an 8 KiB buffer over a
223        // ready slice) already read frame 2 ahead.
224        assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 1));
225        // Re-box exactly like kb does: into_inner -> Box<dyn AsyncRead> -> new reader.
226        let boxed: Box<dyn AsyncRead + Unpin + Send> = Box::new(r.into_inner());
227        let mut r2 = FrameReader::new(boxed, 1024);
228        assert!(matches!(r2.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 2));
229        assert!(r2.next().await.unwrap().is_none());
230    }
231}