Skip to main content

melin_wire_protocol/
blocking.rs

1//! Blocking (synchronous) frame reader/writer for dedicated I/O threads.
2//!
3//! Same length-prefixed framing as the async TCP/UDS transports, but uses
4//! `std::io::Read`/`Write` directly. Used by the server's reader and
5//! response threads to avoid tokio task scheduling overhead on the hot path.
6
7use std::io::{self, BufReader, BufWriter, Read, Write};
8
9use crate::framing::FrameTooLarge;
10
11/// Bound on one frame's payload, after the 4-byte length prefix: the
12/// same limit on every transport. A reader refuses a frame declaring
13/// more, and [`BlockingFrameWriter`] refuses to write one, so an
14/// oversized frame is an error at its source rather than a dropped
15/// connection at the far end.
16pub const MAX_FRAME_SIZE: usize = 1024;
17
18/// Blocking frame reader. Reads length-prefixed frames from any `Read` source.
19///
20/// Uses `BufReader` to amortize read syscalls — a single recv fills the
21/// buffer with many frames, so subsequent `read_exact` calls hit the
22/// buffer instead of making kernel transitions. This is critical for
23/// round-trip latency: without buffering, each frame requires 2 read
24/// syscalls (4-byte length prefix + payload).
25///
26/// Generic over the reader type so it works with both `std::net::TcpStream`
27/// and `std::os::unix::net::UnixStream`.
28pub struct BlockingFrameReader<R> {
29    reader: BufReader<R>,
30    /// Reusable frame buffer — avoids a heap allocation per frame.
31    /// Fixed at MAX_FRAME_SIZE (1 KiB); the valid slice is `&buf[..len]`.
32    buf: [u8; MAX_FRAME_SIZE],
33    /// Length of the last successfully read frame (valid bytes in `buf`).
34    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    /// Read the next complete frame into the internal buffer.
47    /// Returns a borrowed slice of the frame payload, or `None` on clean
48    /// disconnect. The slice is valid until the next `read_frame()` call.
49    ///
50    /// A prefix declaring more than [`MAX_FRAME_SIZE`] is
51    /// [`io::ErrorKind::InvalidData`] wrapping a [`FrameTooLarge`]
52    /// (reachable through [`io::Error::get_ref`]); the stream has no frame
53    /// boundary left to resume from.
54    pub fn read_frame(&mut self) -> io::Result<Option<&[u8]>> {
55        // Read the 4-byte length prefix.
56        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            // Carries the typed error, so a caller can tell an oversized
67            // frame from a failing socket by downcasting.
68            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    /// The frame the last successful [`read_frame`](Self::read_frame)
82    /// returned, for callers that decided what it was and now want it
83    /// back without holding the borrow across that decision. Empty
84    /// before the first read.
85    pub fn frame(&self) -> &[u8] {
86        &self.buf[..self.frame_len]
87    }
88
89    /// Borrow the underlying reader. Mirrors `BufReader::get_ref` — used
90    /// by callers that need to reach the raw stream for socket-level
91    /// configuration (`set_read_timeout`, `set_nodelay`, …) without
92    /// going through the framed layer.
93    pub fn get_ref(&self) -> &R {
94        self.reader.get_ref()
95    }
96}
97
98/// Blocking frame writer. Writes length-prefixed frames to any `Write` sink.
99///
100/// Uses `BufWriter` to batch small writes (length prefix + payload) into
101/// fewer syscalls. Flushed explicitly after each batch.
102pub 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    /// Write a complete frame (prepends the 4-byte LE length prefix).
114    pub fn write_frame(&mut self, data: &[u8]) -> io::Result<()> {
115        self.write_frame_parts(&[data])
116    }
117
118    /// Write one frame whose payload is `parts` laid end to end, with
119    /// the length prefix of their total. Spares a caller that holds a
120    /// frame in pieces — a header it computes and a body it was given —
121    /// from assembling them in a buffer of its own: each part goes
122    /// straight into the write buffer, and the frame reaches the socket
123    /// whole at the next flush. A part that arrives while the buffer is
124    /// too full to take it still lands on the socket in order, only in a
125    /// separate write.
126    ///
127    /// A frame over [`MAX_FRAME_SIZE`] is refused with
128    /// [`io::ErrorKind::InvalidInput`] before any of it is written: the
129    /// far end would drop the connection on it, and this is the error
130    /// that says why.
131    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    /// Flush buffered data to the underlying writer.
147    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        // One byte over the cap, split across parts: the total is what
210        // counts.
211        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        // Nothing of it went out: the frame at the cap that follows is the
216        // first thing the reader sees, and it is whole.
217        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        // Send a length prefix claiming 2 MiB.
246        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        // The typed error rides inside, with the declared length.
257        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}