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