Skip to main content

rskit_codec/framing/
frame.rs

1//! Raw byte-frame transport: a 4-byte big-endian length prefix plus payload.
2
3use std::io::{Read, Write};
4
5use rskit_errors::{AppError, AppResult, ErrorCode};
6
7/// Default maximum accepted payload size for a single frame (16 MiB).
8///
9/// Generous enough for large structured payloads yet bounded so a corrupt
10/// length prefix cannot trigger an unbounded allocation. Callers may pass a
11/// tighter cap to [`read_frame`].
12pub const DEFAULT_MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
13
14/// Width of the big-endian length prefix that precedes every payload.
15const LEN_PREFIX_BYTES: usize = 4;
16
17/// Write one length-delimited frame carrying `payload`, flushing on completion.
18///
19/// # Errors
20///
21/// Returns a typed [`AppError`] if `payload` exceeds `max_bytes` or the
22/// underlying writer fails (cause preserved).
23pub fn write_frame<W: Write>(writer: &mut W, payload: &[u8], max_bytes: usize) -> AppResult<()> {
24    if payload.len() > max_bytes {
25        return Err(AppError::invalid_input(
26            "frame",
27            format!(
28                "payload of {} bytes exceeds the {max_bytes}-byte frame limit",
29                payload.len()
30            ),
31        ));
32    }
33    let len = u32::try_from(payload.len())
34        .map_err(|_| AppError::invalid_input("frame", "payload length exceeds u32 range"))?;
35    writer
36        .write_all(&len.to_be_bytes())
37        .map_err(|error| transport_error("write frame length", &error))?;
38    writer
39        .write_all(payload)
40        .map_err(|error| transport_error("write frame payload", &error))?;
41    writer
42        .flush()
43        .map_err(|error| transport_error("flush frame", &error))?;
44    Ok(())
45}
46
47/// Read one length-delimited frame, bounded by `max_bytes`.
48///
49/// Returns `Ok(None)` on a clean end-of-stream observed *before* any length
50/// byte (the peer closed the connection between frames). A partial prefix or
51/// payload is a hard transport error.
52///
53/// # Errors
54///
55/// Returns a typed [`AppError`] on a truncated frame, a length above
56/// `max_bytes`, or any underlying read failure (cause preserved).
57pub fn read_frame<R: Read>(reader: &mut R, max_bytes: usize) -> AppResult<Option<Vec<u8>>> {
58    let mut prefix = [0u8; LEN_PREFIX_BYTES];
59    match read_exact_or_eof(reader, &mut prefix)? {
60        ReadEnd::Eof => return Ok(None),
61        ReadEnd::Filled => {}
62    }
63    let len = u32::from_be_bytes(prefix) as usize;
64    if len > max_bytes {
65        return Err(AppError::invalid_input(
66            "frame",
67            format!("incoming frame length {len} exceeds the {max_bytes}-byte limit"),
68        ));
69    }
70    let mut payload = vec![0u8; len];
71    reader
72        .read_exact(&mut payload)
73        .map_err(|error| transport_error("read frame payload", &error))?;
74    Ok(Some(payload))
75}
76
77/// Whether a fixed-size read filled the buffer or hit a clean EOF first.
78enum ReadEnd {
79    /// The buffer was completely filled.
80    Filled,
81    /// End-of-stream was reached before any byte was read.
82    Eof,
83}
84
85/// Fill `buf` exactly, distinguishing a clean leading EOF from a truncated read.
86fn read_exact_or_eof<R: Read>(reader: &mut R, buf: &mut [u8]) -> AppResult<ReadEnd> {
87    let mut read = 0;
88    while read < buf.len() {
89        match reader.read(&mut buf[read..]) {
90            Ok(0) => {
91                if read == 0 {
92                    return Ok(ReadEnd::Eof);
93                }
94                return Err(AppError::new(
95                    ErrorCode::ServiceUnavailable,
96                    "framed transport: stream ended mid-frame (truncated length prefix)",
97                ));
98            }
99            Ok(count) => read += count,
100            Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {}
101            Err(error) => return Err(transport_error("read frame length", &error)),
102        }
103    }
104    Ok(ReadEnd::Filled)
105}
106
107/// Build a typed transport error preserving the underlying I/O cause.
108fn transport_error(context: &str, error: &std::io::Error) -> AppError {
109    AppError::new(
110        ErrorCode::ServiceUnavailable,
111        format!("framed transport: {context}"),
112    )
113    .with_cause(std::io::Error::new(error.kind(), error.to_string()))
114}
115
116#[cfg(test)]
117mod tests {
118    use super::{DEFAULT_MAX_FRAME_BYTES, read_frame, write_frame};
119
120    #[test]
121    fn round_trips_a_frame() {
122        let mut buffer = Vec::new();
123        write_frame(&mut buffer, b"hello", DEFAULT_MAX_FRAME_BYTES).expect("write");
124        let mut cursor = std::io::Cursor::new(buffer);
125        let frame = read_frame(&mut cursor, DEFAULT_MAX_FRAME_BYTES)
126            .expect("read")
127            .expect("frame present");
128        assert_eq!(frame, b"hello");
129    }
130
131    #[test]
132    fn clean_eof_between_frames_is_none() {
133        let mut cursor = std::io::Cursor::new(Vec::new());
134        assert!(
135            read_frame(&mut cursor, DEFAULT_MAX_FRAME_BYTES)
136                .expect("read")
137                .is_none()
138        );
139    }
140
141    #[test]
142    fn truncated_prefix_is_a_transport_error() {
143        let mut cursor = std::io::Cursor::new(vec![0u8, 0u8]);
144        let error = read_frame(&mut cursor, DEFAULT_MAX_FRAME_BYTES).expect_err("truncated errors");
145        assert_eq!(error.code(), rskit_errors::ErrorCode::ServiceUnavailable);
146    }
147
148    #[test]
149    fn oversized_frame_is_rejected() {
150        let mut buffer = Vec::new();
151        write_frame(&mut buffer, &[0u8; 8], DEFAULT_MAX_FRAME_BYTES).expect("write");
152        assert!(read_frame(&mut std::io::Cursor::new(buffer), 4).is_err());
153    }
154}