waos_dev_protocol/
connection.rs

1//! One connection, as frames over a stream: each message as postcard encodes it, behind a
2//! four-byte little-endian count of its bytes, so the reader knows how much to wait for before
3//! it allocates anything, and refuses a frame over [`MAX_FRAME_BYTES`] on the count alone. The
4//! first frame each way is the protocol's number, [`PROTOCOL`], outside the messages' encoding,
5//! so two sides of different builds find out before either reads a message it would misread.
6//!
7//! The invariants a later editor will break:
8//! - After `send` or `receive` fails, the connection is done. A frame cut short, or one refused,
9//!   leaves the stream somewhere inside a frame, and nothing says where the next one begins; the
10//!   side that saw the error closes.
11//! - The protocol frame's shape never changes: a count of four, then the number, little-endian.
12//!   A side of any build reads it, and reads nothing else until the numbers agree.
13
14use std::io::{self, ErrorKind};
15
16use serde::Serialize;
17use serde::de::DeserializeOwned;
18use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
19
20/// The protocol's number, which the first frame each way says: the messages and their encoding
21/// as this build of the crate has them. Raised with any change a side of the build before would
22/// misread, which, with postcard, is any change to a message but a variant or a field added at
23/// the end.
24pub const PROTOCOL: u32 = 2;
25
26/// The largest frame either side accepts. A component is megabytes, a debug build tens of
27/// them; a count beyond this is a stream that is not ours, refused before a byte of it is read.
28pub const MAX_FRAME_BYTES: usize = 256 << 20;
29
30const PREFIX_BYTES: usize = 4;
31
32/// The protocol frame's body: the number's four bytes.
33const PROTOCOL_BYTES: usize = 4;
34
35/// One connection: length-prefixed postcard frames over `S`, the Unix socket on this machine.
36/// Split into its halves, each is a `Connection` that only sends or only receives, so a WaOS
37/// sends a program's output on one task while another watches for the connection closing.
38#[derive(Debug)]
39pub struct Connection<S> {
40    pub(crate) stream: S,
41}
42
43impl<S> Connection<S> {
44    pub fn new(stream: S) -> Connection<S> {
45        Connection { stream }
46    }
47}
48
49impl<S: AsyncRead + AsyncWrite + Unpin> Connection<S> {
50    /// The first frame each way: says this build's [`PROTOCOL`], and reads the other side's
51    /// before anything else. When they differ the connection is done, and the error, of kind
52    /// `Unsupported`, names both numbers; `cargo waos` prints it. An other side that closed
53    /// before saying its number is `UnexpectedEof`, or `BrokenPipe` when it was gone before this
54    /// side's frame went out.
55    pub async fn agree_on_protocol(&mut self) -> io::Result<()> {
56        self.agree_on(PROTOCOL).await
57    }
58
59    async fn agree_on(&mut self, ours: u32) -> io::Result<()> {
60        self.stream.write_all(&protocol_frame(ours)).await?;
61        let theirs = self.read_protocol().await?;
62        if theirs != ours {
63            return Err(io::Error::new(
64                ErrorKind::Unsupported,
65                format!(
66                    "the other side speaks protocol {theirs} and this side protocol {ours}: cargo waos and the WaOS it reaches must be built from the same waos-dev-protocol"
67                ),
68            ));
69        }
70        Ok(())
71    }
72
73    /// The other side's protocol number, from its first frame.
74    async fn read_protocol(&mut self) -> io::Result<u32> {
75        let len = self.read_prefix().await?.ok_or_else(|| {
76            io::Error::new(
77                ErrorKind::UnexpectedEof,
78                "the other side closed before saying its protocol",
79            )
80        })?;
81        if len != PROTOCOL_BYTES {
82            return Err(invalid_data(format!(
83                "the first frame is {len} bytes long, not a protocol number"
84            )));
85        }
86        let mut number = [0; PROTOCOL_BYTES];
87        self.stream.read_exact(&mut number).await?;
88        Ok(u32::from_le_bytes(number))
89    }
90}
91
92impl<S: AsyncWrite + Unpin> Connection<S> {
93    /// Sends `message` as one frame; refused, unsent, when it is over the limit.
94    pub async fn send(&mut self, message: &impl Serialize) -> io::Result<()> {
95        let frame = frame(message)?;
96        self.stream.write_all(&frame).await
97    }
98}
99
100impl<S: AsyncRead + Unpin> Connection<S> {
101    /// The next message; none once the other side closed between frames.
102    pub async fn receive<M: DeserializeOwned>(&mut self) -> io::Result<Option<M>> {
103        let Some(len) = self.read_prefix().await? else {
104            return Ok(None);
105        };
106        let mut body = vec![0; len];
107        self.stream.read_exact(&mut body).await?;
108        postcard::from_bytes(&body).map(Some).map_err(invalid_data)
109    }
110
111    /// The next frame's byte count, checked against the limit; none when the other side closed
112    /// between frames.
113    async fn read_prefix(&mut self) -> io::Result<Option<usize>> {
114        let mut prefix = [0; PREFIX_BYTES];
115        let mut filled = 0;
116        while filled < PREFIX_BYTES {
117            match self.stream.read(&mut prefix[filled..]).await? {
118                0 if filled == 0 => return Ok(None),
119                0 => return Err(cut_short()),
120                read => filled += read,
121            }
122        }
123        accepted(u32::from_le_bytes(prefix) as usize).map(Some)
124    }
125}
126
127/// The protocol frame: a count of four, then `protocol` little-endian.
128fn protocol_frame(protocol: u32) -> [u8; PREFIX_BYTES + PROTOCOL_BYTES] {
129    let mut frame = [0; PREFIX_BYTES + PROTOCOL_BYTES];
130    frame[..PREFIX_BYTES].copy_from_slice(&(PROTOCOL_BYTES as u32).to_le_bytes());
131    frame[PREFIX_BYTES..].copy_from_slice(&protocol.to_le_bytes());
132    frame
133}
134
135/// `message` as one frame: its byte count, then its bytes.
136fn frame(message: &impl Serialize) -> io::Result<Vec<u8>> {
137    let mut frame = postcard::to_extend(message, vec![0; PREFIX_BYTES]).map_err(invalid_data)?;
138    let len = accepted(frame.len() - PREFIX_BYTES)?;
139    // Within the limit, so within a `u32`.
140    frame[..PREFIX_BYTES].copy_from_slice(&(len as u32).to_le_bytes());
141    Ok(frame)
142}
143
144/// `len` when a frame that long is accepted; the error names the limit otherwise.
145fn accepted(len: usize) -> io::Result<usize> {
146    if len > MAX_FRAME_BYTES {
147        return Err(invalid_data(format!(
148            "a frame of {len} bytes, over the {MAX_FRAME_BYTES} a connection accepts"
149        )));
150    }
151    Ok(len)
152}
153
154fn cut_short() -> io::Error {
155    io::Error::new(
156        ErrorKind::UnexpectedEof,
157        "the other side closed inside a frame",
158    )
159}
160
161fn invalid_data(error: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> io::Error {
162    io::Error::new(ErrorKind::InvalidData, error)
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168
169    use tokio::io::DuplexStream;
170
171    fn pair() -> (Connection<DuplexStream>, Connection<DuplexStream>) {
172        let (one, other) = tokio::io::duplex(64);
173        (Connection::new(one), Connection::new(other))
174    }
175
176    #[tokio::test]
177    async fn a_message_comes_back_as_sent() {
178        let (mut sender, mut receiver) = pair();
179        tokio::spawn(async move { sender.send(&(7u32, "seven".to_string())).await });
180        let received: Option<(u32, String)> = receiver.receive().await.unwrap();
181        assert_eq!(received, Some((7, "seven".into())));
182    }
183
184    #[tokio::test]
185    async fn two_sides_of_one_build_agree_and_go_on_to_the_messages() {
186        let (mut one, mut other) = pair();
187        let (agreed, agreed_too) = tokio::join!(one.agree_on_protocol(), other.agree_on_protocol());
188        agreed.unwrap();
189        agreed_too.unwrap();
190        tokio::spawn(async move { one.send(&7u32).await });
191        assert_eq!(other.receive::<u32>().await.unwrap(), Some(7));
192    }
193
194    #[test]
195    fn the_protocol_frame_is_a_count_of_four_then_the_number() {
196        assert_eq!(protocol_frame(1), [4, 0, 0, 0, 1, 0, 0, 0]);
197        assert_eq!(protocol_frame(258), [4, 0, 0, 0, 2, 1, 0, 0]);
198    }
199
200    #[tokio::test]
201    async fn two_sides_of_different_builds_each_refuse_naming_both_numbers() {
202        let (mut one, mut other) = pair();
203        let (refused, refused_too) = tokio::join!(one.agree_on(1), other.agree_on(2));
204        let (refused, refused_too) = (refused.unwrap_err(), refused_too.unwrap_err());
205        assert_eq!(refused.kind(), ErrorKind::Unsupported);
206        assert_eq!(refused_too.kind(), ErrorKind::Unsupported);
207        assert!(
208            refused.to_string().contains("protocol 2")
209                && refused.to_string().contains("protocol 1"),
210            "{refused}"
211        );
212        assert!(
213            refused_too.to_string().contains("protocol 1")
214                && refused_too.to_string().contains("protocol 2"),
215            "{refused_too}"
216        );
217    }
218
219    #[tokio::test]
220    async fn closing_before_the_protocol_frame_is_an_error() {
221        let (one, mut other) = pair();
222        drop(one);
223        let error = other.agree_on_protocol().await.unwrap_err();
224        assert!(
225            matches!(
226                error.kind(),
227                ErrorKind::UnexpectedEof | ErrorKind::BrokenPipe
228            ),
229            "{error}"
230        );
231    }
232
233    #[tokio::test]
234    async fn a_first_frame_that_is_not_a_number_is_refused() {
235        let (mut one, mut other) = pair();
236        tokio::spawn(async move { one.send(&"a message before the protocol").await });
237        let error = other.agree_on_protocol().await.unwrap_err();
238        assert_eq!(error.kind(), ErrorKind::InvalidData);
239    }
240
241    #[test]
242    fn the_count_is_four_bytes_little_endian_before_the_message() {
243        let frame = frame(&vec![1u8, 2, 3]).unwrap();
244        // postcard writes a `Vec` as its length, then its bytes.
245        assert_eq!(frame, [4, 0, 0, 0, 3, 1, 2, 3]);
246    }
247
248    #[tokio::test]
249    async fn a_count_over_the_limit_is_refused_before_a_body_is_read() {
250        let (mut sender, mut receiver) = pair();
251        sender
252            .stream
253            .write_all(&u32::MAX.to_le_bytes())
254            .await
255            .unwrap();
256        // No body follows: a receive that tried to read one would wait here forever.
257        let error = receiver.receive::<Vec<u8>>().await.unwrap_err();
258        assert_eq!(error.kind(), ErrorKind::InvalidData);
259        assert!(
260            error.to_string().contains(&MAX_FRAME_BYTES.to_string()),
261            "the error names the limit: {error}"
262        );
263    }
264
265    #[test]
266    fn a_frame_at_the_limit_is_accepted_and_one_over_it_refused() {
267        assert_eq!(accepted(MAX_FRAME_BYTES).unwrap(), MAX_FRAME_BYTES);
268        assert_eq!(
269            accepted(MAX_FRAME_BYTES + 1).unwrap_err().kind(),
270            ErrorKind::InvalidData
271        );
272    }
273
274    #[tokio::test]
275    async fn closing_between_frames_is_none() {
276        let (sender, mut receiver) = pair();
277        drop(sender);
278        assert_eq!(receiver.receive::<u32>().await.unwrap(), None);
279    }
280
281    #[tokio::test]
282    async fn closing_inside_a_frame_is_an_error() {
283        let (mut sender, mut receiver) = pair();
284        sender.stream.write_all(&[8, 0]).await.unwrap();
285        drop(sender);
286        let error = receiver.receive::<u32>().await.unwrap_err();
287        assert_eq!(error.kind(), ErrorKind::UnexpectedEof);
288    }
289}
waos_dev_protocol/connection.rs

waos_dev_protocol/
connection.rs
