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
10/// so a corrupt length prefix cannot trigger an unbounded allocation.
11/// Callers may pass a 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`
22/// or the 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 byte (the peer closed the connection between frames).
50/// A partial prefix or payload is a hard transport error.
51///
52/// # Errors
53///
54/// Returns a typed [`AppError`] on a truncated frame, a length above `max_bytes`,
55/// or any underlying read failure (cause preserved).
56pub fn read_frame<R: Read>(reader: &mut R, max_bytes: usize) -> AppResult<Option<Vec<u8>>> {
57    let mut prefix = [0u8; LEN_PREFIX_BYTES];
58    match read_exact_or_eof(reader, &mut prefix)? {
59        ReadEnd::Eof => return Ok(None),
60        ReadEnd::Filled => {}
61    }
62    let len = u32::from_be_bytes(prefix) as usize;
63    if len > max_bytes {
64        return Err(AppError::invalid_input(
65            "frame",
66            format!("incoming frame length {len} exceeds the {max_bytes}-byte limit"),
67        ));
68    }
69    let mut payload = vec![0u8; len];
70    reader
71        .read_exact(&mut payload)
72        .map_err(|error| transport_error("read frame payload", &error))?;
73    Ok(Some(payload))
74}
75
76/// Whether a fixed-size read filled the buffer or hit a clean EOF first.
77enum ReadEnd {
78    /// The buffer was completely filled.
79    Filled,
80    /// End-of-stream was reached before any byte was read.
81    Eof,
82}
83
84/// Fill `buf` exactly, distinguishing a clean leading EOF from a truncated read.
85fn read_exact_or_eof<R: Read>(reader: &mut R, buf: &mut [u8]) -> AppResult<ReadEnd> {
86    let mut read = 0;
87    while read < buf.len() {
88        match reader.read(&mut buf[read..]) {
89            Ok(0) => {
90                if read == 0 {
91                    return Ok(ReadEnd::Eof);
92                }
93                return Err(AppError::new(
94                    ErrorCode::ServiceUnavailable,
95                    "framed transport: stream ended mid-frame (truncated length prefix)",
96                ));
97            }
98            Ok(count) => read += count,
99            Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {}
100            Err(error) => return Err(transport_error("read frame length", &error)),
101        }
102    }
103    Ok(ReadEnd::Filled)
104}
105
106/// Build a typed transport error preserving the underlying I/O cause.
107fn transport_error(context: &str, error: &std::io::Error) -> AppError {
108    AppError::new(
109        ErrorCode::ServiceUnavailable,
110        format!("framed transport: {context}"),
111    )
112    .with_cause(std::io::Error::new(error.kind(), error.to_string()))
113}
114
115#[cfg(test)]
116mod tests {
117    use std::io::{Read, Write};
118
119    use super::{DEFAULT_MAX_FRAME_BYTES, read_frame, write_frame};
120
121    #[test]
122    fn round_trips_a_frame() {
123        let mut buffer = Vec::new();
124        write_frame(&mut buffer, b"hello", DEFAULT_MAX_FRAME_BYTES).expect("write");
125        let mut cursor = std::io::Cursor::new(buffer);
126        let frame = read_frame(&mut cursor, DEFAULT_MAX_FRAME_BYTES)
127            .expect("read")
128            .expect("frame present");
129        assert_eq!(frame, b"hello");
130    }
131
132    #[test]
133    fn clean_eof_between_frames_is_none() {
134        let mut cursor = std::io::Cursor::new(Vec::new());
135        assert!(
136            read_frame(&mut cursor, DEFAULT_MAX_FRAME_BYTES)
137                .expect("read")
138                .is_none()
139        );
140    }
141
142    #[test]
143    fn truncated_prefix_is_a_transport_error() {
144        let mut cursor = std::io::Cursor::new(vec![0u8, 0u8]);
145        let error = read_frame(&mut cursor, DEFAULT_MAX_FRAME_BYTES).expect_err("truncated errors");
146        assert_eq!(error.code(), rskit_errors::ErrorCode::ServiceUnavailable);
147    }
148
149    #[test]
150    fn oversized_frame_is_rejected() {
151        let mut buffer = Vec::new();
152        write_frame(&mut buffer, &[0u8; 8], DEFAULT_MAX_FRAME_BYTES).expect("write");
153        assert!(read_frame(&mut std::io::Cursor::new(buffer), 4).is_err());
154    }
155
156    #[test]
157    fn write_frame_rejects_oversized_payload_and_preserves_io_errors() {
158        struct FailingWriter;
159
160        impl Write for FailingWriter {
161            fn write(&mut self, _buf: &[u8]) -> std::io::Result<usize> {
162                Err(std::io::Error::other("boom"))
163            }
164
165            fn flush(&mut self) -> std::io::Result<()> {
166                Ok(())
167            }
168        }
169
170        let mut buffer = Vec::new();
171        assert!(write_frame(&mut buffer, b"too large", 1).is_err());
172        assert!(write_frame(&mut FailingWriter, b"ok", DEFAULT_MAX_FRAME_BYTES).is_err());
173    }
174
175    #[test]
176    fn read_frame_retries_interrupted_prefix_reads_and_reports_io_errors() {
177        struct InterruptedThenData {
178            data: std::io::Cursor<Vec<u8>>,
179            interrupted: bool,
180        }
181
182        impl Read for InterruptedThenData {
183            fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
184                if !self.interrupted {
185                    self.interrupted = true;
186                    return Err(std::io::Error::from(std::io::ErrorKind::Interrupted));
187                }
188                self.data.read(buf)
189            }
190        }
191
192        struct FailingReader;
193        impl Read for FailingReader {
194            fn read(&mut self, _buf: &mut [u8]) -> std::io::Result<usize> {
195                Err(std::io::Error::other("boom"))
196            }
197        }
198
199        let mut framed = Vec::new();
200        write_frame(&mut framed, b"x", DEFAULT_MAX_FRAME_BYTES).unwrap();
201        let mut reader = InterruptedThenData {
202            data: std::io::Cursor::new(framed),
203            interrupted: false,
204        };
205        assert_eq!(
206            read_frame(&mut reader, DEFAULT_MAX_FRAME_BYTES)
207                .unwrap()
208                .unwrap(),
209            b"x"
210        );
211
212        assert!(read_frame(&mut FailingReader, DEFAULT_MAX_FRAME_BYTES).is_err());
213    }
214}