AWG2_harness/reboot_module/src/main.rs

use std::{
    fs,
    io::{ErrorKind, Write},
    path::PathBuf,
    process::Command,
    thread,
    time::Duration as StdDuration,
};

use anyhow::Context;
use chrono::{Duration, NaiveDateTime, NaiveTime, Timelike, Utc};
use clap::Parser;
use rand::Rng;
use serde::Deserialize;

const DEFAULT_CONFIG_PATH: &str = "/etc/awg2-harness/config.toml";
const RUNTIME_DIR: &str = "/run/awg2-harness";
const REBOOT_PLAN_FILE: &str = "/run/awg2-harness/reboot.plan";

#[derive(Debug, Parser)]
#[command(name = "awg2-reboot-scheduler")]
struct Cli {
    #[arg(long)]
    config: Option<PathBuf>,
}

#[derive(Debug, Clone, Deserialize)]
struct RootConfig {
    reboot: RebootConfig,
}

#[derive(Debug, Clone, Deserialize)]
struct RebootConfig {
    enabled: bool,
    window_start_utc: String,
    window_end_utc: String,
}

fn main() -> anyhow::Result<()> {
    let cli = Cli::parse();
    let config_path = cli
        .config
        .unwrap_or_else(|| PathBuf::from(DEFAULT_CONFIG_PATH));

    let cfg = load_config(&config_path)?;
    if !cfg.reboot.enabled {
        // A plan left over from when the module was enabled would keep showing up in
        // `status` as a scheduled reboot that will never happen.
        clear_plan()?;
        println!("disabled by config");
        return Ok(());
    }

    let planned = next_reboot_time(&cfg.reboot.window_start_utc, &cfg.reboot.window_end_utc)?;
    persist_plan(planned)?;
    println!("reboot planned at UTC {}", planned);

    loop {
        let now = Utc::now().naive_utc();
        if now >= planned {
            break;
        }
        thread::sleep(StdDuration::from_secs(30));
    }

    loop {
        let out = Command::new("systemctl")
            .args(["--no-block", "reboot"])
            .output()
            .context("failed to run systemctl reboot")?;
        if out.status.success() {
            println!("reboot request accepted by systemctl");
            // The plan has been carried out. Leaving it behind would make `status` show a
            // past timestamp as if a reboot were still pending.
            clear_plan()?;
            return Ok(());
        }
        eprintln!(
            "reboot request failed: {}",
            String::from_utf8_lossy(&out.stderr)
        );
        thread::sleep(StdDuration::from_secs(60));
    }
}

fn load_config(path: &PathBuf) -> anyhow::Result<RootConfig> {
    let raw = fs::read_to_string(path)
        .with_context(|| format!("failed to read config: {}", path.display()))?;
    let cfg: RootConfig = toml::from_str(&raw)
        .with_context(|| format!("failed to parse config TOML: {}", path.display()))?;
    Ok(cfg)
}

/// Writes the plan atomically, and reports failure.
///
/// This is the only thing that tells an operator when the host is going to reboot. It used
/// to be entirely best-effort - the directory creation, the write and any error all
/// discarded - so a failure produced a scheduler that would reboot at a time nothing on
/// the system could report.
fn persist_plan(planned: NaiveDateTime) -> anyhow::Result<()> {
    fs::create_dir_all(RUNTIME_DIR).with_context(|| format!("failed to create {RUNTIME_DIR}"))?;

    let content = format!(
        "planned_reboot_utc={}\nplanned_reboot_unix={}\n",
        planned.format("%Y-%m-%dT%H:%M:%SZ"),
        planned.and_utc().timestamp()
    );

    let temp = PathBuf::from(format!("{REBOOT_PLAN_FILE}.tmp"));
    match fs::remove_file(&temp) {
        Ok(()) => {}
        Err(e) if e.kind() == ErrorKind::NotFound => {}
        Err(e) => return Err(e).with_context(|| format!("failed to clear {}", temp.display())),
    }

    let mut f = fs::OpenOptions::new()
        .create_new(true)
        .write(true)
        .open(&temp)
        .with_context(|| format!("failed to create {}", temp.display()))?;
    f.write_all(content.as_bytes())
        .with_context(|| format!("failed to write {}", temp.display()))?;
    f.sync_all()
        .with_context(|| format!("failed to flush {}", temp.display()))?;
    drop(f);

    fs::rename(&temp, REBOOT_PLAN_FILE)
        .with_context(|| format!("failed to install {REBOOT_PLAN_FILE}"))?;
    Ok(())
}

/// Removes the plan file. An absent plan is the desired end state, not an error.
fn clear_plan() -> anyhow::Result<()> {
    match fs::remove_file(REBOOT_PLAN_FILE) {
        Ok(()) => Ok(()),
        Err(e) if e.kind() == ErrorKind::NotFound => Ok(()),
        Err(e) => Err(e).with_context(|| format!("failed to remove {REBOOT_PLAN_FILE}")),
    }
}

