Initial commit

src/config.rs

@@ -0,0 +1,85 @@
+use std::{env, fs, io::Read, path::PathBuf};
+
+use serde_derive::Deserialize;
+
+#[derive(Clone, Deserialize, Debug)]
+pub struct Config {
+    /// Bind on address address. eg. `127.0.0.1:1080`
+    pub listen_addr: String,
+    /// Request timeout
+    pub request_timeout: u64,
+    /// Authentication
+    #[serde(default)]
+    pub auth: Option<PasswordAuth>,
+    /// Avoid useless roundtrips if we don't need the Authentication layer
+    #[serde(default)]
+    pub skip_auth: bool,
+    /// Enable dns-resolving
+    #[serde(default = "default_true")]
+    pub dns_resolve: bool,
+    /// Enable command execution
+    #[serde(default = "default_true")]
+    pub execute_command: bool,
+    /// Enable UDP support
+    #[serde(default = "default_true")]
+    pub allow_udp: bool,
+}
+
+fn default_true() -> bool {
+    true
+}
+
+/// Password authentication data
+#[derive(Clone, Deserialize, Debug)]
+pub struct PasswordAuth {
+    pub username: String,
+    pub password: String,
+}
+
+impl Default for Config {
+    fn default() -> Self {
+        Config {
+            listen_addr: "127.0.0.1:1080".into(),
+            request_timeout: 120,
+            auth: None,
+            skip_auth: false,
+            dns_resolve: true,
+            execute_command: true,
+            allow_udp: true,
+        }
+    }
+}
+
+impl Config {
+    const FILENAME: &'static str = "config.toml";
+    fn read(filename: &str) -> Result<Self, &str> {
+        let mut file = fs::File::open(filename).map_err(|_| "can't open config")?;
+        let mut data = vec![];
+        file.read_to_end(&mut data)
+            .map_err(|_| "can't read config")?;
+        toml::from_slice(&data).map_err(|_| "can't parse config")
+    }
+    fn file_location() -> Result<PathBuf, &'static str> {
+        let mut res = env::current_exe().map_err(|_| "can't get current exe path")?;
+        res.pop();
+        res.push(Config::FILENAME);
+        Ok(res)
+    }
+    pub fn get() -> Self {
+        let path = Config::file_location();
+        if path.is_err() {
+            log::error!("Error: {}, using default config", path.err().unwrap());
+            return Config::default();
+        }
+
+        let path = path.unwrap();
+        let cfg = Config::read(path.to_str().unwrap());
+        match cfg {
+            Err(e) => {
+                log::error!("Error: {e}, using default config");
+                Config::default()
+            }
+            Ok(cfg) => cfg,
+        }
+    }
+}

src/main.rs

@@ -0,0 +1,45 @@
+extern crate flexi_logger;
+
+use tokio_util::sync::CancellationToken;
+
+mod config;
+mod server;
+use crate::config::Config;
+use crate::server::spawn_socks5_server;
+
+use flexi_logger::{AdaptiveFormat, Age, Cleanup, Criterion, Duplicate, FileSpec, Logger, Naming};
+
+#[tokio::main]
+async fn main() {
+    Logger::try_with_str("info")
+        .unwrap()
+        .log_to_file(FileSpec::default())
+        .rotate(
+            Criterion::Age(Age::Day),
+            Naming::Timestamps,
+            Cleanup::KeepLogFiles(4),
+        )
+        .adaptive_format_for_stderr(AdaptiveFormat::Detailed)
+        .print_message()
+        .duplicate_to_stderr(Duplicate::Warn)
+        .start_with_specfile("logspec.toml")
+        .unwrap();
+
+    let cfg = tokio::task::spawn_blocking(Config::get)
+        .await
+        .expect("get config");
+    log::info!("cfg: {:#?}", cfg);
+
+    let token = CancellationToken::new();
+    let child_token = token.child_token();
+
+    let (r, _) = tokio::join!(
+        spawn_socks5_server(cfg, child_token),
+        tokio::spawn(async move {
+            tokio::time::sleep(std::time::Duration::from_secs(10)).await;
+            token.cancel();
+        })
+    );
+
+    r.unwrap();
+}

src/server.rs

@@ -0,0 +1,87 @@
+use fast_socks5::{
+    server::{SimpleUserPassword, Socks5Server, Socks5Socket},
+    Result,
+};
+use std::future::Future;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::select;
+use tokio::task;
+use tokio_stream::StreamExt;
+use tokio_util::sync::CancellationToken;
+
+use crate::config::Config;
+use crate::config::PasswordAuth;
+
+pub async fn spawn_socks5_server(cfg: Config, token: CancellationToken) -> Result<()> {
+    let mut server_config = fast_socks5::server::Config::default();
+    server_config.set_request_timeout(cfg.request_timeout);
+    server_config.set_skip_auth(cfg.skip_auth);
+    server_config.set_dns_resolve(cfg.dns_resolve);
+    server_config.set_execute_command(cfg.execute_command);
+    server_config.set_udp_support(cfg.allow_udp);
+
+    if let Some(PasswordAuth { username, password }) = cfg.auth {
+        server_config.set_authentication(SimpleUserPassword { username, password });
+        log::info!("Simple auth system has been set.");
+    } else {
+        log::warn!("No authentication has been set!");
+    }
+
+    let mut listener = Socks5Server::bind(&cfg.listen_addr).await?;
+    listener.set_config(server_config);
+
+    let mut incoming = listener.incoming();
+
+    log::info!("Listen for socks connections @ {}", &cfg.listen_addr);
+
+    // Standard TCP loop
+    while let Some(socket_res) = or_chancel(incoming.next(), token.child_token()).await {
+        match socket_res {
+            Ok(socket) => {
+                let child_token = token.child_token();
+                spawn_and_log_error(socket.upgrade_to_socks5(), child_token);
+            }
+            Err(err) => {
+                log::error!("accept error = {:?}", err);
+            }
+        }
+    }
+
+    Ok(())
+}
+
+async fn or_chancel<F, R>(future: F, token: CancellationToken) -> Option<R>
+where
+    F: Future<Output = Option<R>>,
+{
+    select! {
+        _ = token.cancelled() => {
+            log::error!("canceled");
+            None
+        }
+        res = future => {
+            res
+        }
+    }
+}
+
+fn spawn_and_log_error<F, T>(future: F, token: CancellationToken) -> task::JoinHandle<()>
+where
+    F: Future<Output = Result<Socks5Socket<T>>> + Send + 'static,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    tokio::spawn(async move {
+        // Wait for either cancellation or a very long time
+        let result = select! {
+            _ = token.cancelled() => {
+                Err("Client connection canceled".to_string())
+            }
+            res = future => {
+                res.map_err(|e| format!("{:#}", &e))
+            }
+        };
+        if let Err(e) = result {
+            log::error!("{}", &e);
+        }
+    })
+}