OpenPXE/crates/core/src/host_bindings.rs

//! Per-MAC host bindings.
//!
//! Inspired by the Tinkerbell `smee` "MAC-prepended URL" pattern: an
//! operator can attach a preferred boot target (a `BootEntry::id`) to a
//! specific MAC address. When a client with that MAC arrives, the
//! top-level boot script chains straight to that target instead of
//! showing the interactive menu.
//!
//! Use cases:
//! - "This rack of Dell servers always images with Ubuntu Server 24.04"
//! - "Tom's laptop always boots from local disk"
//! - "Bench QA machines always boot Memtest until released"
//!
//! Persisted to `<work_dir>/hosts.json`. Like the SettingsStore, on-disk
//! corruption falls back to an empty registry rather than failing
//! startup — a bad hosts file should never block PXE for the network.

use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use time::OffsetDateTime;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HostBinding {
    /// Lowercase, colon-separated MAC (e.g. `aa:bb:cc:dd:ee:ff`). The
    /// HTTP layer normalizes incoming MACs before lookup so callers
    /// don't have to worry about case.
    pub mac: String,
    /// Preferred boot entry id (matches a `BootEntry::id` in the iso
    /// store) OR one of the reserved menu names: `_local`, `_gate`,
    /// `_tools_menu`. Empty string falls back to the menu.
    pub target: String,
    /// Optional human-readable label shown in the UI (`"Tom's laptop"`,
    /// `"rack-3 spine"`). Empty if unset.
    #[serde(default)]
    pub label: String,
    #[serde(with = "time::serde::rfc3339")]
    pub created_at: OffsetDateTime,
    #[serde(with = "time::serde::rfc3339")]
    pub updated_at: OffsetDateTime,
}

#[derive(Debug, Default)]
struct Inner {
    by_mac: HashMap<String, HostBinding>,
}

/// Registry of per-MAC bindings. Cheap to clone; locks are held
/// briefly. Persistence is best-effort and mirrors `SettingsStore`'s
/// "in-memory authoritative, disk is a cache" policy.
#[derive(Debug, Clone)]
pub struct HostBindings {
    path: Arc<PathBuf>,
    inner: Arc<RwLock<Inner>>,
}

impl HostBindings {
    /// Load from `work_dir/hosts.json`, or start empty if absent /
    /// unreadable.
    #[must_use]
    pub fn load_or_default(work_dir: &std::path::Path) -> Self {
        let path = work_dir.join("hosts.json");
        let inner = match std::fs::read_to_string(&path) {
            Ok(text) => match serde_json::from_str::<Vec<HostBinding>>(&text) {
                Ok(items) => {
                    let mut by_mac = HashMap::new();
                    for b in items {
                        by_mac.insert(normalize_mac(&b.mac), b);
                    }
                    Inner { by_mac }
                }
                Err(e) => {
                    tracing::warn!(
                        target: "pxeforge::hosts",
                        "hosts.json present but unreadable ({e}); starting empty"
                    );
                    Inner::default()
                }
            },
            Err(_) => Inner::default(),
        };
        Self {
            path: Arc::new(path),
            inner: Arc::new(RwLock::new(inner)),
        }
    }

    /// Look up a binding by MAC. Match is case-insensitive and tolerates
    /// `-` or `:` separators.
    #[must_use]
    pub fn lookup(&self, mac: &str) -> Option<HostBinding> {
        self.inner.read().by_mac.get(&normalize_mac(mac)).cloned()
    }

    /// Insert or update. Returns the resulting binding (with timestamps).
    pub fn upsert(&self, mac: &str, target: &str, label: &str) -> HostBinding {
        let key = normalize_mac(mac);
        let now = OffsetDateTime::now_utc();
        let binding = {
            let mut g = self.inner.write();
            let entry = g.by_mac.entry(key.clone()).or_insert_with(|| HostBinding {
                mac: key.clone(),
                target: target.to_string(),
                label: label.to_string(),
                created_at: now,
                updated_at: now,
            });
            entry.target = target.to_string();
            entry.label = label.to_string();
            entry.updated_at = now;
            entry.clone()
        };
        self.persist();
        binding
    }

