socks5_impl/server/connection/
bind.rs

1use crate::protocol::{Address, AsyncStreamOperation, Reply, Response};
2use std::{
3    marker::PhantomData,
4    net::SocketAddr,
5    pin::Pin,
6    task::{Context, Poll},
7    time::Duration,
8};
9use tokio::{
10    io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf},
11    net::{
12        TcpStream,
13        tcp::{ReadHalf, WriteHalf},
14    },
15};
16
17/// Socks5 command type `Bind`
18///
19/// By [`wait_request`](crate::server::connection::Authenticated::wait_request)
20/// on an [`Authenticated`](crate::server::connection::Authenticated) from SOCKS5 client,
21/// you may get a `Bind<NeedFirstReply>`. After replying the client 2 times
22/// using [`reply()`](crate::server::connection::Bind::reply),
23/// you will get a `Bind<Ready>`, which can be used as a regular async TCP stream.
24///
25/// A `Bind<S>` can be converted to a regular tokio [`TcpStream`](https://docs.rs/tokio/latest/tokio/net/struct.TcpStream.html) by using the `From` trait.
26#[derive(Debug)]
27pub struct Bind<S> {
28    stream: TcpStream,
29    _state: PhantomData<S>,
30}
31
32/// Marker type indicating that the connection needs its first reply.
33#[derive(Debug, Default)]
34pub struct NeedFirstReply;
35
36/// Marker type indicating that the connection needs its second reply.
37#[derive(Debug, Default)]
38pub struct NeedSecondReply;
39
40/// Marker type indicating that the connection is ready to use as a regular TCP stream.
41#[derive(Debug, Default)]
42pub struct Ready;
43
44impl Bind<NeedFirstReply> {
45    #[inline]
46    pub(super) fn new(stream: TcpStream) -> Self {
47        Self {
48            stream,
49            _state: PhantomData,
50        }
51    }
52
53    /// Reply to the SOCKS5 client with the given reply and address.
54    ///
55    /// If encountered an error while writing the reply, the error alongside the original `TcpStream` is returned.
56    pub async fn reply(mut self, reply: Reply, addr: Address) -> std::io::Result<Bind<NeedSecondReply>> {
57        let resp = Response::new(reply, addr);
58        resp.write_to_async_stream(&mut self.stream).await?;
59        Ok(Bind::<NeedSecondReply>::new(self.stream))
60    }
61
62    /// 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.
63    #[inline]
64    pub async fn shutdown(&mut self) -> std::io::Result<()> {
65        self.stream.shutdown().await
66    }
67
68    /// Returns the local address that this stream is bound to.
69    #[inline]
70    pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
71        self.stream.local_addr()
72    }
73
74    /// Returns the remote address that this stream is connected to.
75    #[inline]
76    pub fn peer_addr(&self) -> std::io::Result<SocketAddr> {
77        self.stream.peer_addr()
78    }
79
80    /// Gets the value of the `TCP_NODELAY` option on this socket.
81    ///
82    /// For more information about this option, see [`set_nodelay`](crate::server::connection::Bind::set_nodelay).
83    #[inline]
84    pub fn nodelay(&self) -> std::io::Result<bool> {
85        self.stream.nodelay()
86    }
87
88    /// Sets the value of the `TCP_NODELAY` option on this socket.
89    ///
90    /// If set, this option disables the Nagle algorithm. This means that segments are always sent as soon as possible,
91    /// 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,
92    /// thereby avoiding the frequent sending of small packets.
93    pub fn set_nodelay(&self, nodelay: bool) -> std::io::Result<()> {
94        self.stream.set_nodelay(nodelay)
95    }
96
97    /// Gets the value of the `IP_TTL` option for this socket.
98    ///
99    /// For more information about this option, see [`set_ttl`](crate::server::connection::Bind::set_ttl).
100    pub fn ttl(&self) -> std::io::Result<u32> {
101        self.stream.ttl()
102    }
103
104    /// Sets the value for the `IP_TTL` option on this socket.
105    ///
106    /// This value sets the time-to-live field that is used in every packet sent from this socket.
107    pub fn set_ttl(&self, ttl: u32) -> std::io::Result<()> {
108        self.stream.set_ttl(ttl)
109    }
110}
111
112impl Bind<NeedSecondReply> {
113    #[inline]
114    fn new(stream: TcpStream) -> Self {
115        Self {
116            stream,
117            _state: PhantomData,
118        }
119    }
120
121    /// Reply to the SOCKS5 client with the given reply and address.
122    ///
123    /// If encountered an error while writing the reply, the error alongside the original `TcpStream` is returned.
124    pub async fn reply(mut self, reply: Reply, addr: Address) -> Result<Bind<Ready>, (std::io::Error, TcpStream)> {
125        let resp = Response::new(reply, addr);
126
127        if let Err(err) = resp.write_to_async_stream(&mut self.stream).await {
128            return Err((err, self.stream));
129        }
130
131        Ok(Bind::<Ready>::new(self.stream))
132    }
133
134    /// 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.
135    #[inline]
136    pub async fn shutdown(&mut self) -> std::io::Result<()> {
137        self.stream.shutdown().await
138    }
139
140    /// Returns the local address that this stream is bound to.
141    #[inline]
142    pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
143        self.stream.local_addr()
144    }
145
146    /// Returns the remote address that this stream is connected to.
147    #[inline]
148    pub fn peer_addr(&self) -> std::io::Result<SocketAddr> {
149        self.stream.peer_addr()
150    }
151
152    /// Gets the value of the `TCP_NODELAY` option on this socket.
153    ///
154    /// For more information about this option, see
155    /// [`set_nodelay`](crate::server::connection::Bind::set_nodelay).
156    #[inline]
157    pub fn nodelay(&self) -> std::io::Result<bool> {
158        self.stream.nodelay()
159    }
160
161    /// Sets the value of the `TCP_NODELAY` option on this socket.
162    ///
163    /// If set, this option disables the Nagle algorithm. This means that segments are always sent as soon as possible,
164    /// 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,
165    /// thereby avoiding the frequent sending of small packets.
166    pub fn set_nodelay(&self, nodelay: bool) -> std::io::Result<()> {
167        self.stream.set_nodelay(nodelay)
168    }
169
170    /// Gets the value of the `IP_TTL` option for this socket.
171    ///
172    /// For more information about this option, see [`set_ttl`](crate::server::connection::Bind::set_ttl).
173    pub fn ttl(&self) -> std::io::Result<u32> {
174        self.stream.ttl()
175    }
176
177    /// Sets the value for the `IP_TTL` option on this socket.
178    ///
179    /// This value sets the time-to-live field that is used in every packet sent from this socket.
180    pub fn set_ttl(&self, ttl: u32) -> std::io::Result<()> {
181        self.stream.set_ttl(ttl)
182    }
183}
184
185impl Bind<Ready> {
186    #[inline]
187    fn new(stream: TcpStream) -> Self {
188        Self {
189            stream,
190            _state: PhantomData,
191        }
192    }
193
194    /// Split the connection into a read and a write half.
195    #[inline]
196    pub fn split(&mut self) -> (ReadHalf<'_>, WriteHalf<'_>) {
197        self.stream.split()
198    }
199}
200
201impl std::ops::Deref for Bind<Ready> {
202    type Target = TcpStream;
203
204    #[inline]
205    fn deref(&self) -> &Self::Target {
206        &self.stream
207    }
208}
209
210impl std::ops::DerefMut for Bind<Ready> {
211    #[inline]
212    fn deref_mut(&mut self) -> &mut Self::Target {
213        &mut self.stream
214    }
215}
216
217impl AsyncRead for Bind<Ready> {
218    #[inline]
219    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
220        Pin::new(&mut self.stream).poll_read(cx, buf)
221    }
222}
223
224impl AsyncWrite for Bind<Ready> {
225    #[inline]
226    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<std::io::Result<usize>> {
227        Pin::new(&mut self.stream).poll_write(cx, buf)
228    }
229
230    #[inline]
231    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
232        Pin::new(&mut self.stream).poll_flush(cx)
233    }
234
235    #[inline]
236    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
237        Pin::new(&mut self.stream).poll_shutdown(cx)
238    }
239}
240
241impl<S> From<Bind<S>> for TcpStream {
242    #[inline]
243    fn from(conn: Bind<S>) -> Self {
244        conn.stream
245    }
246}
socks5_impl/server/connection/bind.rs

socks5_impl/server/connection/
bind.rs
