rlmctl 0.2.0

//! Where `rlm guard enable` finds the guard binary and which systemd user
//! unit it uses. A distro package ships a unit under a system directory; a
//! `cargo install` does not, so `enable` writes a user unit whose
//! `ExecStart` points at the real binary.

use std::ffi::OsStr;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};

pub const UNIT_TEMPLATE: &str = include_str!("../assets/rlm-guard.service");

/// First line of a unit written by `rlm guard enable`. A user unit without
/// it is the user's own and is never overwritten.
pub const GENERATED_MARKER: &str = "# Generated by `rlm guard enable`. Rerunning it refreshes ExecStart; delete this line to keep your own edits.";

/// Directories where a packaged `rlm-guard.service` may live.
pub const SYSTEM_UNIT_DIRS: &[&str] = &[
    "/etc/systemd/user",
    "/usr/local/lib/systemd/user",
    "/usr/lib/systemd/user",
    "/lib/systemd/user",
];

/// `path` as an `ExecStart` argument, double-quoted when it contains
/// whitespace.
pub fn exec_arg(path: &Path) -> String {
    let s = path.display().to_string();
    if s.chars().any(char::is_whitespace) {
        format!("\"{s}\"")
    } else {
        s
    }
}

/// `template` with its `ExecStart=` line pointing at `exec`, prefixed by
/// [`GENERATED_MARKER`].
pub fn render_user_unit(template: &str, exec: &Path) -> String {
    let mut out = String::with_capacity(template.len() + GENERATED_MARKER.len() + 64);
    out.push_str(GENERATED_MARKER);
    out.push('\n');
    for line in template.split_inclusive('\n') {
        if line.starts_with("ExecStart=") {
            out.push_str(&format!("ExecStart={}", exec_arg(exec)));
            if line.ends_with('\n') {
                out.push('\n');
            }
        } else {
            out.push_str(line);
        }
    }
    out
}

fn is_executable(p: &Path) -> bool {
    p.is_file()
        && std::fs::metadata(p)
            .map(|m| m.permissions().mode() & 0o111 != 0)
            .unwrap_or(false)
}

/// The `rlm-guard` next to the running `rlm` first, then the first one on
/// `path_env`.
pub fn find_guard_binary(current_exe: Option<&Path>, path_env: Option<&OsStr>) -> Option<PathBuf> {
    let sibling = current_exe
        .and_then(Path::parent)
        .map(|dir| dir.join("rlm-guard"));
    let on_path = path_env
        .into_iter()
        .flat_map(std::env::split_paths)
        .map(|dir| dir.join("rlm-guard"));
    sibling
        .into_iter()
        .chain(on_path)
        .find(|p| is_executable(p))
}

/// Whether any of `dirs` holds an `rlm-guard.service`.
pub fn system_unit_installed(dirs: &[&Path]) -> bool {
    dirs.iter().any(|d| d.join("rlm-guard.service").is_file())
}

pub fn user_unit_path(config_dir: &Path) -> PathBuf {
    config_dir.join("systemd/user/rlm-guard.service")
}

/// Description line of the unit 0.1 shipped in `dist/`. Its README told users
/// to copy that file into `~/.config/systemd/user/` and point ExecStart at
/// `~/.cargo/bin`, so it has no generated marker but is still ours to replace.
const LEGACY_DESCRIPTION: &str =
    "Description=rlm freeze guard - proactively prevents system freezes";

fn is_ours(unit: &str) -> bool {
    unit.starts_with(GENERATED_MARKER) || unit.lines().any(|l| l.trim() == LEGACY_DESCRIPTION)
}

#[derive(Debug, PartialEq, Eq)]
pub enum EnablePlan {
    /// A packaged unit exists and no user unit shadows it.
    UseSystemUnit,
    /// The generated user unit already points at the right binary.
    UserUnitCurrent,
    /// The user unit was written or edited by the user; leave it alone.
    UserUnitCustom,
    WriteUserUnit {
        path: PathBuf,
        contents: String,
    },
    /// No packaged unit and no guard binary to point a user unit at.
    NoBinary,
}

/// Decide what `rlm guard enable` does before it runs `systemctl`. A
/// generated user unit takes precedence over a system one (systemd prefers
/// it too), so a stale one is refreshed even when a package is installed.
pub fn plan_enable(
    system_unit: bool,
    existing: Option<&str>,
    bin: Option<&Path>,
    unit_path: &Path,
) -> EnablePlan {
    match (existing, bin) {
        (Some(text), _) if !is_ours(text) => EnablePlan::UserUnitCustom,
        (Some(text), Some(bin)) => {
            let want = render_user_unit(UNIT_TEMPLATE, bin);
            if text == want {
                EnablePlan::UserUnitCurrent
            } else {
                EnablePlan::WriteUserUnit {
                    path: unit_path.to_path_buf(),
                    contents: want,
                }
            }
        }
        (Some(_), None) => EnablePlan::UserUnitCurrent,
        (None, _) if system_unit => EnablePlan::UseSystemUnit,
        (None, Some(bin)) => EnablePlan::WriteUserUnit {
            path: unit_path.to_path_buf(),
            contents: render_user_unit(UNIT_TEMPLATE, bin),
        },
        (None, None) => EnablePlan::NoBinary,
    }
}