    /// Remove a binding. Returns true if something was removed.
    pub fn remove(&self, mac: &str) -> bool {
        let key = normalize_mac(mac);
        let removed = self.inner.write().by_mac.remove(&key).is_some();
        if removed {
            self.persist();
        }
        removed
    }

    #[must_use]
    pub fn list(&self) -> Vec<HostBinding> {
        let g = self.inner.read();
        let mut v: Vec<_> = g.by_mac.values().cloned().collect();
        v.sort_by(|a, b| a.mac.cmp(&b.mac));
        v
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.inner.read().by_mac.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn persist(&self) {
        let items: Vec<HostBinding> = self.inner.read().by_mac.values().cloned().collect();
        let body = match serde_json::to_vec_pretty(&items) {
            Ok(b) => b,
            Err(e) => {
                tracing::warn!(target: "pxeforge::hosts", "serialize hosts.json: {e}");
                return;
            }
        };
        let tmp = self.path.with_extension("json.tmp");
        if let Some(parent) = self.path.parent() {
            let _ = std::fs::create_dir_all(parent);
        }
        if let Err(e) = std::fs::write(&tmp, body) {
            tracing::warn!(target: "pxeforge::hosts", "write hosts.json tmp: {e}");
            return;
        }
        if let Err(e) = std::fs::rename(&tmp, self.path.as_path()) {
            tracing::warn!(target: "pxeforge::hosts", "rename hosts.json: {e}");
        }
    }
}

/// Lowercase a MAC and normalise `-` separators to `:`. We never strip
/// the separator entirely — `aabbccddeeff` formats are rejected at the
/// HTTP layer because they're ambiguous (could be a device id).
#[must_use]
pub fn normalize_mac(mac: &str) -> String {
    mac.trim().to_ascii_lowercase().replace('-', ":")
}

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

    #[test]
    fn normalize_handles_case_and_dashes() {
        assert_eq!(normalize_mac("AA:BB:CC:DD:EE:FF"), "aa:bb:cc:dd:ee:ff");
        assert_eq!(normalize_mac("aa-bb-cc-dd-ee-ff"), "aa:bb:cc:dd:ee:ff");
        assert_eq!(normalize_mac("  aA-Bb-CC:DD-ee:fF "), "aa:bb:cc:dd:ee:ff");
    }

    #[test]
    fn upsert_then_lookup() {
        let dir = tempdir().unwrap();
        let h = HostBindings::load_or_default(dir.path());
        assert!(h.is_empty());
        h.upsert("AA:BB:CC:00:00:01", "ubuntu-24-04-linux", "rack-3 spine");
        let found = h.lookup("aa-bb-cc-00-00-01").expect("lookup");
        assert_eq!(found.target, "ubuntu-24-04-linux");
        assert_eq!(found.label, "rack-3 spine");
        assert_eq!(found.mac, "aa:bb:cc:00:00:01");
    }

    #[test]
    fn upsert_replaces_existing_target() {
        let dir = tempdir().unwrap();
        let h = HostBindings::load_or_default(dir.path());
        h.upsert("aa:bb:cc:00:00:01", "old-target", "label1");
        h.upsert("aa:bb:cc:00:00:01", "new-target", "label2");
        assert_eq!(h.len(), 1);
        let b = h.lookup("aa:bb:cc:00:00:01").unwrap();
        assert_eq!(b.target, "new-target");
        assert_eq!(b.label, "label2");
    }

    #[test]
    fn remove_works_and_reports_outcome() {
        let dir = tempdir().unwrap();
        let h = HostBindings::load_or_default(dir.path());
        h.upsert("aa:bb:cc:00:00:01", "x", "");
        assert!(h.remove("AA:BB:CC:00:00:01"));
        assert!(!h.remove("aa:bb:cc:00:00:01")); // already gone
        assert!(h.is_empty());
    }

    #[test]
    fn round_trip_persists_to_disk() {
        let dir = tempdir().unwrap();
        let h = HostBindings::load_or_default(dir.path());
        h.upsert("aa:bb:cc:00:00:01", "ubuntu-linux", "rack-3");
        h.upsert("aa:bb:cc:00:00:02", "_local", "tom-laptop");
        drop(h);
        let h2 = HostBindings::load_or_default(dir.path());
        assert_eq!(h2.len(), 2);
        assert_eq!(h2.lookup("aa:bb:cc:00:00:02").unwrap().target, "_local");
    }
}