openlatch-client 0.5.4

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

use crate::error::OlError;

use super::{
    Supervisor, SupervisorKind, SupervisorStatus, ERR_SUPERVISION_CONTROL_FAILED,
    ERR_SUPERVISION_INSTALL_FAILED,
};

const TASK_NAME: &str = "OpenLatch\\Client";

pub struct TaskSchedulerSupervisor;

impl Default for TaskSchedulerSupervisor {
    fn default() -> Self {
        Self::new()
    }
}

impl TaskSchedulerSupervisor {
    pub fn new() -> Self {
        Self
    }

    fn generate_xml(&self, binary_path: &Path, user: &str) -> String {
        let bin = binary_path.display();
        // schtasks opens XML files as UTF-16LE regardless of declared
        // encoding — we therefore emit the declaration as UTF-16 and later
        // write the bytes as UTF-16LE with BOM (see `encode_utf16le_bom`).
        //
        // <Principal> MUST carry a <UserId>; without it schtasks defaults the
        // task to the SYSTEM principal, which then requires admin to register.
        // Naming the current user keeps registration inside the normal user
        // scope so no UAC prompt is required. Actions Context="Author" wires
        // the <Exec> block to that same principal.
        //
        // The version marker rides in <Description>, never in an XML comment.
        // Task Scheduler parses the XML into its own store and re-serialises it
        // on `/Query /XML`, and comments do not survive the trip — a marker
        // written as one made every Windows install read as outdated, from the
        // first `init` onwards, with a remedy that could never clear it.
        let marker = super::unit_version_marker();
        format!(
            r#"<?xml version="1.0" encoding="UTF-16"?>
<Task version="1.2" xmlns="http://schemas.microsoft.com/windows/2004/02/mit/task">
  <RegistrationInfo>
    <Description>OpenLatch client daemon ({marker})</Description>
  </RegistrationInfo>
  <Triggers>
    <LogonTrigger>
      <Enabled>true</Enabled>
      <UserId>{user}</UserId>
    </LogonTrigger>
  </Triggers>
  <Settings>
    <MultipleInstancesPolicy>IgnoreNew</MultipleInstancesPolicy>
    <DisallowStartIfOnBatteries>false</DisallowStartIfOnBatteries>
    <StopIfGoingOnBatteries>false</StopIfGoingOnBatteries>
    <ExecutionTimeLimit>PT0S</ExecutionTimeLimit>
    <Hidden>true</Hidden>
    <!--
      PT1M is Task Scheduler's floor: the schema rejects any RestartOnFailure
      Interval below one minute, so Windows cannot match systemd's 2s or
      launchd's 2s cadence. This is the honest bound documented in
      src/model_relay/mod.rs — a sub-minute restart on Windows needs the model relay
      split into its own Service with SCM failure actions, which is out of
      scope. Count is raised 5 → 10 so a transient crash loop (port in
      TIME_WAIT, disk momentarily full) does not exhaust the retries and leave
      the host with no daemon at all.
    -->
    <RestartOnFailure>
      <Interval>PT1M</Interval>
      <Count>10</Count>
    </RestartOnFailure>
  </Settings>
  <Principals>
    <Principal id="Author">
      <UserId>{user}</UserId>
      <LogonType>InteractiveToken</LogonType>
      <RunLevel>LeastPrivilege</RunLevel>
    </Principal>
  </Principals>
  <Actions Context="Author">
    <Exec>
      <Command>{bin}</Command>
      <Arguments>daemon start --foreground</Arguments>
    </Exec>
  </Actions>
</Task>"#
        )
    }
}

/// Encode a string as UTF-16LE with a byte-order mark.
///
/// schtasks always reads XML files as UTF-16 and rejects any other encoding
/// with "unable to switch the encoding", regardless of what the XML prolog
/// claims. Writing the bytes as UTF-16LE with a BOM is the only format it
/// reliably accepts — plain UTF-8 (with or without BOM) triggers the
/// encoding switch error.
fn encode_utf16le_bom(s: &str) -> Vec<u8> {
    let units: Vec<u16> = s.encode_utf16().collect();
    let mut bytes = Vec::with_capacity(2 + units.len() * 2);
    bytes.extend_from_slice(&[0xFF, 0xFE]); // UTF-16LE BOM
    for u in units {
        bytes.extend_from_slice(&u.to_le_bytes());
    }
    bytes
}

