socks5_protocol_async/protocol/
reply.rs

1use super::address::Address;
2use super::constant;
3use futures_io::AsyncWrite;
4use futures_util::AsyncWriteExt;
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct Reply {
8    pub reply_code: u8,
9    pub address: Address,
10    pub port: u16,
11}
12
13impl Reply {
14    pub async fn write<AW>(&self, writer: &mut AW) -> std::io::Result<()>
15    where
16        AW: AsyncWrite + Unpin,
17    {
18        let mut buf = Vec::with_capacity(3 + (2 + 0xFF) + 2);
19        buf.extend(&[constant::protocol_version::SOCKS5, self.reply_code, 0x00]);
20        self.address.write(&mut buf).await?;
21        buf.extend(&self.port.to_be_bytes());
22        writer.write_all(&buf).await?;
23        Ok(())
24    }
25}
26
27#[cfg(test)]
28mod test {
29    use super::*;
30    use crate::test_util::*;
31    use futures_util::io::Cursor;
32
33    #[test]
34    fn happy_path() {
35        let mut writer = [0u8; 10];
36        let mut writer = Cursor::new(&mut writer[..]);
37        let res = Reply {
38            reply_code: 0x01,
39            address: Address::IPv4([127, 0, 0, 1]),
40            port: 80,
41        };
42        let future = res.write(&mut writer);
43        let result = extract_future_output(future);
44        assert_eq!(result.unwrap(), ());
45        assert_eq!(
46            writer.into_inner(),
47            [5u8, 0x01, 0x00, 0x01, 127, 0, 0, 1, 00, 80]
48        )
49    }
50
51    #[test]
52    fn unable_to_write_whole_thing() {
53        let mut writer = [0u8; 2];
54        let mut writer = Cursor::new(&mut writer[..]);
55        let res = Reply {
56            reply_code: 0x01,
57            address: Address::IPv4([127, 0, 0, 1]),
58            port: 80,
59        };
60        let future = res.write(&mut writer);
61        let result = extract_future_output(future);
62        let err = result.unwrap_err();
63        assert_eq!(err.kind(), std::io::ErrorKind::WriteZero);
64        assert_eq!(writer.into_inner(), [5u8, 0x01])
65    }
66}