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::{io::AsyncRead, 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"]; const MAX_COMMAND_OUTPUT_BYTES: usize = 1024 * 1024; pub struct RealExecutor; async fn read_bounded(reader: R) -> std::io::Result> where R: AsyncRead + Unpin, { use tokio::io::AsyncReadExt; let mut output = Vec::new(); reader .take((MAX_COMMAND_OUTPUT_BYTES + 1) as u64) .read_to_end(&mut output) .await?; if output.len() > MAX_COMMAND_OUTPUT_BYTES { output.truncate(MAX_COMMAND_OUTPUT_BYTES); } Ok(output) } async fn collect_output( stdout: tokio::process::ChildStdout, stderr: tokio::process::ChildStderr, ) -> Result<(Vec, Vec), ProcessManagerError> { let (stdout_result, stderr_result) = tokio::join!(read_bounded(stdout), read_bounded(stderr)); let stdout = stdout_result.map_err(|error| { ProcessManagerError::new( ProcessManagerErrorKind::CommandFailed, format!("stdout read failed: {error}"), ) })?; let stderr = stderr_result.map_err(|error| { ProcessManagerError::new( ProcessManagerErrorKind::CommandFailed, format!("stderr read failed: {error}"), ) })?; Ok((stdout, stderr)) } impl RealExecutor { async fn output_with_timeout( &self, program: &str, args: &[&str], process_timeout: std::time::Duration, ) -> Result { let mut child = Command::new(program) .args(args) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .kill_on_drop(true) .spawn() .map_err(|e| { ProcessManagerError::new( ProcessManagerErrorKind::CommandUnavailable, format!("Could not run {program}: {e}"), ) })?; let stdout = child.stdout.take().ok_or_else(|| { ProcessManagerError::new( ProcessManagerErrorKind::CommandFailed, "command stdout pipe was not created", ) })?; let stderr = child.stderr.take().ok_or_else(|| { ProcessManagerError::new( ProcessManagerErrorKind::CommandFailed, "command stderr pipe was not created", ) })?; let output_task = tokio::spawn(collect_output(stdout, stderr)); let status = match timeout(process_timeout, child.wait()).await { Ok(result) => result.map_err(|e| { ProcessManagerError::new(ProcessManagerErrorKind::CommandFailed, e.to_string()) })?, Err(_) => { // Keep the Child alive across the timeout. Explicitly // terminate it and await wait() so the OS child is // reaped before reporting the timeout. let kill_error = child.start_kill().err(); let wait_error = child.wait().await.err(); output_task.abort(); let _ = output_task.await; if let Some(error) = wait_error { return Err(ProcessManagerError::new( ProcessManagerErrorKind::CommandFailed, format!("{program} timed out and could not be reaped: {error}"), )); } let termination_detail = kill_error .map(|error| format!("; termination request reported: {error}")) .unwrap_or_default(); return Err(ProcessManagerError::new( ProcessManagerErrorKind::TimedOut, format!( "{program} timed out after {} seconds{termination_detail}", process_timeout.as_secs(), ), )); } }; let (stdout, stderr) = output_task.await.map_err(|error| { ProcessManagerError::new( ProcessManagerErrorKind::CommandFailed, format!("command output task failed: {error}"), ) })??; Ok(CommandOutput { success: status.success(), stdout: String::from_utf8_lossy(&stdout).into_owned(), stderr: String::from_utf8_lossy(&stderr).into_owned(), }) } } #[async_trait] impl CommandExecutor for RealExecutor { async fn output( &self, program: &str, args: &[&str], ) -> Result { self.output_with_timeout(program, args, PROCESS_MANAGER_TIMEOUT) .await } } 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())); } #[tokio::test] async fn timed_out_real_child_is_terminated_and_reaped() { use std::fs; use std::time::Duration; let directory = tempfile::tempdir().unwrap(); let pid_file = directory.path().join("child.pid"); let script = format!( "printf '%s' \"$$\" > '{}'; exec sleep 60", pid_file.display() ); let executor = RealExecutor; let task = tokio::spawn(async move { executor .output_with_timeout("sh", &["-c", &script], Duration::from_millis(50)) .await }); let deadline = tokio::time::Instant::now() + Duration::from_secs(2); let pid = loop { if let Ok(contents) = fs::read_to_string(&pid_file) { if let Ok(pid) = contents.parse::() { break pid; } } assert!(tokio::time::Instant::now() < deadline); tokio::task::yield_now().await; }; let result = task.await.unwrap(); assert_eq!( result.unwrap_err().kind(), ProcessManagerErrorKind::TimedOut ); assert!(!std::path::Path::new(&format!("/proc/{pid}")).exists()); let mut status = 0; let wait_result = unsafe { libc::waitpid(pid, &mut status, libc::WNOHANG) }; assert_eq!(wait_result, -1); assert_eq!( std::io::Error::last_os_error().raw_os_error(), Some(libc::ECHILD) ); } } } #[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()); } }