1use std::io::Result;
2use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
3
4use bytes::{BufMut, BytesMut};
5use tokio::io::AsyncReadExt;
6
7use crate::{consts::*, Streamable, ToBytes};
8
9#[rustfmt::skip]
11pub mod consts {
12 pub const SOCKS5_VERSION: u8 = 0x05;
13
14 pub const SOCKS5_CMD_CONNECT: u8 = 0x01;
15 pub const SOCKS5_CMD_BIND: u8 = 0x02;
16 pub const SOCKS5_CMD_ASSOCIATE: u8 = 0x03;
17
18 pub const SOCKS5_ADDRESS_TYPE_IPV4: u8 = 0x01;
19 pub const SOCKS5_ADDRESS_TYPE_DOMAIN_NAME: u8 = 0x03;
20 pub const SOCKS5_ADDRESS_TYPE_IPV6: u8 = 0x04;
21
22 pub const SOCKS5_REPLY_SUCCEEDED: u8 = 0x00;
23 pub const SOCKS5_REPLY_GENERAL_FAILURE: u8 = 0x01;
24 pub const SOCKS5_REPLY_CONNECTION_NOT_ALLOWED: u8 = 0x02;
25 pub const SOCKS5_REPLY_NETWORK_UNREACHABLE: u8 = 0x03;
26 pub const SOCKS5_REPLY_HOST_UNREACHABLE: u8 = 0x04;
27 pub const SOCKS5_REPLY_CONNECTION_REFUSED: u8 = 0x05;
28 pub const SOCKS5_REPLY_TTL_EXPIRED: u8 = 0x06;
29 pub const SOCKS5_REPLY_COMMAND_NOT_SUPPORTED: u8 = 0x07;
30 pub const SOCKS5_REPLY_ADDRESS_TYPE_NOT_SUPPORTED: u8 = 0x08;
31}
32
33#[derive(Debug, Clone, PartialEq)]
34pub enum Method {
35 NoAuthentication,
36 GSSAPI,
37 UsernamePassword,
38 IanaAssigned(u8),
39 ReservedPrivate(u8),
40 NoAcceptableMethod,
41}
42
43impl Method {
44 #[rustfmt::skip]
45 pub fn as_u8(&self) -> u8 {
46 match self {
47 Self::NoAuthentication => 0x00,
48 Self::GSSAPI => 0x01,
49 Self::UsernamePassword => 0x03,
50 Self::IanaAssigned(value) => *value,
51 Self::ReservedPrivate(value) => *value,
52 Self::NoAcceptableMethod => 0xFF,
53 }
54 }
55
56 #[rustfmt::skip]
57 pub fn from_u8(value: u8) -> Self {
58 match value {
59 0x00 => Self::NoAuthentication,
60 0x01 => Self::GSSAPI,
61 0x02 => Self::UsernamePassword,
62 0x03..=0x7F => Self::IanaAssigned(value),
63 0x80..=0xFE => Self::ReservedPrivate(value),
64 0xFF => Self::NoAcceptableMethod,
65 }
66 }
67}
68
69impl Streamable for Vec<Method> {
70 async fn read<T>(stream: &mut T) -> Result<Self>
81 where
82 T: AsyncReadExt + Unpin + Send,
83 {
84 let mut buffer = [0u8; 2];
85 stream.read_exact(&mut buffer).await?;
86
87 let method_num = buffer[1];
88 if method_num == 1 {
89 let method = stream.read_u8().await?;
90 return Ok(vec![Method::from_u8(method)]);
91 }
92
93 let mut methods = vec![0u8; method_num as usize];
94 stream.read_exact(&mut methods).await?;
95
96 let result = methods.into_iter().map(|e| Method::from_u8(e)).collect();
97
98 Ok(result)
99 }
100}
101
102impl ToBytes for Vec<Method> {
103 fn to_bytes(&self) -> BytesMut {
104 let mut bytes = BytesMut::new();
105
106 bytes.put_u8(consts::SOCKS5_VERSION);
107 bytes.put_u8(self.len() as u8);
108
109 for e in self.iter() {
110 bytes.put_u8(e.as_u8());
111 }
112
113 bytes
114 }
115}
116
117impl Streamable for Method {
118 async fn read<T>(stream: &mut T) -> Result<Self>
129 where
130 T: AsyncReadExt + Unpin + Send,
131 {
132 let mut buffer = [0u8; 2];
133 stream.read_exact(&mut buffer).await?;
134
135 Ok(Self::from_u8(buffer[1]))
136 }
137}
138
139impl ToBytes for Method {
140 fn to_bytes(&self) -> BytesMut {
141 let mut bytes = BytesMut::new();
142
143 bytes.put_u8(consts::SOCKS5_VERSION);
144 bytes.put_u8(self.as_u8());
145
146 bytes
147 }
148}
149
150#[derive(Debug, Clone, PartialEq)]
151pub enum Request {
152 Bind(Address),
153 Connect(Address),
154 Associate(Address),
155}
156
157impl Streamable for Request {
158 async fn read<T>(stream: &mut T) -> Result<Self>
185 where
186 T: AsyncReadExt + Unpin + Send,
187 {
188 use std::io::{Error, ErrorKind};
189
190 let mut buffer = [0u8; 3];
191 stream.read_exact(&mut buffer).await?;
192
193 let command = buffer[1];
194 let address = Address::read(stream).await?;
195
196 let result = match command {
197 consts::SOCKS5_CMD_BIND => Request::Bind(address),
198 consts::SOCKS5_CMD_CONNECT => Request::Connect(address),
199 consts::SOCKS5_CMD_ASSOCIATE => Request::Associate(address),
200 _command => {
201 return Err(Error::new(
202 ErrorKind::InvalidData,
203 format!("unsupported socks request command {}", _command),
204 ))
205 }
206 };
207
208 Ok(result)
209 }
210}
211
212impl ToBytes for Request {
213 fn to_bytes(&self) -> BytesMut {
214 let mut bytes = BytesMut::new();
215
216 bytes.put_u8(consts::SOCKS5_VERSION);
217
218 let address_bytes = match self {
219 Self::Connect(address) => {
220 bytes.put_u8(consts::SOCKS5_CMD_CONNECT);
221 address.to_bytes()
222 }
223 Self::Bind(address) => {
224 bytes.put_u8(consts::SOCKS5_CMD_BIND);
225 address.to_bytes()
226 }
227 Self::Associate(address) => {
228 bytes.put_u8(consts::SOCKS5_CMD_ASSOCIATE);
229 address.to_bytes()
230 }
231 };
232
233 bytes.put_u8(0x00);
234 bytes.extend(address_bytes);
235
236 bytes
237 }
238}
239
240#[derive(Debug, Clone, PartialEq)]
241pub enum Address {
242 IPv4(SocketAddrV4),
243 IPv6(SocketAddrV6),
244 Domain(String, u16),
245}
246
247impl Address {
248 pub fn from_socket_address(address: SocketAddr) -> Self {
249 match address {
250 SocketAddr::V4(addr) => Self::IPv4(addr),
251 SocketAddr::V6(addr) => Self::IPv6(addr),
252 }
253 }
254}
255
256impl Streamable for Address {
257 async fn read<T>(stream: &mut T) -> Result<Self>
283 where
284 T: AsyncReadExt + Unpin + Send,
285 {
286 use std::io::{Error, ErrorKind};
287 use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
288
289 let result = match stream.read_u8().await? {
290 consts::SOCKS5_ADDRESS_TYPE_IPV4 => {
291 let mut buffer = [0u8; IPV4_ADDRESS_LENGTH + PORT_LENGTH];
292 stream.read_exact(&mut buffer).await?;
293
294 let ip = Ipv4Addr::new(buffer[0], buffer[1], buffer[2], buffer[3]);
295 let port = ((buffer[4] as u16) << 8) | (buffer[5] as u16);
296
297 Address::IPv4(SocketAddrV4::new(ip, port))
298 }
299
300 consts::SOCKS5_ADDRESS_TYPE_IPV6 => {
301 let mut buffer = [0u8; IPV6_ADDRESS_LENGTH + PORT_LENGTH];
302 stream.read_exact(&mut buffer).await?;
303
304 let ip = Ipv6Addr::new(
305 (buffer[0] as u16) << 8 | buffer[1] as u16,
306 (buffer[2] as u16) << 8 | buffer[3] as u16,
307 (buffer[4] as u16) << 8 | buffer[5] as u16,
308 (buffer[6] as u16) << 8 | buffer[7] as u16,
309 (buffer[8] as u16) << 8 | buffer[9] as u16,
310 (buffer[10] as u16) << 8 | buffer[11] as u16,
311 (buffer[12] as u16) << 8 | buffer[13] as u16,
312 (buffer[14] as u16) << 8 | buffer[15] as u16,
313 );
314 let port = ((buffer[16] as u16) << 8) | (buffer[17] as u16);
315
316 Address::IPv6(SocketAddrV6::new(ip, port, 0, 0))
317 }
318
319 consts::SOCKS5_ADDRESS_TYPE_DOMAIN_NAME => {
320 let domain_len = stream.read_u8().await? as usize;
321
322 let mut buffer = vec![0u8; domain_len + PORT_LENGTH];
323 stream.read_exact(&mut buffer).await?;
324
325 let domain = std::str::from_utf8(&buffer[0..domain_len])
326 .map_err(|_| Error::new(ErrorKind::InvalidData, "invalid domain name"))?;
327
328 let port = ((buffer[domain_len] as u16) << 8) | (buffer[domain_len + 1] as u16);
329
330 Address::Domain(domain.to_string(), port)
331 }
332
333 _address_type => {
334 return Err(Error::new(
335 ErrorKind::InvalidData,
336 format!("unsupported socks request address type {}", _address_type),
337 ))
338 }
339 };
340
341 Ok(result)
342 }
343}
344
345impl ToBytes for Address {
346 fn to_bytes(&self) -> BytesMut {
347 let mut bytes = BytesMut::new();
348
349 match self {
350 Self::Domain(domain, port) => {
351 let domain_bytes = domain.as_bytes();
352 bytes.put_u8(consts::SOCKS5_ADDRESS_TYPE_DOMAIN_NAME);
353 bytes.put_u8(domain_bytes.len() as u8);
354 bytes.extend_from_slice(domain_bytes);
355 bytes.extend_from_slice(&port.to_be_bytes());
356 }
357 Self::IPv4(addr) => {
358 bytes.put_u8(consts::SOCKS5_ADDRESS_TYPE_IPV4);
359 bytes.extend_from_slice(&addr.ip().octets());
360 bytes.extend_from_slice(&addr.port().to_be_bytes());
361 }
362 Self::IPv6(addr) => {
363 bytes.put_u8(consts::SOCKS5_ADDRESS_TYPE_IPV6);
364 bytes.extend_from_slice(&addr.ip().octets());
365 bytes.extend_from_slice(&addr.port().to_be_bytes());
366 }
367 }
368
369 bytes
370 }
371}
372
373#[derive(Debug, Clone)]
374pub enum Response {
375 Success(Address),
376 GeneralFailure,
377 ConnectionNotAllowed,
378 NetworkUnreachable,
379 HostUnreachable,
380 ConnectionRefused,
381 TTLExpired,
382 CommandNotSupported,
383 AddressTypeNotSupported,
384 Unassigned(u8),
385}
386
387impl Response {
388 pub fn unspecified_success() -> Self {
389 use std::net::{Ipv4Addr, SocketAddrV4};
390 use std::sync::OnceLock;
391
392 static ADDRESS: OnceLock<Address> = OnceLock::new();
393 let unspecified_address =
394 ADDRESS.get_or_init(|| Address::IPv4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)));
395
396 Self::Success(unspecified_address.clone())
397 }
398}
399
400impl Streamable for Response {
401 async fn read<T>(stream: &mut T) -> Result<Self>
413 where
414 T: AsyncReadExt + Unpin + Send,
415 {
416 let mut buffer = [0u8; 3];
417 stream.read_exact(&mut buffer).await?;
418
419 let reply = buffer[1];
420 let address = Address::read(stream).await?;
421
422 #[rustfmt::skip]
423 let result = match reply {
424 consts::SOCKS5_REPLY_SUCCEEDED => Self::Success(address),
425 consts::SOCKS5_REPLY_GENERAL_FAILURE => Self::GeneralFailure,
426 consts::SOCKS5_REPLY_CONNECTION_NOT_ALLOWED => Self::ConnectionNotAllowed,
427 consts::SOCKS5_REPLY_NETWORK_UNREACHABLE => Self::NetworkUnreachable,
428 consts::SOCKS5_REPLY_HOST_UNREACHABLE => Self::HostUnreachable,
429 consts::SOCKS5_REPLY_CONNECTION_REFUSED => Self::ConnectionRefused,
430 consts::SOCKS5_REPLY_TTL_EXPIRED => Self::TTLExpired,
431 consts::SOCKS5_REPLY_COMMAND_NOT_SUPPORTED => Self::CommandNotSupported,
432 consts::SOCKS5_REPLY_ADDRESS_TYPE_NOT_SUPPORTED => Self::AddressTypeNotSupported,
433
434 _code => Self::Unassigned(_code),
435 };
436
437 Ok(result)
438 }
439}
440
441impl ToBytes for Response {
442 fn to_bytes(&self) -> BytesMut {
443 use std::net::{Ipv4Addr, SocketAddrV4};
444 use std::sync::OnceLock;
445
446 static ADDRESS: OnceLock<Address> = OnceLock::new();
447 let unspecified_address =
448 ADDRESS.get_or_init(|| Address::IPv4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)));
449
450 let mut bytes = BytesMut::new();
451
452 let (reply, address) = match &self {
453 Self::GeneralFailure
454 | Self::ConnectionNotAllowed
455 | Self::NetworkUnreachable
456 | Self::HostUnreachable
457 | Self::ConnectionRefused
458 | Self::TTLExpired
459 | Self::CommandNotSupported
460 | Self::AddressTypeNotSupported
461 | Self::Unassigned(_) => (self.as_u8(), unspecified_address),
462 Self::Success(address) => (self.as_u8(), address),
463 };
464
465 bytes.put_u8(consts::SOCKS5_VERSION);
466 bytes.put_u8(reply);
467 bytes.put_u8(0x00);
468 bytes.extend(address.to_bytes());
469
470 bytes
471 }
472}
473
474impl Response {
475 #[rustfmt::skip]
476 fn as_u8(&self) -> u8 {
477 match self {
478 Self::Success(_) => consts::SOCKS5_REPLY_SUCCEEDED,
479 Self::GeneralFailure => consts::SOCKS5_REPLY_GENERAL_FAILURE,
480 Self::ConnectionNotAllowed => consts::SOCKS5_REPLY_CONNECTION_NOT_ALLOWED,
481 Self::NetworkUnreachable => consts::SOCKS5_REPLY_NETWORK_UNREACHABLE,
482 Self::HostUnreachable => consts::SOCKS5_REPLY_HOST_UNREACHABLE,
483 Self::ConnectionRefused => consts::SOCKS5_REPLY_CONNECTION_REFUSED,
484 Self::TTLExpired => consts::SOCKS5_REPLY_TTL_EXPIRED,
485 Self::CommandNotSupported => consts::SOCKS5_REPLY_COMMAND_NOT_SUPPORTED,
486 Self::AddressTypeNotSupported => consts::SOCKS5_REPLY_ADDRESS_TYPE_NOT_SUPPORTED,
487 Self::Unassigned(code) => *code
488 }
489 }
490}
491
492#[derive(Debug)]
493pub struct UdpPacket {
494 pub frag: u8,
495 pub address: Address,
496 pub data: BytesMut,
497}
498
499impl UdpPacket {
500 pub fn un_frag(address: Address, data: BytesMut) -> Self {
501 Self {
502 frag: 0,
503 address,
504 data,
505 }
506 }
507}
508
509impl Streamable for UdpPacket {
510 async fn read<T>(stream: &mut T) -> Result<Self>
522 where
523 T: AsyncReadExt + Unpin + Send,
524 {
525 let mut buffer = [0u8; 3];
526 stream.read_exact(&mut buffer).await?;
527
528 let frag = buffer[2];
529 let address = Address::read(stream).await?;
530
531 let mut data = Vec::new();
532 stream.read_to_end(&mut data).await?;
533
534 let data = BytesMut::from(data.as_slice());
535
536 Ok(Self {
537 frag,
538 address,
539 data,
540 })
541 }
542}
543
544impl ToBytes for UdpPacket {
545 fn to_bytes(&self) -> BytesMut {
546 let mut bytes = BytesMut::new();
547
548 bytes.put_u8(0x00);
549 bytes.put_u8(0x00);
550
551 bytes.put_u8(self.frag);
552 bytes.extend(self.address.to_bytes());
553 bytes.extend_from_slice(&self.data);
554 bytes
555 }
556}