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

socks5_impl/
address.rs