/// Does a `schtasks /Query /XML` export carry the current version marker?
///
/// The export is the scheduler's re-serialisation of what it stored, not the
/// file `install` handed it — which is why the marker lives in `<Description>`
/// (see `generate_xml`). Written to a pipe, schtasks emits the console code
/// page, not the UTF-16 its prolog declares; the UTF-16LE read is kept for a
/// host that does, and drops the high bytes rather than pull in an encoding
/// crate for one ASCII marker.
fn export_is_current(stdout: &[u8]) -> bool {
    let utf16_low_bytes: String = stdout
        .chunks(2)
        .filter_map(|c| c.first().copied())
        .map(char::from)
        .collect();
    super::unit_is_current(&String::from_utf8_lossy(stdout))
        || super::unit_is_current(&utf16_low_bytes)
}

/// Resolve the interactive user identifier that the scheduled task runs as.
///
/// Preference order:
/// 1. `USERDOMAIN\USERNAME` — works for both domain-joined and local accounts
///    (on local machines `USERDOMAIN` is the computer name).
/// 2. Bare `USERNAME` as a fallback when `USERDOMAIN` is empty.
/// 3. `S-1-5-32-545` (BUILTIN\Users) as a last-resort SID that maps to the
///    interactive user without requiring admin to register.
fn current_user_identifier() -> String {
    let username = std::env::var("USERNAME").unwrap_or_default();
    let domain = std::env::var("USERDOMAIN").unwrap_or_default();
    if !domain.is_empty() && !username.is_empty() {
        format!("{domain}\\{username}")
    } else if !username.is_empty() {
        username
    } else {
        "S-1-5-32-545".to_string()
    }
}

impl Supervisor for TaskSchedulerSupervisor {
    fn kind(&self) -> SupervisorKind {
        SupervisorKind::TaskScheduler
    }

    fn install(&self, binary_path: &Path) -> Result<(), OlError> {
        let user = current_user_identifier();
        let xml = self.generate_xml(binary_path, &user);
        let bytes = encode_utf16le_bom(&xml);

        let tmp_dir =
            std::env::var("TEMP").unwrap_or_else(|_| std::env::var("TMP").unwrap_or_default());
        let xml_path = PathBuf::from(&tmp_dir).join("openlatch-task.xml");
        std::fs::write(&xml_path, &bytes).map_err(|e| {
            OlError::new(
                ERR_SUPERVISION_INSTALL_FAILED,
                format!("Cannot write task XML: {e}"),
            )
        })?;

        // /RU names the run-as user explicitly so schtasks never falls back
        // to SYSTEM. /IT marks the task "interactive only" — it fires when
        // the named user is logged in, which is what a per-user LaunchAgent-
        // equivalent wants on Windows.
        let output = std::process::Command::new("schtasks")
            .args([
                "/Create",
                "/TN",
                TASK_NAME,
                "/XML",
                &xml_path.display().to_string(),
                "/RU",
                &user,
                "/IT",
                "/F",
            ])
            .output();

        let _ = std::fs::remove_file(&xml_path);

        match output {
            Ok(o) if o.status.success() => Ok(()),
            Ok(o) => {
                let stderr = String::from_utf8_lossy(&o.stderr);
                Err(OlError::new(
                    ERR_SUPERVISION_INSTALL_FAILED,
                    format!("schtasks /Create failed: {stderr}"),
                ))
            }
            Err(e) => Err(OlError::new(
                ERR_SUPERVISION_INSTALL_FAILED,
                format!("Cannot run schtasks: {e}"),
            )),
        }
    }

    fn uninstall(&self) -> Result<(), OlError> {
        let _ = std::process::Command::new("schtasks")
            .args(["/Delete", "/TN", TASK_NAME, "/F"])
            .output();
        Ok(())
    }

    fn status(&self) -> Result<SupervisorStatus, OlError> {
        let output = std::process::Command::new("schtasks")
            .args(["/Query", "/TN", TASK_NAME])
            .output();

        match output {
            Ok(o) if o.status.success() => {
                // Registered tasks live in the scheduler's own store, not as a
                // file we can re-read, so the marker is recovered from the XML
                // export. A failed export is reported as drift rather than
                // silently as "current" — the point of the flag is to notice
                // when we cannot confirm.
                let unit_current = std::process::Command::new("schtasks")
                    .args(["/Query", "/TN", TASK_NAME, "/XML"])
                    .output()
                    .ok()
                    .filter(|x| x.status.success())
                    .is_some_and(|x| export_is_current(&x.stdout));
                Ok(SupervisorStatus {
                    installed: true,
                    running: true,
                    unit_current,
                    description: if unit_current {
                        "Task Scheduler (RestartOnFailure active)".into()
                    } else {
                        "Task Scheduler (outdated task definition — reinstall to refresh it)".into()
                    },
                })
            }
            _ => Ok(SupervisorStatus {
                installed: false,
                running: false,
                unit_current: false,
                description: "not installed".into(),
            }),
        }
    }

    fn start(&self) -> Result<(), OlError> {
        // `MultipleInstancesPolicy=IgnoreNew` makes this idempotent: running
        // a task that already has an instance is a no-op, not a second daemon.
        run_control(&["/Run", "/TN", TASK_NAME])
    }

