From 1ac2e1c8e3907285353c1479e8e812b634a58f8c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=92=B8=E9=B1=BC=E8=80=8C=E5=B7=B2?= <64669899+hvvvvvvv@users.noreply.github.com> Date: Sat, 26 Oct 2024 01:32:25 +0800 Subject: [PATCH] Support running as a Windows service. (#445) When the program is started by the Service Control Manager (SCM), it needs to call StartServiceCtrlDispatcher (which corresponds to service_dispatcher::start in the windows-service crate) to inform the SCM of the service's entry function. The SCM then calls the service entry function passed to StartServiceCtrlDispatcher. The process calling StartServiceCtrlDispatcher will block until the service's status is set to Stopped. If the current program is not run through the SCM, StartServiceCtrlDispatcher will return the error ERROR_FAILED_SERVICE_CONTROLLER_CONNECT, and the program will run according to its original mechanism. For more details about SCM, please refer to Microsoft's documentation. --- Cargo.lock | 18 ++++++ easytier/Cargo.toml | 1 + easytier/src/easytier-core.rs | 110 +++++++++++++++++++++++++++++++++- 3 files changed, 128 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 9cfb919..e885bd2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1691,6 +1691,7 @@ dependencies = [ "url", "uuid", "wildmatch", + "windows-service", "windows-sys 0.52.0", "winreg 0.52.0", "zerocopy", @@ -7345,6 +7346,12 @@ dependencies = [ "rustix", ] +[[package]] +name = "widestring" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7219d36b6eac893fa81e84ebe06485e7dcbb616177469b142df14f1f4deb1311" + [[package]] name = "wildmatch" version = "2.3.4" @@ -7466,6 +7473,17 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "windows-service" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d24d6bcc7f734a4091ecf8d7a64c5f7d7066f45585c1861eba06449909609c8a" +dependencies = [ + "bitflags 2.6.0", + "widestring", + "windows-sys 0.52.0", +] + [[package]] name = "windows-sys" version = "0.45.0" diff --git a/easytier/Cargo.toml b/easytier/Cargo.toml index 9f3bec5..57e6146 100644 --- a/easytier/Cargo.toml +++ b/easytier/Cargo.toml @@ -193,6 +193,7 @@ windows-sys = { version = "0.52", features = [ ] } encoding = "0.2" winreg = "0.52" +windows-service = "0.7.0" [build-dependencies] tonic-build = "0.12" diff --git a/easytier/src/easytier-core.rs b/easytier/src/easytier-core.rs index 4c11ede..0c22229 100644 --- a/easytier/src/easytier-core.rs +++ b/easytier/src/easytier-core.rs @@ -4,7 +4,7 @@ extern crate rust_i18n; use std::{ - net::{Ipv4Addr, SocketAddr}, + net::{Ipv4Addr, SocketAddr}, path::PathBuf, }; @@ -28,6 +28,9 @@ use easytier::{ web_client, }; +#[cfg(target_os = "windows")] +windows_service::define_windows_service!(ffi_service_main, win_service_main); + #[cfg(feature = "mimalloc")] use mimalloc_rust::GlobalMiMalloc; @@ -649,11 +652,116 @@ pub fn handle_event(mut events: EventBusSubscriber) -> tokio::task::JoinHandle<( }) } +#[cfg(target_os = "windows")] +fn win_service_event_loop( + stop_notify: std::sync::Arc, + inst: launcher::NetworkInstance, + status_handle: windows_service::service_control_handler::ServiceStatusHandle, +) { + use tokio::runtime::Runtime; + use std::time::Duration; + use windows_service::service::*; + + std::thread::spawn(move || { + let rt = Runtime::new().unwrap(); + rt.block_on(async move { + tokio::select! { + res = inst.wait() => { + if let Some(e) = res { + status_handle.set_service_status(ServiceStatus { + service_type: ServiceType::OWN_PROCESS, + current_state: ServiceState::Stopped, + controls_accepted: ServiceControlAccept::empty(), + checkpoint: 0, + wait_hint: Duration::default(), + exit_code: ServiceExitCode::ServiceSpecific(1u32), + process_id: None + }).unwrap(); + panic!("launcher error: {:?}", e); + } + }, + _ = stop_notify.notified() => { + status_handle.set_service_status(ServiceStatus { + service_type: ServiceType::OWN_PROCESS, + current_state: ServiceState::Stopped, + controls_accepted: ServiceControlAccept::empty(), + checkpoint: 0, + wait_hint: Duration::default(), + exit_code: ServiceExitCode::Win32(0), + process_id: None + }).unwrap(); + std::process::exit(0); + } + } + }); + }); +} + +#[cfg(target_os = "windows")] +fn win_service_main(_: Vec) { + use std::time::Duration; + use windows_service::service_control_handler::*; + use windows_service::service::*; + use std::sync::Arc; + use tokio::sync::Notify; + + let cli = Cli::parse(); + let cfg = TomlConfigLoader::from(cli); + + init_logger(&cfg, false).unwrap(); + + let stop_notify_send = Arc::new(Notify::new()); + let stop_notify_recv = Arc::clone(&stop_notify_send); + let event_handler = move |control_event| -> ServiceControlHandlerResult { + match control_event { + ServiceControl::Interrogate => { + ServiceControlHandlerResult::NoError + } + ServiceControl::Stop => + { + stop_notify_send.notify_one(); + ServiceControlHandlerResult::NoError + } + _ => ServiceControlHandlerResult::NotImplemented, + } + }; + let status_handle = register(String::new(), event_handler).expect("register service fail"); + let next_status = ServiceStatus { + service_type: ServiceType::OWN_PROCESS, + current_state: ServiceState::Running, + controls_accepted: ServiceControlAccept::STOP, + exit_code: ServiceExitCode::Win32(0), + checkpoint: 0, + wait_hint: Duration::default(), + process_id: None, + }; + let mut inst = launcher::NetworkInstance::new(cfg).set_fetch_node_info(false); + + inst.start().unwrap(); + status_handle.set_service_status(next_status).expect("set service status fail"); + win_service_event_loop(stop_notify_recv, inst, status_handle); +} + #[tokio::main] async fn main() { let locale = sys_locale::get_locale().unwrap_or_else(|| String::from("en-US")); rust_i18n::set_locale(&locale); + #[cfg(target_os = "windows")] + match windows_service::service_dispatcher::start(String::new(), ffi_service_main) { + Ok(_) => std::thread::park(), + Err(e) => + { + let should_panic = if let windows_service::Error::Winapi(ref io_error) = e { + io_error.raw_os_error() != Some(0x427) // ERROR_FAILED_SERVICE_CONTROLLER_CONNECT + } else { true }; + + if should_panic { + panic!("SCM start an error: {}", e); + } + } + }; + let cli = Cli::parse(); setup_panic_handler();