New explicit and extensible server API

Cargo.toml

@@ -46,4 +46,7 @@ name = "server"
 name = "client"
 
 [[example]]
-name = "simple_tcp_server"
+name = "custom_auth_server"
+
+[[example]]
+name = "router"

README.md

@@ -12,25 +12,23 @@ This library is maintained by [anyip.io](https://anyip.io/) a residential and mo
 - An `async`/`.await` [SOCKS4 Client](https://www.openssh.com/txt/socks4.protocol) implementation.
 - An `async`/`.await` [SOCKS4a Client](https://www.openssh.com/txt/socks4a.protocol) implementation.
 - No **unsafe** code
-- Built on-top of `tokio` library
+- Built on top of the [Tokio](https://tokio.rs/) runtime
 - Ultra lightweight and scalable
 - No system dependencies
 - Cross-platform
+- Infinitely extensible, explicit server API based on typestates for safety
+  - You control the request handling, the library only ensures you follow the proper protocol flow
+  - Can skip DNS resolution
+  - Can skip the authentication/handshake process (not RFC-compliant, for private use, to save on useless round-trips)
+  - Instead of proxying in-process, swap out `run_tcp_proxy` for custom handling to build a router or to use a custom accelerated proxying method
 - Authentication methods:
-  - No-Auth method
-  - Username/Password auth method
-  - Custom auth methods can be implemented via the Authentication Trait
-  - Credentials returned on authentication success
+  - No-Auth method (`0x00`)
+  - Username/Password auth method (`0x02`)
+  - Custom auth methods can be implemented on the server side via the `AuthMethod` Trait
+    - Multiple auth methods with runtime negotiation can be supported, with fast *static* dispatch (enums can be generated with the `auth_method_enums` macro)
 - UDP is supported
 - All SOCKS5 RFC errors (replies) should be mapped
-- `AsyncRead + AsyncWrite` traits are implemented on Socks5Stream & Socks5Socket
 - `IPv4`, `IPv6`, and `Domains` types are supported
-- Config helper for Socks5Server
-- Helpers to run a Socks5Server à la *"std's TcpStream"* via `incoming.next().await`
-- Examples come with real cases commands scenarios
-- Can disable `DNS resolving`
-- Can skip the authentication/handshake process, which will directly handle command's request (useful to save useless round-trips in a current authenticated environment)
-- Can disable command execution (useful if you just want to forward the request to a different server)
 
 
 ## Install

examples/custom_auth_server.rs

@@ -0,0 +1,179 @@
+#[forbid(unsafe_code)]
+#[macro_use]
+extern crate log;
+
+use fast_socks5::{
+    auth_method_enums,
+    server::{
+        run_tcp_proxy, AuthMethod, AuthMethodSuccessState, DnsResolveHelper as _,
+        PasswordAuthentication, PasswordAuthenticationStarted, Socks5ServerProtocol,
+    },
+    ReplyError, Result, Socks5Command, SocksError,
+};
+use std::{future::Future, time::Duration};
+use structopt::StructOpt;
+use tokio::task;
+use tokio::{
+    io::{AsyncRead, AsyncReadExt},
+    net::TcpListener,
+};
+
+/// # How to use it:
+///
+/// Listen on a local address:
+///     `$ RUST_LOG=debug cargo run --example custom_auth_server -- --listen-addr 127.0.0.1:1337`
+///
+/// then try a client to connect to this server:
+///     `$ RUST_LOG=debug cargo run --example client -- --socks-server 127.0.0.1:1337 --username user --password "correct_horse_battery_staple" -a perdu.com -p 80`
+///
+/// or via a cURL command
+///     `curl -v -s --proxy "socks5://user:[email protected]:1337" "https://httpbin.org/get"`
+///
+#[derive(Debug, StructOpt)]
+#[structopt(
+    name = "socks5-server-custom-auth",
+    about = "A socks5 server with a curious secret."
+)]
+struct Opt {
+    /// Bind on address address. eg. `127.0.0.1:1080`
+    #[structopt(short, long)]
+    pub listen_addr: String,
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    env_logger::init();
+
+    spawn_socks_server().await
+}
+
+async fn spawn_socks_server() -> Result<()> {
+    let opt: Opt = Opt::from_args();
+
+    let listener = TcpListener::bind(&opt.listen_addr).await?;
+
+    info!("Listen for socks connections @ {}", &opt.listen_addr);
+
+    // Standard TCP loop
+    loop {
+        match listener.accept().await {
+            Ok((socket, _client_addr)) => {
+                spawn_and_log_error(serve_socks5(socket));
+            }
+            Err(err) => {
+                error!("accept error = {:?}", err);
+            }
+        }
+    }
+}
+
+pub struct BackdoorAuthenticationStarted<T>(T);
+pub struct BackdoorAuthenticationSuccess<T>(T);
+
+impl<T: AsyncRead + Unpin> BackdoorAuthenticationStarted<T> {
+    pub async fn verify_timing(self) -> Result<BackdoorAuthenticationSuccess<T>> {
+        let mut socket = self.0;
+        let mut buf = vec![0u8; 2];
+        if tokio::time::timeout(Duration::from_millis(500), socket.read_exact(&mut buf))
+            .await
+            .is_ok()
+        {
+            debug!("too early!");
+            return Err(SocksError::AuthenticationRejected("nope".to_owned()));
+        }
+        if tokio::time::timeout(Duration::from_millis(500), socket.read_exact(&mut buf))
+            .await
+            .is_err()
+        {
+            debug!("too late!");
+            return Err(SocksError::AuthenticationRejected("nope".to_owned()));
+        }
+        if buf[0] == 0x13 && buf[1] == 0x37 {
+            Ok(BackdoorAuthenticationSuccess(socket))
+        } else {
+            debug!("wrong contents!");
+            Err(SocksError::AuthenticationRejected("nope".to_owned()))
+        }
+    }
+}
+
+impl<T> AuthMethodSuccessState<T> for BackdoorAuthenticationSuccess<T> {
+    fn into_inner(self) -> T {
+        self.0
+    }
+}
+
+/// A silly example of a custom authentication method.
+#[derive(Debug, Clone, Copy)]
+pub struct BackdoorAuthentication;
+
+impl<T> AuthMethod<T> for BackdoorAuthentication {
+    type StartingState = BackdoorAuthenticationStarted<T>;
+
+    fn method_id(self) -> u8 {
+        0xF0 // From the "RESERVED FOR PRIVATE METHODS" range
+    }
+
+    fn new(self, inner: T) -> Self::StartingState {
+        BackdoorAuthenticationStarted(inner)
+    }
+}
+
+auth_method_enums! {
+    pub enum Auth / AuthStarted<T> {
+        PasswordAuthentication(PasswordAuthenticationStarted<T>),
+        BackdoorAuthentication(BackdoorAuthenticationStarted<T>),
+    }
+}
+
+async fn serve_socks5(socket: tokio::net::TcpStream) -> Result<(), SocksError> {
+    let proto = match Socks5ServerProtocol::start(socket)
+        .negotiate_auth(&[
+            // The order of authentication methods can be tested by clients in sequence,
+            // so list more secure or preferred methods first
+            Auth::PasswordAuthentication(PasswordAuthentication),
+            Auth::BackdoorAuthentication(BackdoorAuthentication),
+        ])
+        .await?
+    {
+        AuthStarted::PasswordAuthentication(auth) => {
+            let (user, pass, auth) = auth.read_username_password().await?;
+            if user == "user" && pass == "correct_horse_battery_staple" {
+                // better to not use spaces for trying out with cURL
+                auth.accept().await?.finish_auth()
+            } else {
+                auth.reject().await?;
+                return Err(SocksError::AuthenticationRejected(
+                    "Wrong username/password".to_owned(),
+                ));
+            }
+        }
+        AuthStarted::BackdoorAuthentication(auth) => auth.verify_timing().await?.finish_auth(),
+    };
+
+    let (proto, cmd, target_addr) = proto.read_command().await?.resolve_dns().await?;
+
+    const REQUEST_TIMEOUT: u64 = 10;
+    match cmd {
+        Socks5Command::TCPConnect => {
+            run_tcp_proxy(proto, &target_addr, REQUEST_TIMEOUT, false).await?;
+        }
+        _ => {
+            proto.reply_error(&ReplyError::CommandNotSupported).await?;
+            return Err(ReplyError::CommandNotSupported.into());
+        }
+    };
+    Ok(())
+}
+
+fn spawn_and_log_error<F>(fut: F) -> task::JoinHandle<()>
+where
+    F: Future<Output = Result<()>> + Send + 'static,
+{
+    task::spawn(async move {
+        match fut.await {
+            Ok(()) => {}
+            Err(err) => error!("{:#}", &err),
+        }
+    })
+}