OpenPXE/crates/dhcp-proxy/src/escalation.rs

//! Automatic per-MAC boot-binary escalation (v0.6.1, extended v0.7.x).
//!
//! OpenPXE serves the firmware-net iPXE build (`snponly`/`undionly`) by
//! default — it's the most reliable choice for chainloading because the
//! firmware just proved its network works by downloading the NBP. Two
//! classes of machine can't run it:
//!
//! * a minority of NICs have a missing or buggy firmware UNDI/SNP stack —
//!   they TFTP the binary fine but iPXE can't bring the link up;
//! * Secure-Boot firmware downloads it fine but refuses to *execute* an
//!   unsigned image.
//!
//! Both look identical from here: the tell-tale second DHCP DISCOVER
//! carrying the `iPXE` user-class never arrives and the machine
//! re-PXE-boots. So a fresh firmware DISCOVER from a MAC whose previous
//! attempt was never confirmed climbs one rung:
//! `Firmware → Builtin → Shim` (the signed shim+GRUB chain). The decision
//! is sticky; there is no operator toggle; the default path is unchanged
//! so hardware that already boots never regresses.
//!
//! v0.7.1 — **learned modes persist**. Walking the ladder costs one or
//! two failed boot cycles, so a machine should pay it once *ever*, not
//! once per idle window or server restart. Two events pin a MAC's mode
//! to disk (`<work_dir>/driver_modes.json`):
//!
//! * a confirmed iPXE handoff at a non-default mode (Builtin proved to
//!   work — also Shim, via the GRUB→iPXE same-boot chainload);
//! * reaching the terminal Shim rung (Secure-Boot machines never produce
//!   an iPXE handoff from the signed menu, so escalation itself is the
//!   best knowledge we'll ever have).
//!
//! Pinned entries are immune to the TTL and reload at startup. The
//! operator escape hatch is a rules-level driver-mode pin (which
//! overrides this table entirely) or deleting `driver_modes.json`.

use openpxe_core::DriverMode;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant};

/// Multiple DISCOVERs within this window belong to the *same* boot (DHCP
/// retransmits, plus the :4011 PXE Boot Server query that follows the :67
/// DISCOVER). They must not be mistaken for a failed-and-retried boot.
const SAME_BOOT_DEBOUNCE: Duration = Duration::from_secs(8);

/// Forget an *unpinned* MAC's state after this long with no activity, so
/// a transient mid-walk state doesn't linger and the map stays bounded.
/// Pinned (learned) entries are exempt — that's their whole point.
const ENTRY_TTL: Duration = Duration::from_mins(30);

/// Hard cap on tracked MACs. Past this we evict the least-recently-seen
/// entry (unpinned first) — escalation is best-effort, never a
/// memory-growth vector.
const MAX_ENTRIES: usize = 4096;

/// How often (at most) the whole map is swept for expired entries.
/// Correctness doesn't depend on the sweep — a stale entry is also
/// detected inline when its MAC next appears — so the sweep only bounds
/// memory for MACs that never return, and amortizing it keeps the
/// per-packet path O(1) instead of O(map).
const PRUNE_INTERVAL: Duration = Duration::from_mins(1);

#[derive(Debug, Clone, Copy)]
struct Entry {
    mode: DriverMode,
    /// True once we've served `mode` and are waiting for the iPXE handoff to
    /// confirm it worked. A *new* boot arriving while this is still true means
    /// the previous attempt failed and we should escalate.
    awaiting_confirm: bool,
    /// Learned mode (v0.7.1): persisted to disk, exempt from the TTL.
    pinned: bool,
    last_seen: Instant,
}

#[derive(Debug)]
struct Inner {
    map: HashMap<String, Entry>,
    /// When the last full TTL sweep ran — see [`PRUNE_INTERVAL`].
    last_prune: Instant,
}

