initial commit

src/main.rs

@@ -0,0 +1,125 @@
+mod list;
+mod opt;
+
+use anyhow::Context;
+use fast_socks5::{
+    ReplyError, Result, Socks5Command, SocksError,
+    server::{DnsResolveHelper as _, Socks5ServerProtocol, run_tcp_proxy, run_udp_proxy},
+};
+use log::*;
+use opt::{AuthMode, Opt};
+use std::{future::Future, sync::Arc};
+use structopt::StructOpt;
+use tokio::net::TcpListener;
+use tokio::task;
+
+use crate::list::List;
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    env_logger::init();
+    spawn_socks_server().await
+}
+
+async fn spawn_socks_server() -> Result<()> {
+    let opt: &'static Opt = Box::leak(Box::new(Opt::from_args()));
+    if opt.allow_udp && opt.public_addr.is_none() {
+        return Err(SocksError::ArgumentInputError(
+            "Can't allow UDP if public-addr is not set",
+        ));
+    }
+    if opt.skip_auth && opt.auth != AuthMode::NoAuth {
+        return Err(SocksError::ArgumentInputError(
+            "Can't use skip-auth flag and authentication altogether.",
+        ));
+    }
+
+    let list = Arc::new(List::from_opt(&opt.allow_list).await.map_err(|err| {
+        error!("Can't parse whitelist: `{err}`");
+        SocksError::ArgumentInputError("Can't parse whitelist")
+    })?);
+
+    let listener = TcpListener::bind(&opt.listen_addr).await?;
+
+    info!("Listen for socks connections @ {}", &opt.listen_addr);
+
+    loop {
+        match listener.accept().await {
+            Ok((socket, _client_addr)) => {
+                spawn_and_log_error(serve_socks5(opt, socket, list.clone()));
+            }
+            Err(err) => error!("accept error = {:?}", err),
+        }
+    }
+}
+
+async fn serve_socks5(
+    opt: &Opt,
+    socket: tokio::net::TcpStream,
+    list: Arc<List>,
+) -> Result<(), SocksError> {
+    let request = match &opt.auth {
+        AuthMode::NoAuth if opt.skip_auth => {
+            Socks5ServerProtocol::skip_auth_this_is_not_rfc_compliant(socket)
+        }
+        AuthMode::NoAuth => Socks5ServerProtocol::accept_no_auth(socket).await?,
+        AuthMode::Password { username, password } => {
+            Socks5ServerProtocol::accept_password_auth(socket, |user, pass| {
+                user == *username && pass == *password
+            })
+            .await?
+            .0
+        }
+    }
+    .read_command()
+    .await?;
+
+    let (host, _) = request.2.clone().into_string_and_port(); // @TODO ref
+    let (proto, cmd, addr) = request.resolve_dns().await?;
+
+    if !list.has(&host) && !list.has(&addr.to_string()) {
+        info!("Blocked connection attempt to: {host}");
+        proto.reply_error(&ReplyError::ConnectionNotAllowed).await?;
+        return Err(ReplyError::ConnectionNotAllowed.into());
+    }
+
+    match cmd {
+        Socks5Command::TCPConnect => {
+            run_tcp_proxy(proto, &addr, opt.request_timeout, false).await?;
+        }
+        Socks5Command::UDPAssociate if opt.allow_udp => {
+            run_udp_proxy(
+                proto,
+                &addr,
+                None,
+                opt.public_addr.context("invalid reply ip")?,
+                None,
+            )
+            .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) => match err {
+                SocksError::ReplyError(reply_error) => {
+                    if !matches!(reply_error, ReplyError::ConnectionNotAllowed) {
+                        error!("{reply_error:#}")
+                    }
+                }
+                _ => error!("{err:#}"),
+            },
+        }
+    })
+}