confium_net/
io.rs

1//! Adapter from message-framed [`Transport`]s to byte-stream
2//! `Read`/`Write` callers.
3//!
4//! Length-prefix protocols (the coordinator's message framing) are
5//! written against `Read + Write` streams. The registry transports
6//! are message-framed: one `send` is one `recv` payload.
7//! [`TransportIo`] bridges the two worlds:
8//!
9//! - `write` buffers; `flush` emits exactly one transport message
10//!   (so `write_all(prefix); write_all(body); flush()` becomes a
11//!   single frame, preserving the caller's framing).
12//! - `read` serves from an internal buffer refilled one transport
13//!   message at a time; a clean close yields `Ok(0)` (EOF).
14
15use std::io::Read;
16use std::io::Write;
17
18use crate::Transport;
19
20pub struct TransportIo {
21    transport: Box<dyn Transport>,
22    wbuf: Vec<u8>,
23    rbuf: Vec<u8>,
24    rpos: usize,
25    closed: bool,
26}
27
28impl TransportIo {
29    pub fn new(transport: Box<dyn Transport>) -> Self {
30        Self {
31            transport,
32            wbuf: Vec::new(),
33            rbuf: Vec::new(),
34            rpos: 0,
35            closed: false,
36        }
37    }
38}
39
40impl Write for TransportIo {
41    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
42        self.wbuf.extend_from_slice(buf);
43        Ok(buf.len())
44    }
45
46    fn flush(&mut self) -> std::io::Result<()> {
47        if self.wbuf.is_empty() {
48            return Ok(());
49        }
50        let frame = std::mem::take(&mut self.wbuf);
51        self.transport.send(&frame).map_err(|_| {
52            std::io::Error::new(
53                std::io::ErrorKind::ConnectionAborted,
54                "transport send failed",
55            )
56        })
57    }
58}
59
60impl Read for TransportIo {
61    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
62        if self.rpos >= self.rbuf.len() {
63            if self.closed {
64                return Ok(0);
65            }
66            // Refill with the next whole transport message.
67            let mut next = vec![0u8; 16 * 1024 * 1024];
68            match self.transport.recv(&mut next) {
69                Ok(0) => {
70                    self.closed = true;
71                    return Ok(0);
72                }
73                Ok(n) => {
74                    next.truncate(n);
75                    self.rbuf = next;
76                    self.rpos = 0;
77                }
78                Err(_) => {
79                    self.closed = true;
80                    return Err(std::io::Error::new(
81                        std::io::ErrorKind::ConnectionAborted,
82                        "transport recv failed",
83                    ));
84                }
85            }
86        }
87        let avail = &self.rbuf[self.rpos..];
88        let n = avail.len().min(buf.len());
89        buf[..n].copy_from_slice(&avail[..n]);
90        self.rpos += n;
91        Ok(n)
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use crate::Result;
99    use std::sync::mpsc;
100
101    /// Minimal in-memory Transport: a channel pair.
102    struct ChanTransport {
103        out: mpsc::Sender<Vec<u8>>,
104        inp: mpsc::Receiver<Vec<u8>>,
105        pending: Vec<u8>,
106    }
107
108    fn pair() -> (ChanTransport, ChanTransport) {
109        let (a2b_tx, a2b_rx) = mpsc::channel();
110        let (b2a_tx, b2a_rx) = mpsc::channel();
111        (
112            ChanTransport {
113                out: a2b_tx,
114                inp: b2a_rx,
115                pending: Vec::new(),
116            },
117            ChanTransport {
118                out: b2a_tx,
119                inp: a2b_rx,
120                pending: Vec::new(),
121            },
122        )
123    }
124
125    impl Transport for ChanTransport {
126        fn send(&mut self, data: &[u8]) -> Result<()> {
127            self.out.send(data.to_vec()).expect("channel pair alive");
128            Ok(())
129        }
130        fn recv(&mut self, buf: &mut [u8]) -> Result<usize> {
131            if self.pending.is_empty() {
132                let msg = self
133                    .inp
134                    .recv()
135                    .map_err(|_| crate::error::ClosedSnafu.build())?;
136                self.pending = msg;
137            }
138            let n = self.pending.len().min(buf.len());
139            buf[..n].copy_from_slice(&self.pending[..n]);
140            self.pending.drain(..n);
141            Ok(n)
142        }
143        fn close(&mut self) -> Result<()> {
144            Ok(())
145        }
146    }
147
148    #[test]
149    fn write_flush_is_one_message() {
150        let (a, mut b) = pair();
151        let mut io = TransportIo::new(Box::new(a));
152        io.write_all(&4u32.to_be_bytes()).unwrap();
153        io.write_all(b"body").unwrap();
154        io.flush().unwrap();
155        let mut buf = [0u8; 16];
156        let n = b.recv(&mut buf).unwrap();
157        assert_eq!(&buf[..n], &[0, 0, 0, 4, b'b', b'o', b'd', b'y']);
158    }
159
160    #[test]
161    fn read_serves_partial_then_refills() {
162        let (mut a, b) = pair();
163        a.send(b"hello world").unwrap();
164        let mut io = TransportIo::new(Box::new(b));
165        let mut five = [0u8; 5];
166        io.read_exact(&mut five).unwrap();
167        assert_eq!(&five, b"hello");
168        let mut rest = [0u8; 6];
169        io.read_exact(&mut rest).unwrap();
170        assert_eq!(&rest, b" world");
171    }
172}
confium_net/io.rs

confium_net/
io.rs