impl Default for Inner {
    fn default() -> Self {
        Self {
            map: HashMap::new(),
            last_prune: Instant::now(),
        }
    }
}

/// Tracks per-MAC driver-mode escalation. Cheap to share via `Arc`.
#[derive(Debug, Default)]
pub struct DriverEscalation {
    inner: Mutex<Inner>,
    /// Persistence target for learned modes; `None` = ephemeral (tests).
    path: Option<Arc<PathBuf>>,
}

impl DriverEscalation {
    /// Ephemeral instance (no persistence) — used by tests.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Instance backed by `<work_dir>/driver_modes.json`. Learned modes
    /// from previous runs are reloaded as pinned entries; a missing or
    /// corrupt file starts empty (same crash-cache policy as every other
    /// store — a bad file must never block PXE).
    #[must_use]
    pub fn load_or_default(work_dir: &Path) -> Self {
        let path = work_dir.join("driver_modes.json");
        let mut map = HashMap::new();
        if let Ok(text) = std::fs::read_to_string(&path) {
            match serde_json::from_str::<HashMap<String, DriverMode>>(&text) {
                Ok(loaded) => {
                    let now = Instant::now();
                    for (mac, mode) in loaded {
                        // Firmware is the default — persisting it would be
                        // noise; tolerate it in the file but don't track it.
                        if mode == DriverMode::Firmware {
                            continue;
                        }
                        map.insert(
                            mac,
                            Entry {
                                mode,
                                awaiting_confirm: false,
                                pinned: true,
                                last_seen: now,
                            },
                        );
                    }
                    tracing::info!(
                        target: "openpxe::dhcp",
                        learned = map.len(),
                        "loaded learned driver modes"
                    );
                }
                Err(e) => {
                    tracing::warn!(
                        target: "openpxe::dhcp",
                        "driver_modes.json present but unreadable ({e}); starting empty"
                    );
                }
            }
        }
        Self {
            inner: Mutex::new(Inner {
                map,
                last_prune: Instant::now(),
            }),
            path: Some(Arc::new(path)),
        }
    }

    /// Decide the driver mode for a firmware (PXEClient/HTTPClient) boot from
    /// `mac`. `primary` is true for the main DHCP DISCOVER (:67) and false for
    /// the PXE Boot Server query (:4011); only the primary path drives
    /// escalation, and only when it's clearly a *new* boot (outside the
    /// same-boot debounce). The :4011 path just echoes the current mode.
    pub fn mode_for_firmware_attempt(&self, mac: &str, primary: bool) -> DriverMode {
        self.decide_at(mac, primary, Instant::now())
    }

    /// Record that `mac` completed the iPXE handoff (a DISCOVER carrying the
    /// `iPXE` user-class). The mode we last served worked, so stop awaiting
    /// confirmation, keep it sticky, and — for non-default modes — pin it to
    /// disk so the machine never re-walks the ladder (v0.7.1).
    pub fn mark_ipxe_success(&self, mac: &str) {
        self.confirm_at(mac, Instant::now());
    }

