iota/iota-process-manager/src/lib.rs

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<String>) -> 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<CommandOutput, ProcessManagerError>;
}

#[async_trait]
pub trait ProcessManager: Send + Sync {
    fn name(&self) -> &'static str;
    async fn unit_status(&self, unit: &str) -> Result<UnitStatus, ProcessManagerError>;
    async fn set_iota_startup_mode(
        &self,
        mode: StartupMode,
    ) -> Result<DaemonStartupStatus, ProcessManagerError>;
    async fn iota_startup_status(&self) -> Result<DaemonStartupStatus, ProcessManagerError> {
        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<DaemonStartupStatus, ProcessManagerError> {
        self.set_iota_startup_mode(mode).await
    }
    async fn disable_startup(&self) -> Result<DaemonStartupStatus, ProcessManagerError> {
        self.set_iota_startup_mode(StartupMode::SocketActivated)
            .await
    }
    async fn process_action(
        &self,
        action: ProcessAction,
    ) -> Result<DaemonStartupStatus, ProcessManagerError> {
        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<Arc<dyn ProcessManager>> {
    #[cfg(target_os = "linux")]
    {
        systemd::SystemdManager::detect()
            .await
            .map(|m| Arc::new(m) as Arc<dyn ProcessManager>)
    }
    #[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<CommandOutput, ProcessManagerError> {
            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<dyn CommandExecutor>,
    }
    impl SystemdManager {
        pub async fn detect() -> Option<Self> {
            if !Path::new("/run/systemd/system").is_dir() {
                return None;
            }
            let executor: Arc<dyn CommandExecutor> = Arc::new(RealExecutor);
            executor
                .output("systemctl", &["--version", &COMMON[0], &COMMON[1]])
                .await
                .ok()
                .filter(|r| r.success)
                .map(|_| Self { executor })
        }
        #[cfg(test)]
        pub fn with_executor(executor: Arc<dyn CommandExecutor>) -> Self {
            Self { executor }
        }
        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<UnitStatus, ProcessManagerError> {
            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<DaemonStartupStatus, ProcessManagerError> {
            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<UnitStatus, ProcessManagerError> {
            self.status(unit).await
        }
        async fn set_iota_startup_mode(
            &self,
            mode: StartupMode,
        ) -> Result<DaemonStartupStatus, ProcessManagerError> {
            match mode {
                StartupMode::AlwaysOn => {
                    self.run(&["disable", SOCKET]).await?;
                    self.run(&["enable", "--now", SERVICE]).await?;
                    self.verify(DetectedStartupMode::AlwaysOn).await
                }
                StartupMode::SocketActivated => {
                    self.run(&["disable", "--now", SERVICE]).await?;
                    self.run(&["enable", "--now", SOCKET]).await?;
                    self.verify(DetectedStartupMode::SocketActivated).await
                }
            }
        }
        async fn unit_action(&self, action: &[&str]) -> Result<(), ProcessManagerError> {
            self.run(action).await
        }
        async fn disable_startup(&self) -> Result<DaemonStartupStatus, ProcessManagerError> {
            self.run(&["disable", "--now", SERVICE]).await?;
            self.run(&["disable", "--now", SOCKET]).await?;
            self.verify(DetectedStartupMode::Disabled).await
        }
    }

    #[cfg(test)]
    mod tests {
        use super::*;
        use std::sync::Mutex;

        struct Fake {
            calls: Mutex<Vec<Vec<String>>>,
            results: Mutex<Vec<CommandOutput>>,
        }
        #[async_trait]
        impl CommandExecutor for Fake {
            async fn output(
                &self,
                _: &str,
                args: &[&str],
            ) -> Result<CommandOutput, ProcessManagerError> {
                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<CommandOutput, ProcessManagerError> {
            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());
    }
}