socks5_impl/server/connection/
associate.rs

1use crate::{
2    Address, AsyncStreamOperation, StreamOperation,
3    protocol::{Reply, Response, UdpHeader},
4};
5use bytes::{Bytes, BytesMut};
6use std::{
7    net::SocketAddr,
8    pin::Pin,
9    sync::atomic::{AtomicUsize, Ordering},
10    task::{Context, Poll},
11};
12use tokio::{
13    io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf},
14    net::{TcpStream, ToSocketAddrs, UdpSocket},
15};
16
17/// Socks5 connection type `UdpAssociate`
18#[derive(Debug)]
19pub struct UdpAssociate<S> {
20    stream: TcpStream,
21    _state: S,
22}
23
24impl<S: Default> UdpAssociate<S> {
25    #[inline]
26    pub(super) fn new(stream: TcpStream) -> Self {
27        Self {
28            stream,
29            _state: S::default(),
30        }
31    }
32
33    /// Reply to the SOCKS5 client with the given reply and address.
34    ///
35    /// If encountered an error while writing the reply, the error alongside the original `TcpStream` is returned.
36    pub async fn reply(mut self, reply: Reply, addr: Address) -> std::io::Result<UdpAssociate<Ready>> {
37        let resp = Response::new(reply, addr);
38        resp.write_to_async_stream(&mut self.stream).await?;
39        Ok(UdpAssociate::<Ready>::new(self.stream))
40    }
41
42    /// Causes the other peer to receive a read of length 0, indicating that no more data will be sent. This only closes the stream in one direction.
43    #[inline]
44    pub async fn shutdown(&mut self) -> std::io::Result<()> {
45        self.stream.shutdown().await
46    }
47
48    /// Returns the local address that this stream is bound to.
49    #[inline]
50    pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
51        self.stream.local_addr()
52    }
53
54    /// Returns the remote address that this stream is connected to.
55    #[inline]
56    pub fn peer_addr(&self) -> std::io::Result<SocketAddr> {
57        self.stream.peer_addr()
58    }
59
60    /// Gets the value of the `TCP_NODELAY` option on this socket.
61    ///
62    /// For more information about this option, see [`set_nodelay`](#method.set_nodelay).
63    #[inline]
64    pub fn nodelay(&self) -> std::io::Result<bool> {
65        self.stream.nodelay()
66    }
67
68    /// Sets the value of the `TCP_NODELAY` option on this socket.
69    ///
70    /// If set, this option disables the Nagle algorithm. This means that segments are always sent as soon as possible,
71    /// even if there is only a small amount of data. When not set, data is buffered until there is a sufficient amount to send out,
72    /// thereby avoiding the frequent sending of small packets.
73    pub fn set_nodelay(&self, nodelay: bool) -> std::io::Result<()> {
74        self.stream.set_nodelay(nodelay)
75    }
76
77    /// Gets the value of the `IP_TTL` option for this socket.
78    ///
79    /// For more information about this option, see [`set_ttl`](#method.set_ttl).
80    pub fn ttl(&self) -> std::io::Result<u32> {
81        self.stream.ttl()
82    }
83
84    /// Sets the value for the `IP_TTL` option on this socket.
85    ///
86    /// This value sets the time-to-live field that is used in every packet sent from this socket.
87    pub fn set_ttl(&self, ttl: u32) -> std::io::Result<()> {
88        self.stream.set_ttl(ttl)
89    }
90}
91
92#[derive(Debug, Default)]
93pub struct NeedReply;
94
95#[derive(Debug, Default)]
96pub struct Ready;
97
98impl UdpAssociate<Ready> {
99    /// Wait until the client closes this TCP connection.
100    ///
101    /// Socks5 protocol defines that when the client closes the TCP connection used to send the associate command,
102    /// the server should release the associated UDP socket.
103    pub async fn wait_until_closed(&mut self) -> std::io::Result<()> {
104        loop {
105            match self.stream.read(&mut [0]).await {
106                Ok(0) => break Ok(()),
107                Ok(_) => {}
108                Err(err) => break Err(err),
109            }
110        }
111    }
112}
113
114impl std::ops::Deref for UdpAssociate<Ready> {
115    type Target = TcpStream;
116
117    #[inline]
118    fn deref(&self) -> &Self::Target {
119        &self.stream
120    }
121}
122
123impl std::ops::DerefMut for UdpAssociate<Ready> {
124    #[inline]
125    fn deref_mut(&mut self) -> &mut Self::Target {
126        &mut self.stream
127    }
128}
129
130impl AsyncRead for UdpAssociate<Ready> {
131    #[inline]
132    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
133        Pin::new(&mut self.stream).poll_read(cx, buf)
134    }
135}
136
137impl AsyncWrite for UdpAssociate<Ready> {
138    #[inline]
139    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<std::io::Result<usize>> {
140        Pin::new(&mut self.stream).poll_write(cx, buf)
141    }
142
143    #[inline]
144    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
145        Pin::new(&mut self.stream).poll_flush(cx)
146    }
147
148    #[inline]
149    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
150        Pin::new(&mut self.stream).poll_shutdown(cx)
151    }
152}
153
154impl<S> From<UdpAssociate<S>> for TcpStream {
155    #[inline]
156    fn from(conn: UdpAssociate<S>) -> Self {
157        conn.stream
158    }
159}
160
161/// This is a helper for managing the associated UDP socket.
162///
163/// It will add the socks5 UDP header to every UDP packet it sends, also try to parse the socks5 UDP header from any UDP packet received.
164///
165/// The receiving buffer size for each UDP packet can be set with [`set_recv_buffer_size()`](#method.set_recv_buffer_size),
166/// and be read with [`get_max_packet_size()`](#method.get_recv_buffer_size).
167///
168/// You can create this struct by using [`AssociatedUdpSocket::from::<(UdpSocket, usize)>()`](#impl-From<UdpSocket>),
169/// the first element of the tuple is the UDP socket, the second element is the receiving buffer size.
170///
171/// This struct can also be revert into a raw tokio UDP socket with [`UdpSocket::from::<AssociatedUdpSocket>()`](#impl-From<AssociatedUdpSocket>).
172///
173/// [`AssociatedUdpSocket`] can be used as the associated UDP socket.
174#[derive(Debug)]
175pub struct AssociatedUdpSocket {
176    socket: UdpSocket,
177    buf_size: AtomicUsize,
178}
179
180impl AssociatedUdpSocket {
181    /// Connects the UDP socket setting the default destination for send() and limiting packets that are read via recv from the address specified in addr.
182    #[inline]
183    pub async fn connect<A: ToSocketAddrs>(&self, addr: A) -> std::io::Result<()> {
184        self.socket.connect(addr).await
185    }
186
187    /// Get the maximum UDP packet size, with socks5 UDP header included.
188    pub fn get_max_packet_size(&self) -> usize {
189        self.buf_size.load(Ordering::Relaxed)
190    }
191
192    /// Set the maximum UDP packet size, with socks5 UDP header included, for adjusting the receiving buffer size.
193    pub fn set_max_packet_size(&self, size: usize) {
194        self.buf_size.store(size, Ordering::Release);
195    }
196
197    /// Receives a socks5 UDP relay packet on the socket from the remote address to which it is connected.
198    /// On success, returns the packet itself, the fragment number and the remote target address.
199    ///
200    /// The [`connect`](#method.connect) method will connect this socket to a remote address.
201    /// This method will fail if the socket is not connected.
202    pub async fn recv(&self) -> std::io::Result<(Bytes, u8, Address)> {
203        loop {
204            let max_packet_size = self.buf_size.load(Ordering::Acquire);
205            let mut buf = vec![0; max_packet_size];
206            let len = self.socket.recv(&mut buf).await?;
207            buf.truncate(len);
208            let pkt = Bytes::from(buf);
209
210            if let Ok(header) = UdpHeader::retrieve_from_async_stream(&mut pkt.as_ref()).await {
211                let pkt = pkt.slice(header.len()..);
212                return Ok((pkt, header.frag, header.address));
213            }
214        }
215    }
216
217    /// Receives a socks5 UDP relay packet on the socket from the any remote address.
218    /// On success, returns the packet itself, the fragment number, the remote target address and the source address.
219    pub async fn recv_from(&self) -> std::io::Result<(Bytes, u8, Address, SocketAddr)> {
220        loop {
221            let max_packet_size = self.buf_size.load(Ordering::Acquire);
222            let mut buf = vec![0; max_packet_size];
223            let (len, src_addr) = self.socket.recv_from(&mut buf).await?;
224            buf.truncate(len);
225            let pkt = Bytes::from(buf);
226
227            if let Ok(header) = UdpHeader::retrieve_from_async_stream(&mut pkt.as_ref()).await {
228                let pkt = pkt.slice(header.len()..);
229                return Ok((pkt, header.frag, header.address, src_addr));
230            }
231        }
232    }
233
234    /// Sends a UDP relay packet to the remote address to which it is connected. The socks5 UDP header will be added to the packet.
235    pub async fn send<P: AsRef<[u8]>>(&self, pkt: P, frag: u8, from_addr: Address) -> std::io::Result<usize> {
236        let header = UdpHeader::new(frag, from_addr);
237        let mut buf = BytesMut::with_capacity(header.len() + pkt.as_ref().len());
238        header.write_to_buf(&mut buf);
239        buf.extend_from_slice(pkt.as_ref());
240
241        self.socket.send(&buf).await.map(|len| len - header.len())
242    }
243
244    /// Sends a UDP relay packet to a specified remote address to which it is connected. The socks5 UDP header will be added to the packet.
245    pub async fn send_to<P: AsRef<[u8]>>(&self, pkt: P, frag: u8, from_addr: Address, to_addr: SocketAddr) -> std::io::Result<usize> {
246        let header = UdpHeader::new(frag, from_addr);
247        let mut buf = BytesMut::with_capacity(header.len() + pkt.as_ref().len());
248        header.write_to_buf(&mut buf);
249        buf.extend_from_slice(pkt.as_ref());
250
251        self.socket.send_to(&buf, to_addr).await.map(|len| len - header.len())
252    }
253}
254
255impl From<(UdpSocket, usize)> for AssociatedUdpSocket {
256    #[inline]
257    fn from(from: (UdpSocket, usize)) -> Self {
258        AssociatedUdpSocket {
259            socket: from.0,
260            buf_size: AtomicUsize::new(from.1),
261        }
262    }
263}
264
265impl From<AssociatedUdpSocket> for UdpSocket {
266    #[inline]
267    fn from(from: AssociatedUdpSocket) -> Self {
268        from.socket
269    }
270}
271
272impl AsRef<UdpSocket> for AssociatedUdpSocket {
273    #[inline]
274    fn as_ref(&self) -> &UdpSocket {
275        &self.socket
276    }
277}
278
279impl AsMut<UdpSocket> for AssociatedUdpSocket {
280    #[inline]
281    fn as_mut(&mut self) -> &mut UdpSocket {
282        &mut self.socket
283    }
284}
socks5_impl/server/connection/associate.rs

socks5_impl/server/connection/
associate.rs
