melin_wire_protocol/

blocking.rs

//! Blocking (synchronous) frame reader/writer for dedicated I/O threads.
//!
//! Same length-prefixed framing as the async TCP/UDS transports, but uses
//! `std::io::Read`/`Write` directly. Used by the server's reader and
//! response threads to avoid tokio task scheduling overhead on the hot path.

use std::io::{self, BufReader, BufWriter, Read, Write};

/// Bound on one frame's payload, after the 4-byte length prefix: the
/// same limit on every transport. A reader refuses a frame declaring
/// more, and [`BlockingFrameWriter`] refuses to write one, so an
/// oversized frame is an error at its source rather than a dropped
/// connection at the far end.
pub const MAX_FRAME_SIZE: usize = 1024;

/// Blocking frame reader. Reads length-prefixed frames from any `Read` source.
///
/// Uses `BufReader` to amortize read syscalls — a single recv fills the
/// buffer with many frames, so subsequent `read_exact` calls hit the
/// buffer instead of making kernel transitions. This is critical for
/// round-trip latency: without buffering, each frame requires 2 read
/// syscalls (4-byte length prefix + payload).
///
/// Generic over the reader type so it works with both `std::net::TcpStream`
/// and `std::os::unix::net::UnixStream`.
pub struct BlockingFrameReader<R> {
    reader: BufReader<R>,
    /// Reusable frame buffer — avoids a heap allocation per frame.
    /// Fixed at MAX_FRAME_SIZE (1 KiB); the valid slice is `&buf[..len]`.
    buf: [u8; MAX_FRAME_SIZE],
    /// Length of the last successfully read frame (valid bytes in `buf`).
    frame_len: usize,
}

impl<R: Read> BlockingFrameReader<R> {
    pub fn new(reader: R) -> Self {
        Self {
            reader: BufReader::new(reader),
            buf: [0u8; MAX_FRAME_SIZE],
            frame_len: 0,
        }
    }

    /// Read the next complete frame into the internal buffer.
    /// Returns a borrowed slice of the frame payload, or `None` on clean
    /// disconnect. The slice is valid until the next `read_frame()` call.
    pub fn read_frame(&mut self) -> io::Result<Option<&[u8]>> {
        // Read the 4-byte length prefix.
        let mut len_buf = [0u8; 4];
        match self.reader.read_exact(&mut len_buf) {
            Ok(()) => {}
            Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
            Err(e) => return Err(e),
        }

        let len = u32::from_le_bytes(len_buf) as usize;
        if len > MAX_FRAME_SIZE {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("frame too large: {len} bytes (max {MAX_FRAME_SIZE})"),
            ));
        }

        self.reader.read_exact(&mut self.buf[..len])?;
        self.frame_len = len;

        Ok(Some(&self.buf[..len]))
    }

    /// The frame the last successful [`read_frame`](Self::read_frame)
    /// returned, for callers that decided what it was and now want it
    /// back without holding the borrow across that decision. Empty
    /// before the first read.
    pub fn frame(&self) -> &[u8] {
        &self.buf[..self.frame_len]
    }

    /// Borrow the underlying reader. Mirrors `BufReader::get_ref` — used
    /// by callers that need to reach the raw stream for socket-level
    /// configuration (`set_read_timeout`, `set_nodelay`, …) without
    /// going through the framed layer.
    pub fn get_ref(&self) -> &R {
        self.reader.get_ref()
    }
}

/// Blocking frame writer. Writes length-prefixed frames to any `Write` sink.
///
/// Uses `BufWriter` to batch small writes (length prefix + payload) into
/// fewer syscalls. Flushed explicitly after each batch.
pub struct BlockingFrameWriter<W: Write> {
    writer: BufWriter<W>,
}

impl<W: Write> BlockingFrameWriter<W> {
    pub fn new(writer: W) -> Self {
        Self {
            writer: BufWriter::new(writer),
        }
    }

    /// Write a complete frame (prepends the 4-byte LE length prefix).
    pub fn write_frame(&mut self, data: &[u8]) -> io::Result<()> {
        self.write_frame_parts(&[data])
    }

