sqry_daemon_protocol/
framing.rs

1//! 4-byte little-endian length-prefix frame codec.
2//!
3//! Every IPC message (handshake, JSON-RPC request, JSON-RPC response,
4//! batch array, shim header) travels as:
5//!
6//! ```text
7//! +--------+------------------------------------+
8//! | len:u32|                body                |
9//! +--------+------------------------------------+
10//! ```
11//!
12//! - `len` is little-endian, excludes the prefix itself, and must be
13//!   `<= MAX_FRAME_BYTES` (64 MiB). Larger frames are rejected at the
14//!   codec boundary with [`io::ErrorKind::InvalidData`].
15//! - `body` is UTF-8 JSON; one complete JSON document per frame, no
16//!   chunking, no streaming.
17//! - An EOF at a frame boundary (zero bytes read before the length
18//!   prefix) is a clean disconnect and returns `Ok(None)`.
19//! - An EOF *inside* a length prefix or body is an
20//!   [`io::ErrorKind::UnexpectedEof`] — truncated frames are never
21//!   silently swallowed (this is the Phase 8a iter-1 B3 fix).
22
23use std::io;
24
25use serde::{Serialize, de::DeserializeOwned};
26use thiserror::Error;
27use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
28
29/// Per-frame ceiling. Values above this limit are rejected at
30/// [`read_frame`] / [`write_frame`] boundaries. The bound is generous
31/// enough for Phase 8b's largest expected MCP tool responses
32/// (subgraph exports on monorepo-scale graphs) without forcing the
33/// complexity of a chunked protocol.
34pub const MAX_FRAME_BYTES: usize = 64 * 1024 * 1024;
35
36/// Errors surfaced by the JSON-flavoured codec helpers in this module.
37#[derive(Debug, Error)]
38pub enum FrameError {
39    /// Low-level transport failure (socket EOF mid-frame, permission,
40    /// interrupted syscall, etc.).
41    #[error("frame io: {0}")]
42    Io(#[from] io::Error),
43
44    /// Frame body failed `serde_json` decode. Used by
45    /// [`read_frame_json`] when the caller expects a typed shape.
46    #[error("frame json: {0}")]
47    Json(#[from] serde_json::Error),
48}
49
50/// Write a single framed blob.
51///
52/// # Errors
53///
54/// Returns [`io::ErrorKind::InvalidInput`] if `body.len() > MAX_FRAME_BYTES`.
55/// Propagates write and flush errors from the underlying async writer.
56pub async fn write_frame<W>(w: &mut W, body: &[u8]) -> io::Result<()>
57where
58    W: AsyncWrite + Unpin,
59{
60    if body.len() > MAX_FRAME_BYTES {
61        return Err(io::Error::new(
62            io::ErrorKind::InvalidInput,
63            format!(
64                "frame body length {} exceeds MAX_FRAME_BYTES ({MAX_FRAME_BYTES})",
65                body.len()
66            ),
67        ));
68    }
69    let len = u32::try_from(body.len()).map_err(|_| {
70        io::Error::new(
71            io::ErrorKind::InvalidInput,
72            "frame body length exceeds u32::MAX",
73        )
74    })?;
75    w.write_all(&len.to_le_bytes()).await?;
76    w.write_all(body).await?;
77    w.flush().await?;
78    Ok(())
79}
80
81/// Read a single framed blob.
82///
83/// Returns `Ok(None)` when the peer closes cleanly at a frame boundary
84/// (zero bytes read before any length-prefix byte). Returns
85/// [`io::ErrorKind::UnexpectedEof`] when the stream ends mid-prefix or
86/// mid-body.
87///
88/// # Errors
89///
90/// Returns [`io::ErrorKind::InvalidData`] when the frame length exceeds
91/// [`MAX_FRAME_BYTES`]. Propagates transport read errors from the underlying
92/// async reader.
93pub async fn read_frame<R>(r: &mut R) -> io::Result<Option<Vec<u8>>>
94where
95    R: AsyncRead + Unpin,
96{
97    let mut len_buf = [0u8; 4];
98    let mut filled = 0usize;
99    while filled < 4 {
100        match r.read(&mut len_buf[filled..]).await? {
101            0 if filled == 0 => return Ok(None),
102            0 => {
103                return Err(io::Error::new(
104                    io::ErrorKind::UnexpectedEof,
105                    format!("truncated frame: got {filled}/4 length bytes before EOF"),
106                ));
107            }
108            n => filled += n,
109        }
110    }
111    let len = u32::from_le_bytes(len_buf) as usize;
112    if len > MAX_FRAME_BYTES {
113        return Err(io::Error::new(
114            io::ErrorKind::InvalidData,
115            format!("frame len {len} exceeds MAX_FRAME_BYTES ({MAX_FRAME_BYTES})"),
116        ));
117    }
118    let mut body = vec![0u8; len];
119    match r.read_exact(&mut body).await {
120        Ok(_) => Ok(Some(body)),
121        Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => Err(io::Error::new(
122            io::ErrorKind::UnexpectedEof,
123            format!("truncated frame body: expected {len} bytes"),
124        )),
125        Err(e) => Err(e),
126    }
127}
128
129/// Write a typed value as a framed JSON blob.
130///
131/// # Errors
132///
133/// Returns [`FrameError::Json`] when serialization fails and [`FrameError::Io`]
134/// when writing the frame fails.
135pub async fn write_frame_json<W, T>(w: &mut W, value: &T) -> Result<(), FrameError>
136where
137    W: AsyncWrite + Unpin,
138    T: Serialize + ?Sized,
139{
140    let body = serde_json::to_vec(value)?;
141    write_frame(w, &body).await?;
142    Ok(())
143}
144
145/// Read a framed JSON blob and decode into `T`. Returns `Ok(None)` on a
146/// clean frame-boundary EOF.
147///
148/// # Errors
149///
150/// Returns [`FrameError::Io`] for frame transport failures and
151/// [`FrameError::Json`] when the frame body cannot be decoded as `T`.
152pub async fn read_frame_json<R, T>(r: &mut R) -> Result<Option<T>, FrameError>
153where
154    R: AsyncRead + Unpin,
155    T: DeserializeOwned,
156{
157    let Some(bytes) = read_frame(r).await? else {
158        return Ok(None);
159    };
160    let value = serde_json::from_slice(&bytes)?;
161    Ok(Some(value))
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use tokio::io::duplex;
168
169    #[tokio::test]
170    async fn round_trip_small_frame() {
171        let (mut a, mut b) = duplex(1024);
172        let payload = br#"{"hello":"world"}"#;
173        write_frame(&mut a, payload).await.expect("write");
174        drop(a);
175        let got = read_frame(&mut b).await.expect("read").expect("some");
176        assert_eq!(got, payload);
177    }
178
179    #[tokio::test]
180    async fn rejects_oversize_frame_on_write() {
181        let (mut a, _b) = duplex(1024);
182        let body = vec![0u8; MAX_FRAME_BYTES + 1];
183        let err = write_frame(&mut a, &body)
184            .await
185            .expect_err("oversize must fail");
186        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
187    }
188
189    #[tokio::test]
190    async fn clean_eof_at_frame_boundary_returns_ok_none() {
191        let (a, mut b) = duplex(64);
192        drop(a); // immediate EOF before any prefix byte
193        let got = read_frame(&mut b).await.expect("no error on clean EOF");
194        assert!(got.is_none());
195    }
196
197    #[tokio::test]
198    async fn truncated_prefix_is_error() {
199        let (mut a, mut b) = duplex(64);
200        a.write_all(&[0x01, 0x00]).await.unwrap(); // 2/4 length bytes
201        drop(a);
202        let err = read_frame(&mut b).await.expect_err("truncated prefix");
203        assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
204        assert!(
205            err.to_string().contains("got 2/4 length bytes"),
206            "got unexpected message: {err}"
207        );
208    }
209
210    #[tokio::test]
211    async fn truncated_body_is_error() {
212        let (mut a, mut b) = duplex(64);
213        // Claim a 16-byte body, send only 5 bytes.
214        let len = 16u32.to_le_bytes();
215        a.write_all(&len).await.unwrap();
216        a.write_all(b"short").await.unwrap();
217        drop(a);
218        let err = read_frame(&mut b).await.expect_err("truncated body");
219        assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
220        assert!(
221            err.to_string().contains("truncated frame body"),
222            "got unexpected message: {err}"
223        );
224    }
225
226    #[tokio::test]
227    async fn oversize_read_is_rejected() {
228        let (mut a, mut b) = duplex(64);
229        let bad_len = (MAX_FRAME_BYTES as u32 + 1).to_le_bytes();
230        a.write_all(&bad_len).await.unwrap();
231        let err = read_frame(&mut b).await.expect_err("oversize claim");
232        assert_eq!(err.kind(), io::ErrorKind::InvalidData);
233    }
234}
sqry_daemon_protocol/framing.rs

sqry_daemon_protocol/
framing.rs
