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>>> {
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
76enum ReadEnd {
78 Filled,
80 Eof,
82}
83
84fn 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
106fn 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}