socks_lib/socks5/
mod.rs

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// SOCKS5 const
10#[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    /// # Authentication Methods
71    ///
72    /// ## Stream
73    /// ```text
74    ///          +-----+----------+----------+
75    ///          | VER | NMETHODS | METHODS  |
76    ///          +-----+----------+----------+
77    ///          |  1  |    1     | 1 to 255 |
78    ///          +-----+----------+----------+
79    /// ```
80    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    /// # Method
119    ///
120    /// ## Stream
121    /// ```text
122    ///      +-----+--------+
123    ///      | VER | METHOD |
124    ///      +-----+--------+
125    ///      |  1  |   1    |
126    ///      +-----+--------+
127    /// ```
128    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    /// # Request
159    ///
160    /// ## Stream
161    /// ```text
162    ///      +-----+-----+-------+------+----------+----------+
163    ///      | VER | CMD |  RSV  | ATYP | DST.ADDR | DST.PORT |
164    ///      +-----+-----+-------+------+----------+----------+
165    ///      |  1  |  1  | X'00' |  1   | Variable |    2     |
166    ///      +-----+-----+-------+------+----------+----------+
167    /// ```
168    ///
169    /// ## Where
170    /// o  VER    protocol version: X'05'  
171    /// o  CMD  
172    ///    o  CONNECT X'01'  
173    ///    o  BIND X'02'  
174    ///    o  UDP ASSOCIATE X'03'  
175    /// o  RSV    RESERVED  
176    /// o  ATYP   address type of following address  
177    ///    o  IP V4 address: X'01'  
178    ///    o  DOMAINNAME: X'03'  
179    ///    o  IP V6 address: X'04'  
180    /// o  DST.ADDR       desired destination address  
181    /// o  DST.PORT desired destination port in network octet  
182    ///    order  
183    ///
184    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    /// # Address
258    ///
259    /// ## Stream
260    /// ```text
261    ///      +------+----------+----------+
262    ///      | ATYP | DST.ADDR | DST.PORT |
263    ///      +------+----------+----------+
264    ///      |  1   | Variable |    2     |
265    ///      +------+----------+----------+
266    /// ```
267    /// ## DST.ADDR BND.ADDR
268    ///   In an address field (DST.ADDR, BND.ADDR), the ATYP field specifies
269    ///   the type of address contained within the field:
270    ///   
271    /// o ATYP: X'01'
272    ///   the address is a version-4 IP address, with a length of 4 octets
273    ///   
274    /// o ATYP: X'03'
275    ///   the address field contains a fully-qualified domain name.  The first
276    ///   octet of the address field contains the number of octets of name that
277    ///   follow, there is no terminating NUL octet.
278    ///   
279    /// o ATYP: X'04'  
280    ///   the address is a version-6 IP address, with a length of 16 octets.
281    ///     
282    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    /// # Response
402    ///
403    /// ## Stream
404    ///
405    /// ```text
406    ///      +-----+-----+-------+------+----------+----------+
407    ///      | VER | REP |  RSV  | ATYP | BND.ADDR | BND.PORT |
408    ///      +-----+-----+-------+------+----------+----------+
409    ///      |  1  |  1  | X'00' |  1   | Variable |    2     |
410    ///      +-----+-----+-------+------+----------+----------+
411    /// ```
412    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    /// # UDP Packet
511    ///
512    /// ## Stream
513    ///
514    /// ```text
515    ///      +-----+------+------+----------+----------+----------+
516    ///      | RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
517    ///      +-----+------+------+----------+----------+----------+
518    ///      |  2  |  1   |  1   | Variable |    2     | Variable |
519    ///      +-----+------+------+----------+----------+----------+
520    /// ```
521    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}