    fn decide_at(&self, mac: &str, primary: bool, now: Instant) -> DriverMode {
        let (mode, snapshot) = {
            let mut g = self.inner.lock();
            if now.duration_since(g.last_prune) >= PRUNE_INTERVAL {
                g.map
                    .retain(|_, e| e.pinned || now.duration_since(e.last_seen) < ENTRY_TTL);
                g.last_prune = now;
            }
            // Inline staleness check: an unpinned MAC whose entry outlived
            // the TTL starts fresh even when the amortized sweep above
            // hasn't caught it yet. Pinned entries never go stale.
            if g.map
                .get(mac)
                .is_some_and(|e| !e.pinned && now.duration_since(e.last_seen) >= ENTRY_TTL)
            {
                g.map.remove(mac);
            }

            let mut newly_pinned = false;
            let mode = match g.map.get_mut(mac) {
                None => {
                    g.map.insert(
                        mac.to_owned(),
                        Entry {
                            mode: DriverMode::Firmware,
                            // Only the primary DISCOVER opens a confirmation window.
                            awaiting_confirm: primary,
                            pinned: false,
                            last_seen: now,
                        },
                    );
                    if g.map.len() > MAX_ENTRIES {
                        evict_oldest(&mut g.map);
                    }
                    DriverMode::Firmware
                }
                Some(entry) => {
                    let recent = now.duration_since(entry.last_seen) < SAME_BOOT_DEBOUNCE;
                    if primary && !recent {
                        // A genuinely new boot. If the previous attempt was
                        // never confirmed, the build we served failed → climb
                        // one rung: Firmware (firmware NIC stack) → Builtin
                        // (iPXE's own drivers) → Shim (signed shim+GRUB —
                        // covers Secure Boot firmware that downloads our
                        // unsigned iPXE but refuses to execute it). Shim is
                        // terminal and pins to disk: SB machines never emit
                        // an iPXE handoff from the signed menu, so reaching
                        // the rung *is* the durable knowledge.
                        if entry.awaiting_confirm {
                            entry.mode = match entry.mode {
                                DriverMode::Firmware => DriverMode::Builtin,
                                DriverMode::Builtin | DriverMode::Shim => DriverMode::Shim,
                            };
                            if entry.mode == DriverMode::Shim && !entry.pinned {
                                entry.pinned = true;
                                newly_pinned = true;
                            }
                        }
                        entry.awaiting_confirm = true;
                    }
                    entry.last_seen = now;
                    entry.mode
                }
            };
            (mode, newly_pinned.then(|| pinned_snapshot(&g.map)))
        };
        if let Some(s) = snapshot {
            self.persist(&s);
        }
        mode
    }

    fn confirm_at(&self, mac: &str, now: Instant) {
        let snapshot = {
            let mut g = self.inner.lock();
            let Some(e) = g.map.get_mut(mac) else {
                return;
            };
            e.awaiting_confirm = false;
            e.last_seen = now;
            // A proven non-default mode is worth remembering forever —
            // the machine demonstrably can't use the default path.
            if e.mode != DriverMode::Firmware && !e.pinned {
                e.pinned = true;
                Some(pinned_snapshot(&g.map))
            } else {
                None
            }
        };
        if let Some(s) = snapshot {
            self.persist(&s);
        }
    }

    /// Best-effort atomic write of the learned-mode table. No-op for
    /// ephemeral instances. Failure logs and moves on — persistence is an
    /// optimization, never a correctness requirement.
    fn persist(&self, snapshot: &HashMap<String, DriverMode>) {
        let Some(path) = &self.path else { return };
        let body = match serde_json::to_vec_pretty(snapshot) {
            Ok(b) => b,
            Err(e) => {
                tracing::warn!(target: "openpxe::dhcp", "serialize driver_modes.json: {e}");
                return;
            }
        };
        if let Some(parent) = path.parent() {
            let _ = std::fs::create_dir_all(parent);
        }
        let tmp = path.with_extension("json.tmp");
        if let Err(e) = std::fs::write(&tmp, body) {
            tracing::warn!(target: "openpxe::dhcp", "write driver_modes.json tmp: {e}");
            return;
        }
        if let Err(e) = std::fs::rename(&tmp, path.as_path()) {
            tracing::warn!(target: "openpxe::dhcp", "rename driver_modes.json: {e}");
        }
    }
}

fn pinned_snapshot(map: &HashMap<String, Entry>) -> HashMap<String, DriverMode> {
    map.iter()
        .filter(|(_, e)| e.pinned)
        .map(|(k, e)| (k.clone(), e.mode))
        .collect()
}

