melin_wire_protocol/
blocking.rs1use std::io::{self, BufReader, BufWriter, Read, Write};
8
9use crate::framing::FrameTooLarge;
10
11pub const MAX_FRAME_SIZE: usize = 1024;
17
18pub struct BlockingFrameReader<R> {
29 reader: BufReader<R>,
30 buf: [u8; MAX_FRAME_SIZE],
33 frame_len: usize,
35}
36
37impl<R: Read> BlockingFrameReader<R> {
38 pub fn new(reader: R) -> Self {
39 Self {
40 reader: BufReader::new(reader),
41 buf: [0u8; MAX_FRAME_SIZE],
42 frame_len: 0,
43 }
44 }
45
46 pub fn read_frame(&mut self) -> io::Result<Option<&[u8]>> {
55 let mut len_buf = [0u8; 4];
57 match self.reader.read_exact(&mut len_buf) {
58 Ok(()) => {}
59 Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
60 Err(e) => return Err(e),
61 }
62
63 let declared = u32::from_le_bytes(len_buf);
64 let len = declared as usize;
65 if len > MAX_FRAME_SIZE {
66 return Err(FrameTooLarge {
69 declared,
70 max: MAX_FRAME_SIZE,
71 }
72 .into());
73 }
74
75 self.reader.read_exact(&mut self.buf[..len])?;
76 self.frame_len = len;
77
78 Ok(Some(&self.buf[..len]))
79 }
80
81 pub fn frame(&self) -> &[u8] {
86 &self.buf[..self.frame_len]
87 }
88
89 pub fn get_ref(&self) -> &R {
94 self.reader.get_ref()
95 }
96}
97
98pub struct BlockingFrameWriter<W: Write> {
103 writer: BufWriter<W>,
104}
105
106impl<W: Write> BlockingFrameWriter<W> {
107 pub fn new(writer: W) -> Self {
108 Self {
109 writer: BufWriter::new(writer),
110 }
111 }
112
113 pub fn write_frame(&mut self, data: &[u8]) -> io::Result<()> {
115 self.write_frame_parts(&[data])
116 }
117
118 pub fn write_frame_parts(&mut self, parts: &[&[u8]]) -> io::Result<()> {
132 let len: usize = parts.iter().map(|part| part.len()).sum();
133 if len > MAX_FRAME_SIZE {
134 return Err(io::Error::new(
135 io::ErrorKind::InvalidInput,
136 format!("frame too large: {len} bytes (max {MAX_FRAME_SIZE})"),
137 ));
138 }
139 self.writer.write_all(&(len as u32).to_le_bytes())?;
140 for part in parts {
141 self.writer.write_all(part)?;
142 }
143 Ok(())
144 }
145
146 pub fn flush(&mut self) -> io::Result<()> {
148 self.writer.flush()
149 }
150}
151
152#[cfg(test)]
153mod tests {
154 use super::*;
155 use std::net::{TcpListener, TcpStream};
156
157 #[test]
158 fn frame_round_trip_blocking() {
159 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
160 let addr = listener.local_addr().unwrap();
161
162 let client = TcpStream::connect(addr).unwrap();
163 let (server, _) = listener.accept().unwrap();
164
165 let mut writer = BlockingFrameWriter::new(client);
166 let mut reader = BlockingFrameReader::new(server);
167
168 let data = b"hello melin";
169 writer.write_frame(data).unwrap();
170 writer.flush().unwrap();
171
172 let received = reader.read_frame().unwrap().unwrap();
173 assert_eq!(received, data);
174 }
175
176 #[test]
177 fn frame_parts_arrive_as_one_frame() {
178 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
179 let addr = listener.local_addr().unwrap();
180
181 let client = TcpStream::connect(addr).unwrap();
182 let (server, _) = listener.accept().unwrap();
183
184 let mut writer = BlockingFrameWriter::new(client);
185 let mut reader = BlockingFrameReader::new(server);
186
187 writer
188 .write_frame_parts(&[&7u64.to_le_bytes(), &[0x10], b"body"])
189 .unwrap();
190 writer.write_frame_parts(&[]).unwrap();
191 writer.flush().unwrap();
192
193 let expected = [&7u64.to_le_bytes()[..], &[0x10], b"body"].concat();
194 assert_eq!(reader.read_frame().unwrap().unwrap(), expected);
195 assert_eq!(reader.read_frame().unwrap().unwrap(), b"");
196 }
197
198 #[test]
199 fn an_oversized_frame_is_refused_before_anything_is_written() {
200 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
201 let addr = listener.local_addr().unwrap();
202
203 let client = TcpStream::connect(addr).unwrap();
204 let (server, _) = listener.accept().unwrap();
205
206 let mut writer = BlockingFrameWriter::new(client);
207 let mut reader = BlockingFrameReader::new(server);
208
209 let body = vec![0xAB; MAX_FRAME_SIZE];
212 let err = writer.write_frame_parts(&[&[0x01], &body]).unwrap_err();
213 assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
214
215 writer.write_frame(&body).unwrap();
218 writer.flush().unwrap();
219 assert_eq!(reader.read_frame().unwrap().unwrap(), &body[..]);
220 }
221
222 #[test]
223 fn clean_disconnect_returns_none() {
224 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
225 let addr = listener.local_addr().unwrap();
226
227 let client = TcpStream::connect(addr).unwrap();
228 let (server, _) = listener.accept().unwrap();
229
230 drop(client);
231
232 let mut reader = BlockingFrameReader::new(server);
233 let result = reader.read_frame().unwrap();
234 assert!(result.is_none());
235 }
236
237 #[test]
238 fn oversized_frame_rejected() {
239 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
240 let addr = listener.local_addr().unwrap();
241
242 let mut client = TcpStream::connect(addr).unwrap();
243 let (server, _) = listener.accept().unwrap();
244
245 let fake_len = 2_000_000u32;
247 client.write_all(&fake_len.to_le_bytes()).unwrap();
248
249 let mut reader = BlockingFrameReader::new(server);
250 let err = reader.read_frame().unwrap_err();
251 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
252 assert_eq!(
253 err.to_string(),
254 format!("frame too large: {fake_len} bytes (max {MAX_FRAME_SIZE})")
255 );
256 assert_eq!(
258 err.get_ref()
259 .and_then(|inner| inner.downcast_ref::<FrameTooLarge>()),
260 Some(&FrameTooLarge {
261 declared: fake_len,
262 max: MAX_FRAME_SIZE
263 })
264 );
265 }
266
267 #[test]
268 fn multiple_frames_in_sequence() {
269 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
270 let addr = listener.local_addr().unwrap();
271
272 let client = TcpStream::connect(addr).unwrap();
273 let (server, _) = listener.accept().unwrap();
274
275 let mut writer = BlockingFrameWriter::new(client);
276 let mut reader = BlockingFrameReader::new(server);
277
278 for i in 0u32..100 {
279 writer.write_frame(&i.to_le_bytes()).unwrap();
280 }
281 writer.flush().unwrap();
282
283 for i in 0u32..100 {
284 let frame = reader.read_frame().unwrap().unwrap();
285 assert_eq!(frame, i.to_le_bytes());
286 }
287 }
288}