socks5_impl/protocol/
address.rs

1#[cfg(feature = "tokio")]
2use crate::protocol::AsyncStreamOperation;
3use crate::protocol::StreamOperation;
4use bytes::BufMut;
5use std::{
6    io::Cursor,
7    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, ToSocketAddrs},
8};
9#[cfg(feature = "tokio")]
10use tokio::io::{AsyncRead, AsyncReadExt};
11
12#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Default)]
13#[repr(u8)]
14pub enum AddressType {
15    #[default]
16    IPv4 = 0x01,
17    Domain = 0x03,
18    IPv6 = 0x04,
19}
20
21impl TryFrom<u8> for AddressType {
22    type Error = std::io::Error;
23    fn try_from(code: u8) -> core::result::Result<Self, Self::Error> {
24        let err = format!("Unsupported address type code {code:#x}");
25        match code {
26            0x01 => Ok(AddressType::IPv4),
27            0x03 => Ok(AddressType::Domain),
28            0x04 => Ok(AddressType::IPv6),
29            _ => Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, err)),
30        }
31    }
32}
33
34impl From<AddressType> for u8 {
35    fn from(addr_type: AddressType) -> Self {
36        match addr_type {
37            AddressType::IPv4 => 0x01,
38            AddressType::Domain => 0x03,
39            AddressType::IPv6 => 0x04,
40        }
41    }
42}
43
44/// SOCKS5 Address Format
45///
46/// ```plain
47/// +------+----------+----------+
48/// | ATYP | DST.ADDR | DST.PORT |
49/// +------+----------+----------+
50/// |  1   | Variable |    2     |
51/// +------+----------+----------+
52/// ```
53#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
54#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
55pub enum Address {
56    /// Represents an IPv4 or IPv6 socket address.
57    SocketAddress(SocketAddr),
58    /// Represents a domain name and a port.
59    DomainAddress(Box<str>, u16),
60}
61
62impl Default for Address {
63    fn default() -> Self {
64        Address::unspecified()
65    }
66}
67
68impl Address {
69    /// Returns an unspecified IPv4 address (0.0.0.0:0).
70    pub fn unspecified() -> Self {
71        Address::SocketAddress(SocketAddr::from((Ipv4Addr::UNSPECIFIED, 0)))
72    }
73
74    /// Returns the type of the address (IPv4, IPv6, or Domain).
75    pub fn get_type(&self) -> AddressType {
76        match self {
77            Self::SocketAddress(SocketAddr::V4(_)) => AddressType::IPv4,
78            Self::SocketAddress(SocketAddr::V6(_)) => AddressType::IPv6,
79            Self::DomainAddress(_, _) => AddressType::Domain,
80        }
81    }
82
83    /// Returns the port number.
84    pub fn port(&self) -> u16 {
85        match self {
86            Self::SocketAddress(addr) => addr.port(),
87            Self::DomainAddress(_, port) => *port,
88        }
89    }
90
91    /// Returns the host name or IP address as a string.
92    pub fn host(&self) -> String {
93        match self {
94            Self::SocketAddress(addr) => addr.ip().to_string(),
95            Self::DomainAddress(addr, _) => addr.to_string(),
96        }
97    }
98
99    /// Returns the host name or IP address as a string.
100    #[deprecated(note = "use host() instead")]
101    pub fn domain(&self) -> String {
102        self.host()
103    }
104
105    /// Returns `true` if it is an IPv4 address.
106    pub fn is_ipv4(&self) -> bool {
107        matches!(self, Self::SocketAddress(SocketAddr::V4(_)))
108    }
109
110    /// Returns `true` if it is an IPv6 address.
111    pub fn is_ipv6(&self) -> bool {
112        matches!(self, Self::SocketAddress(SocketAddr::V6(_)))
113    }
114
115    /// Returns `true` if it is a domain address.
116    pub fn is_domain(&self) -> bool {
117        matches!(self, Self::DomainAddress(_, _))
118    }
119
120    /// Returns the maximum possible serialized length of a SOCKS5 address.
121    pub const fn max_serialized_len() -> usize {
122        1 + 1 + u8::MAX as usize + 2
123    }
124}
125
126impl StreamOperation for Address {
127    fn retrieve_from_stream<R: std::io::Read>(stream: &mut R) -> std::io::Result<Self> {
128        let mut atyp_buf = [0; 1];
129        stream.read_exact(&mut atyp_buf)?;
130        match AddressType::try_from(atyp_buf[0])? {
131            AddressType::IPv4 => {
132                let mut buf = [0; 6];
133                stream.read_exact(&mut buf)?;
134                let addr = Ipv4Addr::new(buf[0], buf[1], buf[2], buf[3]);
135                let port = u16::from_be_bytes([buf[4], buf[5]]);
136                Ok(Self::SocketAddress(SocketAddr::from((addr, port))))
137            }
138            AddressType::Domain => {
139                let mut len_buf = [0; 1];
140                stream.read_exact(&mut len_buf)?;
141                let len = len_buf[0] as usize;
142                let mut domain_buf = vec![0; len];
143                stream.read_exact(&mut domain_buf)?;
144
145                let mut port_buf = [0; 2];
146                stream.read_exact(&mut port_buf)?;
147                let port = u16::from_be_bytes(port_buf);
148
149                let addr = String::from_utf8(domain_buf)
150                    .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, format!("Invalid address encoding: {err}")))?;
151                Ok(Self::DomainAddress(addr.into_boxed_str(), port))
152            }
153            AddressType::IPv6 => {
154                let mut buf = [0; 18];
155                stream.read_exact(&mut buf)?;
156                let port = u16::from_be_bytes([buf[16], buf[17]]);
157                let mut addr_bytes = [0; 16];
158                addr_bytes.copy_from_slice(&buf[..16]);
159                Ok(Self::SocketAddress(SocketAddr::from((Ipv6Addr::from(addr_bytes), port))))
160            }
161        }
162    }
163
164    fn write_to_buf<B: BufMut>(&self, buf: &mut B) {
165        match self {
166            Self::SocketAddress(SocketAddr::V4(addr)) => {
167                buf.put_u8(AddressType::IPv4.into());
168                buf.put_slice(&addr.ip().octets());
169                buf.put_u16(addr.port());
170            }
171            Self::SocketAddress(SocketAddr::V6(addr)) => {
172                buf.put_u8(AddressType::IPv6.into());
173                buf.put_slice(&addr.ip().octets());
174                buf.put_u16(addr.port());
175            }
176            Self::DomainAddress(addr, port) => {
177                let addr = addr.as_bytes();
178                buf.put_u8(AddressType::Domain.into());
179                buf.put_u8(u8::try_from(addr.len()).expect("domain address too long for SOCKS5"));
180                buf.put_slice(addr);
181                buf.put_u16(*port);
182            }
183        }
184    }
185
186    fn len(&self) -> usize {
187        match self {
188            Address::SocketAddress(SocketAddr::V4(_)) => 1 + 4 + 2,
189            Address::SocketAddress(SocketAddr::V6(_)) => 1 + 16 + 2,
190            Address::DomainAddress(addr, _) => 1 + 1 + addr.len() + 2,
191        }
192    }
193}
194
195#[cfg(feature = "tokio")]
196#[async_trait::async_trait]
197impl AsyncStreamOperation for Address {
198    async fn retrieve_from_async_stream<R>(stream: &mut R) -> std::io::Result<Self>
199    where
200        R: AsyncRead + Unpin + Send + ?Sized,
201    {
202        let atyp = stream.read_u8().await?;
203        match AddressType::try_from(atyp)? {
204            AddressType::IPv4 => {
205                let mut addr_bytes = [0; 4];
206                stream.read_exact(&mut addr_bytes).await?;
207                let port = stream.read_u16().await?;
208                let addr = Ipv4Addr::from(addr_bytes);
209                Ok(Self::SocketAddress(SocketAddr::from((addr, port))))
210            }
211            AddressType::Domain => {
212                let len = stream.read_u8().await? as usize;
213                let mut domain_buf = vec![0; len];
214                stream.read_exact(&mut domain_buf).await?;
215                let port = stream.read_u16().await?;
216
217                let addr = String::from_utf8(domain_buf)
218                    .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, format!("Invalid address encoding: {err}")))?;
219                Ok(Self::DomainAddress(addr.into_boxed_str(), port))
220            }
221            AddressType::IPv6 => {
222                let mut addr_bytes = [0; 16];
223                stream.read_exact(&mut addr_bytes).await?;
224                let port = stream.read_u16().await?;
225                Ok(Self::SocketAddress(SocketAddr::from((Ipv6Addr::from(addr_bytes), port))))
226            }
227        }
228    }
229}
230
231impl ToSocketAddrs for Address {
232    type Iter = std::vec::IntoIter<SocketAddr>;
233
234    fn to_socket_addrs(&self) -> std::io::Result<Self::Iter> {
235        match self {
236            Address::SocketAddress(addr) => Ok(vec![*addr].into_iter()),
237            Address::DomainAddress(addr, port) => Ok((&**addr, *port).to_socket_addrs()?),
238        }
239    }
240}
241
242impl std::fmt::Display for Address {
243    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
244        match self {
245            Address::DomainAddress(hostname, port) => write!(f, "{hostname}:{port}"),
246            Address::SocketAddress(socket_addr) => write!(f, "{socket_addr}"),
247        }
248    }
249}
250
251impl TryFrom<Address> for SocketAddr {
252    type Error = std::io::Error;
253
254    fn try_from(address: Address) -> std::result::Result<Self, Self::Error> {
255        match address {
256            Address::SocketAddress(addr) => Ok(addr),
257            Address::DomainAddress(addr, port) => {
258                let addr = addr.as_ref();
259                if let Ok(addr) = addr.parse::<IpAddr>() {
260                    Ok(SocketAddr::from((addr, port)))
261                } else {
262                    let err = format!("domain address {addr} cannot be converted to a SocketAddr without DNS, use to_socket_addrs instead");
263                    Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, err))
264                }
265            }
266        }
267    }
268}
269
270impl TryFrom<&Address> for SocketAddr {
271    type Error = std::io::Error;
272
273    fn try_from(address: &Address) -> std::result::Result<Self, Self::Error> {
274        TryFrom::<Address>::try_from(address.clone())
275    }
276}
277
278impl From<Address> for Vec<u8> {
279    fn from(addr: Address) -> Self {
280        let mut buf = Vec::with_capacity(addr.len());
281        addr.write_to_buf(&mut buf);
282        buf
283    }
284}
285
286impl TryFrom<Vec<u8>> for Address {
287    type Error = std::io::Error;
288
289    fn try_from(data: Vec<u8>) -> std::result::Result<Self, Self::Error> {
290        let mut rdr = Cursor::new(data);
291        Self::retrieve_from_stream(&mut rdr)
292    }
293}
294
295impl TryFrom<&[u8]> for Address {
296    type Error = std::io::Error;
297
298    fn try_from(data: &[u8]) -> std::result::Result<Self, Self::Error> {
299        let mut rdr = Cursor::new(data);
300        Self::retrieve_from_stream(&mut rdr)
301    }
302}
303
304impl From<SocketAddr> for Address {
305    fn from(addr: SocketAddr) -> Self {
306        Address::SocketAddress(addr)
307    }
308}
309
310impl From<&SocketAddr> for Address {
311    fn from(addr: &SocketAddr) -> Self {
312        Address::SocketAddress(*addr)
313    }
314}
315
316impl From<(Ipv4Addr, u16)> for Address {
317    fn from((addr, port): (Ipv4Addr, u16)) -> Self {
318        Address::SocketAddress(SocketAddr::from((addr, port)))
319    }
320}
321
322impl From<(Ipv6Addr, u16)> for Address {
323    fn from((addr, port): (Ipv6Addr, u16)) -> Self {
324        Address::SocketAddress(SocketAddr::from((addr, port)))
325    }
326}
327
328impl From<(IpAddr, u16)> for Address {
329    fn from((addr, port): (IpAddr, u16)) -> Self {
330        Address::SocketAddress(SocketAddr::from((addr, port)))
331    }
332}
333
334impl From<(String, u16)> for Address {
335    fn from((addr, port): (String, u16)) -> Self {
336        Address::from((addr.as_str(), port))
337    }
338}
339
340impl From<(&String, u16)> for Address {
341    fn from((addr, port): (&String, u16)) -> Self {
342        Address::from((addr.as_str(), port))
343    }
344}
345
346impl From<(&str, u16)> for Address {
347    fn from((addr, port): (&str, u16)) -> Self {
348        if let Ok(addr) = addr.parse::<IpAddr>() {
349            return Address::SocketAddress(SocketAddr::from((addr, port)));
350        }
351        Address::DomainAddress(addr.into(), port)
352    }
353}
354
355impl From<&Address> for Address {
356    fn from(addr: &Address) -> Self {
357        addr.clone()
358    }
359}
360
361impl TryFrom<&str> for Address {
362    type Error = crate::Error;
363
364    fn try_from(addr: &str) -> std::result::Result<Self, Self::Error> {
365        if let Ok(addr) = addr.parse::<SocketAddr>() {
366            Ok(Address::SocketAddress(addr))
367        } else {
368            if addr.matches(':').count() > 1 {
369                let err = format!("invalid socket address format: {addr}");
370                return Err(crate::Error::InvalidAddress(err));
371            }
372
373            let Some(pos) = addr.rfind(':') else {
374                let err = format!("missing port in address: {addr}");
375                return Err(crate::Error::InvalidAddress(err));
376            };
377
378            let (addr, port) = (&addr[..pos], &addr[pos + 1..]);
379            let port = port.parse::<u16>()?;
380            Ok(Address::from((addr, port)))
381        }
382    }
383}
384
385impl std::str::FromStr for Address {
386    type Err = crate::Error;
387
388    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
389        Self::try_from(s)
390    }
391}
392
393#[test]
394fn test_address() {
395    let addr = Address::from((Ipv4Addr::new(127, 0, 0, 1), 8080));
396    let mut buf = Vec::new();
397    addr.write_to_buf(&mut buf);
398    assert_eq!(buf, vec![0x01, 127, 0, 0, 1, 0x1f, 0x90]);
399    let addr2 = Address::retrieve_from_stream(&mut Cursor::new(&buf)).unwrap();
400    assert_eq!(addr, addr2);
401
402    let addr = Address::from((Ipv6Addr::new(0x45, 0xff89, 0, 0, 0, 0, 0, 1), 8080));
403    let mut buf = Vec::new();
404    addr.write_to_buf(&mut buf);
405    assert_eq!(buf, vec![0x04, 0, 0x45, 0xff, 0x89, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0x1f, 0x90]);
406    let addr2 = Address::retrieve_from_stream(&mut Cursor::new(&buf)).unwrap();
407    assert_eq!(addr, addr2);
408
409    let addr = Address::from(("sex.com", 8080));
410    let mut buf = Vec::new();
411    addr.write_to_buf(&mut buf);
412    assert_eq!(buf, vec![0x03, 0x07, b's', b'e', b'x', b'.', b'c', b'o', b'm', 0x1f, 0x90]);
413    let addr2 = Address::retrieve_from_stream(&mut Cursor::new(&buf)).unwrap();
414    assert_eq!(addr, addr2);
415
416    let addr = Address::try_from("example.com:8080").unwrap();
417    assert_eq!(addr.host(), "example.com");
418    assert!(Address::try_from("example.com").is_err());
419
420    let addr = Address::try_from("123.45.67.89:8080").unwrap();
421    assert!(addr.is_ipv4());
422    assert!(Address::try_from("123.45.67.89").is_err());
423
424    let addr = Address::try_from("[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:8080").unwrap();
425    assert_eq!(addr.host(), "2001:db8:85a3::8a2e:370:7334");
426    assert!(addr.is_ipv6());
427
428    assert!(Address::try_from("2001:0db8:85a3:0000:0000:8a2e:0370:7334:8080").is_err());
429}
430
431#[test]
432fn test_try_from_address_for_socketaddr() {
433    let addr = Address::from((Ipv4Addr::new(127, 0, 0, 1), 8080));
434    let socket_addr = SocketAddr::try_from(addr.clone()).unwrap();
435    assert_eq!(socket_addr, SocketAddr::from(([127, 0, 0, 1], 8080)));
436
437    let addr = Address::from(("192.168.1.1", 8080));
438    let socket_addr = SocketAddr::try_from(addr).unwrap();
439    assert_eq!(socket_addr, SocketAddr::from(([192, 168, 1, 1], 8080)));
440
441    let addr = Address::from(("example.com", 8080));
442    let err = SocketAddr::try_from(addr).unwrap_err();
443    assert!(format!("{err}").contains("cannot be converted to a SocketAddr without DNS"));
444
445    let addr = Address::from(("baidu.com", 8080));
446    let err = SocketAddr::try_from(addr).unwrap_err();
447    assert!(format!("{err}").contains("cannot be converted to a SocketAddr without DNS"));
448
449    let addr = Address::from(("baidu.com", 8080));
450    let socket_addrs: Vec<SocketAddr> = addr.to_socket_addrs().unwrap().collect();
451    assert!(!socket_addrs.is_empty());
452}
453
454#[cfg(feature = "tokio")]
455#[tokio::test]
456async fn test_address_async() {
457    let addr = Address::from((Ipv4Addr::new(127, 0, 0, 1), 8080));
458    let mut buf = Vec::new();
459    addr.write_to_async_stream(&mut buf).await.unwrap();
460    assert_eq!(buf, vec![0x01, 127, 0, 0, 1, 0x1f, 0x90]);
461    let addr2 = Address::retrieve_from_async_stream(&mut Cursor::new(&buf)).await.unwrap();
462    assert_eq!(addr, addr2);
463
464    let addr = Address::from((Ipv6Addr::new(0x45, 0xff89, 0, 0, 0, 0, 0, 1), 8080));
465    let mut buf = Vec::new();
466    addr.write_to_async_stream(&mut buf).await.unwrap();
467    assert_eq!(buf, vec![0x04, 0, 0x45, 0xff, 0x89, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0x1f, 0x90]);
468    let addr2 = Address::retrieve_from_async_stream(&mut Cursor::new(&buf)).await.unwrap();
469    assert_eq!(addr, addr2);
470
471    let addr = Address::from(("sex.com", 8080));
472    let mut buf = Vec::new();
473    addr.write_to_async_stream(&mut buf).await.unwrap();
474    assert_eq!(buf, vec![0x03, 0x07, b's', b'e', b'x', b'.', b'c', b'o', b'm', 0x1f, 0x90]);
475    let addr2 = Address::retrieve_from_async_stream(&mut Cursor::new(&buf)).await.unwrap();
476    assert_eq!(addr, addr2);
477}
socks5_impl/protocol/address.rs

socks5_impl/protocol/
address.rs
