PicoBot/src/service.rs

use std::path::{Path, PathBuf};
use std::process::Stdio;

use anyhow::{Context, Result, bail};

const UNIT_NAME: &str = "picobot.service";

#[derive(Debug, Clone, Copy)]
pub enum ServiceCommand {
    Install,
    Start,
    Stop,
    Status,
    Restart,
    Uninstall,
}

pub async fn execute(command: ServiceCommand) -> Result<()> {
    ensure_systemd_supported()?;

    match command {
        ServiceCommand::Install => install().await,
        ServiceCommand::Start => run_systemctl(&["start", UNIT_NAME]).await,
        ServiceCommand::Stop => run_systemctl(&["stop", UNIT_NAME]).await,
        ServiceCommand::Status => run_systemctl(&["status", "--no-pager", UNIT_NAME]).await,
        ServiceCommand::Restart => run_systemctl(&["restart", UNIT_NAME]).await,
        ServiceCommand::Uninstall => uninstall().await,
    }
}

fn ensure_systemd_supported() -> Result<()> {
    if cfg!(target_os = "linux") {
        Ok(())
    } else {
        bail!("service management currently requires Linux with a systemd user session")
    }
}

async fn install() -> Result<()> {
    let executable = std::env::current_exe()
        .context("failed to determine the PicoBot executable path")?
        .canonicalize()
        .context("failed to canonicalize the PicoBot executable path")?;
    let working_directory = std::env::current_dir()
        .context("failed to determine the current working directory")?
        .canonicalize()
        .context("failed to canonicalize the current working directory")?;
    let service_path = service_path()?;

    let parent = service_path
        .parent()
        .context("systemd user service path has no parent directory")?;
    std::fs::create_dir_all(parent)
        .with_context(|| format!("failed to create {}", parent.display()))?;
    std::fs::write(
        &service_path,
        render_systemd_unit(&executable, &working_directory),
    )
    .with_context(|| format!("failed to write {}", service_path.display()))?;

    if let Err(error) = async {
        run_systemctl(&["daemon-reload"]).await?;
        run_systemctl(&["enable", UNIT_NAME]).await
    }
    .await
    {
        return Err(error).with_context(|| {
            format!(
                "service file was written to {}, but systemd registration failed",
                service_path.display()
            )
        });
    }

    println!(
        "Installed and enabled {UNIT_NAME} at {}",
        service_path.display()
    );
    println!("Start it with: picobot service start");
    Ok(())
}

async fn uninstall() -> Result<()> {
    let service_path = service_path()?;
    if !service_path.exists() {
        println!("{UNIT_NAME} is not installed at {}", service_path.display());
        return Ok(());
    }

    run_systemctl(&["disable", "--now", UNIT_NAME]).await?;
    std::fs::remove_file(&service_path)
        .with_context(|| format!("failed to remove {}", service_path.display()))?;
    run_systemctl(&["daemon-reload"]).await?;
    let _ = run_systemctl_quiet(&["reset-failed", UNIT_NAME]).await;

    println!("Uninstalled {UNIT_NAME}");
    Ok(())
}

fn service_path() -> Result<PathBuf> {
    let home = dirs::home_dir().context("failed to determine the home directory")?;
    Ok(home
        .join(".config")
        .join("systemd")
        .join("user")
        .join(UNIT_NAME))
}

async fn run_systemctl(args: &[&str]) -> Result<()> {
    let status = tokio::process::Command::new("systemctl")
        .arg("--user")
        .args(args)
        .status()
        .await
        .with_context(|| format!("failed to execute systemctl --user {}", args.join(" ")))?;

    if !status.success() {
        bail!(
            "systemctl --user {} failed with status {}",
            args.join(" "),
            status
        );
    }
    Ok(())
}

async fn run_systemctl_quiet(args: &[&str]) -> Result<()> {
    let status = tokio::process::Command::new("systemctl")
        .arg("--user")
        .args(args)
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .status()
        .await
        .with_context(|| format!("failed to execute systemctl --user {}", args.join(" ")))?;

    if !status.success() {
        bail!(
            "systemctl --user {} failed with status {}",
            args.join(" "),
            status
        );
    }
    Ok(())
}

fn render_systemd_unit(executable: &Path, working_directory: &Path) -> String {
    let executable = quote_systemd_value(executable.to_string_lossy().as_ref());
    let working_directory = quote_systemd_value(working_directory.to_string_lossy().as_ref());

    format!(
        "[Unit]\n\
Description=PicoBot Gateway\n\
After=network-online.target\n\
Wants=network-online.target\n\
\n\
[Service]\n\
Type=simple\n\
ExecStart={executable} gateway\n\
WorkingDirectory={working_directory}\n\
Restart=on-failure\n\
RestartSec=5\n\
Environment=RUST_LOG=info\n\
\n\
[Install]\n\
WantedBy=default.target\n"
    )
}

fn quote_systemd_value(value: &str) -> String {
    let escaped = value
        .replace('%', "%%")
        .replace('\\', "\\\\")
        .replace('"', "\\\"");
    format!("\"{escaped}\"")
}

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

    #[test]
    fn unit_starts_current_binary_as_gateway_and_restarts_on_failure() {
        let unit = render_systemd_unit(
            Path::new("/home/alice/bin/picobot"),
            Path::new("/home/alice/picobot"),
        );

        assert!(unit.contains("ExecStart=\"/home/alice/bin/picobot\" gateway"));
        assert!(unit.contains("WorkingDirectory=\"/home/alice/picobot\""));
        assert!(unit.contains("Restart=on-failure"));
        assert!(unit.contains("WantedBy=default.target"));
    }

    #[test]
    fn unit_quotes_spaces_quotes_backslashes_and_specifiers() {
        let unit = render_systemd_unit(
            Path::new("/home/alice/Pico Bot/%i/\"picobot\""),
            Path::new("/home/alice/work\\space"),
        );

        assert!(unit.contains("ExecStart=\"/home/alice/Pico Bot/%%i/\\\"picobot\\\"\" gateway"));
        assert!(unit.contains("WorkingDirectory=\"/home/alice/work\\\\space\""));
    }
}