rskit_codec/framing/
frame.rs1use std::io::{Read, Write};
4
5use rskit_errors::{AppError, AppResult, ErrorCode};
6
7pub const DEFAULT_MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
13
14const LEN_PREFIX_BYTES: usize = 4;
16
17pub 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
47pub 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
77enum ReadEnd {
79 Filled,
81 Eof,
83}
84
85fn 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
107fn 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}