socks5_impl/
address.rs

1#[cfg(feature = "tokio")]
2use crate::stream::AsyncStreamOperation;
3use crate::stream::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))]
55#[cfg_attr(feature = "serde", serde(try_from = "String", into = "String"))]
56pub enum Address {
57    /// Represents an IPv4 or IPv6 socket address.
58    SocketAddress(SocketAddr),
59    /// Represents a domain name and a port.
60    DomainAddress(Box<str>, u16),
61}
62
63impl Default for Address {
64    fn default() -> Self {
65        Address::unspecified_v4()
66    }
67}
68
69impl Address {
70    /// Returns an unspecified IPv4 address (0.0.0.0:0).
71    #[deprecated(since = "0.10.3", note = "Use 'unspecified_v4' instead")]
72    pub fn unspecified() -> Self {
73        Address::unspecified_v4()
74    }
75
76    /// Returns an unspecified IPv4 address (0.0.0.0:0).
77    pub fn unspecified_v4() -> Self {
78        Address::SocketAddress(SocketAddr::from((Ipv4Addr::UNSPECIFIED, 0)))
79    }
80
81    /// Returns an unspecified IPv6 address (::0:0).
82    pub fn unspecified_v6() -> Self {
83        Address::SocketAddress(SocketAddr::from((Ipv6Addr::UNSPECIFIED, 0)))
84    }
85
86    /// Returns the type of the address (IPv4, IPv6, or Domain).
87    pub fn get_type(&self) -> AddressType {
88        match self {
89            Self::SocketAddress(SocketAddr::V4(_)) => AddressType::IPv4,
90            Self::SocketAddress(SocketAddr::V6(_)) => AddressType::IPv6,
91            Self::DomainAddress(_, _) => AddressType::Domain,
92        }
93    }
94
95    /// Returns the port number.
96    pub fn port(&self) -> u16 {
97        match self {
98            Self::SocketAddress(addr) => addr.port(),
99            Self::DomainAddress(_, port) => *port,
100        }
101    }
102
103    /// Returns the host name or IP address as a string.
104    pub fn host(&self) -> String {
105        match self {
106            Self::SocketAddress(addr) => addr.ip().to_string(),
107            Self::DomainAddress(addr, _) => addr.to_string(),
108        }
109    }
110
111    /// Returns the host name or IP address as a string.
112    #[deprecated(note = "use host() instead")]
113    pub fn domain(&self) -> String {
114        self.host()
115    }
116
117    /// Returns `true` if it is an IPv4 address.
118    pub fn is_ipv4(&self) -> bool {
119        matches!(self, Self::SocketAddress(SocketAddr::V4(_)))
120    }
121
122    /// Returns `true` if it is an IPv6 address.
123    pub fn is_ipv6(&self) -> bool {
124        matches!(self, Self::SocketAddress(SocketAddr::V6(_)))
125    }
126
127    /// Returns `true` if it is a domain address.
128    pub fn is_domain(&self) -> bool {
129        matches!(self, Self::DomainAddress(_, _))
130    }
131
132    /// Returns the maximum possible serialized length of a SOCKS5 address.
133    pub const fn max_serialized_len() -> usize {
134        1 + 1 + u8::MAX as usize + 2
135    }
136
137    /// Resolves the address using Tokio's asynchronous DNS resolution.
138    /// Returns an iterator over the resolved socket addresses.
139    #[cfg(feature = "tokio")]
140    pub async fn resolve_using_tokio(&self) -> std::io::Result<std::vec::IntoIter<SocketAddr>> {
141        match self {
142            Address::SocketAddress(addr) => Ok(vec![*addr].into_iter()),
143            Address::DomainAddress(host, port) => {
144                let host = host.strip_prefix('[').and_then(|host| host.strip_suffix(']')).unwrap_or(host);
145                if let Ok(ip) = host.parse::<IpAddr>() {
146                    return Ok(vec![SocketAddr::new(ip, *port)].into_iter());
147                }
148                Ok(tokio::net::lookup_host((host, *port)).await?.collect::<Vec<_>>().into_iter())
149            }
150        }
151    }
152
153    #[cfg(feature = "tokio")]
154    #[deprecated(since = "0.10.3", note = "Use 'resolve_using_tokio' instead.")]
155    #[inline]
156    pub async fn resolve_tokio(&self) -> std::io::Result<std::vec::IntoIter<SocketAddr>> {
157        self.resolve_using_tokio().await
158    }
159
160    /// Resolves the address using the standard library's DNS resolution.
161    /// Returns an iterator over the resolved socket addresses.
162    /// # Errors
163    ///
164    /// Returns an `std::io::Error` if address resolution fails.
165    /// # Examples
166    ///
167    /// ```
168    /// let addr = socks5_impl::Address::from(("www.qq.com", 80));
169    /// for socket_addr in addr.resolve().unwrap() {
170    ///     println!("Resolved address: {}", socket_addr);
171    /// }
172    ///
173    /// let addr = socks5_impl::Address::from(("127.0.0.1", 8080));
174    /// for socket_addr in addr.resolve().unwrap() {
175    ///     println!("Resolved address: {}", socket_addr);
176    /// }
177    /// ```
178    pub fn resolve(&self) -> std::io::Result<std::vec::IntoIter<SocketAddr>> {
179        self.to_socket_addrs()
180    }
181}
182
183impl StreamOperation for Address {
184    fn retrieve_from_stream<R: std::io::Read>(stream: &mut R) -> std::io::Result<Self> {
185        let mut atyp_buf = [0; 1];
186        stream.read_exact(&mut atyp_buf)?;
187        match AddressType::try_from(atyp_buf[0])? {
188            AddressType::IPv4 => {
189                let mut buf = [0; 6];
190                stream.read_exact(&mut buf)?;
191                let addr = Ipv4Addr::new(buf[0], buf[1], buf[2], buf[3]);
192                let port = u16::from_be_bytes([buf[4], buf[5]]);
193                Ok(Self::SocketAddress(SocketAddr::from((addr, port))))
194            }
195            AddressType::Domain => {
196                let mut len_buf = [0; 1];
197                stream.read_exact(&mut len_buf)?;
198                let len = len_buf[0] as usize;
199                let mut domain_buf = vec![0; len];
200                stream.read_exact(&mut domain_buf)?;
201
202                let mut port_buf = [0; 2];
203                stream.read_exact(&mut port_buf)?;
204                let port = u16::from_be_bytes(port_buf);
205
206                let addr = String::from_utf8(domain_buf)
207                    .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, format!("Invalid address encoding: {err}")))?;
208                Ok(Self::DomainAddress(addr.into_boxed_str(), port))
209            }
210            AddressType::IPv6 => {
211                let mut buf = [0; 18];
212                stream.read_exact(&mut buf)?;
213                let port = u16::from_be_bytes([buf[16], buf[17]]);
214                let mut addr_bytes = [0; 16];
215                addr_bytes.copy_from_slice(&buf[..16]);
216                Ok(Self::SocketAddress(SocketAddr::from((Ipv6Addr::from(addr_bytes), port))))
217            }
218        }
219    }
220
221    fn write_to_buf<B: BufMut>(&self, buf: &mut B) {
222        match self {
223            Self::SocketAddress(SocketAddr::V4(addr)) => {
224                buf.put_u8(AddressType::IPv4.into());
225                buf.put_slice(&addr.ip().octets());
226                buf.put_u16(addr.port());
227            }
228            Self::SocketAddress(SocketAddr::V6(addr)) => {
229                buf.put_u8(AddressType::IPv6.into());
230                buf.put_slice(&addr.ip().octets());
231                buf.put_u16(addr.port());
232            }
233            Self::DomainAddress(addr, port) => {
234                let addr = addr.as_bytes();
235                buf.put_u8(AddressType::Domain.into());
236                buf.put_u8(u8::try_from(addr.len()).expect("domain address too long for SOCKS5"));
237                buf.put_slice(addr);
238                buf.put_u16(*port);
239            }
240        }
241    }
242
243    fn len(&self) -> usize {
244        match self {
245            Address::SocketAddress(SocketAddr::V4(_)) => 1 + 4 + 2,
246            Address::SocketAddress(SocketAddr::V6(_)) => 1 + 16 + 2,
247            Address::DomainAddress(addr, _) => 1 + 1 + addr.len() + 2,
248        }
249    }
250}
251
252#[cfg(feature = "tokio")]
253#[async_trait::async_trait]
254impl AsyncStreamOperation for Address {
255    async fn retrieve_from_async_stream<R>(stream: &mut R) -> std::io::Result<Self>
256    where
257        R: AsyncRead + Unpin + Send + ?Sized,
258    {
259        let atyp = stream.read_u8().await?;
260        match AddressType::try_from(atyp)? {
261            AddressType::IPv4 => {
262                let mut addr_bytes = [0; 4];
263                stream.read_exact(&mut addr_bytes).await?;
264                let port = stream.read_u16().await?;
265                let addr = Ipv4Addr::from(addr_bytes);
266                Ok(Self::SocketAddress(SocketAddr::from((addr, port))))
267            }
268            AddressType::Domain => {
269                let len = stream.read_u8().await? as usize;
270                let mut domain_buf = vec![0; len];
271                stream.read_exact(&mut domain_buf).await?;
272                let port = stream.read_u16().await?;
273
274                let addr = String::from_utf8(domain_buf)
275                    .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, format!("Invalid address encoding: {err}")))?;
276                Ok(Self::DomainAddress(addr.into_boxed_str(), port))
277            }
278            AddressType::IPv6 => {
279                let mut addr_bytes = [0; 16];
280                stream.read_exact(&mut addr_bytes).await?;
281                let port = stream.read_u16().await?;
282                Ok(Self::SocketAddress(SocketAddr::from((Ipv6Addr::from(addr_bytes), port))))
283            }
284        }
285    }
286}
287
288impl ToSocketAddrs for Address {
289    type Iter = std::vec::IntoIter<SocketAddr>;
290
291    fn to_socket_addrs(&self) -> std::io::Result<Self::Iter> {
292        match self {
293            Address::SocketAddress(addr) => Ok(vec![*addr].into_iter()),
294            Address::DomainAddress(host, port) => {
295                let host = host.strip_prefix('[').and_then(|host| host.strip_suffix(']')).unwrap_or(host);
296                if let Ok(ip) = host.parse::<IpAddr>() {
297                    return Ok(vec![SocketAddr::new(ip, *port)].into_iter());
298                }
299                Ok((host, *port).to_socket_addrs()?)
300            }
301        }
302    }
303}
304
305impl std::fmt::Display for Address {
306    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
307        match self {
308            Address::DomainAddress(hostname, port) => write!(f, "{hostname}:{port}"),
309            Address::SocketAddress(socket_addr) => write!(f, "{socket_addr}"),
310        }
311    }
312}
313
314impl TryFrom<String> for Address {
315    type Error = crate::Error;
316
317    fn try_from(addr: String) -> std::result::Result<Self, Self::Error> {
318        Self::try_from(addr.as_str())
319    }
320}
321
322impl TryFrom<&String> for Address {
323    type Error = crate::Error;
324
325    fn try_from(addr: &String) -> std::result::Result<Self, Self::Error> {
326        Self::try_from(addr.as_str())
327    }
328}
329
330impl From<Address> for String {
331    fn from(addr: Address) -> Self {
332        addr.to_string()
333    }
334}
335
336impl From<&Address> for String {
337    fn from(addr: &Address) -> Self {
338        addr.to_string()
339    }
340}
341
342impl TryFrom<&Address> for SocketAddr {
343    type Error = std::io::Error;
344
345    fn try_from(address: &Address) -> std::result::Result<Self, Self::Error> {
346        match address {
347            Address::SocketAddress(addr) => Ok(*addr),
348            Address::DomainAddress(host, port) => {
349                let host = host.strip_prefix('[').and_then(|host| host.strip_suffix(']')).unwrap_or(host);
350                if let Ok(addr) = host.parse::<IpAddr>() {
351                    return Ok(SocketAddr::from((addr, *port)));
352                }
353                let err = format!("domain address {host} cannot be converted to a SocketAddr without DNS, use to_socket_addrs instead");
354                Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, err))
355            }
356        }
357    }
358}
359
360impl TryFrom<Address> for SocketAddr {
361    type Error = std::io::Error;
362
363    fn try_from(address: Address) -> std::result::Result<Self, Self::Error> {
364        Self::try_from(&address)
365    }
366}
367
368impl From<Address> for Vec<u8> {
369    fn from(addr: Address) -> Self {
370        let mut buf = Vec::with_capacity(addr.len());
371        addr.write_to_buf(&mut buf);
372        buf
373    }
374}
375
376impl TryFrom<Vec<u8>> for Address {
377    type Error = std::io::Error;
378
379    fn try_from(data: Vec<u8>) -> std::result::Result<Self, Self::Error> {
380        let mut rdr = Cursor::new(data);
381        Self::retrieve_from_stream(&mut rdr)
382    }
383}
384
385impl TryFrom<&[u8]> for Address {
386    type Error = std::io::Error;
387
388    fn try_from(data: &[u8]) -> std::result::Result<Self, Self::Error> {
389        let mut rdr = Cursor::new(data);
390        Self::retrieve_from_stream(&mut rdr)
391    }
392}
393
394impl From<SocketAddr> for Address {
395    fn from(addr: SocketAddr) -> Self {
396        Address::SocketAddress(addr)
397    }
398}
399
400impl From<&SocketAddr> for Address {
401    fn from(addr: &SocketAddr) -> Self {
402        Address::SocketAddress(*addr)
403    }
404}
405
406impl From<(Ipv4Addr, u16)> for Address {
407    fn from((addr, port): (Ipv4Addr, u16)) -> Self {
408        Address::SocketAddress(SocketAddr::from((addr, port)))
409    }
410}
411
412impl From<(Ipv6Addr, u16)> for Address {
413    fn from((addr, port): (Ipv6Addr, u16)) -> Self {
414        Address::SocketAddress(SocketAddr::from((addr, port)))
415    }
416}
417
418impl From<(IpAddr, u16)> for Address {
419    fn from((addr, port): (IpAddr, u16)) -> Self {
420        Address::SocketAddress(SocketAddr::from((addr, port)))
421    }
422}
423
424impl From<(String, u16)> for Address {
425    fn from((addr, port): (String, u16)) -> Self {
426        Address::from((addr.as_str(), port))
427    }
428}
429
430impl From<(&String, u16)> for Address {
431    fn from((addr, port): (&String, u16)) -> Self {
432        Address::from((addr.as_str(), port))
433    }
434}
435
436impl From<(&str, u16)> for Address {
437    fn from((addr, port): (&str, u16)) -> Self {
438        if let Ok(addr) = addr.parse::<IpAddr>() {
439            return Address::SocketAddress(SocketAddr::from((addr, port)));
440        }
441        Address::DomainAddress(addr.into(), port)
442    }
443}
444
445impl From<&Address> for Address {
446    fn from(addr: &Address) -> Self {
447        addr.clone()
448    }
449}
450
451impl TryFrom<&str> for Address {
452    type Error = crate::Error;
453
454    fn try_from(addr: &str) -> std::result::Result<Self, Self::Error> {
455        if let Ok(addr) = addr.parse::<SocketAddr>() {
456            Ok(Address::SocketAddress(addr))
457        } else {
458            if addr.matches(':').count() > 1 {
459                let err = format!("invalid socket address format: {addr}");
460                return Err(crate::Error::InvalidAddress(err));
461            }
462
463            let Some(pos) = addr.rfind(':') else {
464                let err = format!("missing port in address: {addr}");
465                return Err(crate::Error::InvalidAddress(err));
466            };
467
468            let (addr, port) = (&addr[..pos], &addr[pos + 1..]);
469            let port = port.parse::<u16>()?;
470            Ok(Address::from((addr, port)))
471        }
472    }
473}
474
475impl std::str::FromStr for Address {
476    type Err = crate::Error;
477
478    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
479        Self::try_from(s)
480    }
481}
482
483#[test]
484fn test_address() {
485    let addr = Address::from((Ipv4Addr::new(127, 0, 0, 1), 8080));
486    let mut buf = Vec::new();
487    addr.write_to_buf(&mut buf);
488    assert_eq!(buf, vec![0x01, 127, 0, 0, 1, 0x1f, 0x90]);
489    let addr2 = Address::retrieve_from_stream(&mut Cursor::new(&buf)).unwrap();
490    assert_eq!(addr, addr2);
491
492    let addr = Address::from((Ipv6Addr::new(0x45, 0xff89, 0, 0, 0, 0, 0, 1), 8080));
493    let mut buf = Vec::new();
494    addr.write_to_buf(&mut buf);
495    assert_eq!(buf, vec![0x04, 0, 0x45, 0xff, 0x89, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0x1f, 0x90]);
496    let addr2 = Address::retrieve_from_stream(&mut Cursor::new(&buf)).unwrap();
497    assert_eq!(addr, addr2);
498
499    let addr = Address::from(("sex.com", 8080));
500    let mut buf = Vec::new();
501    addr.write_to_buf(&mut buf);
502    assert_eq!(buf, vec![0x03, 0x07, b's', b'e', b'x', b'.', b'c', b'o', b'm', 0x1f, 0x90]);
503    let addr2 = Address::retrieve_from_stream(&mut Cursor::new(&buf)).unwrap();
504    assert_eq!(addr, addr2);
505
506    let addr = Address::try_from("example.com:8080").unwrap();
507    assert_eq!(addr.host(), "example.com");
508    assert!(Address::try_from("example.com").is_err());
509
510    let addr = Address::try_from("123.45.67.89:8080").unwrap();
511    assert!(addr.is_ipv4());
512    assert!(Address::try_from("123.45.67.89").is_err());
513
514    let addr = Address::try_from("[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:8080").unwrap();
515    assert_eq!(addr.host(), "2001:db8:85a3::8a2e:370:7334");
516    assert!(addr.is_ipv6());
517
518    let addr = Address::from((Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1), 8080));
519    assert_eq!(addr.host(), "2001:db8::1");
520
521    assert!(Address::try_from("2001:0db8:85a3:0000:0000:8a2e:0370:7334:8080").is_err());
522}
523
524#[test]
525fn test_address_display_for_ipv6() {
526    let addr = Address::from((Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1), 8080));
527    assert_eq!(addr.to_string(), "[2001:db8::1]:8080");
528}
529
530#[cfg(feature = "serde")]
531#[test]
532fn test_address_serde_string_roundtrip() {
533    for value in ["127.0.0.1:8080", "example.com:8080", "[2001:db8::1]:8080"] {
534        let address = value.parse::<Address>().unwrap();
535        let serialized = serde_json::to_string(&address).unwrap();
536
537        assert_eq!(serialized, serde_json::to_string(&address.to_string()).unwrap());
538        assert_eq!(serde_json::from_str::<Address>(&serialized).unwrap(), address);
539    }
540    assert!(serde_json::from_str::<Address>(r#""example.com""#).is_err());
541}
542
543#[test]
544fn test_try_from_address_for_socketaddr() {
545    let addr = Address::from((Ipv4Addr::new(127, 0, 0, 1), 8080));
546    let socket_addr = SocketAddr::try_from(addr.clone()).unwrap();
547    assert_eq!(socket_addr, SocketAddr::from(([127, 0, 0, 1], 8080)));
548
549    let addr = Address::from(("192.168.1.1", 8080));
550    let socket_addr = SocketAddr::try_from(addr).unwrap();
551    assert_eq!(socket_addr, SocketAddr::from(([192, 168, 1, 1], 8080)));
552
553    let addr = Address::from(("example.com", 8080));
554    let err = SocketAddr::try_from(addr).unwrap_err();
555    assert!(format!("{err}").contains("cannot be converted to a SocketAddr without DNS"));
556
557    let addr = Address::from(("baidu.com", 8080));
558    let err = SocketAddr::try_from(addr).unwrap_err();
559    assert!(format!("{err}").contains("cannot be converted to a SocketAddr without DNS"));
560
561    let addr = Address::from(("baidu.com", 8080));
562    let socket_addrs: Vec<SocketAddr> = addr.to_socket_addrs().unwrap().collect();
563    assert!(!socket_addrs.is_empty());
564}
565
566#[cfg(feature = "tokio")]
567#[tokio::test]
568async fn test_address_async() {
569    let addr = Address::from((Ipv4Addr::new(127, 0, 0, 1), 8080));
570    let mut buf = Vec::new();
571    addr.write_to_async_stream(&mut buf).await.unwrap();
572    assert_eq!(buf, vec![0x01, 127, 0, 0, 1, 0x1f, 0x90]);
573    let addr2 = Address::retrieve_from_async_stream(&mut Cursor::new(&buf)).await.unwrap();
574    assert_eq!(addr, addr2);
575
576    let addr = Address::from((Ipv6Addr::new(0x45, 0xff89, 0, 0, 0, 0, 0, 1), 8080));
577    let mut buf = Vec::new();
578    addr.write_to_async_stream(&mut buf).await.unwrap();
579    assert_eq!(buf, vec![0x04, 0, 0x45, 0xff, 0x89, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0x1f, 0x90]);
580    let addr2 = Address::retrieve_from_async_stream(&mut Cursor::new(&buf)).await.unwrap();
581    assert_eq!(addr, addr2);
582
583    let addr = Address::from(("sex.com", 8080));
584    let mut buf = Vec::new();
585    addr.write_to_async_stream(&mut buf).await.unwrap();
586    assert_eq!(buf, vec![0x03, 0x07, b's', b'e', b'x', b'.', b'c', b'o', b'm', 0x1f, 0x90]);
587    let addr2 = Address::retrieve_from_async_stream(&mut Cursor::new(&buf)).await.unwrap();
588    assert_eq!(addr, addr2);
589}
socks5_impl/address.rs

socks5_impl/
address.rs
