socks5_server/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::{
4    fmt::Debug,
5    io::Error,
6    net::SocketAddr,
7    sync::Arc,
8    task::{Context, Poll},
9};
10use tokio::net::TcpListener;
11
12pub mod auth;
13pub mod connection;
14
15pub use crate::{
16    auth::Auth,
17    connection::{
18        associate::{Associate, AssociatedUdpSocket},
19        bind::Bind,
20        connect::Connect,
21        Command, IncomingConnection,
22    },
23};
24
25pub use socks5_proto as proto;
26
27pub(crate) type AuthAdaptor<A> = Arc<dyn Auth<Output = A> + Send + Sync>;
28
29type ServerAcceptResult<A> = Result<
30    (
31        IncomingConnection<A, connection::state::NeedAuthenticate>,
32        SocketAddr,
33    ),
34    Error,
35>;
36
37/// A SOCKS5 server listener
38///
39/// This server listens on a socket and treats incoming connections as SOCKS5 connections.
40///
41/// Generic `<A>` is the output type of the authentication adapter. See trait [`Auth`].
42///
43/// # Example
44///
45/// ```rust
46/// use socks5_server::{auth::NoAuth, Server};
47/// use std::sync::Arc;
48/// use tokio::net::TcpListener;
49///
50/// async fn listen() {
51///     let listener = TcpListener::bind("127.0.0.1:5000").await.unwrap();
52///     let auth = Arc::new(NoAuth) as Arc<_>;
53///
54///     let server = Server::new(listener, auth);
55///
56///     while let Ok((conn, _)) = server.accept().await {
57///         tokio::spawn(async move {
58///             todo!();
59///         });
60///     }
61/// }
62/// ```
63pub struct Server<A> {
64    listener: TcpListener,
65    auth: AuthAdaptor<A>,
66}
67
68impl<A> Server<A> {
69    /// Creates a new [`Server<A>`] with a [`TcpListener`](tokio::net::TcpListener) and an `Arc<dyn Auth<Output = A> + Send + Sync>`.
70    #[inline]
71    pub fn new(listener: TcpListener, auth: AuthAdaptor<A>) -> Self {
72        Self { listener, auth }
73    }
74
75    /// Accept an [`IncomingConnection`].
76    ///
77    /// The connection is only a freshly created TCP connection and may not be a valid SOCKS5 connection. You should call [`IncomingConnection::authenticate()`] to perform a SOCKS5 authentication handshake.
78    #[inline]
79    pub async fn accept(&self) -> ServerAcceptResult<A> {
80        let (stream, addr) = self.listener.accept().await?;
81        Ok((IncomingConnection::new(stream, self.auth.clone()), addr))
82    }
83
84    /// Polls to accept an [`IncomingConnection`].
85    ///
86    /// The connection is only a freshly created TCP connection and may not be a valid SOCKS5 connection. You should call [`IncomingConnection::authenticate()`] to perform a SOCKS5 authentication handshake.
87    ///
88    /// If there is no connection to accept, Poll::Pending is returned and the current task will be notified by a waker. Note that on multiple calls to poll_accept, only the Waker from the Context passed to the most recent call is scheduled to receive a wakeup.
89    #[inline]
90    pub fn poll_accept(&self, cx: &mut Context<'_>) -> Poll<ServerAcceptResult<A>> {
91        self.listener
92            .poll_accept(cx)
93            .map_ok(|(stream, addr)| (IncomingConnection::new(stream, self.auth.clone()), addr))
94    }
95
96    /// Returns the local address that this server is bound to.
97    ///
98    /// This can be useful, for example, when binding to port 0 to figure out which port was actually bound.
99    #[inline]
100    pub fn local_addr(&self) -> Result<SocketAddr, Error> {
101        self.listener.local_addr()
102    }
103
104    /// Returns a shared reference to the listener.
105    ///
106    /// Note that this may break the encapsulation of the [`Server`] and you should not use this method unless you know what you are doing.
107    #[inline]
108    pub fn get_ref(&self) -> &TcpListener {
109        &self.listener
110    }
111
112    /// Returns a mutable reference to the listener.
113    ///
114    /// Note that this may break the encapsulation of the [`Server`] and you should not use this method unless you know what you are doing.
115    #[inline]
116    pub fn get_mut(&mut self) -> &mut TcpListener {
117        &mut self.listener
118    }
119
120    /// Consumes the [`Server<A>`] and returns the underlying [`TcpListener`](tokio::net::TcpListener) and `Arc<dyn Auth<Output = A> + Send + Sync>`.
121    #[inline]
122    pub fn into_inner(self) -> (TcpListener, AuthAdaptor<A>) {
123        (self.listener, self.auth)
124    }
125}
126
127impl<A> Debug for Server<A> {
128    #[inline]
129    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
130        f.debug_struct("Server")
131            .field("listener", &self.listener)
132            .finish()
133    }
134}