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

socks5_impl/
address.rs
