diff --git a/Cargo.lock b/Cargo.lock index 38f6d26..1ec1070 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -61,6 +61,45 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "clap" +version = "3.2.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86447ad904c7fb335a790c9d7fe3d0d971dc523b8ccd1561a520de9a85302750" +dependencies = [ + "atty", + "bitflags", + "clap_derive", + "clap_lex", + "indexmap", + "once_cell", + "strsim", + "termcolor", + "textwrap", +] + +[[package]] +name = "clap_derive" +version = "3.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea0c8bce528c4be4da13ea6fead8965e95b6073585a2f05204bd8f4119f82a65" +dependencies = [ + "heck", + "proc-macro-error", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2850f2f5a82cbf437dd5af4d49848fbdfc27c157c3d010345776f952765261c5" +dependencies = [ + "os_str_bytes", +] + [[package]] name = "crossbeam-channel" version = "0.5.6" @@ -91,6 +130,20 @@ dependencies = [ "once_cell", ] +[[package]] +name = "err-derive" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c34a887c8df3ed90498c1c437ce21f211c8e27672921a8ffa293cb8d6d4caa9e" +dependencies = [ + "proc-macro-error", + "proc-macro2", + "quote", + "rustversion", + "syn", + "synstructure", +] + [[package]] name = "fast-socks5" version = "0.8.1" @@ -144,6 +197,18 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b919933a397b79c37e33b77bb2aa3dc8eb6e165ad809e58ff75bc7db2e34574" +[[package]] +name = "hashbrown" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" + +[[package]] +name = "heck" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2540771e65fc8cb83cd6e8a237f70c319bd5c29f78ed1084ba5d50eeac86f7f9" + [[package]] name = "hermit-abi" version = "0.1.19" @@ -153,6 +218,16 @@ dependencies = [ "libc", ] +[[package]] +name = "indexmap" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a35a97730320ffe8e2d410b5d3b69279b98d2c14bdb8b70ea89ecf7888d41e" +dependencies = [ + "autocfg", + "hashbrown", +] + [[package]] name = "itoa" version = "1.0.3" @@ -233,6 +308,12 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e82dad04139b71a90c080c8463fe0dc7902db5192d939bd0950f074d014339e1" +[[package]] +name = "os_str_bytes" +version = "6.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ff7415e9ae3fff1225851df9e0d9e4e5479f947619774677a63572e55e80eff" + [[package]] name = "parking_lot" version = "0.12.1" @@ -262,6 +343,30 @@ version = "0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e0a7ae3ac2f1173085d398531c705756c94a4c56843785df85a60c1a0afac116" +[[package]] +name = "proc-macro-error" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" +dependencies = [ + "proc-macro-error-attr", + "proc-macro2", + "quote", + "syn", + "version_check", +] + +[[package]] +name = "proc-macro-error-attr" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" +dependencies = [ + "proc-macro2", + "quote", + "version_check", +] + [[package]] name = "proc-macro2" version = "1.0.43" @@ -367,6 +472,7 @@ dependencies = [ name = "socks5ws" version = "0.1.0" dependencies = [ + "clap", "fast-socks5", "flexi_logger", "log", @@ -376,8 +482,15 @@ dependencies = [ "tokio-stream", "tokio-util", "toml", + "windows-service", ] +[[package]] +name = "strsim" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" + [[package]] name = "syn" version = "1.0.100" @@ -389,6 +502,33 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "synstructure" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f36bdaa60a83aca3921b5259d5400cbf5e90fc51931376a9bd4a0eb79aa7210f" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "unicode-xid", +] + +[[package]] +name = "termcolor" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bab24d30b911b2376f3a13cc2cd443142f0c81dda04c118693e35b3835757755" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "textwrap" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "949517c0cf1bf4ee812e2e07e08ab448e3ae0d23472aee8a06c985f0c8815b16" + [[package]] name = "thiserror" version = "1.0.35" @@ -498,12 +638,30 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dcc811dc4066ac62f84f11307873c4850cb653bfa9b1719cee2bd2204a4bc5dd" +[[package]] +name = "unicode-xid" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f962df74c8c05a667b5ee8bcf162993134c104e96440b663c8daa176dc772d8c" + +[[package]] +name = "version_check" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" + [[package]] name = "wasi" version = "0.11.0+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +[[package]] +name = "widestring" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "653f141f39ec16bba3c5abe400a0c60da7468261cc2cbf36805022876bc721a8" + [[package]] name = "winapi" version = "0.3.9" @@ -520,12 +678,33 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" +[[package]] +name = "winapi-util" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" +dependencies = [ + "winapi", +] + [[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-service" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "917fdb865e7ff03af9dd86609f8767bc88fefba89e8efd569de8e208af8724b3" +dependencies = [ + "bitflags", + "err-derive", + "widestring", + "windows-sys", +] + [[package]] name = "windows-sys" version = "0.36.1" diff --git a/Cargo.toml b/Cargo.toml index c6c9e28..cbdae10 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +clap = { version = "3.2.22", features = ["derive"] } fast-socks5 = "0.8.1" flexi_logger = { version = "0.23.3", features = ["specfile_without_notification", "async"] } log = "0.4.17" @@ -15,3 +16,4 @@ tokio = { version = "1.21.1", features = ["full", "winapi", "mio"] } tokio-stream = "0.1.10" tokio-util = "0.7.4" toml = "0.5.9" +windows-service = "0.5.0" diff --git a/src/main.rs b/src/main.rs index bd2fa26..ea3e364 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,18 +1,55 @@ extern crate flexi_logger; -use tokio_util::sync::CancellationToken; - mod config; mod server; -use crate::config::Config; -use crate::server::server_executor; +mod service; use flexi_logger::{AdaptiveFormat, Age, Cleanup, Criterion, Duplicate, FileSpec, Logger, Naming}; +use clap::{Parser, Subcommand}; + +#[derive(Subcommand, Debug)] +enum Command { + /// install service + Install, + /// uninstall service + Uninstall, + /// start service + Start, + /// stop service + Stop, + /// run service (by Windows) + Run, +} + +/// SOCKS5 proxy windows service +#[derive(Parser, Debug)] +#[clap(author, version, about, long_about = None)] +#[clap(propagate_version = true)] +struct Cli { + #[clap(subcommand)] + command: Command, +} + +macro_rules! handle_error { + ($name:expr, $res:expr) => { + match $res { + Err(e) => { + log::error!("{} error: {:#?}", $name, e) + } + _ => (), + } + }; +} + fn main() { + let args = Cli::parse(); + Logger::try_with_str("info") .unwrap() - .log_to_file(FileSpec::default()) + .log_to_file( + FileSpec::default().directory(std::env::current_exe().unwrap().parent().unwrap()), + ) .rotate( Criterion::Age(Age::Day), Naming::Timestamps, @@ -21,18 +58,18 @@ fn main() { .adaptive_format_for_stderr(AdaptiveFormat::Detailed) .print_message() .duplicate_to_stderr(Duplicate::Warn) - .start_with_specfile("logspec.toml") + .start_with_specfile( + std::env::current_exe() + .unwrap() + .with_file_name("logspec.toml"), + ) .unwrap(); - let cfg = Config::get(); - log::info!("cfg: {:#?}", cfg); - - let token = CancellationToken::new(); - let child_token = token.child_token(); - let handle = std::thread::spawn(move || server_executor(cfg, child_token)); - - std::thread::sleep(std::time::Duration::from_secs(10)); - token.cancel(); - - handle.join().unwrap(); + match args.command { + Command::Install => handle_error!("install", service::install()), + Command::Uninstall => handle_error!("uninstall", service::uninstall()), + Command::Run => handle_error!("run", service::run()), + Command::Start => handle_error!("start", service::start()), + Command::Stop => handle_error!("stop", service::stop()), + } } diff --git a/src/service.rs b/src/service.rs new file mode 100644 index 0000000..571e07f --- /dev/null +++ b/src/service.rs @@ -0,0 +1,183 @@ +use tokio_util::sync::CancellationToken; + +use std::{ffi::OsString, sync::mpsc, thread, time::Duration}; +use windows_service::{ + define_windows_service, + service::{ + ServiceAccess, ServiceControl, ServiceControlAccept, ServiceErrorControl, ServiceExitCode, + ServiceInfo, ServiceStartType, ServiceState, ServiceStatus, ServiceType, + }, + service_control_handler::{self, ServiceControlHandlerResult}, + service_dispatcher, + service_manager::{ServiceManager, ServiceManagerAccess}, + Result, +}; + +use crate::config::Config; +use crate::server::server_executor; + +const SERVICE_NAME: &str = "socks5ws_srv"; +const SERVICE_TYPE: ServiceType = ServiceType::OWN_PROCESS; +const SERVICE_DISPLAY: &str = "socks5ws proxy"; +const SERVICE_DESCRIPTION: &str = "SOCKS5 proxy windows service"; + +pub fn install() -> windows_service::Result<()> { + let manager_access = ServiceManagerAccess::CONNECT | ServiceManagerAccess::CREATE_SERVICE; + let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?; + + let service_binary_path = ::std::env::current_exe().unwrap(); + + let service_info = ServiceInfo { + name: SERVICE_NAME.into(), + display_name: OsString::from(SERVICE_DISPLAY), + service_type: ServiceType::OWN_PROCESS, + start_type: ServiceStartType::OnDemand, + error_control: ServiceErrorControl::Normal, + executable_path: service_binary_path, + launch_arguments: vec!["run".into()], + dependencies: vec![], + account_name: Some(OsString::from(r#"NT AUTHORITY\NetworkService"#)), + account_password: None, + }; + let service = service_manager.create_service(&service_info, ServiceAccess::CHANGE_CONFIG)?; + service.set_description(SERVICE_DESCRIPTION)?; + log::info!("service installed"); + Ok(()) +} + +pub fn uninstall() -> windows_service::Result<()> { + let manager_access = ServiceManagerAccess::CONNECT; + let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?; + + let service_access = ServiceAccess::QUERY_STATUS | ServiceAccess::STOP | ServiceAccess::DELETE; + let service = service_manager.open_service(SERVICE_NAME, service_access)?; + + let service_status = service.query_status()?; + if service_status.current_state != ServiceState::Stopped { + log::warn!("stopping service"); + service.stop()?; + // Wait for service to stop + thread::sleep(Duration::from_secs(5)); + } + + service.delete()?; + log::warn!("service deleted"); + Ok(()) +} + +pub fn stop() -> windows_service::Result<()> { + let manager_access = ServiceManagerAccess::CONNECT; + let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?; + + let service_access = ServiceAccess::QUERY_STATUS | ServiceAccess::STOP; + let service = service_manager.open_service(SERVICE_NAME, service_access)?; + + let service_status = service.query_status()?; + if service_status.current_state != ServiceState::Stopped { + log::info!("stopping service"); + service.stop()?; + } + Ok(()) +} + +pub fn start() -> windows_service::Result<()> { + let manager_access = ServiceManagerAccess::CONNECT; + let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?; + + let service_access = ServiceAccess::QUERY_STATUS | ServiceAccess::START; + let service = service_manager.open_service(SERVICE_NAME, service_access)?; + + let service_status = service.query_status()?; + if service_status.current_state != ServiceState::Running { + log::info!("start service"); + service.start(Vec::<&str>::new().as_slice())?; + } + Ok(()) +} + +pub fn run() -> Result<()> { + // Register generated `ffi_service_main` with the system and start the service, blocking + // this thread until the service is stopped. + log::info!("service run"); + service_dispatcher::start(SERVICE_NAME, ffi_service_main) +} + +// Generate the windows service boilerplate. +// The boilerplate contains the low-level service entry function (ffi_service_main) that parses +// incoming service arguments into Vec and passes them to user defined service +// entry (my_service_main). +define_windows_service!(ffi_service_main, my_service_main); + +// Service entry function which is called on background thread by the system with service +// parameters. There is no stdout or stderr at this point so make sure to configure the log +// output to file if needed. +pub fn my_service_main(_arguments: Vec) { + if let Err(e) = run_service() { + log::error!("{:#?}", e); + } +} + +pub fn run_service() -> Result<()> { + // Create a channel to be able to poll a stop event from the service worker loop. + let (shutdown_tx, shutdown_rx) = mpsc::channel(); + + // Define system service event handler that will be receiving service events. + let event_handler = move |control_event| -> ServiceControlHandlerResult { + match control_event { + // Notifies a service to report its current status information to the service + // control manager. Always return NoError even if not implemented. + ServiceControl::Interrogate => ServiceControlHandlerResult::NoError, + + // Handle stop + ServiceControl::Stop => { + log::debug!("Stop signal from system"); + shutdown_tx.send(()).unwrap(); + ServiceControlHandlerResult::NoError + } + + _ => ServiceControlHandlerResult::NotImplemented, + } + }; + + // Register system service event handler. + // The returned status handle should be used to report service status changes to the system. + let status_handle = service_control_handler::register(SERVICE_NAME, event_handler)?; + + // Tell the system that service is running + status_handle.set_service_status(ServiceStatus { + service_type: SERVICE_TYPE, + current_state: ServiceState::Running, + controls_accepted: ServiceControlAccept::STOP, + exit_code: ServiceExitCode::Win32(0), + checkpoint: 0, + wait_hint: Duration::default(), + process_id: None, + })?; + + let cfg = Config::get(); + log::info!("start with config: {:#?}", cfg); + + let token = CancellationToken::new(); + let child_token = token.child_token(); + let server_handle = std::thread::spawn(move || server_executor(cfg, child_token)); + + shutdown_rx.recv().unwrap(); // wait for shutdown signal + log::info!("service stop"); + + // stop server + token.cancel(); + server_handle.join().unwrap(); + + // Tell the system that service has stopped. + status_handle.set_service_status(ServiceStatus { + service_type: SERVICE_TYPE, + current_state: ServiceState::Stopped, + controls_accepted: ServiceControlAccept::empty(), + exit_code: ServiceExitCode::Win32(0), + checkpoint: 0, + wait_hint: Duration::default(), + process_id: None, + })?; + + Ok(()) +}