socks5_impl/server/connection/
mod.rs

1use self::{associate::UdpAssociate, bind::Bind, connect::Connect};
2use crate::{
3    protocol::{self, Address, AsyncStreamOperation, AuthMethod, Command, handshake},
4    server::AuthAdaptor,
5};
6use std::{net::SocketAddr, time::Duration};
7use tokio::{io::AsyncWriteExt, net::TcpStream};
8
9pub mod associate;
10pub mod bind;
11pub mod connect;
12
13/// An incoming connection. This may not be a valid socks5 connection. You need to call [`authenticate()`](#method.authenticate)
14/// to perform the socks5 handshake. It will be converted to a proper socks5 connection after the handshake succeeds.
15pub struct IncomingConnection {
16    stream: TcpStream,
17    auth: AuthAdaptor,
18}
19
20impl IncomingConnection {
21    #[inline]
22    pub(crate) fn new(stream: TcpStream, auth: AuthAdaptor) -> Self {
23        IncomingConnection { stream, auth }
24    }
25
26    /// Returns the local address that this stream is bound to.
27    #[inline]
28    pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
29        self.stream.local_addr()
30    }
31
32    /// Returns the remote address that this stream is connected to.
33    #[inline]
34    pub fn peer_addr(&self) -> std::io::Result<SocketAddr> {
35        self.stream.peer_addr()
36    }
37
38    /// Shutdown the TCP stream.
39    #[inline]
40    pub async fn shutdown(&mut self) -> std::io::Result<()> {
41        self.stream.shutdown().await
42    }
43
44    /// Gets the value of the `TCP_NODELAY` option on this socket.
45    ///
46    /// For more information about this option, see
47    /// [`set_nodelay`](#method.set_nodelay).
48    #[inline]
49    pub fn nodelay(&self) -> std::io::Result<bool> {
50        self.stream.nodelay()
51    }
52
53    /// Sets the value of the `TCP_NODELAY` option on this socket.
54    ///
55    /// If set, this option disables the Nagle algorithm. This means that segments are always sent as soon as possible,
56    /// even if there is only a small amount of data. When not set, data is buffered until there is a sufficient amount
57    /// to send out, thereby avoiding the frequent sending of small packets.
58    pub fn set_nodelay(&self, nodelay: bool) -> std::io::Result<()> {
59        self.stream.set_nodelay(nodelay)
60    }
61
62    /// Gets the value of the `IP_TTL` option for this socket.
63    ///
64    /// For more information about this option, see
65    /// [`set_ttl`](#method.set_ttl).
66    pub fn ttl(&self) -> std::io::Result<u32> {
67        self.stream.ttl()
68    }
69
70    /// Sets the value for the `IP_TTL` option on this socket.
71    ///
72    /// This value sets the time-to-live field that is used in every packet sent from this socket.
73    pub fn set_ttl(&self, ttl: u32) -> std::io::Result<()> {
74        self.stream.set_ttl(ttl)
75    }
76
77    /// Set a timeout for the SOCKS5 handshake.
78    pub async fn authenticate_with_timeout(self, timeout: Duration) -> crate::Result<Authenticated> {
79        tokio::time::timeout(timeout, self.authenticate())
80            .await
81            .map_err(|_| crate::Error::String("handshake timeout".into()))?
82    }
83
84    /// Perform a SOCKS5 authentication handshake using the given
85    /// [`AuthExecutor`](crate::server::auth::AuthExecutor) adapter.
86    ///
87    /// If the handshake succeeds, an [`Authenticated`] stream is returned.
88    /// Otherwise, the error and the original [`TcpStream`](https://docs.rs/tokio/latest/tokio/net/struct.TcpStream.html) is returned.
89    ///
90    /// Note that this method will not implicitly close the connection even if the handshake failed.
91    pub async fn authenticate(mut self) -> crate::Result<Authenticated> {
92        let request = handshake::Request::retrieve_from_async_stream(&mut self.stream).await?;
93        if let Some(method) = self.evaluate_request(&request) {
94            let response = handshake::Response::new(method);
95            response.write_to_async_stream(&mut self.stream).await?;
96            if !self.auth.execute(&mut self.stream).await? {
97                use std::io::{Error, ErrorKind::PermissionDenied};
98                return Err(crate::Error::Io(Error::new(PermissionDenied, "authentication failed")));
99            }
100            Ok(Authenticated::new(self.stream))
101        } else {
102            let response = handshake::Response::new(AuthMethod::NoAcceptableMethods);
103            response.write_to_async_stream(&mut self.stream).await?;
104            let err = "No available handshake method provided by client";
105            Err(crate::Error::Io(std::io::Error::new(std::io::ErrorKind::Unsupported, err)))
106        }
107    }
108
109    fn evaluate_request(&self, req: &handshake::Request) -> Option<AuthMethod> {
110        let method = self.auth.auth_method();
111        if req.evaluate_method(method) { Some(method) } else { None }
112    }
113}
114
115impl std::fmt::Debug for IncomingConnection {
116    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
117        f.debug_struct("IncomingConnection").field("stream", &self.stream).finish()
118    }
119}
120
121impl From<IncomingConnection> for TcpStream {
122    #[inline]
123    fn from(conn: IncomingConnection) -> Self {
124        conn.stream
125    }
126}
127
128/// A TCP stream that has been authenticated.
129///
130/// To get the command from the SOCKS5 client, use
131/// [`wait_request`](crate::server::connection::Authenticated::wait_request).
132///
133/// It can also be converted back into a raw [`tokio::TcpStream`](https://docs.rs/tokio/latest/tokio/net/struct.TcpStream.html) with `From` trait.
134pub struct Authenticated(TcpStream);
135
136impl Authenticated {
137    #[inline]
138    fn new(stream: TcpStream) -> Self {
139        Self(stream)
140    }
141
142    /// Waits the SOCKS5 client to send a request.
143    ///
144    /// This method will return a [`Command`] if the client sends a valid command.
145    ///
146    /// When encountering an error, the stream will be returned alongside the error.
147    ///
148    /// Note that this method will not implicitly close the connection even if the client sends an invalid request.
149    pub async fn wait_request(mut self) -> crate::Result<ClientConnection> {
150        let req = protocol::Request::retrieve_from_async_stream(&mut self.0).await?;
151
152        match req.command {
153            Command::UdpAssociate => Ok(ClientConnection::UdpAssociate(
154                UdpAssociate::<associate::NeedReply>::new(self.0),
155                req.address,
156            )),
157            Command::Bind => Ok(ClientConnection::Bind(Bind::<bind::NeedFirstReply>::new(self.0), req.address)),
158            Command::Connect => Ok(ClientConnection::Connect(Connect::<connect::NeedReply>::new(self.0), req.address)),
159        }
160    }
161
162    /// 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.
163    #[inline]
164    pub async fn shutdown(&mut self) -> std::io::Result<()> {
165        self.0.shutdown().await
166    }
167
168    /// Returns the local address that this stream is bound to.
169    #[inline]
170    pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
171        self.0.local_addr()
172    }
173
174    /// Returns the remote address that this stream is connected to.
175    #[inline]
176    pub fn peer_addr(&self) -> std::io::Result<SocketAddr> {
177        self.0.peer_addr()
178    }
179
180    /// Gets the value of the `TCP_NODELAY` option on this socket.
181    ///
182    /// For more information about this option, see
183    /// [`set_nodelay`](crate::server::connection::Authenticated::set_nodelay).
184    #[inline]
185    pub fn nodelay(&self) -> std::io::Result<bool> {
186        self.0.nodelay()
187    }
188
189    /// Sets the value of the `TCP_NODELAY` option on this socket.
190    ///
191    /// If set, this option disables the Nagle algorithm. This means that segments are always sent as soon as possible,
192    /// 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,
193    /// thereby avoiding the frequent sending of small packets.
194    pub fn set_nodelay(&self, nodelay: bool) -> std::io::Result<()> {
195        self.0.set_nodelay(nodelay)
196    }
197
198    /// Gets the value of the `IP_TTL` option for this socket.
199    ///
200    /// For more information about this option, see
201    /// [`set_ttl`](crate::server::connection::Authenticated::set_ttl).
202    pub fn ttl(&self) -> std::io::Result<u32> {
203        self.0.ttl()
204    }
205
206    /// Sets the value for the `IP_TTL` option on this socket.
207    ///
208    /// This value sets the time-to-live field that is used in every packet sent from this socket.
209    pub fn set_ttl(&self, ttl: u32) -> std::io::Result<()> {
210        self.0.set_ttl(ttl)
211    }
212}
213
214impl From<Authenticated> for TcpStream {
215    #[inline]
216    fn from(conn: Authenticated) -> Self {
217        conn.0
218    }
219}
220
221/// After the socks5 handshake succeeds, the connection may become:
222///
223/// - Associate
224/// - Bind
225/// - Connect
226#[derive(Debug)]
227pub enum ClientConnection {
228    UdpAssociate(UdpAssociate<associate::NeedReply>, Address),
229    Bind(Bind<bind::NeedFirstReply>, Address),
230    Connect(Connect<connect::NeedReply>, Address),
231}
socks5_impl/server/connection/mod.rs

socks5_impl/server/connection/
mod.rs
