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

socks5_impl/protocol/
address.rs
