Skip to main content

rag_rat_sync/
codec.rs

1//! Length-prefixed framing over an async byte stream (phase D, #406).
2//!
3//! Each frame is a 4-byte big-endian length followed by that many CBOR bytes. The length is capped
4//! so a peer cannot announce a multi-gigabyte frame and force an unbounded read before the frame is
5//! even parsed — the transport-level half of "bounded frames, no amplification". This layer is
6//! transport-agnostic: it runs over an iroh bi-stream in production and over an in-memory duplex in
7//! tests, so the session logic is exercised without the network.
8
9use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
10
11use crate::wire::{Frame, WireError};
12
13/// The largest frame this codec will read. Chosen well above one full [`Frame::Entries`] page of
14/// account entries (each entry is at most the §18a envelope, 64 KiB) plus overhead. A larger
15/// declared length is refused before any allocation.
16pub const MAX_FRAME_BYTES: u32 = 24 * 1024 * 1024;
17
18/// What can go wrong moving a frame over the wire.
19#[derive(Debug)]
20pub enum CodecError {
21    /// The underlying stream failed or closed mid-frame.
22    Io(std::io::Error),
23    /// A frame's declared length exceeded [`MAX_FRAME_BYTES`].
24    FrameTooLarge(u32),
25    /// The frame bytes were not a valid protocol frame.
26    Wire(WireError),
27    /// The stream ended cleanly at a frame boundary — not an error, but distinguished so the
28    /// session can tell "peer hung up" from "peer sent garbage".
29    Eof,
30}
31
32impl std::fmt::Display for CodecError {
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        match self {
35            CodecError::Io(e) => write!(f, "sync stream io: {e}"),
36            CodecError::FrameTooLarge(n) => {
37                write!(f, "sync frame declared {n} bytes, over {MAX_FRAME_BYTES}")
38            },
39            CodecError::Wire(e) => write!(f, "{e}"),
40            CodecError::Eof => write!(f, "sync stream closed at a frame boundary"),
41        }
42    }
43}
44
45impl std::error::Error for CodecError {}
46
47impl From<WireError> for CodecError {
48    fn from(e: WireError) -> Self {
49        CodecError::Wire(e)
50    }
51}
52
53/// Write one frame: a 4-byte big-endian length prefix, then the CBOR body.
54pub async fn write_frame<W: AsyncWrite + Unpin>(
55    w: &mut W,
56    frame: &Frame,
57) -> Result<(), CodecError> {
58    let body = frame.encode();
59    // Local frames are built within the caps, so this only fires on a programmer bug, not a wire
60    // condition — but guard rather than truncate a silently-huge length prefix.
61    let len = u32::try_from(body.len()).map_err(|_| CodecError::FrameTooLarge(u32::MAX))?;
62    w.write_all(&len.to_be_bytes()).await.map_err(CodecError::Io)?;
63    w.write_all(&body).await.map_err(CodecError::Io)?;
64    Ok(())
65}
66
67/// Read one frame, or [`CodecError::Eof`] if the stream ends cleanly before the next length prefix.
68pub async fn read_frame<R: AsyncRead + Unpin>(r: &mut R) -> Result<Frame, CodecError> {
69    read_frame_within(r, MAX_FRAME_BYTES).await
70}
71
72/// [`read_frame`] with a caller-supplied maximum, checked against the length prefix BEFORE any
73/// allocation. The auth phase passes a tight bound so an unauthenticated peer cannot force the full
74/// [`MAX_FRAME_BYTES`] allocation with its first frame (#881) — the frame-level cap is what
75/// actually bounds the pre-auth allocation, since the body is sized from the length prefix.
76pub async fn read_frame_within<R: AsyncRead + Unpin>(
77    r: &mut R,
78    max_bytes: u32,
79) -> Result<Frame, CodecError> {
80    let mut len_buf = [0u8; 4];
81    match r.read_exact(&mut len_buf).await {
82        Ok(_) => {},
83        Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Err(CodecError::Eof),
84        Err(e) => return Err(CodecError::Io(e)),
85    }
86    let len = u32::from_be_bytes(len_buf);
87    if len > max_bytes {
88        return Err(CodecError::FrameTooLarge(len));
89    }
90    let mut body = vec![0u8; len as usize];
91    r.read_exact(&mut body).await.map_err(CodecError::Io)?;
92    Ok(Frame::decode(&body)?)
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[tokio::test]
100    async fn frames_roundtrip_over_a_duplex() {
101        let (mut a, mut b) = tokio::io::duplex(64 * 1024);
102        let sent = vec![
103            Frame::Hello { account_id: [3; 32], have: vec![[1; 32]] },
104            Frame::Entries { entries: vec![vec![9, 9, 9]], more: false },
105            Frame::Done,
106        ];
107        let to_send = sent.clone();
108        let writer = tokio::spawn(async move {
109            for f in &to_send {
110                write_frame(&mut a, f).await.unwrap();
111            }
112            // drop `a` → clean EOF on `b`
113        });
114        let mut got = Vec::new();
115        loop {
116            match read_frame(&mut b).await {
117                Ok(f) => got.push(f),
118                Err(CodecError::Eof) => break,
119                Err(e) => panic!("unexpected {e}"),
120            }
121        }
122        writer.await.unwrap();
123        assert_eq!(got, sent);
124    }
125
126    #[tokio::test]
127    async fn an_oversized_length_prefix_is_refused_before_allocating() {
128        let (mut a, mut b) = tokio::io::duplex(64);
129        a.write_all(&(MAX_FRAME_BYTES + 1).to_be_bytes()).await.unwrap();
130        drop(a);
131        assert!(matches!(read_frame(&mut b).await, Err(CodecError::FrameTooLarge(_))));
132    }
133
134    #[tokio::test]
135    async fn a_capped_read_refuses_a_prefix_over_its_smaller_limit() {
136        // The auth phase reads with a tight cap so an unauthenticated peer cannot force the full
137        // MAX_FRAME_BYTES allocation. A 2 KiB prefix under a 1 KiB cap is refused from the prefix
138        // alone — the 2 KiB body is never allocated.
139        let (mut a, mut b) = tokio::io::duplex(64);
140        a.write_all(&2048u32.to_be_bytes()).await.unwrap();
141        drop(a);
142        assert!(matches!(
143            read_frame_within(&mut b, 1024).await,
144            Err(CodecError::FrameTooLarge(2048)),
145        ));
146    }
147}