rd_std/socks5/
client.rs

1use super::protocol::{
2    AuthMethod, AuthRequest, AuthResponse, CommandRequest, CommandResponse, Version,
3};
4
5use super::common::{pack_udp, parse_udp, Address};
6use rd_interface::{
7    async_trait, error::map_other, impl_async_read_write, INet, ITcpStream, IUdpSocket,
8    IntoAddress, IntoDyn, Net, Result, TcpStream, UdpSocket, NOT_IMPLEMENTED,
9};
10use std::net::SocketAddr;
11use tokio::io::{split, AsyncWriteExt, BufWriter};
12
13pub struct Socks5Client {
14    address: String,
15    port: u16,
16    net: Net,
17}
18
19pub struct Socks5TcpStream(TcpStream);
20
21impl_async_read_write!(Socks5TcpStream, 0);
22
23pub struct Socks5UdpSocket(UdpSocket, TcpStream, SocketAddr);
24
25#[async_trait]
26impl IUdpSocket for Socks5UdpSocket {
27    async fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, SocketAddr)> {
28        // 259 is max size of address, atype 1 + domain len 1 + domain 255 + port 2
29        let bytes_size = 259 + buf.len();
30        let mut bytes = vec![0u8; bytes_size];
31        let recv_len = loop {
32            let (len, addr) = self.0.recv_from(&mut bytes).await?;
33            if addr == self.2 {
34                break len;
35            }
36        };
37        bytes.truncate(recv_len);
38
39        let (addr, payload) = parse_udp(&bytes).await?;
40        let to_copy = payload.len().min(buf.len());
41        buf[..to_copy].copy_from_slice(&payload[..to_copy]);
42
43        Ok((to_copy, addr.to_socket_addr()?))
44    }
45
46    async fn send_to(&self, buf: &[u8], addr: rd_interface::Address) -> Result<usize> {
47        let addr: Address = addr.into();
48
49        let bytes = pack_udp(addr, buf).await?;
50
51        self.0.send_to(&bytes, self.2.into()).await
52    }
53
54    async fn local_addr(&self) -> Result<SocketAddr> {
55        self.0.local_addr().await
56    }
57}
58
59#[async_trait]
60impl ITcpStream for Socks5TcpStream {
61    async fn peer_addr(&self) -> Result<SocketAddr> {
62        Err(NOT_IMPLEMENTED)
63    }
64
65    async fn local_addr(&self) -> Result<SocketAddr> {
66        Err(NOT_IMPLEMENTED)
67    }
68}
69
70#[async_trait]
71impl INet for Socks5Client {
72    async fn udp_bind(
73        &self,
74        ctx: &mut rd_interface::Context,
75        addr: rd_interface::Address,
76    ) -> Result<UdpSocket> {
77        let mut socket = self.net.tcp_connect(ctx, self.server()?).await?;
78
79        let req = CommandRequest::udp_associate(addr.clone().into_address()?.into());
80        let resp = self.send_command(&mut socket, req).await?;
81        let client = self.net.udp_bind(ctx, addr.clone()).await?;
82
83        let addr = resp.address.to_socket_addr()?;
84
85        Ok(Socks5UdpSocket(client, socket, addr).into_dyn())
86    }
87    async fn tcp_connect(
88        &self,
89        ctx: &mut rd_interface::Context,
90        addr: rd_interface::Address,
91    ) -> Result<TcpStream> {
92        let mut socket = self.net.tcp_connect(ctx, self.server()?).await?;
93
94        let req = CommandRequest::connect(addr.into_address()?.into());
95        let _resp = self.send_command(&mut socket, req).await?;
96
97        Ok(Socks5TcpStream(socket).into_dyn())
98    }
99
100    async fn tcp_bind(
101        &self,
102        _ctx: &mut rd_interface::Context,
103        _addr: rd_interface::Address,
104    ) -> Result<rd_interface::TcpListener> {
105        Err(rd_interface::Error::NotImplemented)
106    }
107}
108
109impl Socks5Client {
110    pub fn new(net: Net, address: String, port: u16) -> Self {
111        Self { address, port, net }
112    }
113    fn server(&self) -> Result<rd_interface::Address> {
114        (self.address.as_str(), self.port)
115            .into_address()
116            .map_err(Into::into)
117    }
118    async fn send_command(
119        &self,
120        socket: &mut TcpStream,
121        command_req: CommandRequest,
122    ) -> Result<CommandResponse> {
123        let (mut rx, tx) = split(socket);
124        let mut tx = BufWriter::with_capacity(512, tx);
125
126        let version = Version::V5;
127        let auth_req = AuthRequest::new(vec![AuthMethod::Noauth]);
128        version.write(&mut tx).await.map_err(map_other)?;
129        auth_req.write(&mut tx).await.map_err(map_other)?;
130        tx.flush().await?;
131
132        Version::read(&mut rx).await.map_err(map_other)?;
133        let resp = AuthResponse::read(&mut rx).await.map_err(map_other)?;
134        if resp.method() != AuthMethod::Noauth {
135            return Err(rd_interface::Error::Other("Auth failed".to_string().into()));
136        }
137
138        command_req.write(&mut tx).await.map_err(map_other)?;
139        tx.flush().await?;
140
141        let command_resp = CommandResponse::read(rx).await.map_err(map_other)?;
142
143        Ok(command_resp)
144    }
145}