Skip to main content

lzc_sdk/remotesocks/
address.rs

1use std::fmt;
2use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
3use std::str::FromStr;
4
5use base64::Engine as _;
6use base64::engine::general_purpose::STANDARD_NO_PAD;
7use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
8
9use crate::Error;
10
11const IPV4_ADDRESS: u8 = 0x01;
12const DOMAIN_ADDRESS: u8 = 0x03;
13const IPV6_ADDRESS: u8 = 0x04;
14const CUSTOM_SUFFIX: &str = ".custom";
15
16/// An RFC 1928 SOCKS address or `LazyCat` custom-network address.
17#[derive(Clone, Debug, Eq, Hash, PartialEq)]
18pub enum SocksAddress {
19    /// Numeric IPv4 or IPv6 socket address.
20    Ip(SocketAddr),
21    /// Domain name and port resolved by the SOCKS server.
22    Domain {
23        /// Domain name carried on the wire without DNS resolution.
24        host: String,
25        /// Destination port.
26        port: u16,
27    },
28    /// `LazyCat` custom network encoded into the SOCKS domain field.
29    Custom {
30        /// Custom network identifier, such as `unix`.
31        network: String,
32        /// Address interpreted by the custom network implementation.
33        address: String,
34    },
35}
36
37impl SocksAddress {
38    /// Construct a domain address without resolving it locally.
39    ///
40    /// # Errors
41    ///
42    /// Returns an error for an empty host or a host longer than 255 bytes.
43    pub fn domain(host: impl Into<String>, port: u16) -> Result<Self, Error> {
44        let host = host.into();
45        validate_wire_string(&host)?;
46        if host.contains(':')
47            || host
48                .chars()
49                .any(|character| character.is_whitespace() || character.is_control())
50        {
51            return Err(Error::InvalidSocksAddress);
52        }
53        Ok(Self::Domain { host, port })
54    }
55
56    /// Construct a `LazyCat` custom-network address.
57    ///
58    /// # Errors
59    ///
60    /// Returns an error for an empty network or when the encoded custom domain
61    /// is longer than 255 bytes.
62    pub fn custom(network: impl Into<String>, address: impl Into<String>) -> Result<Self, Error> {
63        let network = network.into();
64        let address = address.into();
65        if network.is_empty() {
66            return Err(Error::InvalidSocksAddress);
67        }
68        let encoded = encode_custom_domain(&network, &address);
69        validate_wire_string(&encoded)?;
70        Ok(Self::Custom { network, address })
71    }
72
73    /// Encode this address using the RFC 1928 wire representation.
74    ///
75    /// # Errors
76    ///
77    /// Returns an error if a string address cannot fit in the one-byte length
78    /// field.
79    pub fn encode(&self) -> Result<Vec<u8>, Error> {
80        let mut encoded = Vec::with_capacity(19);
81        match self {
82            Self::Ip(SocketAddr::V4(address)) => {
83                encoded.push(IPV4_ADDRESS);
84                encoded.extend_from_slice(&address.ip().octets());
85                encoded.extend_from_slice(&address.port().to_be_bytes());
86            }
87            Self::Ip(SocketAddr::V6(address)) => {
88                encoded.push(IPV6_ADDRESS);
89                encoded.extend_from_slice(&address.ip().octets());
90                encoded.extend_from_slice(&address.port().to_be_bytes());
91            }
92            Self::Domain { host, port } => {
93                encode_domain(&mut encoded, host, *port)?;
94            }
95            Self::Custom { network, address } => {
96                encode_domain(&mut encoded, &encode_custom_domain(network, address), 0)?;
97            }
98        }
99        Ok(encoded)
100    }
101
102    /// Decode exactly one RFC 1928 address.
103    ///
104    /// # Errors
105    ///
106    /// Returns an error for truncated data, trailing bytes, invalid UTF-8,
107    /// malformed custom addresses, or unsupported address types.
108    pub fn decode(encoded: &[u8]) -> Result<Self, Error> {
109        let (address, consumed) = Self::decode_prefix(encoded)?;
110        if consumed != encoded.len() {
111            return Err(Error::InvalidSocksAddress);
112        }
113        Ok(address)
114    }
115
116    pub(crate) fn decode_prefix(encoded: &[u8]) -> Result<(Self, usize), Error> {
117        let (&address_type, remainder) = encoded.split_first().ok_or(Error::InvalidSocksAddress)?;
118        match address_type {
119            IPV4_ADDRESS => {
120                if remainder.len() < 6 {
121                    return Err(Error::InvalidSocksAddress);
122                }
123                let ip = Ipv4Addr::new(remainder[0], remainder[1], remainder[2], remainder[3]);
124                let port = u16::from_be_bytes([remainder[4], remainder[5]]);
125                Ok((Self::Ip(SocketAddr::new(IpAddr::V4(ip), port)), 7))
126            }
127            IPV6_ADDRESS => {
128                if remainder.len() < 18 {
129                    return Err(Error::InvalidSocksAddress);
130                }
131                let octets: [u8; 16] = remainder[..16]
132                    .try_into()
133                    .map_err(|_| Error::InvalidSocksAddress)?;
134                let port = u16::from_be_bytes([remainder[16], remainder[17]]);
135                Ok((
136                    Self::Ip(SocketAddr::new(IpAddr::V6(Ipv6Addr::from(octets)), port)),
137                    19,
138                ))
139            }
140            DOMAIN_ADDRESS => {
141                let (&length, _) = remainder.split_first().ok_or(Error::InvalidSocksAddress)?;
142                let consumed = usize::from(length) + 4;
143                if encoded.len() < consumed {
144                    return Err(Error::InvalidSocksAddress);
145                }
146                Ok((decode_domain(&remainder[..consumed - 1])?, consumed))
147            }
148            address_type => Err(Error::UnsupportedSocksAddressType { address_type }),
149        }
150    }
151
152    /// Network name associated with this address.
153    #[must_use]
154    pub fn network(&self) -> &str {
155        match self {
156            Self::Custom { network, .. } => network,
157            Self::Ip(_) | Self::Domain { .. } => "socks5",
158        }
159    }
160
161    /// Destination port, or zero for custom-network addresses.
162    #[must_use]
163    pub const fn port(&self) -> u16 {
164        match self {
165            Self::Ip(address) => address.port(),
166            Self::Domain { port, .. } => *port,
167            Self::Custom { .. } => 0,
168        }
169    }
170
171    pub(crate) async fn read_from(reader: &mut (impl AsyncRead + Unpin)) -> Result<Self, Error> {
172        let address_type = reader.read_u8().await?;
173        let mut encoded = vec![address_type];
174        match address_type {
175            IPV4_ADDRESS => {
176                let mut remainder = [0_u8; 6];
177                reader.read_exact(&mut remainder).await?;
178                encoded.extend_from_slice(&remainder);
179            }
180            IPV6_ADDRESS => {
181                let mut remainder = [0_u8; 18];
182                reader.read_exact(&mut remainder).await?;
183                encoded.extend_from_slice(&remainder);
184            }
185            DOMAIN_ADDRESS => {
186                let length = reader.read_u8().await?;
187                encoded.push(length);
188                let mut remainder = vec![0_u8; usize::from(length) + 2];
189                reader.read_exact(&mut remainder).await?;
190                encoded.extend_from_slice(&remainder);
191            }
192            address_type => return Err(Error::UnsupportedSocksAddressType { address_type }),
193        }
194        Self::decode(&encoded)
195    }
196
197    pub(crate) async fn write_to(
198        &self,
199        writer: &mut (impl AsyncWrite + Unpin),
200    ) -> Result<(), Error> {
201        writer.write_all(&self.encode()?).await?;
202        Ok(())
203    }
204}
205
206impl From<SocketAddr> for SocksAddress {
207    fn from(address: SocketAddr) -> Self {
208        Self::Ip(address)
209    }
210}
211
212impl FromStr for SocksAddress {
213    type Err = Error;
214
215    fn from_str(value: &str) -> Result<Self, Self::Err> {
216        if let Ok(address) = value.parse::<SocketAddr>() {
217            return Ok(Self::Ip(address));
218        }
219        let (host, port) = value.rsplit_once(':').ok_or(Error::InvalidSocksAddress)?;
220        if host.is_empty() || host.starts_with('[') || host.ends_with(']') {
221            return Err(Error::InvalidSocksAddress);
222        }
223        let port = port
224            .parse::<u16>()
225            .map_err(|_| Error::InvalidSocksAddress)?;
226        Self::domain(host, port)
227    }
228}
229
230impl fmt::Display for SocksAddress {
231    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
232        match self {
233            Self::Ip(address) => address.fmt(formatter),
234            Self::Domain { host, port } => write!(formatter, "{host}:{port}"),
235            Self::Custom { address, .. } => address.fmt(formatter),
236        }
237    }
238}
239
240fn validate_wire_string(value: &str) -> Result<(), Error> {
241    if value.is_empty() {
242        return Err(Error::InvalidSocksAddress);
243    }
244    if value.len() > u8::MAX.into() {
245        return Err(Error::SocksAddressTooLong);
246    }
247    Ok(())
248}
249
250fn encode_domain(encoded: &mut Vec<u8>, domain: &str, port: u16) -> Result<(), Error> {
251    validate_wire_string(domain)?;
252    encoded.push(DOMAIN_ADDRESS);
253    encoded.push(u8::try_from(domain.len()).map_err(|_| Error::SocksAddressTooLong)?);
254    encoded.extend_from_slice(domain.as_bytes());
255    encoded.extend_from_slice(&port.to_be_bytes());
256    Ok(())
257}
258
259fn decode_domain(remainder: &[u8]) -> Result<SocksAddress, Error> {
260    let (&length, remainder) = remainder.split_first().ok_or(Error::InvalidSocksAddress)?;
261    let length = usize::from(length);
262    if length == 0 || remainder.len() != length + 2 {
263        return Err(Error::InvalidSocksAddress);
264    }
265    let host = std::str::from_utf8(&remainder[..length]).map_err(|_| Error::InvalidSocksAddress)?;
266    let port = u16::from_be_bytes([remainder[length], remainder[length + 1]]);
267    if port == 0 && host.ends_with(CUSTOM_SUFFIX) {
268        let (network, address) = decode_custom_domain(host)?;
269        SocksAddress::custom(network, address)
270    } else {
271        SocksAddress::domain(host, port)
272    }
273}
274
275fn encode_custom_domain(network: &str, address: &str) -> String {
276    format!(
277        "{}.{}{}",
278        STANDARD_NO_PAD.encode(network),
279        STANDARD_NO_PAD.encode(address),
280        CUSTOM_SUFFIX
281    )
282}
283
284fn decode_custom_domain(domain: &str) -> Result<(String, String), Error> {
285    let encoded = domain
286        .strip_suffix(CUSTOM_SUFFIX)
287        .ok_or(Error::InvalidSocksAddress)?;
288    let mut fields = encoded.split('.');
289    let network = fields.next().ok_or(Error::InvalidSocksAddress)?;
290    let address = fields.next().ok_or(Error::InvalidSocksAddress)?;
291    if network.is_empty() || fields.next().is_some() {
292        return Err(Error::InvalidSocksAddress);
293    }
294    let network = STANDARD_NO_PAD
295        .decode(network)
296        .map_err(|_| Error::InvalidSocksAddress)?;
297    let address = STANDARD_NO_PAD
298        .decode(address)
299        .map_err(|_| Error::InvalidSocksAddress)?;
300    let network = String::from_utf8(network).map_err(|_| Error::InvalidSocksAddress)?;
301    let address = String::from_utf8(address).map_err(|_| Error::InvalidSocksAddress)?;
302    Ok((network, address))
303}