socks5_impl/server/connection/
mod.rs

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

socks5_impl/server/connection/
mod.rs
