socks5_protocol_async/protocol/
address.rs

1use super::constant;
2use super::shared_internal::*;
3use failure::Error;
4use futures_io::{AsyncRead, AsyncWrite};
5use futures_util::{AsyncReadExt, AsyncWriteExt};
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub enum Address {
9    IPv4([u8; 4]),
10    IPv6([u8; 16]),
11    DomainName(Vec<u8>),
12}
13
14impl Address {
15    pub async fn read<AR>(reader: &mut AR) -> Result<Self, Error>
16    where
17        AR: AsyncRead + Unpin,
18    {
19        let address_type = read_u8(reader).await?;
20        match address_type {
21            constant::address_type::IPV4 => {
22                let mut buf = [0u8; 4];
23                reader.read_exact(buf.as_mut()).await?;
24                Ok(Address::IPv4(buf))
25            }
26            constant::address_type::IPV6 => {
27                let mut buf = [0u8; 16];
28                reader.read_exact(buf.as_mut()).await?;
29                Ok(Address::IPv6(buf))
30            }
31            constant::address_type::DOMAIN_NAME => {
32                let len = read_u8(reader).await?;
33                let v = read_vec(reader, len as usize).await?;
34                Ok(Address::DomainName(v))
35            }
36            _ => Err(crate::error::InvalidAddressTypeError(address_type))?,
37        }
38    }
39
40    pub async fn write<AW>(&self, writer: &mut AW) -> std::io::Result<()>
41    where
42        AW: AsyncWrite + Unpin,
43    {
44        match self {
45            Address::IPv4(val) => {
46                let head = [constant::address_type::IPV4];
47                writer.write_all(&head).await?;
48                writer.write_all(val).await?;
49                Ok(())
50            }
51            Address::IPv6(val) => {
52                let head = [constant::address_type::IPV6];
53                writer.write_all(&head).await?;
54                writer.write_all(val).await?;
55                Ok(())
56            }
57            Address::DomainName(val) => {
58                let head = [constant::address_type::DOMAIN_NAME, val.len() as u8];
59                writer.write_all(&head).await?;
60                writer.write_all(val).await?;
61                Ok(())
62            }
63        }
64    }
65}
66
67#[cfg(test)]
68mod test {
69    use super::*;
70    use crate::test_util::*;
71    use futures_util::io::Cursor;
72
73    #[test]
74    fn read_invalid_type() {
75        let mut reader = Cursor::new(&[0x2]);
76        let future = Address::read(&mut reader);
77        let result = extract_future_output(future);
78        let err = result.unwrap_err();
79        let err = err
80            .downcast::<crate::error::InvalidAddressTypeError>()
81            .unwrap();
82        assert_eq!(err, crate::error::InvalidAddressTypeError(0x2));
83    }
84
85    #[test]
86    fn read_ipv4_ok() {
87        let mut reader = Cursor::new(&[constant::address_type::IPV4, 127, 0, 0, 1]);
88        let future = Address::read(&mut reader);
89        let result = extract_future_output(future);
90        assert_eq!(result.unwrap(), Address::IPv4([127, 0, 0, 1]));
91    }
92
93    #[test]
94    fn read_ipv6_ok() {
95        let mut reader = Cursor::new(&[
96            constant::address_type::IPV6,
97            1,
98            0,
99            0,
100            0,
101            0,
102            0,
103            0,
104            0,
105            0,
106            0,
107            0,
108            0,
109            0,
110            0,
111            0,
112            0xFF,
113        ]);
114        let future = Address::read(&mut reader);
115        let result = extract_future_output(future);
116        assert_eq!(
117            result.unwrap(),
118            Address::IPv6([1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xFF])
119        );
120    }
121
122    #[test]
123    fn read_domain_name_ok() {
124        let mut reader = Cursor::new(&[
125            constant::address_type::DOMAIN_NAME,
126            11,
127            'e' as u8,
128            'x' as u8,
129            'a' as u8,
130            'm' as u8,
131            'p' as u8,
132            'l' as u8,
133            'e' as u8,
134            '.' as u8,
135            'c' as u8,
136            'o' as u8,
137            'm' as u8,
138        ]);
139        let future = Address::read(&mut reader);
140        let result = extract_future_output(future);
141        assert_eq!(
142            result.unwrap(),
143            Address::DomainName("example.com".as_bytes().to_vec())
144        );
145    }
146
147    #[test]
148    fn write_ipv4_ok() {
149        let mut writer = [0u8; 1 + 4];
150        let mut writer = Cursor::new(&mut writer[..]);
151        let res = Address::IPv4([127, 0, 0, 1]);
152        let future = res.write(&mut writer);
153        let result = extract_future_output(future);
154        assert_eq!(result.unwrap(), ());
155        assert_eq!(writer.into_inner(), [0x1, 127, 0, 0, 1])
156    }
157
158    #[test]
159    fn write_ipv6_ok() {
160        let mut writer = [0u8; 1 + 16];
161        let mut writer = Cursor::new(&mut writer[..]);
162        let res = Address::IPv6([1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xFF]);
163        let future = res.write(&mut writer);
164        let result = extract_future_output(future);
165        assert_eq!(result.unwrap(), ());
166        assert_eq!(
167            writer.into_inner(),
168            [0x4, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xFF]
169        )
170    }
171
172    #[test]
173    fn write_domain_name_ok() {
174        let mut writer = [0u8; 1 + 1 + 11];
175        let mut writer = Cursor::new(&mut writer[..]);
176        let res = Address::DomainName("example.com".as_bytes().to_vec());
177        let future = res.write(&mut writer);
178        let result = extract_future_output(future);
179        assert_eq!(result.unwrap(), ());
180        assert_eq!(
181            writer.into_inner(),
182            [
183                0x3, 11, 'e' as u8, 'x' as u8, 'a' as u8, 'm' as u8, 'p' as u8, 'l' as u8,
184                'e' as u8, '.' as u8, 'c' as u8, 'o' as u8, 'm' as u8,
185            ]
186        )
187    }
188}
socks5_protocol_async/protocol/address.rs

socks5_protocol_async/protocol/
address.rs