    fn stop(&self) -> Result<(), OlError> {
        // `/End` terminates the running instance. Unlike systemd and launchd
        // there is nothing to suppress afterwards — the task's only trigger is
        // `LogonTrigger`, and `RestartOnFailure` fires on a failed run, not on
        // one Task Scheduler was asked to end. The task stays registered, so
        // the daemon returns at the next logon.
        run_control(&["/End", "/TN", TASK_NAME])
    }

    fn restart(&self) -> Result<(), OlError> {
        // No atomic restart verb. `/End` first; a task with no running
        // instance reports failure, which is benign here — the `/Run` is what
        // has to succeed.
        let _ = run_control(&["/End", "/TN", TASK_NAME]);
        run_control(&["/Run", "/TN", TASK_NAME])
    }
}

/// Invoke a `schtasks` control verb, mapping failure to
/// [`ERR_SUPERVISION_CONTROL_FAILED`] so callers can fall back to acting on
/// the process directly.
fn run_control(args: &[&str]) -> Result<(), OlError> {
    let out = std::process::Command::new("schtasks").args(args).output();
    match out {
        Ok(o) if o.status.success() => Ok(()),
        Ok(o) => Err(OlError::new(
            ERR_SUPERVISION_CONTROL_FAILED,
            format!(
                "schtasks {} failed: {}",
                args.join(" "),
                String::from_utf8_lossy(&o.stderr).trim()
            ),
        )),
        Err(e) => Err(OlError::new(
            ERR_SUPERVISION_CONTROL_FAILED,
            format!("Cannot run schtasks: {e}"),
        )),
    }
}

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

    #[test]
    fn xml_hidden_is_true() {
        // `<Hidden>` keeps the boot-time daemon out of the scheduler UI's
        // default task list. It does NOT suppress the console window — that
        // is `release_unattended_console` in `src/app/openlatch/main.rs`,
        // which frees the console Task Scheduler allocated once it has
        // established that this process is its only member. The two are
        // separate concerns and were conflated while the binary carried
        // `windows_subsystem = "windows"`.
        let sup = TaskSchedulerSupervisor::new();
        let xml = sup.generate_xml(Path::new("C:\\openlatch\\openlatch.exe"), "USER\\me");
        assert!(xml.contains("<Hidden>true</Hidden>"));
    }

    #[test]
    fn xml_has_execution_time_limit_zero() {
        let sup = TaskSchedulerSupervisor::new();
        let xml = sup.generate_xml(Path::new("C:\\openlatch\\openlatch.exe"), "USER\\me");
        assert!(xml.contains("<ExecutionTimeLimit>PT0S</ExecutionTimeLimit>"));
    }

    #[test]
    fn xml_has_logon_trigger() {
        let sup = TaskSchedulerSupervisor::new();
        let xml = sup.generate_xml(Path::new("C:\\openlatch\\openlatch.exe"), "USER\\me");
        assert!(xml.contains("<LogonTrigger>"));
    }

    #[test]
    fn xml_has_restart_on_failure() {
        let sup = TaskSchedulerSupervisor::new();
        let xml = sup.generate_xml(Path::new("C:\\openlatch\\openlatch.exe"), "USER\\me");
        assert!(xml.contains("<RestartOnFailure>"));
        // PT1M is the schema floor — documented, not a choice we can improve.
        assert!(xml.contains("<Interval>PT1M</Interval>"));
        // Raised 5 → 10 so a transient crash loop cannot exhaust the retries.
        assert!(xml.contains("<Count>10</Count>"), "{xml}");
    }

    #[test]
    fn xml_carries_the_version_marker() {
        let sup = TaskSchedulerSupervisor::new();
        let xml = sup.generate_xml(Path::new("C:\\openlatch\\openlatch.exe"), "USER\\me");
        assert!(super::super::unit_is_current(&xml), "{xml}");
    }

    /// Remove every `<!-- … -->`, which is what registration does to the XML.
    fn strip_xml_comments(xml: &str) -> String {
        let mut out = String::new();
        let mut rest = xml;
        while let Some(start) = rest.find("<!--") {
            out.push_str(&rest[..start]);
            let end = rest[start..].find("-->").expect("unterminated comment");
            rest = &rest[start + end + 3..];
        }
        out.push_str(rest);
        out
    }

    #[test]
    fn the_version_marker_survives_registration() {
        // Measured against a live Task Scheduler: a task registered from XML
        // exports without its comments. A marker inside one read as outdated
        // on every Windows install, straight after `init`, and the remedy —
        // reinstall — wrote the same comment and lost it again.
        let sup = TaskSchedulerSupervisor::new();
        let xml = sup.generate_xml(Path::new("C:\\openlatch\\openlatch.exe"), "USER\\me");
        let registered = strip_xml_comments(&xml);
        assert!(export_is_current(registered.as_bytes()), "{registered}");
        assert!(registered.contains(&format!(
            "<Description>OpenLatch client daemon ({})</Description>",
            super::super::unit_version_marker()
        )));
    }

    #[test]
    fn an_export_is_read_in_either_encoding() {
        let sup = TaskSchedulerSupervisor::new();
        let xml = sup.generate_xml(Path::new("C:\\openlatch\\openlatch.exe"), "USER\\me");
        // What a pipe actually receives: the console code page.
        assert!(export_is_current(xml.as_bytes()));
        // What the prolog declares.
        assert!(export_is_current(&encode_utf16le_bom(&xml)));
    }

    #[test]
    fn a_task_registered_with_the_comment_marker_reads_as_outdated() {
        // The shape an install from a release that wrote the marker as a
        // comment exports as. Reporting it stale is correct — nothing in it
        // proves the generation — and `run_start` regenerates it.
        let export = "<?xml version=\"1.0\" encoding=\"UTF-16\"?>\r\r\n\
                      <Task version=\"1.2\" xmlns=\"http://schemas.microsoft.com/windows/2004/02/mit/task\">\r\r\n\
                      <RegistrationInfo><URI>\\OpenLatch\\Client</URI></RegistrationInfo>\r\r\n\
                      </Task>";
        assert!(!export_is_current(export.as_bytes()));
        assert!(!export_is_current(&encode_utf16le_bom(export)));
    }

    #[test]
    fn xml_has_least_privilege() {
        let sup = TaskSchedulerSupervisor::new();
        let xml = sup.generate_xml(Path::new("C:\\openlatch\\openlatch.exe"), "USER\\me");
        assert!(xml.contains("<RunLevel>LeastPrivilege</RunLevel>"));
    }

    #[test]
    fn xml_declares_utf16_to_match_on_disk_bytes() {
        // schtasks always reads the file as UTF-16LE regardless of the
        // declaration; we therefore emit the declaration as UTF-16 and
        // write the bytes via encode_utf16le_bom(). Any other combination
        // triggers "unable to switch the encoding".
        let sup = TaskSchedulerSupervisor::new();
        let xml = sup.generate_xml(Path::new("C:\\openlatch\\openlatch.exe"), "USER\\me");
        assert!(xml.starts_with("<?xml version=\"1.0\" encoding=\"UTF-16\"?>"));
    }

    #[test]
    fn utf16le_bom_encoder_emits_bom_and_little_endian_units() {
        let out = encode_utf16le_bom("A");
        assert_eq!(out, vec![0xFF, 0xFE, 0x41, 0x00]);
    }

    #[test]
    fn utf16le_bom_encoder_roundtrips_ascii_xml_prolog() {
        let prolog = "<?xml version=\"1.0\"?>";
        let bytes = encode_utf16le_bom(prolog);
        // BOM + 21 code units × 2 bytes.
        assert_eq!(bytes.len(), 2 + prolog.len() * 2);
        assert_eq!(&bytes[0..2], &[0xFF, 0xFE]);
        // First character is '<' = 0x3C in UTF-16LE (3C, 00).
        assert_eq!(&bytes[2..4], &[0x3C, 0x00]);
    }

    #[test]
    fn xml_principal_has_user_id() {
        // No UserId → schtasks tries to register against SYSTEM and needs
        // admin. Scoping the task to the current user keeps it in the
        // non-elevated lane.
        let sup = TaskSchedulerSupervisor::new();
        let xml = sup.generate_xml(Path::new("C:\\openlatch\\openlatch.exe"), "MYHOST\\alice");
        assert!(xml.contains("<UserId>MYHOST\\alice</UserId>"));
    }

    #[test]
    fn xml_actions_context_is_author() {
        // Actions inherit the principal via Context="Author"; without it the
        // task has no security context and register fails.
        let sup = TaskSchedulerSupervisor::new();
        let xml = sup.generate_xml(Path::new("C:\\openlatch\\openlatch.exe"), "USER\\me");
        assert!(xml.contains("<Actions Context=\"Author\">"));
    }

    #[test]
    fn current_user_identifier_prefers_domain_backslash_user() {
        // Documents the precedence so future changes to env-var handling
        // don't silently regress to SID fallback.
        let domain = std::env::var("USERDOMAIN").ok();
        let username = std::env::var("USERNAME").ok();
        let id = current_user_identifier();
        if let (Some(d), Some(u)) = (domain.as_deref(), username.as_deref()) {
            if !d.is_empty() && !u.is_empty() {
                assert_eq!(id, format!("{d}\\{u}"));
            }
        }
    }
}