From 268e520a7230da2fd8fb9d014a8202ddbf39bb9e Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Tue, 19 Sep 2023 00:29:35 +0800 Subject: [PATCH] feat(sswinservice): Running ssservice as a Windows Service ref #1292 https://learn.microsoft.com/en-us/windows/win32/services/about-services --- Cargo.lock | 142 +++++++++++++++++++++------- Cargo.toml | 9 ++ bin/sswinservice.rs | 205 +++++++++++++++++++++++++++++++++++++++++ src/service/local.rs | 36 +++++--- src/service/manager.rs | 36 +++++--- src/service/server.rs | 36 +++++--- 6 files changed, 399 insertions(+), 65 deletions(-) create mode 100644 bin/sswinservice.rs diff --git a/Cargo.lock b/Cargo.lock index b6d47714..64986049 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -126,7 +126,7 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5ca11d4be1bab0c8bc8734a9aa7bf4ee8316d462a08c6ac5052f888fef5b494b" dependencies = [ - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -136,7 +136,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "58f54d10c6dfa51283a066ceab3ec1ab78d13fae00aa49243a45e4571fb79dfd" dependencies = [ "anstyle", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -379,7 +379,7 @@ dependencies = [ "js-sys", "num-traits", "wasm-bindgen", - "windows-targets", + "windows-targets 0.48.5", ] [[package]] @@ -633,7 +633,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -715,7 +715,7 @@ checksum = "136526188508e25c6fef639d7927dfb3e0e3084488bf202267829cf7fc23dbdd" dependencies = [ "errno-dragonfly", "libc", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -752,7 +752,7 @@ dependencies = [ "cfg-if", "libc", "redox_syscall 0.3.5", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -1179,7 +1179,7 @@ checksum = "eae7b9aee968036d54dce06cebaefd919e4472e753296daccd6d344e3e2df0c2" dependencies = [ "hermit-abi", "libc", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -1196,7 +1196,7 @@ checksum = "b58db92f96b720de98181bbbe63c831e87005ab460c1bf306eb2622b4707997f" dependencies = [ "socket2 0.5.3", "widestring", - "windows-sys", + "windows-sys 0.48.0", "winreg", ] @@ -1223,7 +1223,7 @@ checksum = "cb0889898416213fab133e1d33a0e5858a48177452750691bde3666d0fdbaf8b" dependencies = [ "hermit-abi", "rustix 0.38.13", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -1458,7 +1458,7 @@ dependencies = [ "libc", "log", "wasi", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -1506,7 +1506,7 @@ dependencies = [ "log", "mio", "walkdir", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -1658,7 +1658,7 @@ dependencies = [ "libc", "redox_syscall 0.3.5", "smallvec", - "windows-targets", + "windows-targets 0.48.5", ] [[package]] @@ -2088,7 +2088,7 @@ dependencies = [ "io-lifetimes", "libc", "linux-raw-sys 0.3.8", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -2101,7 +2101,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.4.7", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -2168,7 +2168,7 @@ version = "0.1.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c3733bf4cf7ea0880754e19cb5a462007c4a8c1914bff372ccc95b464f1df88" dependencies = [ - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -2361,7 +2361,7 @@ dependencies = [ "tokio-tfo", "trust-dns-resolver", "url", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -2421,6 +2421,7 @@ dependencies = [ "tcmalloc", "thiserror", "tokio", + "windows-service", "xdg", ] @@ -2465,7 +2466,7 @@ dependencies = [ "trust-dns-resolver", "tun", "webpki-roots", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -2557,7 +2558,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2538b18701741680e0322a2302176d3253a35388e2e62f172f64f4f16605f877" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -2656,7 +2657,7 @@ dependencies = [ "fastrand", "redox_syscall 0.3.5", "rustix 0.38.13", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -2675,7 +2676,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e6bf6f19e9f8ed8d4048dc22981458ebcf406d67e94cd422e5ecd73d63b3237" dependencies = [ "rustix 0.37.23", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -2740,7 +2741,7 @@ dependencies = [ "signal-hook-registry", "socket2 0.5.3", "tokio-macros", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -2788,7 +2789,7 @@ dependencies = [ "pin-project", "socket2 0.5.3", "tokio", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] @@ -3206,7 +3207,27 @@ version = "0.48.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e686886bc078bc1b0b600cac0147aadb815089b6e4da64016cbd754b6342700f" dependencies = [ - "windows-targets", + "windows-targets 0.48.5", +] + +[[package]] +name = "windows-service" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd9db37ecb5b13762d95468a2fc6009d4b2c62801243223aabd44fca13ad13c8" +dependencies = [ + "bitflags 1.3.2", + "widestring", + "windows-sys 0.45.0", +] + +[[package]] +name = "windows-sys" +version = "0.45.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0" +dependencies = [ + "windows-targets 0.42.2", ] [[package]] @@ -3215,7 +3236,22 @@ version = "0.48.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" dependencies = [ - "windows-targets", + "windows-targets 0.48.5", +] + +[[package]] +name = "windows-targets" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071" +dependencies = [ + "windows_aarch64_gnullvm 0.42.2", + "windows_aarch64_msvc 0.42.2", + "windows_i686_gnu 0.42.2", + "windows_i686_msvc 0.42.2", + "windows_x86_64_gnu 0.42.2", + "windows_x86_64_gnullvm 0.42.2", + "windows_x86_64_msvc 0.42.2", ] [[package]] @@ -3224,51 +3260,93 @@ version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", + "windows_aarch64_gnullvm 0.48.5", + "windows_aarch64_msvc 0.48.5", + "windows_i686_gnu 0.48.5", + "windows_i686_msvc 0.48.5", + "windows_x86_64_gnu 0.48.5", + "windows_x86_64_gnullvm 0.48.5", + "windows_x86_64_msvc 0.48.5", ] +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" + [[package]] name = "windows_aarch64_gnullvm" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" +[[package]] +name = "windows_aarch64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" + [[package]] name = "windows_aarch64_msvc" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" +[[package]] +name = "windows_i686_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" + [[package]] name = "windows_i686_gnu" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" +[[package]] +name = "windows_i686_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" + [[package]] name = "windows_i686_msvc" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" +[[package]] +name = "windows_x86_64_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" + [[package]] name = "windows_x86_64_gnu" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" + [[package]] name = "windows_x86_64_gnullvm" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" +[[package]] +name = "windows_x86_64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" + [[package]] name = "windows_x86_64_msvc" version = "0.48.5" @@ -3282,7 +3360,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "524e57b2c537c0f9b1e69f1965311ec12182b4122e45035b1508cd24d2adadb1" dependencies = [ "cfg-if", - "windows-sys", + "windows-sys 0.48.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index e2ffbc21..5da8636d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,11 @@ name = "ssservice" path = "bin/ssservice.rs" required-features = ["service"] +[[bin]] +name = "sswinservice" +path = "bin/sswinservice.rs" +required-features = ["winservice"] + [workspace] members = ["crates/shadowsocks", "crates/shadowsocks-service"] @@ -75,6 +80,8 @@ manager = ["shadowsocks-service/manager"] utility = ["qrcode"] # Enable service service = ["local", "server", "manager"] +# Enable Windows Service +winservice = ["service", "windows-service"] # Enables trust-dns for replacing tokio's builtin DNS resolver trust-dns = ["shadowsocks-service/trust-dns"] @@ -184,6 +191,8 @@ rpmalloc = { version = "0.2", optional = true } shadowsocks-service = { version = "1.16.1", path = "./crates/shadowsocks-service" } +windows-service = { version = "0.6", optional = true } + [target.'cfg(unix)'.dependencies] daemonize = "0.5" diff --git a/bin/sswinservice.rs b/bin/sswinservice.rs new file mode 100644 index 00000000..07aef218 --- /dev/null +++ b/bin/sswinservice.rs @@ -0,0 +1,205 @@ +use std::{ + ffi::OsString, + future::Future, + process::ExitCode, + sync::atomic::{AtomicU32, Ordering}, + time::Duration, +}; + +use clap::Command; +use log::{error, info}; +use shadowsocks_rust::service::{local, manager, server}; +use tokio::{runtime::Runtime, sync::oneshot}; +use windows_service::{ + define_windows_service, + service::{ServiceControl, ServiceControlAccept, ServiceExitCode, ServiceState, ServiceStatus, ServiceType}, + service_control_handler::{self, ServiceControlHandlerResult, ServiceStatusHandle}, + service_dispatcher, +}; + +const SERVICE_NAME: &str = "ssservice"; +const SERVICE_EXIT_CODE_ARGUMENT_ERROR: u32 = 100; +const SERVICE_EXIT_CODE_EXITED_UNEXPECTLY: u32 = 101; +const SERVICE_EXIT_CODE_CREATE_FAILED: u32 = 102; + +#[inline] +fn set_service_status( + handle: &ServiceStatusHandle, + current_state: ServiceState, + exit_code: ServiceExitCode, + wait_hint: Duration, +) -> Result<(), windows_service::Error> { + static SERVICE_STATE_CHECKPOINT: AtomicU32 = AtomicU32::new(0); + + let next_status = ServiceStatus { + service_type: ServiceType::OWN_PROCESS, + current_state, + controls_accepted: if current_state == ServiceState::StartPending { + ServiceControlAccept::empty() + } else { + ServiceControlAccept::STOP + }, + exit_code, + checkpoint: if matches!(current_state, ServiceState::Running | ServiceState::Stopped) { + SERVICE_STATE_CHECKPOINT.fetch_add(1, Ordering::AcqRel) + } else { + 0 + }, + wait_hint, + process_id: None, + }; + handle.set_service_status(next_status) +} + +fn handle_create_service_result( + status_handle: ServiceStatusHandle, + create_service_result: Result<(Runtime, F), ExitCode>, + stop_receiver: oneshot::Receiver<()>, +) -> Result<(), windows_service::Error> +where + F: Future, +{ + match create_service_result { + Ok((runtime, main_fut)) => { + // Successfully create runtime and future + + // Report running state + set_service_status( + &status_handle, + ServiceState::Running, + ServiceExitCode::Win32(0), + Duration::default(), + )?; + + // Run it right now. + let exited_by_ctrl = runtime.block_on(async move { + tokio::pin!(main_fut); + tokio::pin!(stop_receiver); + + loop { + tokio::select! { + _ = stop_receiver => { + break true; + } + exit_code = main_fut => { + info!("service exited unexpectly with code: {:?}", exit_code); + break false; + } + } + } + }); + + // Report stopped state + set_service_status( + &status_handle, + ServiceState::Stopped, + if exited_by_ctrl { + ServiceExitCode::Win32(0) + } else { + ServiceExitCode::ServiceSpecific(SERVICE_EXIT_CODE_EXITED_UNEXPECTLY) + }, + Duration::default(), + )?; + } + Err(exit_code) => { + error!("failed to create service, exit code: {:?}", exit_code); + + // Report running state + set_service_status( + &status_handle, + ServiceState::Stopped, + ServiceExitCode::ServiceSpecific(SERVICE_EXIT_CODE_CREATE_FAILED), + Duration::default(), + )?; + } + } + + Ok(()) +} + +fn service_main(arguments: Vec) -> Result<(), windows_service::Error> { + // Create a Oneshot channel for receiving Stop event + let (stop_sender, stop_receiver) = oneshot::channel(); + + let mut stop_sender_opt = Some(stop_sender); + let event_handler = move |control_event| -> ServiceControlHandlerResult { + match control_event { + ServiceControl::Stop => { + if let Some(stop_sender) = stop_sender_opt.take() { + let _ = stop_sender.send(()); + } + ServiceControlHandlerResult::NoError + } + ServiceControl::Interrogate => ServiceControlHandlerResult::NoError, + _ => ServiceControlHandlerResult::NotImplemented, + } + }; + + // Register system service event handler + let status_handle = service_control_handler::register(SERVICE_NAME, event_handler)?; + + // Report SERVICE_START_PENDING + // https://learn.microsoft.com/en-us/windows/win32/services/writing-a-servicemain-function + set_service_status( + &status_handle, + ServiceState::StartPending, + ServiceExitCode::Win32(0), + Duration::from_secs(30), + )?; + + let app = Command::new("shadowsocks service") + .version(shadowsocks_rust::VERSION) + .about("A fast tunnel proxy that helps you bypass firewalls. (https://shadowsocks.org)"); + + let app = app + .subcommand_required(true) + .subcommand(local::define_command_line_options(Command::new("local")).about("Shadowsocks Local service")) + .subcommand(server::define_command_line_options(Command::new("server")).about("Shadowsocks Server service")) + .subcommand( + manager::define_command_line_options(Command::new("manager")).about("Shadowsocks Server Manager service"), + ); + + let matches_result = if arguments.len() <= 1 { + app.try_get_matches() + } else { + app.try_get_matches_from(arguments) + }; + + let matches = match matches_result { + Ok(m) => m, + Err(err) => { + error!("failed to parse command line arguments, error: {}", err); + set_service_status( + &status_handle, + ServiceState::Stopped, + ServiceExitCode::ServiceSpecific(SERVICE_EXIT_CODE_ARGUMENT_ERROR), + Duration::default(), + )?; + return Err(windows_service::Error::LaunchArgumentsNotSupported); + } + }; + + match matches.subcommand() { + Some(("local", matches)) => handle_create_service_result(status_handle, local::create(matches), stop_receiver), + Some(("server", matches)) => { + handle_create_service_result(status_handle, server::create(matches), stop_receiver) + } + Some(("manager", matches)) => { + handle_create_service_result(status_handle, manager::create(matches), stop_receiver) + } + _ => Err(windows_service::Error::LaunchArgumentsNotSupported), + } +} + +fn service_entry(arguments: Vec) { + if let Err(err) = service_main(arguments) { + error!("service main exited with error: {}", err); + } +} + +define_windows_service!(ffi_service_entry, service_entry); + +fn main() -> Result<(), windows_service::Error> { + service_dispatcher::start(SERVICE_NAME, ffi_service_entry)?; + Ok(()) +} diff --git a/src/service/local.rs b/src/service/local.rs index 9780f559..6bfdd154 100644 --- a/src/service/local.rs +++ b/src/service/local.rs @@ -1,11 +1,14 @@ //! Local server launchers -use std::{net::IpAddr, path::PathBuf, process::ExitCode, time::Duration}; +use std::{future::Future, net::IpAddr, path::PathBuf, process::ExitCode, time::Duration}; use clap::{builder::PossibleValuesParser, Arg, ArgAction, ArgGroup, ArgMatches, Command, ValueHint}; use futures::future::{self, Either}; use log::{info, trace}; -use tokio::{self, runtime::Builder}; +use tokio::{ + self, + runtime::{Builder, Runtime}, +}; #[cfg(feature = "local-redir")] use shadowsocks_service::config::RedirType; @@ -496,8 +499,8 @@ pub fn define_command_line_options(mut app: Command) -> Command { app } -/// Program entrance `main` -pub fn main(matches: &ArgMatches) -> ExitCode { +/// Create `Runtime` and `main` entry +pub fn create(matches: &ArgMatches) -> Result<(Runtime, impl Future), ExitCode> { let (config, runtime) = { let config_path_opt = matches.get_one::("CONFIG").cloned().or_else(|| { if !matches.contains_id("SERVER_CONFIG") { @@ -518,7 +521,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { Ok(c) => c, Err(err) => { eprintln!("loading config {config_path:?}, {err}"); - return crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into()); } }, None => ServiceConfig::default(), @@ -542,7 +545,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { Ok(cfg) => cfg, Err(err) => { eprintln!("loading config {cpath:?}, {err}"); - return crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into()); } }, None => Config::new(ConfigType::Local), @@ -781,7 +784,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { Ok(acl) => acl, Err(err) => { eprintln!("loading ACL \"{acl_file}\", {err}"); - return crate::EXIT_CODE_LOAD_ACL_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_ACL_FAILURE.into()); } }; config.acl = Some(acl); @@ -831,12 +834,12 @@ pub fn main(matches: &ArgMatches) -> ExitCode { "missing `local_address`, consider specifying it by --local-addr command line option, \ or \"local_address\" and \"local_port\" in configuration file" ); - return crate::EXIT_CODE_INSUFFICIENT_PARAMS.into(); + return Err(crate::EXIT_CODE_INSUFFICIENT_PARAMS.into()); } if let Err(err) = config.check_integrity() { eprintln!("config integrity check failed, {err}"); - return crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into()); } #[cfg(unix)] @@ -870,7 +873,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { (config, runtime) }; - runtime.block_on(async move { + let main_fut = async move { let config_path = config.config_path.clone(); let instance = Server::new(config).await.expect("create local"); @@ -899,7 +902,18 @@ pub fn main(matches: &ArgMatches) -> ExitCode { // The abort signal future resolved. Means we should just exit. Either::Right(_) => ExitCode::SUCCESS, } - }) + }; + + Ok((runtime, main_fut)) +} + +/// Program entrance `main` +#[inline] +pub fn main(matches: &ArgMatches) -> ExitCode { + match create(matches) { + Ok((runtime, main_fut)) => runtime.block_on(main_fut), + Err(code) => code, + } } #[cfg(unix)] diff --git a/src/service/manager.rs b/src/service/manager.rs index afb5171e..0ec8e3d6 100644 --- a/src/service/manager.rs +++ b/src/service/manager.rs @@ -1,11 +1,14 @@ //! Server Manager launchers -use std::{net::IpAddr, path::PathBuf, process::ExitCode, time::Duration}; +use std::{future::Future, net::IpAddr, path::PathBuf, process::ExitCode, time::Duration}; use clap::{builder::PossibleValuesParser, Arg, ArgAction, ArgGroup, ArgMatches, Command, ValueHint}; use futures::future::{self, Either}; use log::{info, trace}; -use tokio::{self, runtime::Builder}; +use tokio::{ + self, + runtime::{Builder, Runtime}, +}; #[cfg(unix)] use shadowsocks_service::config::ManagerServerMode; @@ -256,8 +259,8 @@ pub fn define_command_line_options(mut app: Command) -> Command { app } -/// Program entrance `main` -pub fn main(matches: &ArgMatches) -> ExitCode { +/// Create `Runtime` and `main` entry +pub fn create(matches: &ArgMatches) -> Result<(Runtime, impl Future), ExitCode> { let (config, runtime) = { let config_path_opt = matches.get_one::("CONFIG").cloned().or_else(|| { if !matches.contains_id("SERVER_CONFIG") { @@ -278,7 +281,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { Ok(c) => c, Err(err) => { eprintln!("loading config {config_path:?}, {err}"); - return crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into()); } }, None => ServiceConfig::default(), @@ -302,7 +305,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { Ok(cfg) => cfg, Err(err) => { eprintln!("loading config {cpath:?}, {err}"); - return crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into()); } }, None => Config::new(ConfigType::Manager), @@ -409,7 +412,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { Ok(acl) => acl, Err(err) => { eprintln!("loading ACL \"{acl_file}\", {err}"); - return crate::EXIT_CODE_LOAD_ACL_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_ACL_FAILURE.into()); } }; config.acl = Some(acl); @@ -459,12 +462,12 @@ pub fn main(matches: &ArgMatches) -> ExitCode { "missing `manager_address`, consider specifying it by --manager-address command line option, \ or \"manager_address\" and \"manager_port\" keys in configuration file" ); - return crate::EXIT_CODE_INSUFFICIENT_PARAMS.into(); + return Err(crate::EXIT_CODE_INSUFFICIENT_PARAMS.into()); } if let Err(err) = config.check_integrity() { eprintln!("config integrity check failed, {err}"); - return crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into()); } #[cfg(unix)] @@ -503,7 +506,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { (config, runtime) }; - runtime.block_on(async move { + let main_fut = async move { let abort_signal = monitor::create_signal_monitor(); let server = run_manager(config); @@ -524,7 +527,18 @@ pub fn main(matches: &ArgMatches) -> ExitCode { // The abort signal future resolved. Means we should just exit. Either::Right(_) => ExitCode::SUCCESS, } - }) + }; + + Ok((runtime, main_fut)) +} + +/// Program entrance `main` +#[inline] +pub fn main(matches: &ArgMatches) -> ExitCode { + match create(matches) { + Ok((runtime, main_fut)) => runtime.block_on(main_fut), + Err(code) => code, + } } #[cfg(test)] diff --git a/src/service/server.rs b/src/service/server.rs index 0b9213f7..27505eb8 100644 --- a/src/service/server.rs +++ b/src/service/server.rs @@ -1,11 +1,14 @@ //! Server launchers -use std::{net::IpAddr, path::PathBuf, process::ExitCode, time::Duration}; +use std::{future::Future, net::IpAddr, path::PathBuf, process::ExitCode, time::Duration}; use clap::{builder::PossibleValuesParser, Arg, ArgAction, ArgGroup, ArgMatches, Command, ValueHint}; use futures::future::{self, Either}; use log::{info, trace}; -use tokio::{self, runtime::Builder}; +use tokio::{ + self, + runtime::{Builder, Runtime}, +}; use shadowsocks_service::{ acl::AccessControl, @@ -268,8 +271,8 @@ pub fn define_command_line_options(mut app: Command) -> Command { app } -/// Program entrance `main` -pub fn main(matches: &ArgMatches) -> ExitCode { +/// Create `Runtime` and `main` entry +pub fn create(matches: &ArgMatches) -> Result<(Runtime, impl Future), ExitCode> { let (config, runtime) = { let config_path_opt = matches.get_one::("CONFIG").cloned().or_else(|| { if !matches.contains_id("SERVER_CONFIG") { @@ -290,7 +293,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { Ok(c) => c, Err(err) => { eprintln!("loading config {config_path:?}, {err}"); - return crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into()); } }, None => ServiceConfig::default(), @@ -314,7 +317,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { Ok(cfg) => cfg, Err(err) => { eprintln!("loading config {cpath:?}, {err}"); - return crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into()); } }, None => Config::new(ConfigType::Server), @@ -428,7 +431,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { Ok(acl) => acl, Err(err) => { eprintln!("loading ACL \"{acl_file}\", {err}"); - return crate::EXIT_CODE_LOAD_ACL_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_ACL_FAILURE.into()); } }; config.acl = Some(acl); @@ -479,12 +482,12 @@ pub fn main(matches: &ArgMatches) -> ExitCode { --server-addr, --encrypt-method, --password command line option, \ or configuration file, check more details in https://shadowsocks.org/guide/configs.html" ); - return crate::EXIT_CODE_INSUFFICIENT_PARAMS.into(); + return Err(crate::EXIT_CODE_INSUFFICIENT_PARAMS.into()); } if let Err(err) = config.check_integrity() { eprintln!("config integrity check failed, {err}"); - return crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into(); + return Err(crate::EXIT_CODE_LOAD_CONFIG_FAILURE.into()); } #[cfg(unix)] @@ -523,7 +526,7 @@ pub fn main(matches: &ArgMatches) -> ExitCode { (config, runtime) }; - runtime.block_on(async move { + let main_fut = async move { let abort_signal = monitor::create_signal_monitor(); let server = run_server(config); @@ -544,7 +547,18 @@ pub fn main(matches: &ArgMatches) -> ExitCode { // The abort signal future resolved. Means we should just exit. Either::Right(_) => ExitCode::SUCCESS, } - }) + }; + + Ok((runtime, main_fut)) +} + +/// Program entrance `main` +#[inline] +pub fn main(matches: &ArgMatches) -> ExitCode { + match create(matches) { + Ok((runtime, main_fut)) => runtime.block_on(main_fut), + Err(code) => code, + } } #[cfg(test)]