use async_trait::async_trait; use std::{ fmt::{Display, Formatter}, sync::Arc, }; pub const PROCESS_MANAGER_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(15); #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct UnitStatus { pub active: bool, pub enabled: bool, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum StartupMode { AlwaysOn, SocketActivated, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum ProcessAction { Start, Stop, Restart, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum DetectedStartupMode { AlwaysOn, SocketActivated, Disabled, Conflicting, } #[derive(Clone, Debug, PartialEq, Eq)] pub struct DaemonStartupStatus { pub service: UnitStatus, pub socket: UnitStatus, pub detected: DetectedStartupMode, } impl DaemonStartupStatus { pub fn classify(service: UnitStatus, socket: UnitStatus) -> Self { let detected = match (service.enabled, socket.enabled) { (true, false) => DetectedStartupMode::AlwaysOn, (false, true) => DetectedStartupMode::SocketActivated, (false, false) => DetectedStartupMode::Disabled, (true, true) => DetectedStartupMode::Conflicting, }; Self { service, socket, detected, } } } #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum ProcessManagerErrorKind { CommandUnavailable, PermissionDenied, UnitMissing, CommandFailed, ParseFailed, VerificationFailed, TimedOut, } #[derive(Clone, Debug, PartialEq, Eq)] pub struct ProcessManagerError { pub kind: ProcessManagerErrorKind, message: String, } impl ProcessManagerError { pub fn new(kind: ProcessManagerErrorKind, message: impl Into) -> Self { Self { kind, message: message.into(), } } pub fn kind(&self) -> ProcessManagerErrorKind { self.kind } } impl Display for ProcessManagerError { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.write_str(&self.message) } } impl std::error::Error for ProcessManagerError {} #[derive(Clone, Debug, PartialEq, Eq)] pub struct CommandOutput { pub success: bool, pub stdout: String, pub stderr: String, } #[async_trait] pub trait CommandExecutor: Send + Sync { async fn output( &self, program: &str, args: &[&str], ) -> Result; } #[async_trait] pub trait ProcessManager: Send + Sync { fn name(&self) -> &'static str; async fn unit_status(&self, unit: &str) -> Result; async fn set_iota_startup_mode( &self, mode: StartupMode, ) -> Result; async fn iota_startup_status(&self) -> Result { Ok(DaemonStartupStatus::classify( self.unit_status("iota-daemon.service").await?, self.unit_status("iota-daemon.socket").await?, )) } async fn enable_startup( &self, mode: StartupMode, ) -> Result { self.set_iota_startup_mode(mode).await } async fn disable_startup(&self) -> Result { self.set_iota_startup_mode(StartupMode::SocketActivated) .await } async fn process_action( &self, action: ProcessAction, ) -> Result { let unit = "iota-daemon.service"; match action { ProcessAction::Start => self.unit_action(&["start", unit]).await?, ProcessAction::Stop => self.unit_action(&["stop", unit]).await?, ProcessAction::Restart => self.unit_action(&["restart", unit]).await?, } self.iota_startup_status().await } async fn unit_action(&self, _action: &[&str]) -> Result<(), ProcessManagerError> { Err(ProcessManagerError::new( ProcessManagerErrorKind::CommandFailed, "process actions are unsupported", )) } } pub async fn detect() -> Option> { #[cfg(target_os = "linux")] { systemd::SystemdManager::detect() .await .map(|m| Arc::new(m) as Arc) } #[cfg(not(target_os = "linux"))] { None } } #[cfg(target_os = "linux")] mod systemd { use super::*; use std::{path::Path, process::Stdio}; use tokio::{process::Command, time::timeout}; const SERVICE: &str = "iota-daemon.service"; const SOCKET: &str = "iota-daemon.socket"; const COMMON: [&str; 2] = ["--no-pager", "--no-ask-password"]; pub struct RealExecutor; #[async_trait] impl CommandExecutor for RealExecutor { async fn output( &self, program: &str, args: &[&str], ) -> Result { let child = Command::new(program) .args(args) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .kill_on_drop(false) .spawn() .map_err(|e| { ProcessManagerError::new( ProcessManagerErrorKind::CommandUnavailable, format!("Could not run {program}: {e}"), ) })?; let output = timeout(PROCESS_MANAGER_TIMEOUT, child.wait_with_output()) .await .map_err(|_| { ProcessManagerError::new( ProcessManagerErrorKind::TimedOut, format!( "{program} timed out after {} seconds", PROCESS_MANAGER_TIMEOUT.as_secs() ), ) })? .map_err(|e| { ProcessManagerError::new(ProcessManagerErrorKind::CommandFailed, e.to_string()) })?; Ok(CommandOutput { success: output.status.success(), stdout: String::from_utf8_lossy(&output.stdout).into_owned(), stderr: String::from_utf8_lossy(&output.stderr).into_owned(), }) } } pub struct SystemdManager { executor: Arc, service: &'static str, socket: &'static str, } impl SystemdManager { pub async fn detect() -> Option { if !Path::new("/run/systemd/system").is_dir() { return None; } let executor: Arc = Arc::new(RealExecutor); let mut manager = executor .output("systemctl", &["--version", &COMMON[0], &COMMON[1]]) .await .ok() .filter(|r| r.success) .map(|_| Self { executor, service: SERVICE, socket: SOCKET, })?; if manager.status("iota.service").await.is_ok() { manager.service = "iota.service"; manager.socket = "iota.socket"; } Some(manager) } #[cfg(test)] pub fn with_executor(executor: Arc) -> Self { Self { executor, service: SERVICE, socket: SOCKET, } } async fn run(&self, action: &[&str]) -> Result<(), ProcessManagerError> { let mut args = COMMON.to_vec(); args.extend_from_slice(action); let output = self.executor.output("systemctl", &args).await?; if output.success { return Ok(()); } let detail = if output.stderr.trim().is_empty() { output.stdout.trim() } else { output.stderr.trim() }; let kind = if detail.to_ascii_lowercase().contains("access denied") || detail.to_ascii_lowercase().contains("permission denied") { ProcessManagerErrorKind::PermissionDenied } else { ProcessManagerErrorKind::CommandFailed }; Err(ProcessManagerError::new( kind, if detail.is_empty() { format!("systemctl {} failed", action.join(" ")) } else { detail.to_owned() }, )) } async fn status(&self, unit: &str) -> Result { let mut args = COMMON.to_vec(); args.extend_from_slice(&[ "show", "--property=LoadState", "--property=ActiveState", "--property=UnitFileState", "--value", unit, ]); let output = self.executor.output("systemctl", &args).await?; if !output.success { let detail = if output.stderr.trim().is_empty() { output.stdout.trim() } else { output.stderr.trim() }; let kind = if detail.to_ascii_lowercase().contains("denied") { ProcessManagerErrorKind::PermissionDenied } else { ProcessManagerErrorKind::CommandFailed }; return Err(ProcessManagerError::new( kind, format!("systemctl could not inspect {unit}: {detail}"), )); } let values: Vec<_> = output.stdout.lines().map(str::trim).collect(); if values.len() < 3 { return Err(ProcessManagerError::new( ProcessManagerErrorKind::ParseFailed, format!("systemctl returned incomplete state for {unit}"), )); } if values[0] == "not-found" { return Err(ProcessManagerError::new( ProcessManagerErrorKind::UnitMissing, format!("systemd unit {unit} was not found"), )); } if values[0] != "loaded" { return Err(ProcessManagerError::new( ProcessManagerErrorKind::ParseFailed, format!("unsupported LoadState `{}` for {unit}", values[0]), )); } let active = match values[1] { "active" => true, "inactive" | "failed" | "activating" | "deactivating" | "reloading" => false, v => { return Err(ProcessManagerError::new( ProcessManagerErrorKind::ParseFailed, format!("unsupported ActiveState `{v}` for {unit}"), )); } }; let enabled = match values[2] { "enabled" | "enabled-runtime" => true, "disabled" | "static" | "indirect" | "masked" | "generated" | "transient" => false, v => { return Err(ProcessManagerError::new( ProcessManagerErrorKind::ParseFailed, format!("unsupported UnitFileState `{v}` for {unit}"), )); } }; Ok(UnitStatus { active, enabled }) } async fn verify( &self, expected: DetectedStartupMode, ) -> Result { let status = self.iota_startup_status().await?; if status.detected == expected { Ok(status) } else { Err(ProcessManagerError::new( ProcessManagerErrorKind::VerificationFailed, format!( "systemd reported {:?} after applying {:?}", status.detected, expected ), )) } } } #[async_trait] impl ProcessManager for SystemdManager { fn name(&self) -> &'static str { "systemd" } async fn unit_status(&self, unit: &str) -> Result { self.status(unit).await } async fn iota_startup_status(&self) -> Result { Ok(DaemonStartupStatus::classify( self.status(self.service).await?, self.status(self.socket).await?, )) } async fn set_iota_startup_mode( &self, mode: StartupMode, ) -> Result { match mode { StartupMode::AlwaysOn => { self.run(&["disable", self.socket]).await?; self.run(&["enable", "--now", self.service]).await?; self.verify(DetectedStartupMode::AlwaysOn).await } StartupMode::SocketActivated => { self.run(&["disable", "--now", self.service]).await?; self.run(&["enable", "--now", self.socket]).await?; self.verify(DetectedStartupMode::SocketActivated).await } } } async fn unit_action(&self, action: &[&str]) -> Result<(), ProcessManagerError> { self.run(action).await } async fn process_action( &self, action: ProcessAction, ) -> Result { let verb = match action { ProcessAction::Start => "start", ProcessAction::Stop => "stop", ProcessAction::Restart => "restart", }; self.run(&[verb, self.service]).await?; self.iota_startup_status().await } async fn disable_startup(&self) -> Result { self.run(&["disable", "--now", self.service]).await?; self.run(&["disable", "--now", self.socket]).await?; self.verify(DetectedStartupMode::Disabled).await } } #[cfg(test)] mod tests { use super::*; use std::sync::Mutex; struct Fake { calls: Mutex>>, results: Mutex>, } #[async_trait] impl CommandExecutor for Fake { async fn output( &self, _: &str, args: &[&str], ) -> Result { self.calls .lock() .unwrap() .push(args.iter().map(|arg| (*arg).to_owned()).collect()); Ok(self.results.lock().unwrap().remove(0)) } } fn ok(stdout: &str) -> CommandOutput { CommandOutput { success: true, stdout: stdout.into(), stderr: String::new(), } } #[tokio::test] async fn every_systemctl_operation_disables_interactive_features() { let fake = Arc::new(Fake { calls: Mutex::new(Vec::new()), results: Mutex::new(vec![ok("loaded\nactive\nenabled\n")]), }); let manager = SystemdManager::with_executor(fake.clone()); manager.unit_status(SERVICE).await.unwrap(); let call = &fake.calls.lock().unwrap()[0]; assert!(call.contains(&"--no-pager".into())); assert!(call.contains(&"--no-ask-password".into())); } } } #[cfg(test)] mod tests { use super::*; #[test] fn modes_distinct() { assert_ne!(StartupMode::AlwaysOn, StartupMode::SocketActivated); } struct BlockingExecutor; #[async_trait::async_trait] impl CommandExecutor for BlockingExecutor { async fn output(&self, _: &str, _: &[&str]) -> Result { std::future::pending().await } } #[tokio::test] async fn executor_future_can_be_cancelled_without_blocking_runtime() { let result = tokio::time::timeout( std::time::Duration::from_millis(20), BlockingExecutor.output("systemctl", &["show"]), ) .await; assert!(result.is_err()); } }