/// Note for `rlm guard enable` when the service was already running before
/// the command. `systemctl enable --now` does not restart a running unit,
/// so the old process keeps going until the user restarts it.
pub fn restart_note(was_active: bool, unit_written: bool) -> Option<&'static str> {
    match (was_active, unit_written) {
        (false, _) => None,
        (true, true) => Some(
            "note: rlm-guard is still running with the old unit. Restart it to use the new one: systemctl --user restart rlm-guard",
        ),
        (true, false) => Some(
            "note: rlm-guard was already running. If you upgraded rlm, restart it so the new version runs: systemctl --user restart rlm-guard",
        ),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::os::unix::fs::PermissionsExt;

    #[test]
    fn restart_note_only_when_already_running() {
        assert_eq!(restart_note(false, true), None);
        assert_eq!(restart_note(false, false), None);
        let written = restart_note(true, true).unwrap();
        assert!(written.contains("systemctl --user restart rlm-guard"));
        assert!(written.contains("old unit"));
        assert!(restart_note(true, false)
            .unwrap()
            .contains("systemctl --user restart rlm-guard"));
    }

    #[test]
    fn render_replaces_only_exec_start_and_marks_the_file() {
        let out = render_user_unit(UNIT_TEMPLATE, Path::new("/home/u/.cargo/bin/rlm-guard"));
        assert!(out.starts_with(GENERATED_MARKER));
        assert!(out.contains("\nExecStart=/home/u/.cargo/bin/rlm-guard\n"));
        assert_eq!(out.matches("ExecStart=").count(), 1);
        assert!(out.contains("RestartPreventExitStatus=78"));
    }

    #[test]
    fn paths_with_spaces_are_quoted() {
        assert_eq!(
            exec_arg(Path::new("/opt/my tools/rlm-guard")),
            "\"/opt/my tools/rlm-guard\""
        );
        assert_eq!(
            exec_arg(Path::new("/usr/bin/rlm-guard")),
            "/usr/bin/rlm-guard"
        );
    }

    fn exe(dir: &Path) -> PathBuf {
        let p = dir.join("rlm-guard");
        std::fs::write(&p, "#!/bin/sh\n").unwrap();
        std::fs::set_permissions(&p, std::fs::Permissions::from_mode(0o755)).unwrap();
        p
    }

    #[test]
    fn guard_next_to_rlm_wins_over_path() {
        let a = tempfile::tempdir().unwrap();
        let b = tempfile::tempdir().unwrap();
        let sibling = exe(a.path());
        exe(b.path());
        let rlm = a.path().join("rlm");
        assert_eq!(
            find_guard_binary(Some(&rlm), Some(b.path().as_os_str())),
            Some(sibling)
        );
        let none = tempfile::tempdir().unwrap();
        let on_path = find_guard_binary(Some(&none.path().join("rlm")), Some(b.path().as_os_str()));
        assert_eq!(on_path, Some(b.path().join("rlm-guard")));
        assert_eq!(find_guard_binary(None, None), None);
    }

    #[test]
    fn non_executable_candidates_are_skipped() {
        let a = tempfile::tempdir().unwrap();
        let p = a.path().join("rlm-guard");
        std::fs::write(&p, "").unwrap();
        std::fs::set_permissions(&p, std::fs::Permissions::from_mode(0o644)).unwrap();
        assert_eq!(find_guard_binary(Some(&a.path().join("rlm")), None), None);
    }

    #[test]
    fn system_unit_detection_and_user_unit_path() {
        let d = tempfile::tempdir().unwrap();
        assert!(!system_unit_installed(&[d.path()]));
        std::fs::write(d.path().join("rlm-guard.service"), UNIT_TEMPLATE).unwrap();
        assert!(system_unit_installed(&[
            Path::new("/nonexistent"),
            d.path()
        ]));
        assert_eq!(
            user_unit_path(Path::new("/h/.config")),
            PathBuf::from("/h/.config/systemd/user/rlm-guard.service")
        );
    }

    #[test]
    fn enable_plans() {
        let unit = Path::new("/h/.config/systemd/user/rlm-guard.service");
        let bin = Path::new("/h/.cargo/bin/rlm-guard");
        assert_eq!(
            plan_enable(true, None, Some(bin), unit),
            EnablePlan::UseSystemUnit
        );
        assert_eq!(plan_enable(false, None, None, unit), EnablePlan::NoBinary);
        let fresh = render_user_unit(UNIT_TEMPLATE, bin);
        assert_eq!(
            plan_enable(false, None, Some(bin), unit),
            EnablePlan::WriteUserUnit {
                path: unit.to_path_buf(),
                contents: fresh.clone()
            }
        );
        assert_eq!(
            plan_enable(false, Some(&fresh), Some(bin), unit),
            EnablePlan::UserUnitCurrent
        );
        let stale = render_user_unit(UNIT_TEMPLATE, Path::new("/old/rlm-guard"));
        assert!(matches!(
            plan_enable(true, Some(&stale), Some(bin), unit),
            EnablePlan::WriteUserUnit { .. }
        ));
        assert_eq!(
            plan_enable(false, Some("[Service]\nExecStart=/mine\n"), Some(bin), unit),
            EnablePlan::UserUnitCustom
        );
    }

    #[test]
    fn plan_enable_replaces_the_0_1_unit() {
        let legacy = "[Unit]\nDescription=rlm freeze guard - proactively prevents system freezes\n\n[Service]\nExecStart=%h/.cargo/bin/rlm-guard\n";
        let bin = Path::new("/home/u/.cargo/bin/rlm-guard");
        let unit = Path::new("/home/u/.config/systemd/user/rlm-guard.service");
        assert_eq!(
            plan_enable(false, Some(legacy), Some(bin), unit),
            EnablePlan::WriteUserUnit {
                path: unit.to_path_buf(),
                contents: render_user_unit(UNIT_TEMPLATE, bin),
            }
        );
    }
}