tobytcp/
lib.rs

1#![feature(async_await)]
2
3//! `tobytcp` is a library used when sending messages over a buffer, typically an async `TcpStream`.
4//!
5//! It uses length-prefixing to allow the receiver to differentiate different messages
6//!
7//! # Examples
8//! Here is a tiny example of what it looks like to use `TobyTcp`'s built-in `send` and `receive` fns. Also look at the `/examples` directory
9//! and unit tests in the source code for concrete uses of this library.
10//!
11//! ```no_run
12//! #![feature(async_await)]
13//! # use romio::TcpStream;
14//! # use tobytcp::{send, receive};
15//!
16//! # async fn toby() -> Result<u64, std::io::Error> { // For some reason when I do Result<(), std::io::Error> it complains a ton..
17//! # let mut stream = TcpStream::connect(&"127.0.0.1:7070".parse().unwrap()).await?;
18//! let mut buf = vec![1, 2, 3];
19//! send(&mut buf, &mut stream).await?;
20//!
21//! // Pretend we're connected to an echo server..
22//! let received = receive(&mut stream).await?;
23//!
24//! assert_eq!(buf, received);
25//! # Ok(8)
26//! # };
27//! ```
28//!
29
30/// Includes methods for generating the `tobytcp` prefix, or attempting to decode the length of
31/// the encoded data i.e. payload, from a buffer.
32pub mod protocol;
33
34use futures::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
35use std::io;
36use std::marker::Unpin;
37
38/// Writes the data, encoded as `tobytcp`, to the `Write`. Returns the total number of bytes
39/// written, equal to `data.len() + 8`. See the `examples/` dir in the source code, or the tests
40/// of this fn in the source code for some examples of this being used.
41///
42/// Note: Do *not* perform any IO on the `Write` outside of calling `send` or `receive`! It can corrupt the tobytcp stream
43pub async fn send<'w, W>(data: &'w [u8], write: &'w mut W) -> Result<usize, io::Error>
44where
45    W: AsyncWrite + Unpin,
46{
47    let prefix = protocol::tobytcp_prefix(data.len());
48    write.write_all(&prefix).await?;
49    write.write_all(data).await?;
50    Ok(data.len() + 8)
51}
52
53/// Wait for data, which was encoded as `tobytcp`, to be received from this `Read`. Returns the data or any error.
54/// See the `examples/` dir in the source code, or the tests of this fn in the source code for some examples of this being used.
55///
56/// Note: Do *not* perform any IO on the `Read` outside of calling `send` or `receive`! It can corrupt the tobytcp stream
57pub async fn receive<R>(read: &mut R) -> Result<Vec<u8>, io::Error>
58where
59    R: AsyncRead + Unpin,
60{
61    let mut len_buf: [u8; 8] = [0; 8];
62    read.read_exact(&mut len_buf).await?;
63
64    let len = protocol::tobytcp_len(len_buf);
65
66    let mut buf = vec![0; len as usize];
67    read.read_exact(&mut buf).await?;
68    Ok(buf)
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    #[runtime::test]
76    async fn simple_test() {
77        let to_send: Vec<u8> = vec![13, 58, 2, 4];
78        let mut output = Vec::new();
79
80        let size: usize = 12;
81        assert_eq!(size, send(&to_send, &mut output).await.unwrap());
82
83        // 'manually' check the buffer
84        let mut len_bytes: [u8; 8] = [0; 8];
85        len_bytes.clone_from_slice(&output[0..8]);
86        assert_eq!(4, u64::from_be_bytes(len_bytes));
87        assert_eq!(to_send[0..], output[8..12]);
88
89        // hacky..
90        let x = output.clone();
91        let mut y = x.as_slice();
92
93        let received = receive(&mut y).await.unwrap();
94        assert_eq!(received, to_send);
95    }
96
97    // Tests that if we buffer up a bunch of data that we can separate them, very important!
98    #[runtime::test]
99    async fn many_sends_then_receive_test() {
100        let mut output = Vec::new();
101
102        let num = 20;
103        for i in 0..num {
104            let to_send: Vec<u8> = vec![i, i, i, i];
105            send(&to_send, &mut output).await.unwrap();
106        }
107
108        // hacky..
109        let x = output.clone();
110        let mut y = x.as_slice();
111
112        for i in 0..num {
113            let received = receive(&mut y).await.unwrap();
114            assert_eq!(vec![i, i, i, i], received);
115        }
116    }
117}