From 80b5fb67b927fae0af1bd80f0275fd3253cebff5 Mon Sep 17 00:00:00 2001 From: Dmitry Date: Sat, 24 Sep 2022 19:23:53 +0300 Subject: [PATCH] Initial commit --- .gitignore | 4 + Cargo.lock | 570 ++++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 17 ++ src/config.rs | 85 ++++++++ src/main.rs | 45 ++++ src/server.rs | 87 ++++++++ 6 files changed, 808 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 src/config.rs create mode 100644 src/main.rs create mode 100644 src/server.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3cc4c8a --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +/target +socks5ws_*.log +logspec.toml +config.toml diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..38f6d26 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,570 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "aho-corasick" +version = "0.7.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4f55bd91a0978cbfd91c457a164bab8b4001c833b7f323132c0a4e1922dd44e" +dependencies = [ + "memchr", +] + +[[package]] +name = "ansi_term" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d52a9bb7ec0cf484c551830a7ce27bd20d67eac647e1befb56b0be4ee39a55d2" +dependencies = [ + "winapi", +] + +[[package]] +name = "anyhow" +version = "1.0.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98161a4e3e2184da77bb14f02184cdd111e83bbbcc9979dfee3c44b9a85f5602" + +[[package]] +name = "atty" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" +dependencies = [ + "hermit-abi", + "libc", + "winapi", +] + +[[package]] +name = "autocfg" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "bytes" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec8a7b6a70fde80372154c65702f00a0f56f3e1c36abbc6c440484be248856db" + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "crossbeam-channel" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2dd04ddaf88237dc3b8d8f9a3c1004b506b54b3313403944054d23c0870c521" +dependencies = [ + "cfg-if", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-queue" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1cd42583b04998a5363558e5f9291ee5a5ff6b49944332103f251e7479a82aa7" +dependencies = [ + "cfg-if", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51887d4adc7b564537b15adcfb307936f8075dfcd5f00dde9a9f1d29383682bc" +dependencies = [ + "cfg-if", + "once_cell", +] + +[[package]] +name = "fast-socks5" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2687b5a6108f18ba8621e0e618a3be1dcc2768632dad24b7cea1f87975375a9" +dependencies = [ + "anyhow", + "log", + "thiserror", + "tokio", + "tokio-stream", +] + +[[package]] +name = "flexi_logger" +version = "0.23.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4a12e3b5a8775259ee83ac38aea8cdf9c3a1667c02178d207378c0837808fa9" +dependencies = [ + "ansi_term", + "atty", + "crossbeam-channel", + "crossbeam-queue", + "glob", + "lazy_static", + "log", + "regex", + "rustversion", + "serde", + "serde_derive", + "thiserror", + "time", + "toml", +] + +[[package]] +name = "futures-core" +version = "0.3.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e5aa3de05362c3fb88de6531e6296e85cde7739cccad4b9dfeeb7f6ebce56bf" + +[[package]] +name = "futures-sink" +version = "0.3.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21b20ba5a92e727ba30e72834706623d94ac93a725410b6a6b6fbc1b07f7ba56" + +[[package]] +name = "glob" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b919933a397b79c37e33b77bb2aa3dc8eb6e165ad809e58ff75bc7db2e34574" + +[[package]] +name = "hermit-abi" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62b467343b94ba476dcb2500d242dadbb39557df889310ac77c5d99100aaac33" +dependencies = [ + "libc", +] + +[[package]] +name = "itoa" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c8af84674fe1f223a982c933a0ee1086ac4d4052aa0fb8060c12c6ad838e754" + +[[package]] +name = "lazy_static" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" + +[[package]] +name = "libc" +version = "0.2.133" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f80d65747a3e43d1596c7c5492d95d5edddaabd45a7fcdb02b95f644164966" + +[[package]] +name = "lock_api" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "435011366fe56583b16cf956f9df0095b405b82d76425bc8981c0e22e60ec4df" +dependencies = [ + "autocfg", + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "abb12e687cfb44aa40f41fc3978ef76448f9b6038cad6aef4259d3c095a2382e" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "memchr" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d" + +[[package]] +name = "mio" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57ee1c23c7c63b0c9250c339ffdc69255f110b298b901b9f6c82547b7b87caaf" +dependencies = [ + "libc", + "log", + "wasi", + "windows-sys", +] + +[[package]] +name = "num_cpus" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19e64526ebdee182341572e50e9ad03965aa510cd94427a4549448f285e957a1" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "num_threads" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2819ce041d2ee131036f4fc9d6ae7ae125a3a40e97ba64d04fe799ad9dabbb44" +dependencies = [ + "libc", +] + +[[package]] +name = "once_cell" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e82dad04139b71a90c080c8463fe0dc7902db5192d939bd0950f074d014339e1" + +[[package]] +name = "parking_lot" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09a279cbf25cb0757810394fbc1e359949b59e348145c643a939a525692e6929" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-sys", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0a7ae3ac2f1173085d398531c705756c94a4c56843785df85a60c1a0afac116" + +[[package]] +name = "proc-macro2" +version = "1.0.43" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a2ca2c61bc9f3d74d2886294ab7b9853abd9c1ad903a3ac7815c58989bb7bab" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbe448f377a7d6961e30f5955f9b8d106c3f5e449d493ee1b125c1d43c2b5179" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c4eb3267174b8c6c2f654116623910a0fef09c4753f8dd83db29c48a0df988b" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.6.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3f87b73ce11b1619a3c6332f45341e0047173771e8b8b73f87bfeefb7b56244" + +[[package]] +name = "rustversion" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97477e48b4cf8603ad5f7aaf897467cf42ab4218a38ef76fb14c2d6773a6d6a8" + +[[package]] +name = "scopeguard" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" + +[[package]] +name = "serde" +version = "1.0.145" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728eb6351430bccb993660dfffc5a72f91ccc1295abaa8ce19b27ebe4f75568b" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.145" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81fa1584d3d1bcacd84c277a0dfe21f5b0f6accf4a23d04d4c6d61f1af522b4c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "signal-hook-registry" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e51e73328dc4ac0c7ccbda3a494dfa03df1de2f46018127f60c693f2648455b0" +dependencies = [ + "libc", +] + +[[package]] +name = "smallvec" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fd0db749597d91ff862fd1d55ea87f7855a744a8425a64695b6fca237d1dad1" + +[[package]] +name = "socket2" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02e2d2db9033d13a1567121ddd7a095ee144db4e1ca1b1bda3419bc0da294ebd" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "socks5ws" +version = "0.1.0" +dependencies = [ + "fast-socks5", + "flexi_logger", + "log", + "serde", + "serde_derive", + "tokio", + "tokio-stream", + "tokio-util", + "toml", +] + +[[package]] +name = "syn" +version = "1.0.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52205623b1b0f064a4e71182c3b18ae902267282930c6d5462c91b859668426e" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c53f98874615aea268107765aa1ed8f6116782501d18e53d08b471733bea6c85" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8b463991b4eab2d801e724172285ec4195c650e8ec79b149e6c2a8e6dd3f783" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "time" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c3f9a28b618c3a6b9251b6908e9c99e04b9e5c02e6581ccbb67d59c34ef7f9b" +dependencies = [ + "itoa", + "libc", + "num_threads", + "time-macros", +] + +[[package]] +name = "time-macros" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42657b1a6f4d817cda8e7a0ace261fe0cc946cf3a80314390b22cc61ae080792" + +[[package]] +name = "tokio" +version = "1.21.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0020c875007ad96677dcc890298f4b942882c5d4eb7cc8f439fc3bf813dc9c95" +dependencies = [ + "autocfg", + "bytes", + "libc", + "memchr", + "mio", + "num_cpus", + "once_cell", + "parking_lot", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "winapi", +] + +[[package]] +name = "tokio-macros" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9724f9a975fb987ef7a3cd9be0350edcbe130698af5b8f7a631e23d42d052484" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio-stream" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6edf2d6bc038a43d31353570e27270603f4648d18f5ed10c0e179abe43255af" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tokio-util" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bb2e075f03b3d66d8d8785356224ba688d2906a371015e225beeb65ca92c740" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "toml" +version = "0.5.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d82e1a7758622a465f8cee077614c73484dac5b836c02ff6a40d5d1010324d7" +dependencies = [ + "serde", +] + +[[package]] +name = "unicode-ident" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcc811dc4066ac62f84f11307873c4850cb653bfa9b1719cee2bd2204a4bc5dd" + +[[package]] +name = "wasi" +version = "0.11.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows-sys" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea04155a16a59f9eab786fe12a4a450e75cdb175f9e0d80da1e17db09f55b8d2" +dependencies = [ + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_msvc" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bb8c3fd39ade2d67e9874ac4f3db21f0d710bee00fe7cab16949ec184eeaa47" + +[[package]] +name = "windows_i686_gnu" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "180e6ccf01daf4c426b846dfc66db1fc518f074baa793aa7d9b9aaeffad6a3b6" + +[[package]] +name = "windows_i686_msvc" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2e7917148b2812d1eeafaeb22a97e4813dfa60a3f8f78ebe204bcc88f12f024" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4dcd171b8776c41b97521e5da127a2d86ad280114807d0b2ab1e462bc764d9e1" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c811ca4a8c853ef420abd8592ba53ddbbac90410fab6903b3e79972a631f7680" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..c6c9e28 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "socks5ws" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +fast-socks5 = "0.8.1" +flexi_logger = { version = "0.23.3", features = ["specfile_without_notification", "async"] } +log = "0.4.17" +serde = { version = "1.0.145", features = ["derive"] } +serde_derive = "1.0.145" +tokio = { version = "1.21.1", features = ["full", "winapi", "mio"] } +tokio-stream = "0.1.10" +tokio-util = "0.7.4" +toml = "0.5.9" diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..5175474 --- /dev/null +++ b/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, + /// 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 { + 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 { + 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, + } + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..7260c2e --- /dev/null +++ b/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(); +} diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..41a17fc --- /dev/null +++ b/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(future: F, token: CancellationToken) -> Option +where + F: Future>, +{ + select! { + _ = token.cancelled() => { + log::error!("canceled"); + None + } + res = future => { + res + } + } +} + +fn spawn_and_log_error(future: F, token: CancellationToken) -> task::JoinHandle<()> +where + F: Future>> + 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); + } + }) +}