1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

#![doc = include_str!("../README.md")]

use std::{
    io::Error,
    net::SocketAddr,
    sync::Arc,
    task::{Context, Poll},
};
use tokio::net::TcpListener;

pub mod auth;
pub mod connection;

pub use crate::{
    auth::Auth,
    connection::{
        associate::{Associate, AssociatedUdpSocket},
        bind::Bind,
        connect::Connect,
        Authenticated, Command, IncomingConnection,
    },
};

pub(crate) type AuthAdaptor<O> = Arc<dyn Auth<Output = O> + Send + Sync>;

/// A SOCKS5 server listener
///
/// This server listens on a socket and treats incoming connections as SOCKS5 connections.
///
/// A `(TcpListener, Arc<dyn Auth<Output = O> + Send + Sync>)` can be converted into a `Server<O>` with `From` trait. Also, a `Server<O>` can be converted back.
///
/// Generic type `<O>` is the output type of the authentication adapter. See trait [`Auth`](https://docs.rs/socks5-server/latest/socks5_server/auth/trait.Auth.html).
///
/// # Example
///
/// ```rust
/// use socks5_server::{auth::NoAuth, Server};
/// use std::sync::Arc;
/// use tokio::net::TcpListener;
///
/// async fn listen() {
///     let listener = TcpListener::bind("127.0.0.1:5000").await.unwrap();
///     let auth = Arc::new(NoAuth) as Arc<_>;
///
///     let server = Server::from((listener, auth));
///
///     while let Ok((conn, _)) = server.accept().await {
///         tokio::spawn(async move {
///             todo!();
///         });
///     }
/// }
/// ```
pub struct Server<O> {
    listener: TcpListener,
    auth: AuthAdaptor<O>,
}

impl<O> Server<O> {
    /// Accept an [`IncomingConnection<O>`](https://docs.rs/socks5-server/latest/socks5_server/connection/struct.IncomingConnection.html).
    ///
    /// The connection is only a freshly created TCP connection and may not be a valid SOCKS5 connection. You should call [`IncomingConnection::authenticate()`](https://docs.rs/socks5-server/latest/socks5_server/connection/struct.IncomingConnection.html#method.authenticate) to perform a SOCKS5 authentication handshake.
    #[inline]
    pub async fn accept(&self) -> Result<(IncomingConnection<O>, SocketAddr), Error> {
        let (stream, addr) = self.listener.accept().await?;
        Ok((IncomingConnection::new(stream, self.auth.clone()), addr))
    }

    /// Polls to accept an [`IncomingConnection<O>`](https://docs.rs/socks5-server/latest/socks5_server/connection/struct.IncomingConnection.html).
    ///
    /// The connection is only a freshly created TCP connection and may not be a valid SOCKS5 connection. You should call [`IncomingConnection::authenticate()`](https://docs.rs/socks5-server/latest/socks5_server/connection/struct.IncomingConnection.html#method.authenticate) to perform a SOCKS5 authentication handshake.
    ///
    /// 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.
    #[inline]
    pub fn poll_accept(
        &self,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(IncomingConnection<O>, SocketAddr), Error>> {
        self.listener
            .poll_accept(cx)
            .map_ok(|(stream, addr)| (IncomingConnection::new(stream, self.auth.clone()), addr))
    }

    /// Returns the local address that this server is bound to.
    ///
    /// This can be useful, for example, when binding to port 0 to figure out which port was actually bound.
    #[inline]
    pub fn local_addr(&self) -> Result<SocketAddr, Error> {
        self.listener.local_addr()
    }

    /// Sets the value for the `IP_TTL` option on this socket.
    ///
    /// This value sets the time-to-live field that is used in every packet sent from this socket.
    #[inline]
    pub fn set_ttl(&self, ttl: u32) -> Result<(), Error> {
        self.listener.set_ttl(ttl)
    }

    /// Gets the value of the `IP_TTL` option for this socket.
    ///
    /// For more information about this option, see [set_ttl](https://docs.rs/socks5-server/latest/socks5_server/struct.Server.html#method.set_ttl).
    #[inline]
    pub fn ttl(&self) -> Result<u32, Error> {
        self.listener.ttl()
    }
}

impl<O> From<(TcpListener, AuthAdaptor<O>)> for Server<O> {
    #[inline]
    fn from((listener, auth): (TcpListener, AuthAdaptor<O>)) -> Self {
        Self { listener, auth }
    }
}

impl<O> From<Server<O>> for (TcpListener, AuthAdaptor<O>) {
    #[inline]
    fn from(server: Server<O>) -> Self {
        (server.listener, server.auth)
    }
}