OpenPXE/crates/iso-store/src/introspect.rs

//! ISO introspection — identify the distro family and locate kernel/initrd.
//!
//! We avoid a full ISO9660/Joliet/Rock-Ridge parser by reading a small number
//! of well-known files via `isoinfo` (from cdrtools/genisoimage) when it's on
//! the path. As a pure-Rust fallback we do a crude scan: read the volume
//! descriptor at offset 0x8000 to grab the volume label, and grep for known
//! filenames by scanning raw sectors — good enough to tell Debian from RHEL
//! most of the time, without shelling out.
//!
//! The returned `IntrospectionReport` is what `BootEntry`s get generated from.

use serde::{Deserialize, Serialize};
use std::io::{Read, Seek, SeekFrom};
use std::path::Path;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DistroFamily {
    DebianUbuntu,
    RhelFedora,
    OpenSuse,
    Arch,
    Alpine,
    WindowsPe,
    Unknown,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IntrospectionReport {
    pub family: DistroFamily,
    pub volume_label: Option<String>,
    /// Kernel path inside the ISO (e.g. `/casper/vmlinuz`, `/isolinux/vmlinuz`).
    pub kernel_path: Option<String>,
    /// Initrd path(s) inside the ISO. May be multiple for multi-initrd setups.
    pub initrd_paths: Vec<String>,
    /// True if `sources/boot.wim` present — Windows install media.
    pub has_boot_wim: bool,
}

/// Probe an ISO file on disk. Never fails — on unrecoverable IO error we log
/// and return an `Unknown` family so the uploader still sees a record.
pub fn introspect(path: &Path) -> IntrospectionReport {
    let mut report = IntrospectionReport {
        family: DistroFamily::Unknown,
        volume_label: None,
        kernel_path: None,
        initrd_paths: Vec::new(),
        has_boot_wim: false,
    };

    let Ok(mut f) = std::fs::File::open(path) else {
        tracing::warn!(target: "openpxe::iso", "cannot open ISO for introspection: {}", path.display());
        return report;
    };

    // ISO9660 Primary Volume Descriptor at LBA 16 (offset 0x8000), 2048 bytes.
    // Bytes 40..72 are the Volume Identifier (space-padded, d-characters).
    let mut pvd = [0u8; 2048];
    if f.seek(SeekFrom::Start(0x8000)).is_ok() && f.read_exact(&mut pvd).is_ok() {
        // Byte 0 must be 0x01 (primary descriptor), bytes 1..6 = "CD001".
        if pvd[0] == 0x01 && &pvd[1..6] == b"CD001" {
            let label_raw = &pvd[40..72];
            let label = String::from_utf8_lossy(label_raw).trim().to_string();
            if !label.is_empty() {
                report.volume_label = Some(label.clone());
                report.family = family_from_label(&label);
            }
        }
    }

    // Cheap content scan: read the first ~64 MiB, look for signature filenames.
    // This is enough to identify `sources/boot.wim` (Windows) and common
    // kernel/initrd paths for the major Linux distros.
    let _ = f.seek(SeekFrom::Start(0));
    let scan_bytes = 64 * 1024 * 1024;
    let mut buf = vec![0u8; 1024 * 1024];
    let mut read_total = 0usize;
    let mut haystack = Vec::with_capacity(scan_bytes.min(32 * 1024 * 1024));
    while read_total < scan_bytes {
        let n = f.read(&mut buf).unwrap_or(0);
        if n == 0 {
            break;
        }
        haystack.extend_from_slice(&buf[..n]);
        read_total += n;
    }

    // `sources/boot.wim` is the definitive Windows-install-media marker
    // when the ISO exposes ASCII (ISO9660/Joliet) names. `contains_ascii`
    // is case-insensitive, so one form covers BOOT.WIM / boot.wim and the
    // backslash variant.
    if contains_ascii(&haystack, b"sources/boot.wim")
        || contains_ascii(&haystack, b"sources\\boot.wim")
    {
        report.has_boot_wim = true;
        report.family = DistroFamily::WindowsPe;
    }

    // v0.5.8: broaden Windows detection. Modern Windows 10/11 ISOs are
    // UDF — filenames are stored as UTF-16 (so the ASCII scan above misses
    // them) and the volume label is a cryptic Microsoft string (so
    // `family_from_label` misses it too). Booting is via HTTP sanboot of
    // the raw ISO (no boot.wim extraction), so we only need the *family*.
    // Catch the common cases: well-known Windows markers in either ASCII
    // or UTF-16LE within the first 16 MiB, plus a filename hint.
    if report.family == DistroFamily::Unknown {
        let head = &haystack[..haystack.len().min(16 * 1024 * 1024)];
        let ascii_markers: [&[u8]; 4] = [
            b"bootmgr",
            b"sources/install.wim",
            b"sources/install.esd",
            b"efi/microsoft",
        ];
        let utf16_markers = ["bootmgr", "boot.wim", "install.wim", "microsoft"];
        let looks_windows = ascii_markers.iter().any(|m| contains_ascii(head, m))
            || utf16_markers.iter().any(|m| contains_utf16le_ci(head, m))
            || filename_looks_windows(path);
        if looks_windows {
            report.family = DistroFamily::WindowsPe;
        }
    }

    // Best-effort kernel/initrd path guess from family. These paths are what
    // distro ISOs conventionally ship at — we don't verify extraction here;
    // that happens in the store after introspection.
    let (k, i) = guess_kernel_initrd(report.family);
    report.kernel_path = k.map(str::to_string);
    report.initrd_paths = i.iter().map(std::string::ToString::to_string).collect();

    report
}

fn family_from_label(label: &str) -> DistroFamily {
    let l = label.to_ascii_lowercase();
    if l.contains("ubuntu") || l.contains("debian") || l.contains("mint") {
        DistroFamily::DebianUbuntu
    } else if l.contains("rhel")
        || l.contains("centos")
        || l.contains("fedora")
        || l.contains("rocky")
        || l.contains("alma")
    {
        DistroFamily::RhelFedora
    } else if l.contains("suse") || l.contains("opensuse") {
        DistroFamily::OpenSuse
    } else if l.contains("arch") {
        DistroFamily::Arch
    } else if l.contains("alpine") {
        DistroFamily::Alpine
    } else if l.contains("windows") || l.contains("winpe") {
        DistroFamily::WindowsPe
    } else {
        DistroFamily::Unknown
    }
}

fn guess_kernel_initrd(family: DistroFamily) -> (Option<&'static str>, Vec<&'static str>) {
    match family {
        DistroFamily::DebianUbuntu => (Some("/casper/vmlinuz"), vec!["/casper/initrd"]),
        DistroFamily::RhelFedora => (
            Some("/images/pxeboot/vmlinuz"),
            vec!["/images/pxeboot/initrd.img"],
        ),
        DistroFamily::OpenSuse => (
            Some("/boot/x86_64/loader/linux"),
            vec!["/boot/x86_64/loader/initrd"],
        ),
        DistroFamily::Arch => (
            Some("/arch/boot/x86_64/vmlinuz-linux"),
            vec!["/arch/boot/x86_64/initramfs-linux.img"],
        ),
        DistroFamily::Alpine => (Some("/boot/vmlinuz-lts"), vec!["/boot/initramfs-lts"]),
        DistroFamily::WindowsPe | DistroFamily::Unknown => (None, Vec::new()),
    }
}

fn contains_ascii(haystack: &[u8], needle: &[u8]) -> bool {
    if needle.is_empty() || haystack.len() < needle.len() {
        return false;
    }
    haystack
        .windows(needle.len())
        .any(|w| w.eq_ignore_ascii_case(needle))
}

/// Case-insensitive search for an ASCII string encoded as UTF-16LE — the
/// way UDF (and thus modern Windows ISOs) store filenames. Each character
/// is two bytes: the ASCII low byte (compared case-insensitively) followed
/// by a 0 high byte. v0.5.8.
fn contains_utf16le_ci(haystack: &[u8], ascii: &str) -> bool {
    let n = ascii.len();
    if n == 0 || haystack.len() < n * 2 {
        return false;
    }
    let lower: Vec<u8> = ascii.bytes().map(|b| b.to_ascii_lowercase()).collect();
    haystack.windows(n * 2).any(|w| {
        lower
            .iter()
            .enumerate()
            .all(|(i, &c)| w[i * 2 + 1] == 0 && w[i * 2].to_ascii_lowercase() == c)
    })
}

/// Filename heuristic: a stock Windows ISO almost always carries an obvious
/// token in its name (e.g. `..._windows_11_...`, `Win10`, `winserver`).
/// Used only as a last-resort family hint when the content scan and volume
/// label are inconclusive. v0.5.8.
fn filename_looks_windows(path: &Path) -> bool {
    let name = path
        .file_name()
        .and_then(|s| s.to_str())
        .unwrap_or("")
        .to_ascii_lowercase();
    const TOKENS: [&str; 6] = [
        "windows",
        "winpe",
        "win10",
        "win11",
        "winserver",
        "win-server",
    ];
    TOKENS.iter().any(|t| name.contains(t))
}

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

    #[test]
    fn label_matching() {
        assert_eq!(
            family_from_label("Ubuntu 24.04"),
            DistroFamily::DebianUbuntu
        );
        assert_eq!(
            family_from_label("Rocky-9-x86_64-dvd"),
            DistroFamily::RhelFedora
        );
        assert_eq!(
            family_from_label("openSUSE-Leap-15.6"),
            DistroFamily::OpenSuse
        );
        assert_eq!(family_from_label("ARCH_202604"), DistroFamily::Arch);
        assert_eq!(family_from_label("weird-custom"), DistroFamily::Unknown);
    }

    #[test]
    fn utf16le_marker_matches_case_insensitively() {
        // "boot.wim" encoded UTF-16LE, mixed case — UDF stores Windows
        // filenames this way, which the ASCII scan can't see.
        let s = "BOOT.WIM";
        let utf16: Vec<u8> = s.bytes().flat_map(|b| [b, 0]).collect();
        let mut hay = vec![0u8; 8];
        hay.extend_from_slice(&utf16);
        hay.extend_from_slice(&[1, 2, 3]);
        assert!(contains_utf16le_ci(&hay, "boot.wim"));
        assert!(contains_utf16le_ci(&hay, "Boot.Wim"));
        assert!(!contains_utf16le_ci(&hay, "install.wim"));
        // An ASCII (not UTF-16) occurrence must NOT match the UTF-16 scan.
        assert!(!contains_utf16le_ci(b"boot.wim plain ascii", "boot.wim"));
    }

    #[test]
    fn filename_hint_catches_windows_isos() {
        use std::path::Path;
        assert!(filename_looks_windows(Path::new(
            "en-us_windows_11_iot_enterprise_ltsc_2024_x64_dvd.iso"
        )));
        assert!(filename_looks_windows(Path::new(
            "Win10_22H2_English_x64.iso"
        )));
        assert!(filename_looks_windows(Path::new("winserver2022.iso")));
        assert!(!filename_looks_windows(Path::new(
            "ubuntu-24.04-desktop.iso"
        )));
        assert!(!filename_looks_windows(Path::new(
            "Rocky-9.4-x86_64-dvd.iso"
        )));
    }
}