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 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}