rlmctl 0.2.2

//! 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 the program of an `ExecStart` line. systemd expands `%`
/// specifiers there, so `%` is doubled. `$` is left alone: systemd does not
/// expand variables in the program path. The result is double-quoted when the
/// path has whitespace. Paths systemd refuses as a program (see
/// [`unit_path_problem`]) must be rejected before calling this.
pub fn exec_arg(path: &Path) -> String {
    let s = path.display().to_string();
    let mut out = String::with_capacity(s.len() + 2);
    for c in s.chars() {
        match c {
            '%' => out.push_str("%%"),
            '\\' => out.push_str("\\\\"),
            '"' => out.push_str("\\\""),
            '\'' => out.push_str("\\'"),
            c => out.push(c),
        }
    }
    if s.chars()
        .any(|c| c.is_whitespace() || c == '"' || c == '\'')
    {
        format!("\"{out}\"")
    } else {
        out
    }
}

/// Why systemd would refuse `path` as an `ExecStart` program, if it would:
/// it rejects quotes and backslashes in the executable path, control
/// characters would split the unit line, and the unit file must be UTF-8.
pub fn unit_path_problem(path: &Path) -> Option<&'static str> {
    let Some(s) = path.to_str() else {
        return Some("it is not valid UTF-8");
    };
    if s.chars().any(|c| c.is_control()) {
        return Some("it contains a control character");
    }
    if s.contains(['"', '\'', '\\']) {
        return Some("systemd does not accept quotes or backslashes in an ExecStart path");
    }
    None
}

/// `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`. Relative `PATH` entries (including empty ones and `.`) are
/// skipped: they depend on the working directory, and the result ends up in
/// a unit's `ExecStart`.
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)
        .filter(|dir| dir.is_absolute())
        .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,
    }
}

/// Error for `rlm guard enable` when `systemctl --user is-enabled` reports
/// the unit as masked (`masked` or `masked-runtime`). A masked user unit is
/// a symlink to `/dev/null` at the user unit path, which would otherwise read
/// as an empty user-written unit and be kept as "your own".
pub fn masked_error(enabled: &str) -> Option<String> {
    let runtime = match enabled {
        "masked" => false,
        "masked-runtime" => true,
        _ => return None,
    };
    let unmask = if runtime {
        "systemctl --user unmask --runtime rlm-guard"
    } else {
        "systemctl --user unmask rlm-guard"
    };
    Some(format!(
        "the rlm-guard unit is masked, so systemd will not start it. Unmask it with: {unmask} and then rerun: rlm guard enable"
    ))
}

/// Write `contents` to `path` through a temporary file in the same
/// directory and a rename, so an interrupted write never leaves a truncated
/// unit behind. The temporary name does not end in `.service`, so systemd
/// ignores it if it is left over.
pub fn write_atomically(path: &Path, contents: &str) -> std::io::Result<()> {
    use std::io::Write;
    let dir = path.parent().unwrap_or(Path::new("."));
    let name = path
        .file_name()
        .map(|n| n.to_string_lossy().into_owned())
        .unwrap_or_default();
    let tmp = dir.join(format!(".{name}.tmp-{}", std::process::id()));
    let result = std::fs::File::create(&tmp)
        .and_then(|mut f| {
            f.write_all(contents.as_bytes())?;
            f.sync_all()
        })
        .and_then(|()| std::fs::rename(&tmp, path));
    if result.is_err() {
        let _ = std::fs::remove_file(&tmp);
    }
    result
}

/// 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 masked_units_get_an_unmask_hint() {
        assert_eq!(masked_error("enabled"), None);
        assert_eq!(masked_error("disabled"), None);
        assert_eq!(masked_error("not-found"), None);
        let m = masked_error("masked").unwrap();
        assert!(m.contains("masked"));
        assert!(m.contains("systemctl --user unmask rlm-guard"));
        assert!(masked_error("masked-runtime")
            .unwrap()
            .contains("systemctl --user unmask --runtime rlm-guard"));
    }

    #[test]
    fn write_atomically_replaces_the_file_and_leaves_no_temp() {
        let d = tempfile::tempdir().unwrap();
        let p = d.path().join("rlm-guard.service");
        std::fs::write(&p, "old contents that are longer than the new ones\n").unwrap();
        write_atomically(&p, "new\n").unwrap();
        assert_eq!(std::fs::read_to_string(&p).unwrap(), "new\n");
        let names: Vec<_> = std::fs::read_dir(d.path())
            .unwrap()
            .map(|e| e.unwrap().file_name())
            .collect();
        assert_eq!(names, vec![std::ffi::OsString::from("rlm-guard.service")]);
        assert!(write_atomically(&d.path().join("missing/rlm-guard.service"), "x").is_err());
    }

    #[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"
        );
    }

    #[test]
    fn exec_arg_escapes_systemd_syntax() {
        assert_eq!(
            exec_arg(Path::new("/opt/100%/rlm-guard")),
            "/opt/100%%/rlm-guard"
        );
        assert_eq!(
            exec_arg(Path::new("/opt/a$b/rlm-guard")),
            "/opt/a$b/rlm-guard"
        );
        assert_eq!(
            unit_path_problem(Path::new("/opt/it's/rlm-guard")),
            Some("systemd does not accept quotes or backslashes in an ExecStart path")
        );
        assert!(unit_path_problem(Path::new("/opt/a\nb/rlm-guard")).is_some());
        assert_eq!(
            unit_path_problem(Path::new("/opt/100% sure/rlm-guard")),
            None
        );
        assert_eq!(
            exec_arg(Path::new(r"/opt/a\b/rlm-guard")),
            r"/opt/a\\b/rlm-guard"
        );
        assert_eq!(
            exec_arg(Path::new(r#"/opt/say "hi"/rlm-guard"#)),
            r#""/opt/say \"hi\"/rlm-guard""#
        );
        assert_eq!(
            exec_arg(Path::new("/opt/it's/rlm-guard")),
            r#""/opt/it\'s/rlm-guard""#
        );
        assert_eq!(
            exec_arg(Path::new(r"/opt/my 50% \dir/rlm-guard")),
            r#""/opt/my 50%% \\dir/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 relative_path_entries_are_skipped() {
        let d = tempfile::tempdir().unwrap();
        exe(d.path());
        // The same directory spelled relative to the working directory.
        let cwd = std::env::current_dir().unwrap();
        let mut rel = PathBuf::new();
        for _ in cwd.components().skip(1) {
            rel.push("..");
        }
        rel.push(d.path().strip_prefix("/").unwrap());
        assert!(rel.join("rlm-guard").is_file());
        let path = std::env::join_paths([Path::new(""), Path::new("."), &rel]).unwrap();
        assert_eq!(find_guard_binary(None, Some(&path)), None);
        let path = std::env::join_paths([rel.as_path(), d.path()]).unwrap();
        assert_eq!(
            find_guard_binary(None, Some(&path)),
            Some(d.path().join("rlm-guard"))
        );
    }

    #[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),
            }
        );
    }
}