    /// Write one frame whose payload is `parts` laid end to end, with
    /// the length prefix of their total. Spares a caller that holds a
    /// frame in pieces — a header it computes and a body it was given —
    /// from assembling them in a buffer of its own: each part goes
    /// straight into the write buffer, and the frame reaches the socket
    /// whole at the next flush. A part that arrives while the buffer is
    /// too full to take it still lands on the socket in order, only in a
    /// separate write.
    ///
    /// A frame over [`MAX_FRAME_SIZE`] is refused with
    /// [`io::ErrorKind::InvalidInput`] before any of it is written: the
    /// far end would drop the connection on it, and this is the error
    /// that says why.
    pub fn write_frame_parts(&mut self, parts: &[&[u8]]) -> io::Result<()> {
        let len: usize = parts.iter().map(|part| part.len()).sum();
        if len > MAX_FRAME_SIZE {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("frame too large: {len} bytes (max {MAX_FRAME_SIZE})"),
            ));
        }
        self.writer.write_all(&(len as u32).to_le_bytes())?;
        for part in parts {
            self.writer.write_all(part)?;
        }
        Ok(())
    }

    /// Flush buffered data to the underlying writer.
    pub fn flush(&mut self) -> io::Result<()> {
        self.writer.flush()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{TcpListener, TcpStream};

    #[test]
    fn frame_round_trip_blocking() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();

        let client = TcpStream::connect(addr).unwrap();
        let (server, _) = listener.accept().unwrap();

        let mut writer = BlockingFrameWriter::new(client);
        let mut reader = BlockingFrameReader::new(server);

        let data = b"hello melin";
        writer.write_frame(data).unwrap();
        writer.flush().unwrap();

        let received = reader.read_frame().unwrap().unwrap();
        assert_eq!(received, data);
    }

    #[test]
    fn frame_parts_arrive_as_one_frame() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();

        let client = TcpStream::connect(addr).unwrap();
        let (server, _) = listener.accept().unwrap();

        let mut writer = BlockingFrameWriter::new(client);
        let mut reader = BlockingFrameReader::new(server);

        writer
            .write_frame_parts(&[&7u64.to_le_bytes(), &[0x10], b"body"])
            .unwrap();
        writer.write_frame_parts(&[]).unwrap();
        writer.flush().unwrap();

        let expected = [&7u64.to_le_bytes()[..], &[0x10], b"body"].concat();
        assert_eq!(reader.read_frame().unwrap().unwrap(), expected);
        assert_eq!(reader.read_frame().unwrap().unwrap(), b"");
    }

    #[test]
    fn an_oversized_frame_is_refused_before_anything_is_written() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();

        let client = TcpStream::connect(addr).unwrap();
        let (server, _) = listener.accept().unwrap();

        let mut writer = BlockingFrameWriter::new(client);
        let mut reader = BlockingFrameReader::new(server);

        // One byte over the cap, split across parts: the total is what
        // counts.
        let body = vec![0xAB; MAX_FRAME_SIZE];
        let err = writer.write_frame_parts(&[&[0x01], &body]).unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);

        // Nothing of it went out: the frame at the cap that follows is the
        // first thing the reader sees, and it is whole.
        writer.write_frame(&body).unwrap();
        writer.flush().unwrap();
        assert_eq!(reader.read_frame().unwrap().unwrap(), &body[..]);
    }

    #[test]
    fn clean_disconnect_returns_none() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();

        let client = TcpStream::connect(addr).unwrap();
        let (server, _) = listener.accept().unwrap();

        drop(client);

        let mut reader = BlockingFrameReader::new(server);
        let result = reader.read_frame().unwrap();
        assert!(result.is_none());
    }

    #[test]
    fn oversized_frame_rejected() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();

        let mut client = TcpStream::connect(addr).unwrap();
        let (server, _) = listener.accept().unwrap();

        // Send a length prefix claiming 2 MiB.
        let fake_len = 2_000_000u32;
        client.write_all(&fake_len.to_le_bytes()).unwrap();

        let mut reader = BlockingFrameReader::new(server);
        let result = reader.read_frame();
        assert!(result.is_err());
    }

    #[test]
    fn multiple_frames_in_sequence() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();

        let client = TcpStream::connect(addr).unwrap();
        let (server, _) = listener.accept().unwrap();

        let mut writer = BlockingFrameWriter::new(client);
        let mut reader = BlockingFrameReader::new(server);

        for i in 0u32..100 {
            writer.write_frame(&i.to_le_bytes()).unwrap();
        }
        writer.flush().unwrap();

        for i in 0u32..100 {
            let frame = reader.read_frame().unwrap().unwrap();
            assert_eq!(frame, i.to_le_bytes());
        }
    }
}