socks5_impl/client/
mod.rs

1use crate::{
2    Address, AddressType, AsyncStreamOperation, StreamOperation, UserKey,
3    error::{Error, Result},
4    protocol::{AuthMethod, Command, Reply, Version},
5};
6use std::{
7    fmt::Debug,
8    io::Cursor,
9    net::{SocketAddr, ToSocketAddrs},
10    time::Duration,
11};
12use tokio::{
13    io::{AsyncReadExt, AsyncWriteExt, BufStream},
14    net::{TcpStream, UdpSocket},
15};
16
17#[async_trait::async_trait]
18pub trait Socks5Reader: AsyncReadExt + Unpin {
19    async fn read_version(&mut self) -> Result<()> {
20        let value = Version::try_from(self.read_u8().await?)?;
21        match value {
22            Version::V4 => Err(Error::WrongVersion),
23            Version::V5 => Ok(()),
24        }
25    }
26
27    async fn read_method(&mut self) -> Result<AuthMethod> {
28        let value = AuthMethod::from(self.read_u8().await?);
29        match value {
30            AuthMethod::NoAuth | AuthMethod::UserPass => Ok(value),
31            _ => Err(Error::InvalidAuthMethod(value)),
32        }
33    }
34
35    async fn read_command(&mut self) -> Result<Command> {
36        let value = self.read_u8().await?;
37        Ok(Command::try_from(value)?)
38    }
39
40    async fn read_atyp(&mut self) -> Result<AddressType> {
41        let value = self.read_u8().await?;
42        Ok(AddressType::try_from(value)?)
43    }
44
45    async fn read_reserved(&mut self) -> Result<()> {
46        let value = self.read_u8().await?;
47        match value {
48            0x00 => Ok(()),
49            _ => Err(Error::InvalidReserved(value)),
50        }
51    }
52
53    async fn read_fragment_id(&mut self) -> Result<()> {
54        let value = self.read_u8().await?;
55        if value == 0x00 {
56            Ok(())
57        } else {
58            Err(Error::InvalidFragmentId(value))
59        }
60    }
61
62    async fn read_reply(&mut self) -> Result<()> {
63        let value = self.read_u8().await?;
64        match Reply::try_from(value)? {
65            Reply::Succeeded => Ok(()),
66            reply => Err(format!("{reply}").into()),
67        }
68    }
69
70    async fn read_address(&mut self) -> Result<Address> {
71        Ok(Address::retrieve_from_async_stream(self).await?)
72    }
73
74    async fn read_string(&mut self) -> Result<String> {
75        let len = self.read_u8().await? as usize;
76        let mut str = vec![0; len];
77        self.read_exact(&mut str).await?;
78        let str = String::from_utf8(str)?;
79        Ok(str)
80    }
81
82    async fn read_auth_version(&mut self) -> Result<()> {
83        let value = self.read_u8().await?;
84        if value != 0x01 {
85            return Err(Error::InvalidAuthSubnegotiation(value));
86        }
87        Ok(())
88    }
89
90    async fn read_auth_status(&mut self) -> Result<()> {
91        let value = self.read_u8().await?;
92        if value != 0x00 {
93            return Err(Error::InvalidAuthStatus(value));
94        }
95        Ok(())
96    }
97
98    async fn read_selection_msg(&mut self) -> Result<AuthMethod> {
99        self.read_version().await?;
100        self.read_method().await
101    }
102
103    async fn read_final(&mut self) -> Result<Address> {
104        self.read_version().await?;
105        self.read_reply().await?;
106        self.read_reserved().await?;
107        let addr = self.read_address().await?;
108        Ok(addr)
109    }
110}
111
112#[async_trait::async_trait]
113impl<T: AsyncReadExt + Unpin> Socks5Reader for T {}
114
115#[async_trait::async_trait]
116pub trait Socks5Writer: AsyncWriteExt + Unpin {
117    async fn write_version(&mut self) -> Result<()> {
118        self.write_u8(0x05).await?;
119        Ok(())
120    }
121
122    async fn write_method(&mut self, method: AuthMethod) -> Result<()> {
123        self.write_u8(u8::from(method)).await?;
124        Ok(())
125    }
126
127    async fn write_command(&mut self, command: Command) -> Result<()> {
128        self.write_u8(u8::from(command)).await?;
129        Ok(())
130    }
131
132    async fn write_atyp(&mut self, atyp: AddressType) -> Result<()> {
133        self.write_u8(u8::from(atyp)).await?;
134        Ok(())
135    }
136
137    async fn write_reserved(&mut self) -> Result<()> {
138        self.write_u8(0x00).await?;
139        Ok(())
140    }
141
142    async fn write_fragment_id(&mut self, id: u8) -> Result<()> {
143        self.write_u8(id).await?;
144        Ok(())
145    }
146
147    async fn write_address(&mut self, address: &Address) -> Result<()> {
148        address.write_to_async_stream(self).await?;
149        Ok(())
150    }
151
152    async fn write_string(&mut self, string: &str) -> Result<()> {
153        let bytes = string.as_bytes();
154        if bytes.len() > 255 {
155            return Err("Too long string".into());
156        }
157        self.write_u8(bytes.len() as u8).await?;
158        self.write_all(bytes).await?;
159        Ok(())
160    }
161
162    async fn write_auth_version(&mut self) -> Result<()> {
163        self.write_u8(0x01).await?;
164        Ok(())
165    }
166
167    async fn write_methods(&mut self, methods: &[AuthMethod]) -> Result<()> {
168        let method_count = u8::try_from(methods.len()).map_err(|_| "Too many authentication methods")?;
169        self.write_u8(method_count).await?;
170        for method in methods {
171            self.write_method(*method).await?;
172        }
173        Ok(())
174    }
175
176    async fn write_selection_msg(&mut self, methods: &[AuthMethod]) -> Result<()> {
177        self.write_version().await?;
178        self.write_methods(methods).await?;
179        self.flush().await?;
180        Ok(())
181    }
182
183    async fn write_final(&mut self, command: Command, addr: &Address) -> Result<()> {
184        self.write_version().await?;
185        self.write_command(command).await?;
186        self.write_reserved().await?;
187        self.write_address(addr).await?;
188        self.flush().await?;
189        Ok(())
190    }
191}
192
193#[async_trait::async_trait]
194impl<T: AsyncWriteExt + Unpin> Socks5Writer for T {}
195
196async fn username_password_auth<S>(stream: &mut S, auth: &UserKey) -> Result<()>
197where
198    S: Socks5Writer + Socks5Reader + Send,
199{
200    stream.write_auth_version().await?;
201    stream.write_string(&auth.username).await?;
202    stream.write_string(&auth.password).await?;
203    stream.flush().await?;
204
205    stream.read_auth_version().await?;
206    stream.read_auth_status().await
207}
208
209async fn init<S, A>(stream: &mut S, command: Command, addr: A, auth: Option<UserKey>) -> Result<Address>
210where
211    S: Socks5Writer + Socks5Reader + Send,
212    A: Into<Address>,
213{
214    let addr: Address = addr.into();
215
216    let mut methods = Vec::with_capacity(2);
217    methods.push(AuthMethod::NoAuth);
218    if auth.is_some() {
219        methods.push(AuthMethod::UserPass);
220    }
221    stream.write_selection_msg(&methods).await?;
222    stream.flush().await?;
223
224    let method: AuthMethod = stream.read_selection_msg().await?;
225    match method {
226        AuthMethod::NoAuth => {}
227        AuthMethod::UserPass if auth.is_some() => {
228            username_password_auth(stream, auth.as_ref().unwrap()).await?;
229        }
230        _ => return Err(Error::InvalidAuthMethod(method)),
231    }
232
233    stream.write_final(command, &addr).await?;
234    stream.read_final().await
235}
236
237/// Proxifies a TCP connection. Performs the [`CONNECT`] command under the hood.
238///
239/// [`CONNECT`]: https://tools.ietf.org/html/rfc1928#page-6
240///
241/// ```no_run
242/// # use socks5_impl::Result;
243/// # #[tokio::main(flavor = "current_thread")]
244/// # async fn main() -> Result<()> {
245/// use socks5_impl::client;
246/// use tokio::{io::BufStream, net::TcpStream};
247///
248/// let stream = TcpStream::connect("my-proxy-server.com:54321").await?;
249/// let mut stream = BufStream::new(stream);
250/// client::connect(&mut stream, ("google.com", 80), None).await?;
251///
252/// # Ok(())
253/// # }
254/// ```
255pub async fn connect<S, A>(socket: &mut S, addr: A, auth: Option<UserKey>) -> Result<Address>
256where
257    S: AsyncWriteExt + AsyncReadExt + Send + Unpin,
258    A: Into<Address>,
259{
260    init(socket, Command::Connect, addr, auth).await
261}
262
263/// A listener that accepts TCP connections through a proxy.
264///
265/// ```no_run
266/// # use socks5_impl::Result;
267/// # #[tokio::main(flavor = "current_thread")]
268/// # async fn main() -> Result<()> {
269/// use socks5_impl::client::SocksListener;
270/// use tokio::{io::BufStream, net::TcpStream};
271///
272/// let stream = TcpStream::connect("my-proxy-server.com:54321").await?;
273/// let mut stream = BufStream::new(stream);
274/// let (stream, addr) = SocksListener::bind(stream, ("ftp-server.org", 21), None)
275///     .await?
276///     .accept()
277///     .await?;
278///
279/// # Ok(())
280/// # }
281/// ```
282#[derive(Debug)]
283pub struct SocksListener<S> {
284    stream: S,
285    proxy_addr: Address,
286}
287
288impl<S> SocksListener<S>
289where
290    S: AsyncWriteExt + AsyncReadExt + Send + Unpin,
291{
292    /// Creates `SocksListener`. Performs the [`BIND`] command under the hood.
293    ///
294    /// [`BIND`]: https://tools.ietf.org/html/rfc1928#page-6
295    pub async fn bind<A>(mut stream: S, addr: A, auth: Option<UserKey>) -> Result<Self>
296    where
297        A: Into<Address>,
298    {
299        let addr = init(&mut stream, Command::Bind, addr, auth).await?;
300        Ok(Self { stream, proxy_addr: addr })
301    }
302
303    pub fn proxy_addr(&self) -> &Address {
304        &self.proxy_addr
305    }
306
307    pub async fn accept(mut self) -> Result<(S, Address)> {
308        let addr = self.stream.read_final().await?;
309        Ok((self.stream, addr))
310    }
311}
312
313/// A UDP socket that sends packets through a proxy.
314#[derive(Debug)]
315pub struct SocksDatagram<S> {
316    socket: UdpSocket,
317    proxy_addr: Address,
318    stream: S,
319}
320
321impl<S> SocksDatagram<S>
322where
323    S: AsyncWriteExt + AsyncReadExt + Send + Unpin,
324{
325    /// Creates `SocksDatagram`. Performs [`UDP ASSOCIATE`] under the hood.
326    ///
327    /// [`UDP ASSOCIATE`]: https://tools.ietf.org/html/rfc1928#page-7
328    pub async fn udp_associate(mut stream: S, socket: UdpSocket, auth: Option<UserKey>) -> Result<Self> {
329        let addr = if socket.local_addr()?.is_ipv4() { "0.0.0.0:0" } else { "[::]:0" };
330        let addr = addr.parse::<SocketAddr>()?;
331        let proxy_addr = init(&mut stream, Command::UdpAssociate, addr, auth).await?;
332        let addr = proxy_addr.to_socket_addrs()?.next().ok_or("InvalidAddress")?;
333        socket.connect(addr).await?;
334        Ok(Self {
335            socket,
336            proxy_addr,
337            stream,
338        })
339    }
340
341    /// Returns the address of the associated udp address.
342    pub fn proxy_addr(&self) -> &Address {
343        &self.proxy_addr
344    }
345
346    /// Returns a reference to the underlying udp socket.
347    pub fn get_ref(&self) -> &UdpSocket {
348        &self.socket
349    }
350
351    /// Returns a mutable reference to the underlying udp socket.
352    pub fn get_mut(&mut self) -> &mut UdpSocket {
353        &mut self.socket
354    }
355
356    /// Returns the associated stream and udp socket.
357    pub fn into_inner(self) -> (S, UdpSocket) {
358        (self.stream, self.socket)
359    }
360
361    //  Builds a udp-based client request packet, the format is as follows:
362    //  +----+------+------+----------+----------+----------+
363    //  |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
364    //  +----+------+------+----------+----------+----------+
365    //  | 2  |  1   |  1   | Variable |    2     | Variable |
366    //  +----+------+------+----------+----------+----------+
367    //  The reference link is as follows:
368    //  https://tools.ietf.org/html/rfc1928#page-8
369    //
370    pub async fn build_socks5_udp_datagram(buf: &[u8], addr: &Address) -> Result<Vec<u8>> {
371        let bytes_size = Self::get_buf_size(addr.len(), buf.len());
372        let bytes = Vec::with_capacity(bytes_size);
373
374        let mut cursor = Cursor::new(bytes);
375        cursor.write_reserved().await?;
376        cursor.write_reserved().await?;
377        cursor.write_fragment_id(0x00).await?;
378        cursor.write_address(addr).await?;
379        cursor.write_all(buf).await?;
380
381        let bytes = cursor.into_inner();
382        Ok(bytes)
383    }
384
385    /// Sends data via the udp socket to the given address.
386    pub async fn send_to<A>(&self, buf: &[u8], addr: A) -> Result<usize>
387    where
388        A: Into<Address>,
389    {
390        let addr: Address = addr.into();
391        let bytes = Self::build_socks5_udp_datagram(buf, &addr).await?;
392        Ok(self.socket.send(&bytes).await?)
393    }
394
395    /// Parses the udp-based server response packet, the format is same as the client request packet.
396    async fn parse_socks5_udp_response(bytes: &mut [u8], buf: &mut Vec<u8>) -> Result<(usize, Address)> {
397        let len = bytes.len();
398        let mut cursor = Cursor::new(bytes);
399        cursor.read_reserved().await?;
400        cursor.read_reserved().await?;
401        cursor.read_fragment_id().await?;
402        let addr = cursor.read_address().await?;
403        let header_len = cursor.position() as usize;
404        buf.resize(len - header_len, 0);
405        _ = cursor.read_exact(buf).await?;
406        Ok((len - header_len, addr))
407    }
408
409    /// Receives data from the udp socket and returns the number of bytes read and the origin of the data.
410    pub async fn recv_from(&self, timeout: Duration, buf: &mut Vec<u8>) -> Result<(usize, Address)> {
411        const UDP_MTU: usize = 1500;
412        // let bytes_size = Self::get_buf_size(Address::max_serialized_len(), buf.len());
413        let bytes_size = UDP_MTU;
414        let mut bytes = vec![0; bytes_size];
415        let len = tokio::time::timeout(timeout, self.socket.recv(&mut bytes)).await??;
416        bytes.truncate(len);
417        let (read, addr) = Self::parse_socks5_udp_response(&mut bytes, buf).await?;
418        Ok((read, addr))
419    }
420
421    fn get_buf_size(addr_size: usize, buf_len: usize) -> usize {
422        // reserved + fragment id + addr_size + buf_len
423        2 + 1 + addr_size + buf_len
424    }
425}
426
427pub type GuardTcpStream = BufStream<TcpStream>;
428pub type SocksUdpClient = SocksDatagram<GuardTcpStream>;
429
430#[async_trait::async_trait]
431pub trait UdpClientTrait {
432    async fn send_to<A>(&mut self, buf: &[u8], addr: A) -> Result<usize>
433    where
434        A: Into<Address> + Send + Unpin;
435
436    async fn recv_from(&mut self, timeout: Duration, buf: &mut Vec<u8>) -> Result<(usize, Address)>;
437}
438
439#[async_trait::async_trait]
440impl UdpClientTrait for SocksUdpClient {
441    async fn send_to<A>(&mut self, buf: &[u8], addr: A) -> Result<usize, Error>
442    where
443        A: Into<Address> + Send + Unpin,
444    {
445        SocksDatagram::send_to(self, buf, addr).await
446    }
447
448    async fn recv_from(&mut self, timeout: Duration, buf: &mut Vec<u8>) -> Result<(usize, Address), Error> {
449        SocksDatagram::recv_from(self, timeout, buf).await
450    }
451}
452
453pub async fn create_udp_client<A: Into<SocketAddr>>(proxy_addr: A, auth: Option<UserKey>) -> Result<SocksUdpClient> {
454    let proxy_addr = proxy_addr.into();
455    let client_addr = if proxy_addr.is_ipv4() { "0.0.0.0:0" } else { "[::]:0" };
456    let proxy = TcpStream::connect(proxy_addr).await?;
457    let proxy = BufStream::new(proxy);
458    let client = UdpSocket::bind(client_addr).await?;
459    SocksDatagram::udp_associate(proxy, client, auth).await
460}
461
462pub struct ClientWrapper<C> {
463    client: C,
464}
465
466impl ClientWrapper<SocksUdpClient> {
467    pub async fn transfer_data<A2: Into<Address>>(&self, server_addr: A2, data: &[u8], timeout: Duration) -> Result<Vec<u8>> {
468        let _len = self.send_to(data, server_addr).await?;
469        let mut buf = Vec::with_capacity(data.len());
470        let (_len, _) = self.recv(timeout, &mut buf).await?;
471        Ok(buf)
472    }
473
474    pub async fn send_to<A2: Into<Address>>(&self, data: &[u8], server_addr: A2) -> Result<usize> {
475        self.client.send_to(data, server_addr).await
476    }
477
478    pub async fn recv(&self, timeout: Duration, buf: &mut Vec<u8>) -> Result<(usize, Address)> {
479        self.client.recv_from(timeout, buf).await
480    }
481
482    pub async fn datagram<A1: Into<SocketAddr>>(proxy_addr: A1, auth: Option<UserKey>) -> Result<Self> {
483        let client = create_udp_client(proxy_addr, auth).await?;
484        Ok(Self { client })
485    }
486}
487
488#[cfg(test)]
489mod tests {
490    use crate::{Address, UserKey};
491    use crate::{
492        Error, Result,
493        client::{self, SocksListener, SocksUdpClient, UdpClientTrait},
494    };
495    use std::{
496        net::{SocketAddr, ToSocketAddrs},
497        sync::Arc,
498        time::Duration,
499    };
500    use tokio::{
501        io::{AsyncReadExt, AsyncWriteExt, BufStream},
502        net::{TcpStream, UdpSocket},
503    };
504
505    const PROXY_ADDR: &str = "127.0.0.1:1080";
506    const PROXY_AUTH_ADDR: &str = "127.0.0.1:1081";
507    const DATA: &[u8] = b"Hello, world!";
508
509    async fn connect(addr: &str, auth: Option<UserKey>) {
510        let socket = TcpStream::connect(addr).await.unwrap();
511        let mut socket = BufStream::new(socket);
512        client::connect(&mut socket, Address::from(("baidu.com", 80)), auth).await.unwrap();
513    }
514
515    #[ignore]
516    #[tokio::test]
517    async fn connect_auth() {
518        connect(PROXY_AUTH_ADDR, Some(UserKey::new("hyper", "proxy"))).await;
519    }
520
521    #[ignore]
522    #[tokio::test]
523    async fn connect_no_auth() {
524        connect(PROXY_ADDR, None).await;
525    }
526
527    #[ignore]
528    #[should_panic = "InvalidAuthMethod(NoAcceptableMethods)"]
529    #[tokio::test]
530    async fn connect_no_auth_panic() {
531        connect(PROXY_AUTH_ADDR, None).await;
532    }
533
534    #[ignore]
535    #[tokio::test]
536    async fn bind() {
537        let run_block = async {
538            let server_addr = Address::from(("127.0.0.1", 8000));
539
540            let client = TcpStream::connect(PROXY_ADDR).await?;
541            let client = BufStream::new(client);
542            let client = SocksListener::bind(client, server_addr, None).await?;
543
544            let server_addr = client.proxy_addr.to_socket_addrs()?.next().ok_or("Invalid address")?;
545            let mut server = TcpStream::connect(&server_addr).await?;
546
547            let (mut client, _) = client.accept().await?;
548
549            server.write_all(DATA).await?;
550
551            let mut buf = [0; DATA.len()];
552            client.read_exact(&mut buf).await?;
553            assert_eq!(buf, DATA);
554            Ok::<_, Error>(())
555        };
556        if let Err(e) = run_block.await {
557            println!("{e:?}");
558        }
559    }
560
561    type TestHalves = (Arc<SocksUdpClient>, Arc<SocksUdpClient>);
562
563    #[async_trait::async_trait]
564    impl UdpClientTrait for TestHalves {
565        async fn send_to<A>(&mut self, buf: &[u8], addr: A) -> Result<usize, Error>
566        where
567            A: Into<Address> + Send,
568        {
569            self.1.send_to(buf, addr).await
570        }
571
572        async fn recv_from(&mut self, timeout: Duration, buf: &mut Vec<u8>) -> Result<(usize, Address), Error> {
573            self.0.recv_from(timeout, buf).await
574        }
575    }
576
577    const SERVER_ADDR: &str = "127.0.0.1:23456";
578
579    struct UdpTest<C> {
580        client: C,
581        server: UdpSocket,
582        server_addr: Address,
583    }
584
585    impl<C: UdpClientTrait> UdpTest<C> {
586        async fn test(mut self) {
587            let mut buf = vec![0; DATA.len()];
588            self.client.send_to(DATA, self.server_addr).await.unwrap();
589            let (len, addr) = self.server.recv_from(&mut buf).await.unwrap();
590            assert_eq!(len, buf.len());
591            assert_eq!(buf.as_slice(), DATA);
592
593            let mut buf = vec![0; DATA.len()];
594            self.server.send_to(DATA, addr).await.unwrap();
595            let timeout = Duration::from_secs(5);
596            let (len, _) = self.client.recv_from(timeout, &mut buf).await.unwrap();
597            assert_eq!(len, buf.len());
598            assert_eq!(buf.as_slice(), DATA);
599        }
600    }
601
602    impl UdpTest<SocksUdpClient> {
603        async fn datagram() -> Self {
604            let addr = PROXY_ADDR.parse::<SocketAddr>().unwrap();
605            let client = client::create_udp_client(addr, None).await.unwrap();
606
607            let server_addr: SocketAddr = SERVER_ADDR.parse().unwrap();
608            let server = UdpSocket::bind(server_addr).await.unwrap();
609            let server_addr = Address::from(server_addr);
610
611            Self {
612                client,
613                server,
614                server_addr,
615            }
616        }
617    }
618
619    impl UdpTest<TestHalves> {
620        async fn halves() -> Self {
621            let this = UdpTest::<SocksUdpClient>::datagram().await;
622            let client = Arc::new(this.client);
623            Self {
624                client: (client.clone(), client),
625                server: this.server,
626                server_addr: this.server_addr,
627            }
628        }
629    }
630
631    #[ignore]
632    #[tokio::test]
633    async fn udp_datagram_halves() {
634        UdpTest::halves().await.test().await
635    }
636}
socks5_impl/client/mod.rs

socks5_impl/client/
mod.rs
