add proxy router example

examples/router.rs

@@ -0,0 +1,181 @@
+#[forbid(unsafe_code)]
+#[macro_use]
+extern crate log;
+
+use fast_socks5::{
+    server::{Config, SimpleUserPassword, Socks5Server, Socks5Socket},
+    client::{self, Socks5Stream},
+    Result, SocksError,
+};
+use anyhow::{anyhow, Context};
+use structopt::StructOpt;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::task;
+use tokio_stream::StreamExt;
+use std::net::ToSocketAddrs;
+use std::io::ErrorKind;
+
+/// # How to use it:
+///
+/// Listen on a local address, authentication-free:
+///     `$ RUST_LOG=debug cargo run --example router -- --listen-addr 127.0.0.1:1337 --proxy-addr 127.0.0.1:1338 no-auth`
+///
+/// Listen on a local address, with basic username/password requirement:
+///     `$ RUST_LOG=debug cargo run --example router -- --listen-addr 127.0.0.1:1337 --proxy-addr 127.0.0.1:1338 password --username admin --password password`
+///
+#[derive(Debug, StructOpt)]
+#[structopt(
+    name = "socks5-server",
+    about = "A simple implementation of a socks5-server."
+)]
+struct Opt {
+    /// Bind on address for incoming traffic. eg. `127.0.0.1:1080`
+    #[structopt(short, long)]
+    pub listen_addr: String,
+
+    /// Bind on proxy address. eg. `127.0.0.1:1081`
+    #[structopt(short, long)]
+    pub proxy_addr: String,
+
+    /// Request timeout
+    #[structopt(short = "t", long, default_value = "10")]
+    pub request_timeout: u64,
+
+    /// Choose authentication type
+    #[structopt(subcommand, name = "auth")] // Note that we mark a field as a subcommand
+    pub auth: AuthMode,
+
+    /// Don't perform the auth handshake, send directly the command request
+    #[structopt(short = "k", long)]
+    pub skip_auth: bool,
+}
+
+/// Choose the authentication type
+#[derive(StructOpt, Debug)]
+enum AuthMode {
+    NoAuth,
+    Password {
+        #[structopt(short, long)]
+        username: String,
+
+        #[structopt(short, long)]
+        password: String,
+    },
+}
+
+/// Useful read 1. https://blog.yoshuawuyts.com/rust-streams/
+/// Useful read 2. https://blog.yoshuawuyts.com/futures-concurrency/
+/// Useful read 3. https://blog.yoshuawuyts.com/streams-concurrency/
+/// error-libs benchmark: https://blog.yoshuawuyts.com/error-handling-survey/
+///
+/// TODO: Write functional tests: https://github.com/ark0f/async-socks5/blob/master/src/lib.rs#L762
+/// TODO: Write functional tests with cURL?
+#[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 mut config = Config::default();
+    config.set_request_timeout(opt.request_timeout);
+    config.set_skip_auth(opt.skip_auth);
+
+    match opt.auth {
+        AuthMode::NoAuth => warn!("No authentication has been set!"),
+        AuthMode::Password { username, password } => {
+            if opt.skip_auth {
+                return Err(SocksError::ArgumentInputError(
+                    "Can't use skip-auth flag and authentication altogether.",
+                ));
+            }
+
+            config.set_authentication(SimpleUserPassword { username, password });
+            info!("Simple auth system has been set.");
+        }
+    }
+
+    let mut listener = Socks5Server::bind(&opt.listen_addr).await?;
+    listener.set_config(config);
+
+    let mut incoming = listener.incoming();
+
+    info!("Listen for socks connections @ {}, using proxy @ {}", &opt.listen_addr, &opt.proxy_addr);
+
+    // Standard TCP loop
+    while let Some(socket_res) = incoming.next().await {
+        match socket_res {
+            Ok(socket) => {
+                let proxy_addr = opt.proxy_addr.clone();
+                task::spawn(async move {
+                    if let Err(err) = handle_socket(socket, proxy_addr).await {
+                        error!("socket handle error = {:#}", err);
+                    }
+                });
+            }
+            Err(err) => {
+                error!("accept error = {:#}", err);
+            }
+        }
+    }
+
+    Ok(())
+}
+
+async fn handle_socket<T>(socket: Socks5Socket<T>, proxy_addr: String) -> Result<()>
+where
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    // upgrade socket to SOCKS5 proxy
+    let mut socks5_socket = socket
+        .upgrade_to_socks5()
+        .await
+        .context("upgrade incoming socket to socks5")?;
+
+    // get resolved target addr
+    socks5_socket
+        .resolve_dns()
+        .await
+        .context("resolve target dns for incoming socket")?;
+    let socket_addr = socks5_socket
+        .target_addr()
+        .context("find target address for incoming socket")?
+        .to_socket_addrs()
+        .context("convert target address of incoming socket to socket addresses")?
+        .next()
+        .context("reach out to target of incoming socket")?;
+
+    // connect to downstream proxy
+    let mut stream = Socks5Stream::connect(
+        proxy_addr,
+        socket_addr.ip().to_string(),
+        socket_addr.port(),
+        client::Config::default(),
+    )
+    .await
+    .context("connect to downstream proxy for incoming socket")?;
+
+    // copy data between our incoming client and the used downstream proxy
+    match tokio::io::copy_bidirectional(&mut stream, &mut socks5_socket).await {
+        Ok(res) => {
+            info!("socket transfer closed ({}, {})", res.0, res.1);
+            Ok(())
+        }
+        Err(err) => match err.kind() {
+            ErrorKind::NotConnected => {
+                info!("socket transfer closed by client");
+                Ok(())
+            },
+            ErrorKind::ConnectionReset => {
+                info!("socket transfer closed by downstream proxy");
+                Ok(())
+            },
+            _ => Err(SocksError::Other(anyhow!(
+                "socket transfer error: {:#}",
+                err
+            )))
+        },
+    }
+}