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#[derive(Clone, Debug, Eq, Hash, PartialEq)]
18pub enum SocksAddress {
19 Ip(SocketAddr),
21 Domain {
23 host: String,
25 port: u16,
27 },
28 Custom {
30 network: String,
32 address: String,
34 },
35}
36
37impl SocksAddress {
38 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 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 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 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 #[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 #[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}