fn evict_oldest(map: &mut HashMap<String, Entry>) {
    // Prefer evicting an unpinned entry; only touch learned modes when
    // the whole table is pinned (4096 learned machines — at that point
    // the operator has bigger questions than our memory bound).
    let pick = |pinned: bool| {
        map.iter()
            .filter(|(_, e)| e.pinned == pinned)
            .min_by_key(|(_, e)| e.last_seen)
            .map(|(k, _)| k.clone())
    };
    if let Some(oldest) = pick(false).or_else(|| pick(true)) {
        map.remove(&oldest);
    }
}

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

    #[test]
    fn firmware_first_then_escalates_on_unconfirmed_retry() {
        let e = DriverEscalation::new();
        let t0 = Instant::now();
        // Boot 1, primary DISCOVER: firmware.
        assert_eq!(e.decide_at("aa", true, t0), DriverMode::Firmware);
        // Same boot's :4011 query (+1s, within debounce): still firmware, no escalation.
        assert_eq!(
            e.decide_at("aa", false, t0 + Duration::from_secs(1)),
            DriverMode::Firmware
        );
        // Firmware net failed → no iPXE handoff → machine re-PXE-boots much
        // later: escalate to builtin drivers.
        assert_eq!(
            e.decide_at("aa", true, t0 + Duration::from_mins(1)),
            DriverMode::Builtin
        );
    }

    #[test]
    fn builtin_is_sticky_after_success() {
        let e = DriverEscalation::new();
        let t0 = Instant::now();
        assert_eq!(e.decide_at("bb", true, t0), DriverMode::Firmware);
        assert_eq!(
            e.decide_at("bb", true, t0 + Duration::from_mins(1)),
            DriverMode::Builtin
        );
        // Builtin worked this time — confirm the handoff.
        e.confirm_at("bb", t0 + Duration::from_secs(61));
        // Next cold boot goes straight to builtin (no wasted firmware attempt).
        assert_eq!(
            e.decide_at("bb", true, t0 + Duration::from_mins(2)),
            DriverMode::Builtin
        );
    }

    #[test]
    fn confirmed_firmware_never_escalates() {
        let e = DriverEscalation::new();
        let t0 = Instant::now();
        assert_eq!(e.decide_at("cc", true, t0), DriverMode::Firmware);
        // snponly worked: handoff confirmed.
        e.confirm_at("cc", t0 + Duration::from_secs(2));
        // A later boot stays on firmware — no spurious escalation.
        assert_eq!(
            e.decide_at("cc", true, t0 + Duration::from_mins(5)),
            DriverMode::Firmware
        );
    }

    #[test]
    fn third_unconfirmed_attempt_escalates_to_shim_and_stays() {
        // v0.7.0: a Secure-Boot client downloads-but-refuses both unsigned
        // iPXE builds; the third boot gets the signed shim chain, and the
        // MAC stays there for subsequent boots.
        let e = DriverEscalation::new();
        let t0 = Instant::now();
        assert_eq!(e.decide_at("ee", true, t0), DriverMode::Firmware);
        assert_eq!(
            e.decide_at("ee", true, t0 + Duration::from_mins(1)),
            DriverMode::Builtin
        );
        assert_eq!(
            e.decide_at("ee", true, t0 + Duration::from_mins(2)),
            DriverMode::Shim
        );
        // Shim is terminal — a fourth unconfirmed boot stays on Shim.
        assert_eq!(
            e.decide_at("ee", true, t0 + Duration::from_mins(3)),
            DriverMode::Shim
        );
    }

    #[test]
    fn stale_unpinned_entry_is_forgotten_and_resets_to_firmware() {
        let e = DriverEscalation::new();
        let t0 = Instant::now();
        assert_eq!(e.decide_at("dd", true, t0), DriverMode::Firmware);
        assert_eq!(
            e.decide_at("dd", true, t0 + Duration::from_mins(1)),
            DriverMode::Builtin
        );
        // After the TTL with no activity the (unpinned) Builtin walk is
        // pruned → fresh firmware. (A *confirmed* Builtin would be pinned
        // and survive — see learned_builtin_survives_ttl.)
        let later = t0 + Duration::from_mins(1) + ENTRY_TTL + Duration::from_secs(1);
        assert_eq!(e.decide_at("dd", true, later), DriverMode::Firmware);
    }

    #[test]
    fn shim_pin_survives_ttl() {
        // v0.7.1: reaching the Shim rung is durable knowledge — the
        // machine must NOT re-walk the ladder after an idle period.
        let e = DriverEscalation::new();
        let t0 = Instant::now();
        let _ = e.decide_at("ff", true, t0);
        let _ = e.decide_at("ff", true, t0 + Duration::from_mins(1));
        assert_eq!(
            e.decide_at("ff", true, t0 + Duration::from_mins(2)),
            DriverMode::Shim
        );
        let much_later = t0 + Duration::from_mins(2) + ENTRY_TTL + Duration::from_mins(5);
        assert_eq!(e.decide_at("ff", true, much_later), DriverMode::Shim);
    }

    #[test]
    fn learned_builtin_survives_ttl() {
        let e = DriverEscalation::new();
        let t0 = Instant::now();
        let _ = e.decide_at("gg", true, t0);
        assert_eq!(
            e.decide_at("gg", true, t0 + Duration::from_mins(1)),
            DriverMode::Builtin
        );
        // The handoff confirms Builtin → pinned.
        e.confirm_at("gg", t0 + Duration::from_secs(61));
        let much_later = t0 + ENTRY_TTL + Duration::from_mins(10);
        assert_eq!(e.decide_at("gg", true, much_later), DriverMode::Builtin);
    }

    #[test]
    fn learned_modes_persist_across_restart() {
        let dir = tempdir().unwrap();
        let t0 = Instant::now();
        {
            let e = DriverEscalation::load_or_default(dir.path());
            // Walk one MAC to Shim (pins on escalation)...
            let _ = e.decide_at("aa:01", true, t0);
            let _ = e.decide_at("aa:01", true, t0 + Duration::from_mins(1));
            assert_eq!(
                e.decide_at("aa:01", true, t0 + Duration::from_mins(2)),
                DriverMode::Shim
            );
            // ...and another to a confirmed Builtin (pins on handoff).
            let _ = e.decide_at("aa:02", true, t0);
            let _ = e.decide_at("aa:02", true, t0 + Duration::from_mins(1));
            e.confirm_at("aa:02", t0 + Duration::from_secs(61));
        }
        // "Restart": a fresh instance from the same work_dir knows both.
        let e2 = DriverEscalation::load_or_default(dir.path());
        assert_eq!(e2.decide_at("aa:01", true, t0), DriverMode::Shim);
        assert_eq!(e2.decide_at("aa:02", true, t0), DriverMode::Builtin);
        // Unlearned MACs still start at the default.
        assert_eq!(e2.decide_at("aa:03", true, t0), DriverMode::Firmware);
    }

    #[test]
    fn corrupt_persistence_file_starts_empty() {
        let dir = tempdir().unwrap();
        std::fs::write(dir.path().join("driver_modes.json"), b"{broken").unwrap();
        let e = DriverEscalation::load_or_default(dir.path());
        assert_eq!(
            e.decide_at("aa:bb", true, Instant::now()),
            DriverMode::Firmware
        );
    }

    #[test]
    fn confirmed_firmware_is_not_persisted() {
        // The default mode is never written — the file only carries
        // exceptions, so a healthy fleet leaves it absent/empty.
        let dir = tempdir().unwrap();
        let t0 = Instant::now();
        {
            let e = DriverEscalation::load_or_default(dir.path());
            let _ = e.decide_at("aa:09", true, t0);
            e.confirm_at("aa:09", t0 + Duration::from_secs(2));
        }
        assert!(!dir.path().join("driver_modes.json").exists());
    }
}