initial commit

src/list.rs

@@ -0,0 +1,51 @@
+use anyhow::Result;
+use log::*;
+use std::collections::HashSet;
+
+#[derive(PartialEq, Eq, Hash)]
+pub enum Item {
+    Ending(String),
+    Exact(String),
+}
+
+pub enum List {
+    Allow(HashSet<Item>),
+}
+
+impl List {
+    pub async fn from_opt(list: &Vec<String>) -> Result<Self> {
+        let mut this = HashSet::new();
+        for i in list {
+            for line in if i.contains("://") {
+                reqwest::get(i).await?.text().await?
+            } else {
+                std::fs::read_to_string(i)?
+            }
+            .lines()
+            {
+                if line.starts_with("/") || line.starts_with("#") || line.is_empty() {
+                    continue;
+                }
+                if !this.insert(if let Some(item) = line.strip_prefix(".") {
+                    debug!("Add `{line}` to the whitelist");
+                    Item::Ending(item.to_string())
+                } else {
+                    debug!("Add `{line}` exact match to the whitelist");
+                    Item::Exact(line.to_string())
+                }) {
+                    warn!("Duplicated whitelist record: `{line}`")
+                }
+            }
+        }
+        info!("Total whitelist entries parsed: {}", this.len());
+        Ok(Self::Allow(this))
+    }
+    pub fn has(&self, value: &str) -> bool {
+        match self {
+            Self::Allow(list) => list.iter().any(|item| match item {
+                Item::Exact(s) => s == value,
+                Item::Ending(s) => value.ends_with(s),
+            }),
+        }
+    }
+}

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:#}"),
+            },
+        }
+    })
+}

src/opt.rs

@@ -0,0 +1,67 @@
+use std::{num::ParseFloatError, time::Duration};
+use structopt::StructOpt;
+
+/// # How to use it:
+///
+/// Listen on a local address, authentication-free:
+///     `$ RUST_LOG=debug cargo run --example server -- --listen-addr 127.0.0.1:1337 no-auth`
+///
+/// Listen on a local address, with basic username/password requirement:
+///     `$ RUST_LOG=debug cargo run --example server -- --listen-addr 127.0.0.1:1337 password --username admin --password password`
+///
+/// Same as above but with UDP support
+///     `$ RUST_LOG=debug cargo run --example server -- --listen-addr 127.0.0.1:1337 --allow-udp --public-addr 127.0.0.1 password --username admin --password password`
+#[derive(Debug, StructOpt)]
+#[structopt(
+    name = "socks5-server",
+    about = "A simple implementation of a socks5-server."
+)]
+pub struct Opt {
+    /// Bind on address address. eg. `127.0.0.1:1080`
+    #[structopt(short, long, default_value = "127.0.0.1:1080")]
+    pub listen_addr: String,
+
+    /// Our external IP address to be sent in reply packets (required for UDP)
+    #[structopt(long)]
+    pub public_addr: Option<std::net::IpAddr>,
+
+    /// Request timeout
+    #[structopt(short = "t", long, default_value = "10", parse(try_from_str=parse_duration))]
+    pub request_timeout: Duration,
+
+    /// 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,
+
+    /// Allow UDP proxying, requires public-addr to be set
+    #[structopt(short = "U", long)]
+    pub allow_udp: bool,
+
+    /// Allow list:
+    /// * local filename
+    /// * remote URL
+    #[structopt(short = "a", long)]
+    pub allow_list: Vec<String>,
+}
+
+/// Choose the authentication type
+#[derive(StructOpt, Debug, PartialEq)]
+pub enum AuthMode {
+    NoAuth,
+    Password {
+        #[structopt(short, long)]
+        username: String,
+
+        #[structopt(short, long)]
+        password: String,
+    },
+}
+
+fn parse_duration(s: &str) -> Result<Duration, ParseFloatError> {
+    let seconds = s.parse()?;
+    Ok(Duration::from_secs_f64(seconds))
+}