ghast_protocol/
io.rs

1//! Async framed reader and writer over tokio's `AsyncRead` / `AsyncWrite`
2//! (feature `tokio`).
3//!
4//! Use them on the two halves of a connection (e.g. `TcpStream::into_split`). Each half
5//! has its own compression/encryption switches; the connection logic flips both at the
6//! protocol-defined points.
7
8use bytes::Buf;
9use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
10
11use crate::error::{EncodeError, FrameError};
12use crate::frame::{FrameDecoder, FrameEncoder};
13use crate::packet::{Packet, PacketSet, RawPacket};
14
15/// Minimum free space requested from the socket per read.
16const READ_CHUNK: usize = 16 * 1024;
17
18/// Reads packets from an `AsyncRead`.
19#[derive(Debug)]
20pub struct FramedReader<R> {
21    inner: R,
22    decoder: FrameDecoder,
23}
24
25impl<R: AsyncRead + Unpin> FramedReader<R> {
26    /// Wraps a reader.
27    pub fn new(inner: R) -> Self {
28        FramedReader {
29            inner,
30            decoder: FrameDecoder::new(),
31        }
32    }
33
34    /// The frame decoder (compression, encryption, filter, limits, stats).
35    pub fn decoder(&mut self) -> &mut FrameDecoder {
36        &mut self.decoder
37    }
38
39    /// The underlying reader.
40    pub fn get_mut(&mut self) -> &mut R {
41        &mut self.inner
42    }
43
44    /// Reads the next packet that passes the filter. `Ok(None)` on a clean end of
45    /// stream (no partial frame buffered).
46    ///
47    /// Cancel safe: if the future is dropped, no data is lost.
48    pub async fn read_packet(&mut self) -> Result<Option<RawPacket>, FrameError> {
49        loop {
50            if let Some(p) = self.decoder.next_packet()? {
51                return Ok(Some(p));
52            }
53            let want = self.decoder.needed().unwrap_or(0).max(READ_CHUNK);
54            let buf = self.decoder.buffer_mut();
55            buf.reserve(want);
56            let old = buf.len();
57            let n = self.inner.read_buf(buf).await?;
58            if n == 0 {
59                if self.decoder.buffered() == 0 {
60                    return Ok(None);
61                }
62                return Err(FrameError::Io(std::io::ErrorKind::UnexpectedEof.into()));
63            }
64            self.decoder.commit(old);
65        }
66    }
67}
68
69/// Writes packets to an `AsyncWrite`, batching until [`FramedWriter::flush`].
70#[derive(Debug)]
71pub struct FramedWriter<W> {
72    inner: W,
73    encoder: FrameEncoder,
74}
75
76impl<W: AsyncWrite + Unpin> FramedWriter<W> {
77    /// Wraps a writer.
78    pub fn new(inner: W) -> Self {
79        FramedWriter {
80            inner,
81            encoder: FrameEncoder::new(),
82        }
83    }
84
85    /// The frame encoder (compression, encryption).
86    pub fn encoder(&mut self) -> &mut FrameEncoder {
87        &mut self.encoder
88    }
89
90    /// The underlying writer.
91    pub fn get_mut(&mut self) -> &mut W {
92        &mut self.inner
93    }
94
95    /// Queues a typed packet (nothing is written until [`FramedWriter::flush`]).
96    pub fn queue<P: Packet>(&mut self, p: &P) -> Result<(), EncodeError> {
97        self.encoder.write_packet(p)
98    }
99
100    /// Queues a packet-set value.
101    pub fn queue_set<S: PacketSet>(&mut self, p: &S) -> Result<(), EncodeError> {
102        self.encoder.write_set(p)
103    }
104
105    /// Queues a raw packet.
106    pub fn queue_raw(&mut self, p: &RawPacket) -> Result<(), EncodeError> {
107        self.encoder.write_raw(p)
108    }
109
110    /// Writes everything queued and flushes the writer.
111    ///
112    /// Cancel safe in the sense that bytes already written are removed from the queue;
113    /// a dropped flush can be resumed by calling it again.
114    pub async fn flush(&mut self) -> std::io::Result<()> {
115        let out = self.encoder.out_mut();
116        while out.has_remaining() {
117            let n = self.inner.write_buf(out).await?;
118            if n == 0 {
119                return Err(std::io::ErrorKind::WriteZero.into());
120            }
121        }
122        self.inner.flush().await
123    }
124
125    /// Queues a typed packet and flushes.
126    pub async fn send<P: Packet>(&mut self, p: &P) -> Result<(), FrameError> {
127        self.queue(p)
128            .map_err(|e| FrameError::Io(std::io::Error::other(e)))?;
129        self.flush().await?;
130        Ok(())
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use bytes::Bytes;
138
139    #[tokio::test]
140    async fn duplex_roundtrip() {
141        let (a, b) = tokio::io::duplex(64);
142        let mut w = FramedWriter::new(a);
143        let mut r = FramedReader::new(b);
144        let key = [5u8; 16];
145        w.encoder().set_compression(32);
146        w.encoder().enable_encryption(&key);
147        r.decoder().set_compression(32);
148        r.decoder().enable_encryption(&key);
149        let packets: Vec<RawPacket> = (0..20)
150            .map(|i| RawPacket::new(i, Bytes::from(vec![i as u8; (i * 13) as usize])))
151            .collect();
152        let send = packets.clone();
153        let writer = tokio::spawn(async move {
154            for p in &send {
155                w.queue_raw(p).unwrap();
156            }
157            w.flush().await.unwrap();
158            drop(w);
159        });
160        let mut got = Vec::new();
161        while let Some(p) = r.read_packet().await.unwrap() {
162            got.push(p);
163        }
164        writer.await.unwrap();
165        assert_eq!(got, packets);
166    }
167}
ghast_protocol/io.rs

ghast_protocol/
io.rs
