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.
This commit is contained in:
咸鱼而已 2024-10-26 01:32:25 +08:00 committed by GitHub
parent a78b759741
commit 1ac2e1c8e3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 128 additions and 1 deletions

18
Cargo.lock generated
View File

@ -1691,6 +1691,7 @@ dependencies = [
"url", "url",
"uuid", "uuid",
"wildmatch", "wildmatch",
"windows-service",
"windows-sys 0.52.0", "windows-sys 0.52.0",
"winreg 0.52.0", "winreg 0.52.0",
"zerocopy", "zerocopy",
@ -7345,6 +7346,12 @@ dependencies = [
"rustix", "rustix",
] ]
[[package]]
name = "widestring"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7219d36b6eac893fa81e84ebe06485e7dcbb616177469b142df14f1f4deb1311"
[[package]] [[package]]
name = "wildmatch" name = "wildmatch"
version = "2.3.4" version = "2.3.4"
@ -7466,6 +7473,17 @@ dependencies = [
"windows-targets 0.52.6", "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]] [[package]]
name = "windows-sys" name = "windows-sys"
version = "0.45.0" version = "0.45.0"

View File

@ -193,6 +193,7 @@ windows-sys = { version = "0.52", features = [
] } ] }
encoding = "0.2" encoding = "0.2"
winreg = "0.52" winreg = "0.52"
windows-service = "0.7.0"
[build-dependencies] [build-dependencies]
tonic-build = "0.12" tonic-build = "0.12"

View File

@ -4,7 +4,7 @@
extern crate rust_i18n; extern crate rust_i18n;
use std::{ use std::{
net::{Ipv4Addr, SocketAddr}, net::{Ipv4Addr, SocketAddr},
path::PathBuf, path::PathBuf,
}; };
@ -28,6 +28,9 @@ use easytier::{
web_client, web_client,
}; };
#[cfg(target_os = "windows")]
windows_service::define_windows_service!(ffi_service_main, win_service_main);
#[cfg(feature = "mimalloc")] #[cfg(feature = "mimalloc")]
use mimalloc_rust::GlobalMiMalloc; 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<tokio::sync::Notify>,
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<std::ffi::OsString>) {
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] #[tokio::main]
async fn main() { async fn main() {
let locale = sys_locale::get_locale().unwrap_or_else(|| String::from("en-US")); let locale = sys_locale::get_locale().unwrap_or_else(|| String::from("en-US"));
rust_i18n::set_locale(&locale); 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(); let cli = Cli::parse();
setup_panic_handler(); setup_panic_handler();