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}