fn next_reboot_time(start_hhmm: &str, end_hhmm: &str) -> anyhow::Result<NaiveDateTime> {
    let start = parse_hhmm(start_hhmm)?;
    let end = parse_hhmm(end_hhmm)?;
    if start == end {
        anyhow::bail!("reboot window start and end must differ");
    }

    let now = Utc::now().naive_utc();
    let now_t = NaiveTime::from_hms_opt(now.hour(), now.minute(), now.second())
        .ok_or_else(|| anyhow::anyhow!("bad current time"))?;

    let (range_start, range_end) = current_or_next_window(now, now_t, start, end)?;
    let span = (range_end - range_start).num_seconds().max(1);
    let offset = rand::thread_rng().gen_range(0..span);
    Ok(range_start + Duration::seconds(offset))
}

fn current_or_next_window(
    now: NaiveDateTime,
    now_t: NaiveTime,
    start: NaiveTime,
    end: NaiveTime,
) -> anyhow::Result<(NaiveDateTime, NaiveDateTime)> {
    let today = now.date();
    let tomorrow = today
        .succ_opt()
        .ok_or_else(|| anyhow::anyhow!("date overflow"))?;

    if start < end {
        if now_t < start {
            let s = today
                .and_hms_opt(start.hour(), start.minute(), 0)
                .ok_or_else(|| anyhow::anyhow!("invalid start dt"))?;
            let e = today
                .and_hms_opt(end.hour(), end.minute(), 0)
                .ok_or_else(|| anyhow::anyhow!("invalid end dt"))?;
            return Ok((s, e));
        }
        if now_t < end {
            let e = today
                .and_hms_opt(end.hour(), end.minute(), 0)
                .ok_or_else(|| anyhow::anyhow!("invalid end dt"))?;
            return Ok((now, e));
        }
        let s = tomorrow
            .and_hms_opt(start.hour(), start.minute(), 0)
            .ok_or_else(|| anyhow::anyhow!("invalid start dt"))?;
        let e = tomorrow
            .and_hms_opt(end.hour(), end.minute(), 0)
            .ok_or_else(|| anyhow::anyhow!("invalid end dt"))?;
        return Ok((s, e));
    }

    // cross-midnight window, e.g. 23:00..02:00
    if now_t >= start {
        let end_day = tomorrow;
        let e = end_day
            .and_hms_opt(end.hour(), end.minute(), 0)
            .ok_or_else(|| anyhow::anyhow!("invalid end dt"))?;
        return Ok((now, e));
    }
    if now_t < end {
        let e = today
            .and_hms_opt(end.hour(), end.minute(), 0)
            .ok_or_else(|| anyhow::anyhow!("invalid end dt"))?;
        return Ok((now, e));
    }
    let s = today
        .and_hms_opt(start.hour(), start.minute(), 0)
        .ok_or_else(|| anyhow::anyhow!("invalid start dt"))?;
    let e = tomorrow
        .and_hms_opt(end.hour(), end.minute(), 0)
        .ok_or_else(|| anyhow::anyhow!("invalid end dt"))?;
    Ok((s, e))
}

fn parse_hhmm(v: &str) -> anyhow::Result<NaiveTime> {
    let mut s = v.split(':');
    let h = s
        .next()
        .ok_or_else(|| anyhow::anyhow!("bad time"))?
        .parse()?;
    let m = s
        .next()
        .ok_or_else(|| anyhow::anyhow!("bad time"))?
        .parse()?;
    if s.next().is_some() {
        anyhow::bail!("bad time")
    }
    if h > 23 || m > 59 {
        anyhow::bail!("bad time")
    }
    NaiveTime::from_hms_opt(h, m, 0).ok_or_else(|| anyhow::anyhow!("bad time"))
}

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

    #[test]
    fn parse_time_rejects_invalid() {
        assert!(parse_hhmm("24:00").is_err());
        assert!(parse_hhmm("10:60").is_err());
        assert!(parse_hhmm("ab:cd").is_err());
    }

    #[test]
    fn window_inside_day_now_inside_window_uses_now_to_end() {
        let now =
            NaiveDateTime::parse_from_str("2026-01-10 01:30:00", "%Y-%m-%d %H:%M:%S").expect("now");
        let (s, e) = current_or_next_window(
            now,
            parse_hhmm("01:30").expect("time"),
            parse_hhmm("01:00").expect("start"),
            parse_hhmm("02:00").expect("end"),
        )
        .expect("window");
        assert_eq!(s, now);
        assert_eq!(
            e.format("%Y-%m-%d %H:%M:%S").to_string(),
            "2026-01-10 02:00:00"
        );
    }

    #[test]
    fn cross_midnight_window_supported() {
        let now =
            NaiveDateTime::parse_from_str("2026-01-10 23:30:00", "%Y-%m-%d %H:%M:%S").expect("now");
        let (s, e) = current_or_next_window(
            now,
            parse_hhmm("23:30").expect("time"),
            parse_hhmm("23:00").expect("start"),
            parse_hhmm("02:00").expect("end"),
        )
        .expect("window");
        assert_eq!(s, now);
        assert_eq!(
            e.format("%Y-%m-%d %H:%M:%S").to_string(),
            "2026-01-11 02:00:00"
        );
    }

    #[test]
    fn equal_window_boundaries_rejected() {
        let err = next_reboot_time("01:00", "01:00").err();
        assert!(err.is_some());
    }
}