From 1d37f7c21f86a17ea323b9b223473d11e9aa8667 Mon Sep 17 00:00:00 2001 From: chao wan <1013448513@qq.com> Date: Fri, 15 Nov 2024 20:58:46 +0800 Subject: [PATCH] Add more service installation options. --- easytier/locales/app.yml | 5 +- easytier/src/common/constants.rs | 2 + easytier/src/easytier-cli.rs | 267 +++++++++++++++++++++++++++---- easytier/src/easytier-core.rs | 36 +++-- 4 files changed, 259 insertions(+), 51 deletions(-) diff --git a/easytier/locales/app.yml b/easytier/locales/app.yml index bb4cc58..4941bf9 100644 --- a/easytier/locales/app.yml +++ b/easytier/locales/app.yml @@ -120,10 +120,7 @@ core_clap: ipv6_listener: en: "the url of the ipv6 listener, e.g.: tcp://[::]:11010, if not set, will listen on random udp port" zh-CN: "IPv6 监听器的URL，例如：tcp://[::]:11010，如果未设置，将在随机UDP端口上监听" - work_dir: - en: "Specify the working directory for the program. If not specified, the current directory will be used." - zh-CN: "指定程序的工作目录。如果未指定，将使用当前目录。" - + core_app: panic_backtrace_save: en: "backtrace saved to easytier-panic.log" diff --git a/easytier/src/common/constants.rs b/easytier/src/common/constants.rs index 3b86dba..fe561f2 100644 --- a/easytier/src/common/constants.rs +++ b/easytier/src/common/constants.rs @@ -25,6 +25,8 @@ define_global_var!(OSPF_UPDATE_MY_GLOBAL_FOREIGN_NETWORK_INTERVAL_SEC, u64, 10); pub const UDP_HOLE_PUNCH_CONNECTOR_SERVICE_ID: u32 = 2; +pub const WIN_SERVICE_WORK_DIR_REG_KEY: &str = "SOFTWARE\\EasyTier\\Service\\WorkDir"; + pub const EASYTIER_VERSION: &str = git_version::git_version!( args = ["--abbrev=8", "--always", "--dirty=~"], prefix = concat!(env!("CARGO_PKG_VERSION"), "-"), diff --git a/easytier/src/easytier-cli.rs b/easytier/src/easytier-cli.rs index aefa5d7..cab0a86 100644 --- a/easytier/src/easytier-cli.rs +++ b/easytier/src/easytier-cli.rs @@ -1,5 +1,6 @@ use std::{ ffi::OsString, + fmt::Write, net::SocketAddr, path::PathBuf, sync::Mutex, @@ -131,6 +132,9 @@ struct NodeArgs { #[derive(Args, Debug)] struct ServiceArgs{ + #[arg(short, long, default_value = env!("CARGO_PKG_NAME"), help = "service name")] + name: String, + #[command(subcommand)] sub_command: ServiceSubCommand } @@ -146,8 +150,22 @@ enum ServiceSubCommand { #[derive(Args, Debug)] struct InstallArgs { + #[arg(long, default_value = env!("CARGO_PKG_DESCRIPTION"), help = "service description")] + description: String, + + #[cfg(target_os = "windows")] + #[arg(long)] + display_name: Option, + + #[arg(long, default_value = "false")] + disable_autostart: bool, + #[arg(long)] core_path: Option, + + #[arg(long)] + service_work_dir: Option, + #[arg(trailing_var_arg = true, allow_hyphen_values = true)] core_args: Option> } @@ -508,43 +526,49 @@ impl CommandHandler { } } +pub struct ServiceInstallOptions { + pub program: PathBuf, + pub args: Vec, + pub work_directory: PathBuf, + pub disable_autostart: bool, + pub description: Option, + pub display_name: Option, +} pub struct Service{ lable: ServiceLabel, + kind: ServiceManagerKind, service_manager: Box } impl Service { - pub fn new() -> Result { + pub fn new(name: String) -> Result { #[cfg(target_os = "windows")] let service_manager = Box::new( - crate::win_service_manager::WinServiceManager::new( - Some(OsString::from("EasyTier Service")), - Some(OsString::from(env!("CARGO_PKG_DESCRIPTION"))), - vec![OsString::from("dnscache"), OsString::from("rpcss")], - )? + crate::win_service_manager::WinServiceManager::new()? ); #[cfg(not(target_os = "windows"))] let service_manager = ::native()?; + let kind = ServiceManagerKind::native()?; Ok(Self { - lable: env!("CARGO_PKG_NAME").parse().unwrap(), + lable: name.parse()?, + kind, service_manager }) } - pub fn install(&self, bin_path: std::path::PathBuf, bin_args: Vec) -> Result<(), Error> { + pub fn install(&self, options: &ServiceInstallOptions) -> Result<(), Error> { let ctx = ServiceInstallCtx { label: self.lable.clone(), - contents: None, - program: bin_path, - args: bin_args, - autostart: true, - username: None, - working_directory: None, + program: options.program.clone(), + args: options.args.clone(), + contents: self.make_install_content_option(options), + autostart: !options.disable_autostart, + username: None, + working_directory: Some(options.work_directory.clone()), environment: None, }; - if self.status()? != ServiceStatus::NotInstalled { return Err(anyhow::anyhow!("Service is already installed")); } @@ -620,8 +644,130 @@ impl Service { } } + fn make_install_content_option(&self, options: &ServiceInstallOptions) -> Option { + match self.kind { + ServiceManagerKind::Systemd => { + Some(self.make_systemd_unit(options).unwrap()) + } + ServiceManagerKind::Rcd => { + Some(self.make_rcd_script(options).unwrap()) + } + ServiceManagerKind::OpenRc => { + Some(self.make_open_rc_script(options).unwrap()) + } + _ => { + #[cfg(target_os = "windows")] { + let win_options = win_service_manager::WinServiceInstallOptions{ + description: options.description.clone(), + display_name: options.display_name.clone(), + dependencies: Some(vec!["rpcss".to_string(), "dnscache".to_string()]) + }; + + Some(serde_json::to_string(&win_options).unwrap()) + } + + #[cfg(not(target_os = "windows"))] + None + } + } + } + + fn make_systemd_unit(&self, options: &ServiceInstallOptions) -> Result { + let args = options.args.iter() + .map(|a| a.to_string_lossy()) + .collect::>() + .join(" "); + let target_app = options.program.display().to_string(); + let work_dir = options.work_directory.display().to_string(); + let mut unit_content = String::new(); + + writeln!(unit_content, "[Unit]")?; + writeln!(unit_content, "After=network.target syslog.target")?; + if let Some(ref d) = options.description { + writeln!(unit_content, "Description={d}")?; + } + writeln!(unit_content, "[Service]")?; + writeln!(unit_content, "Type=simple")?; + writeln!(unit_content, "WorkingDirectory={work_dir}")?; + writeln!(unit_content, "ExecStart={target_app} {args}")?; + writeln!(unit_content, "LimitNOFILE=infinity")?; + writeln!(unit_content)?; + writeln!(unit_content, "[Install]")?; + writeln!(unit_content, "WantedBy=multi-user.target")?; + + std::result::Result::Ok(unit_content) + } + + fn make_rcd_script(&self, options: &ServiceInstallOptions) -> Result { + let name = self.lable.to_qualified_name(); + let args = options.args.iter() + .map(|a| a.to_string_lossy()) + .collect::>() + .join(" "); + let target_app = options.program.display().to_string(); + let work_dir = options.work_directory.display().to_string(); + let mut script = String::new(); + + writeln!(script, "#!/bin/sh")?; + writeln!(script, "#")?; + writeln!(script, "# PROVIDE: {name}")?; + writeln!(script, "# REQUIRE: LOGIN FILESYSTEMS NETWORKING ")?; + writeln!(script, "# KEYWORD: shutdown")?; + writeln!(script)?; + writeln!(script, ". /etc/rc.subr")?; + writeln!(script)?; + writeln!(script, "name=\"{name}\"")?; + if let Some(ref d) = options.description { + writeln!(script, "desc=\"{d}\"")?; + } + writeln!(script, "rcvar=\"{name}_enable\"")?; + writeln!(script)?; + writeln!(script, "load_rc_config ${{name}}")?; + writeln!(script)?; + writeln!(script, ": ${{{name}_options=\"{args}\"}}")?; + writeln!(script)?; + writeln!(script, "{name}_chdir=\"{work_dir}\"")?; + writeln!(script, "pidfile=\"/var/run/${{name}}.pid\"")?; + writeln!(script, "procname=\"{target_app}\"")?; + writeln!(script, "command=\"/usr/sbin/daemon\"")?; + writeln!(script, "command_args=\"-c -S -T ${{name}} -p ${{pidfile}} ${{procname}} ${{{name}_options}}\"")?; + writeln!(script)?; + writeln!(script, "run_rc_command \"$1\"")?; + + std::result::Result::Ok(script) + } + + fn make_open_rc_script(&self, options: &ServiceInstallOptions) -> Result { + let args = options.args.iter() + .map(|a| a.to_string_lossy()) + .collect::>() + .join(" "); + let target_app = options.program.display().to_string(); + let work_dir = options.work_directory.display().to_string(); + let mut script = String::new(); + + writeln!(script, "#!/sbin/openrc-run")?; + writeln!(script)?; + if let Some(ref d) = options.description { + writeln!(script, "description=\"{d}\"")?; + } + writeln!(script, "command=\"{target_app}\"")?; + writeln!(script, "command_args=\"{args}\"")?; + writeln!(script, "pidfile=\"/run/${{RC_SVCNAME}}.pid\"")?; + writeln!(script, "command_background=\"yes\"")?; + writeln!(script, "directory=\"{work_dir}\"")?; + writeln!(script)?; + writeln!(script, "depend() {{")?; + writeln!(script, " need net")?; + writeln!(script, " use looger")?; + writeln!(script, "}}")?; + + std::result::Result::Ok(script) + + } } + #[tokio::main] #[tracing::instrument] async fn main() -> Result<(), Error> { @@ -785,7 +931,7 @@ async fn main() -> Result<(), Error> { } } SubCommand::Service(service_args) => { - let service = Service::new()?; + let service = Service::new(service_args.name)?; match service_args.sub_command { ServiceSubCommand::Install(install_args) => { let bin_path = install_args.core_path.unwrap_or_else(|| { @@ -805,7 +951,32 @@ async fn main() -> Result<(), Error> { anyhow::anyhow!("failed to get easytier core application: {}", e) })?; let bin_args = install_args.core_args.unwrap_or_default(); - service.install(bin_path, bin_args)?; + let work_dir = install_args.service_work_dir.unwrap_or_else(|| { + if cfg!(target_os = "windows") { + bin_path.parent().unwrap().to_path_buf() + } else { + std::env::temp_dir() + } + }); + + let work_dir = std::fs::canonicalize(&work_dir).map_err(|e| { + anyhow::anyhow!("failed to get service work directory[{}]: {}", work_dir.display(), e) + })?; + + if !work_dir.is_dir() { + return Err(anyhow::anyhow!("work directory is not a directory")); + } + + let install_options = ServiceInstallOptions { + program: bin_path, + args: bin_args, + work_directory: work_dir, + disable_autostart: install_args.disable_autostart, + description: Some(install_args.description), + #[cfg(target_os = "windows")] + display_name: install_args.display_name, + }; + service.install(&install_options)?; } ServiceSubCommand::Uninstall => { service.uninstall()?; @@ -850,7 +1021,8 @@ mod win_service_manager { use std::{ io, ffi::OsString, - ffi::OsStr + ffi::OsStr, + path::PathBuf }; use service_manager::{ @@ -863,24 +1035,34 @@ mod win_service_manager { ServiceStopCtx }; + use winreg::{enums::*, RegKey}; + + use easytier::common::constants::WIN_SERVICE_WORK_DIR_REG_KEY; + + use serde::{ + Serialize, + Deserialize + }; + + #[derive(Serialize, Deserialize)] + pub struct WinServiceInstallOptions { + pub dependencies: Option>, + pub description: Option, + pub display_name: Option, + } + pub struct WinServiceManager { - service_manager: ServiceManager, - display_name: Option, - description: Option, - dependencies: Vec + service_manager: ServiceManager } impl WinServiceManager { - pub fn new(display_name: Option, description: Option, dependencies: Vec,) -> Result { + pub fn new() -> Result { let service_manager = ServiceManager::local_computer( None::<&str>, ServiceManagerAccess::ALL_ACCESS, )?; Ok(Self { - service_manager, - display_name, - description, - dependencies, + service_manager }) } } @@ -892,8 +1074,23 @@ mod win_service_manager { fn install(&self, ctx: ServiceInstallCtx) -> io::Result<()> { let start_type_ = if ctx.autostart { ServiceStartType::AutoStart } else { ServiceStartType::OnDemand }; let srv_name = OsString::from(ctx.label.to_qualified_name()); - let dis_name = self.display_name.clone().unwrap_or_else(|| srv_name.clone()); - let dependencies = self.dependencies.iter().map(|dep| ServiceDependency::Service(dep.clone())).collect::>(); + let mut dis_name = srv_name.clone(); + let mut description: Option = None; + let mut dependencies = Vec::::new(); + + if let Some(s) = ctx.contents.as_ref() { + let options: WinServiceInstallOptions = serde_json::from_str(s.as_str()).unwrap(); + if let Some(d) = options.dependencies { + dependencies = d.iter().map(|dep| ServiceDependency::Service(OsString::from(dep.clone()))).collect::>(); + } + if let Some(d) = options.description { + description = Some(OsString::from(d)); + } + if let Some(d) = options.display_name { + dis_name = OsString::from(d); + } + } + let service_info = ServiceInfo { name: srv_name, display_name: dis_name, @@ -911,11 +1108,15 @@ mod win_service_manager { io::Error::new(io::ErrorKind::Other, e) })?; - if let Some(s) = &self.description { + if let Some(s) = description { service.set_description(s.clone()).map_err(|e| { io::Error::new(io::ErrorKind::Other, e) })?; } + + if let Some(work_dir) = ctx.working_directory { + set_service_work_directory(&ctx.label.to_qualified_name(), work_dir)?; + } Ok(()) } @@ -986,4 +1187,10 @@ mod win_service_manager { } } } + + fn set_service_work_directory(service_name: &str, work_directory: PathBuf) -> io::Result<()> { + let (reg_key, _) = RegKey::predef(HKEY_LOCAL_MACHINE).create_subkey(WIN_SERVICE_WORK_DIR_REG_KEY)?; + reg_key.set_value::(service_name, &work_directory.as_os_str().to_os_string())?; + Ok(()) + } } \ No newline at end of file diff --git a/easytier/src/easytier-core.rs b/easytier/src/easytier-core.rs index bd5b6bb..6fea439 100644 --- a/easytier/src/easytier-core.rs +++ b/easytier/src/easytier-core.rs @@ -289,12 +289,6 @@ struct Cli { help = t!("core_clap.ipv6_listener").to_string() )] ipv6_listener: Option, - - #[arg( - long, - help = t!("core_clap.work_dir").to_string() - )] - work_dir: Option, } rust_i18n::i18n!("locales", fallback = "en"); @@ -640,6 +634,22 @@ pub fn handle_event(mut events: EventBusSubscriber) -> tokio::task::JoinHandle<( }) } +#[cfg(target_os = "windows")] +fn win_service_set_work_dir(service_name: &std::ffi::OsString) -> anyhow::Result<()> { + use winreg::enums::*; + use winreg::RegKey; + use easytier::common::constants::WIN_SERVICE_WORK_DIR_REG_KEY; + + let hklm = RegKey::predef(HKEY_LOCAL_MACHINE); + let key = hklm.open_subkey_with_flags(WIN_SERVICE_WORK_DIR_REG_KEY, KEY_READ)?; + let dir_pat_str = key.get_value::(service_name)?; + let dir_path = std::fs::canonicalize(dir_pat_str)?; + + std::env::set_current_dir(dir_path)?; + + Ok(()) +} + #[cfg(target_os = "windows")] fn win_service_event_loop( stop_notify: std::sync::Arc, @@ -669,12 +679,6 @@ fn win_service_event_loop( process_id: None, }; - if cli.work_dir == None { - let mut path = std::env::current_exe().unwrap(); - path.pop(); - std::env::set_current_dir(path).unwrap(); - } - std::thread::spawn(move || { let rt = Runtime::new().unwrap(); rt.block_on(async move { @@ -701,13 +705,15 @@ fn win_service_event_loop( } #[cfg(target_os = "windows")] -fn win_service_main(_: Vec) { +fn win_service_main(arg: Vec) { use std::sync::Arc; use std::time::Duration; use tokio::sync::Notify; use windows_service::service::*; use windows_service::service_control_handler::*; + _ = win_service_set_work_dir(&arg[0]); + let cli = Cli::parse(); let stop_notify_send = Arc::new(Notify::new()); @@ -738,10 +744,6 @@ fn win_service_main(_: Vec) { } async fn run_main(cli: Cli) -> anyhow::Result<()> { - if let Some(dir) = cli.work_dir.as_ref() { - std::env::set_current_dir(dir).map_err(|e| anyhow::anyhow!("failed to set work dir: {}", e))?; - } - if cli.config_server.is_some() { let config_server_url_s = cli.config_server.clone().unwrap(); let config_server_url = match url::Url::parse(&config_server_url_s) {