blockdev.rs

use std::collections::HashSet;
use std::env;
use std::path::Path;
use std::process::{Command, Stdio};

use anyhow::{Context, Result, anyhow};
use camino::{Utf8Path, Utf8PathBuf};
use cap_std_ext::cap_std::fs::Dir;
use fn_error_context::context;
use serde::Deserialize;

use bootc_utils::CommandRunExt;

/// MBR partition type IDs that indicate an EFI System Partition.
/// 0x06 is FAT16 (used as ESP on some MBR systems), 0xEF is the
/// explicit EFI System Partition type.
/// Refer to <https://en.wikipedia.org/wiki/Partition_type>
pub const ESP_ID_MBR: &[u8] = &[0x06, 0xEF];

/// EFI System Partition (ESP) for UEFI boot on GPT
pub const ESP: &str = "c12a7328-f81f-11d2-ba4b-00a0c93ec93b";

/// BIOS boot partition type GUID for GPT
pub const BIOS_BOOT: &str = "21686148-6449-6e6f-744e-656564454649";

#[derive(Debug, Deserialize)]
struct DevicesOutput {
    blockdevices: Vec<Device>,
}

#[allow(dead_code)]
#[derive(Debug, Clone, Deserialize)]
pub struct Device {
    pub name: String,
    pub serial: Option<String>,
    pub model: Option<String>,
    pub partlabel: Option<String>,
    pub parttype: Option<String>,
    pub partuuid: Option<String>,
    /// Partition number (1-indexed). None for whole disk devices.
    pub partn: Option<u32>,
    pub children: Option<Vec<Device>>,
    pub size: u64,
    #[serde(rename = "maj:min")]
    pub maj_min: Option<String>,
    // NOTE this one is not available on older util-linux, and
    // will also not exist for whole blockdevs (as opposed to partitions).
    pub start: Option<u64>,

    // Filesystem-related properties
    pub label: Option<String>,
    pub fstype: Option<String>,
    pub uuid: Option<String>,
    pub path: Option<String>,
    /// Partition table type (e.g., "gpt", "dos"). Only present on whole disk devices.
    pub pttype: Option<String>,
}

impl Device {
    // RHEL8's lsblk doesn't have PATH, so we do it
    pub fn path(&self) -> String {
        self.path.clone().unwrap_or(format!("/dev/{}", &self.name))
    }

    /// Alias for path() for compatibility
    #[allow(dead_code)]
    pub fn node(&self) -> String {
        self.path()
    }

    #[allow(dead_code)]
    pub fn has_children(&self) -> bool {
        self.children.as_ref().is_some_and(|v| !v.is_empty())
    }

    // Check if the device is mpath
    pub fn is_mpath(&self) -> Result<bool> {
        let dm_path = Utf8PathBuf::from_path_buf(std::fs::canonicalize(self.path())?)
            .map_err(|_| anyhow::anyhow!("Non-UTF8 path"))?;
        let dm_name = dm_path.file_name().unwrap_or("");
        let uuid_path = Utf8PathBuf::from(format!("/sys/class/block/{dm_name}/dm/uuid"));

        if uuid_path.exists() {
            let uuid = std::fs::read_to_string(&uuid_path)
                .with_context(|| format!("Failed to read {uuid_path}"))?;
            if uuid.trim_start().starts_with("mpath-") {
                return Ok(true);
            }
        }
        Ok(false)
    }

    /// Get the numeric partition index of the ESP (e.g. "1", "2").
    ///
    /// We read `/sys/class/block/<name>/partition` rather than parsing device
    /// names because naming conventions vary across disk types (sd, nvme, dm, etc.).
    /// On multipath devices the sysfs `partition` attribute doesn't exist, so we
    /// fall back to the `partn` field reported by lsblk.
    pub fn get_esp_partition_number(&self) -> Result<String> {
        let esp_device = self.find_partition_of_esp()?;
        let devname = &esp_device.name;

        let partition_path = Utf8PathBuf::from(format!("/sys/class/block/{devname}/partition"));
        if partition_path.exists() {
            return std::fs::read_to_string(&partition_path)
                .with_context(|| format!("Failed to read {partition_path}"));
        }

        // On multipath the partition attribute is not existing
        if self.is_mpath()? {
            if let Some(partn) = esp_device.partn {
                return Ok(partn.to_string());
            }
        }
        anyhow::bail!("Not supported for {devname}")
    }

    /// Find BIOS boot partition among children.
    pub fn find_partition_of_bios_boot(&self) -> Option<&Device> {
        self.find_partition_of_type(BIOS_BOOT)
    }

    /// Find all ESP partitions across all root devices backing this device.
    /// Calls find_all_roots() to discover physical disks, then searches each for an ESP.
    /// Returns None if no ESPs are found.
    pub fn find_colocated_esps(&self) -> Result<Option<Vec<Device>>> {
        let mut esps = Vec::new();
        for root in &self.find_all_roots()? {
            if let Some(esp) = root.find_partition_of_esp_optional()? {
                esps.push(esp.clone());
            }
        }
        Ok((!esps.is_empty()).then_some(esps))
    }

    /// Find a single ESP partition among all root devices backing this device.
    ///
    /// Walks the parent chain to find all backing disks, then looks for ESP
    /// partitions on each. Returns the first ESP found. This is the common
    /// case for composefs/UKI boot paths where exactly one ESP is expected.
    pub fn find_first_colocated_esp(&self) -> Result<Device> {
        self.find_colocated_esps()?
            .and_then(|mut v| Some(v.remove(0)))
            .ok_or_else(|| anyhow!("No ESP partition found among backing devices"))
    }

    /// Find all BIOS boot partitions across all root devices backing this device.
    /// Calls find_all_roots() to discover physical disks, then searches each for a BIOS boot partition.
    /// Returns None if no BIOS boot partitions are found.
    pub fn find_colocated_bios_boot(&self) -> Result<Option<Vec<Device>>> {
        let bios_boots: Vec<_> = self
            .find_all_roots()?
            .iter()
            .filter_map(|root| root.find_partition_of_bios_boot())
            .cloned()
            .collect();
        Ok((!bios_boots.is_empty()).then_some(bios_boots))
    }

    /// Find a child partition by partition type (case-insensitive).
    pub fn find_partition_of_type(&self, parttype: &str) -> Option<&Device> {
        self.children.as_ref()?.iter().find(|child| {
            child
                .parttype
                .as_ref()
                .is_some_and(|pt| pt.eq_ignore_ascii_case(parttype))
        })
    }

    /// Find the EFI System Partition (ESP) among children.
    ///
    /// For GPT disks, this matches by the ESP partition type GUID.
    /// For MBR (dos) disks, this matches by the MBR partition type IDs (0x06 or 0xEF).
    ///
    /// If no ESP is found among direct children, this recurses into children
    /// that have their own partition table (e.g. firmware RAID arrays where the
    /// hierarchy is disk → md array → partitions).
    ///
    /// Returns `Ok(None)` when there are no children or no ESP partition
    /// is present. Returns `Err` only for genuinely unexpected conditions
    /// (e.g. an unsupported partition table type).
    pub fn find_partition_of_esp_optional(&self) -> Result<Option<&Device>> {
        let Some(children) = self.children.as_ref() else {
            return Ok(None);
        };
        let direct = match self.pttype.as_deref() {
            Some("dos") => children.iter().find(|child| {
                child
                    .parttype
                    .as_ref()
                    .and_then(|pt| {
                        let pt = pt.strip_prefix("0x").unwrap_or(pt);
                        u8::from_str_radix(pt, 16).ok()
                    })
                    .is_some_and(|pt| ESP_ID_MBR.contains(&pt))
            }),
            // When pttype is None (e.g. older lsblk or partition devices), default
            // to GPT UUID matching which will simply not match MBR hex types.
            Some("gpt") | None => self.find_partition_of_type(ESP),
            Some(other) => return Err(anyhow!("Unsupported partition table type: {other}")),
        };
        if direct.is_some() {
            return Ok(direct);
        }
        // Recurse into children that carry their own partition table, such as
        // firmware RAID arrays (disk → md array → partitions).
        for child in children {
            if child.pttype.is_some() {
                if let Some(esp) = child.find_partition_of_esp_optional()? {
                    return Ok(Some(esp));
                }
            }
        }
        Ok(None)
    }

    /// Find the EFI System Partition (ESP) among children, or error if absent.
    ///
    /// This is a convenience wrapper around [`Self::find_partition_of_esp_optional`]
    /// for callers that require an ESP to be present.
    pub fn find_partition_of_esp(&self) -> Result<&Device> {
        self.find_partition_of_esp_optional()?
            .ok_or_else(|| anyhow!("ESP partition not found on {}", self.path()))
    }

    /// Find a child partition by partition number (1-indexed).
    pub fn find_device_by_partno(&self, partno: u32) -> Result<&Device> {
        self.children
            .as_ref()
            .ok_or_else(|| anyhow!("Device has no children"))?
            .iter()
            .find(|child| child.partn == Some(partno))
            .ok_or_else(|| anyhow!("Missing partition for index {partno}"))
    }

    /// Re-query this device's information from lsblk, updating all fields.
    /// This is useful after partitioning when the device's children have changed.
    pub fn refresh(&mut self) -> Result<()> {
        let path = self.path();
        let new_device = list_dev(Utf8Path::new(&path))?;
        *self = new_device;
        Ok(())
    }

    /// Read a sysfs property for this device and parse it as the target type.
    fn read_sysfs_property<T>(&self, property: &str) -> Result<Option<T>>
    where
        T: std::str::FromStr,
        T::Err: std::error::Error + Send + Sync + 'static,
    {
        let Some(majmin) = self.maj_min.as_deref() else {
            return Ok(None);
        };
        let sysfs_path = format!("/sys/dev/block/{majmin}/{property}");
        if !Utf8Path::new(&sysfs_path).try_exists()? {
            return Ok(None);
        }
        let value = std::fs::read_to_string(&sysfs_path)
            .with_context(|| format!("Reading {sysfs_path}"))?;
        let parsed = value
            .trim()
            .parse()
            .with_context(|| format!("Parsing sysfs {property} property"))?;
        tracing::debug!("backfilled {property} to {value}");
        Ok(Some(parsed))
    }

    /// Older versions of util-linux may be missing some properties. Backfill them if they're missing.
    pub fn backfill_missing(&mut self) -> Result<()> {
        // The "start" parameter was only added in a version of util-linux that's only
        // in Fedora 40 as of this writing.
        if self.start.is_none() {
            self.start = self.read_sysfs_property("start")?;
        }
        // The "partn" column was added in util-linux 2.39, which is newer than
        // what CentOS 9 / RHEL 9 ship (2.37). Note: sysfs uses "partition" not "partn".
        if self.partn.is_none() {
            self.partn = self.read_sysfs_property("partition")?;
        }
        // Recurse to child devices
        for child in self.children.iter_mut().flatten() {
            child.backfill_missing()?;
        }
        Ok(())
    }

    /// Query parent devices via `lsblk --inverse`.
    ///
    /// Returns `Ok(None)` if this device is already a root device (no parents).
    /// In the returned `Vec<Device>`, each device's `children` field contains
    /// *its own* parents (grandparents, etc.), forming the full chain to the
    /// root device(s). A device can have multiple parents (e.g. RAID, LVM).
    pub fn list_parents(&self) -> Result<Option<Vec<Device>>> {
        let path = self.path();
        let output: DevicesOutput = Command::new("lsblk")
            .args(["-J", "-b", "-O", "--inverse"])
            .arg(&path)
            .log_debug()
            .run_and_parse_json()?;

        let device = output
            .blockdevices
            .into_iter()
            .next()
            .ok_or_else(|| anyhow!("no device output from lsblk --inverse for {path}"))?;

        match device.children {
            Some(mut children) if !children.is_empty() => {
                for child in &mut children {
                    child.backfill_missing()?;
                }
                Ok(Some(children))
            }
            _ => Ok(None),
        }
    }

    /// Walk the parent chain to find all root (whole disk) devices,
    /// and fail if more than one root is found.
    ///
    /// This is a convenience wrapper around `find_all_roots` for callers
    /// that expect exactly one backing device (e.g. non-RAID setups).
    pub fn require_single_root(&self) -> Result<Device> {
        let mut roots = self.find_all_roots()?;
        match roots.len() {
            1 => Ok(roots.remove(0)),
            n => anyhow::bail!(
                "Expected a single root device for {}, but found {n}",
                self.path()
            ),
        }
    }

    /// Walk the parent chain to find all root (whole disk) devices.
    ///
    /// Returns all root devices with their children (partitions) populated.
    /// This handles devices backed by multiple parents (e.g. RAID arrays)
    /// by following all branches of the parent tree.
    /// If this device is already a root device, returns a single-element list.
    pub fn find_all_roots(&self) -> Result<Vec<Device>> {
        let Some(parents) = self.list_parents()? else {
            // Already a root device; re-query to ensure children are populated
            return Ok(vec![list_dev(Utf8Path::new(&self.path()))?]);
        };

        let mut roots = Vec::new();
        let mut seen = HashSet::new();
        let mut queue = parents;
        while let Some(mut device) = queue.pop() {
            match device.children.take() {
                Some(grandparents) if !grandparents.is_empty() => {
                    queue.extend(grandparents);
                }
                _ => {
                    // Deduplicate: in complex topologies (e.g. multipath)
                    // multiple branches can converge on the same physical disk.
                    let name = device.name.clone();
                    if seen.insert(name) {
                        // Found a new root; re-query to populate its actual children
                        roots.push(list_dev(Utf8Path::new(&device.path()))?);
                    }
                }
            }
        }
        Ok(roots)
    }
}

#[context("Listing device {dev}")]
pub fn list_dev(dev: &Utf8Path) -> Result<Device> {
    let mut devs: DevicesOutput = Command::new("lsblk")
        .args(["-J", "-b", "-O"])
        .arg(dev)
        .log_debug()
        .run_and_parse_json()?;
    for dev in devs.blockdevices.iter_mut() {
        dev.backfill_missing()?;
    }
    devs.blockdevices
        .into_iter()
        .next()
        .ok_or_else(|| anyhow!("no device output from lsblk for {dev}"))
}

#[context("Finding block device for ZFS dataset {dataset}")]
fn list_dev_for_zfs_dataset(dataset: &str) -> Result<Device> {
    let dataset = dataset.strip_prefix("ZFS=").unwrap_or(dataset);
    let pool = dataset
        .split('/')
        .next()
        .ok_or_else(|| anyhow!("Invalid ZFS dataset: {dataset}"))?;

    let output = Command::new("zpool")
        .args(["list", "-H", "-v", "-P", pool])
        .run_get_string()
        .with_context(|| format!("Querying ZFS pool {pool}"))?;

    for line in output.lines() {
        if line.starts_with('\t') || line.starts_with(' ') {
            let dev_path = line.trim_start().split('\t').next().unwrap_or("").trim();
            if dev_path.starts_with('/') {
                return list_dev(Utf8Path::new(dev_path));
            }
        }
    }

    anyhow::bail!("Could not find a block device backing ZFS pool {pool}")
}

/// List the device containing the filesystem mounted at the given directory.
pub fn list_dev_by_dir(dir: &Dir) -> Result<Device> {
    let fsinfo = bootc_mount::inspect_filesystem_of_dir(dir)?;
    let source = &fsinfo.source;
    if fsinfo.fstype == "zfs" || (!source.starts_with('/') && source.contains('/')) {
        return list_dev_for_zfs_dataset(source);
    }
    list_dev(&Utf8PathBuf::from(source))
}

pub struct LoopbackDevice {
    pub dev: Option<Utf8PathBuf>,
    // Handle to the cleanup helper process
    cleanup_handle: Option<LoopbackCleanupHandle>,
}

/// Handle to manage the cleanup helper process for loopback devices
struct LoopbackCleanupHandle {
    /// Child process handle
    child: std::process::Child,
}

impl LoopbackDevice {
    // Create a new loopback block device targeting the provided file path.
    pub fn new(path: &Path) -> Result<Self> {
        let direct_io = match env::var("BOOTC_DIRECT_IO") {
            Ok(val) => {
                if val == "on" {
                    "on"
                } else {
                    "off"
                }
            }
            Err(_e) => "off",
        };

        let dev = Command::new("losetup")
            .args([
                "--show",
                format!("--direct-io={direct_io}").as_str(),
                "-P",
                "--find",
            ])
            .arg(path)
            .run_get_string()?;
        let dev = Utf8PathBuf::from(dev.trim());
        tracing::debug!("Allocated loopback {dev}");

        // Try to spawn cleanup helper, but don't fail if it doesn't work
        let cleanup_handle = match Self::spawn_cleanup_helper(dev.as_str()) {
            Ok(handle) => Some(handle),
            Err(e) => {
                tracing::warn!(
                    "Failed to spawn loopback cleanup helper for {}: {}. \
                     Loopback device may not be cleaned up if process is interrupted.",
                    dev,
                    e
                );
                None
            }
        };

        Ok(Self {
            dev: Some(dev),
            cleanup_handle,
        })
    }

    // Access the path to the loopback block device.
    pub fn path(&self) -> &Utf8Path {
        // SAFETY: The option cannot be destructured until we are dropped
        self.dev.as_deref().unwrap()
    }

    /// Spawn a cleanup helper process that will clean up the loopback device
    /// if the parent process dies unexpectedly
    fn spawn_cleanup_helper(device_path: &str) -> Result<LoopbackCleanupHandle> {
        // Try multiple strategies to find the bootc binary
        let bootc_path = bootc_utils::reexec::executable_path()
            .context("Failed to locate bootc binary for cleanup helper")?;

        // Create the helper process
        let mut cmd = Command::new(bootc_path);
        cmd.args([
            "internals",
            "loopback-cleanup-helper",
            "--device",
            device_path,
        ]);

        // Set environment variable to indicate this is a cleanup helper
        cmd.env("BOOTC_LOOPBACK_CLEANUP_HELPER", "1");

        // Set up stdio to redirect to /dev/null
        cmd.stdin(Stdio::null());
        cmd.stdout(Stdio::null());
        // Don't redirect stderr so we can see error messages

        // Spawn the process
        let child = cmd
            .spawn()
            .context("Failed to spawn loopback cleanup helper")?;

        Ok(LoopbackCleanupHandle { child })
    }

    // Shared backend for our `close` and `drop` implementations.
    fn impl_close(&mut self) -> Result<()> {
        // SAFETY: This is the only place we take the option
        let Some(dev) = self.dev.take() else {
            tracing::trace!("loopback device already deallocated");
            return Ok(());
        };

        // Kill the cleanup helper since we're cleaning up normally
        if let Some(mut cleanup_handle) = self.cleanup_handle.take() {
            // Send SIGTERM to the child process and let it do the cleanup
            let _ = cleanup_handle.child.kill();
        }

        Command::new("losetup")
            .args(["-d", dev.as_str()])
            .run_capture_stderr()
    }

    /// Consume this device, unmounting it.
    pub fn close(mut self) -> Result<()> {
        self.impl_close()
    }
}

impl Drop for LoopbackDevice {
    fn drop(&mut self) {
        // Best effort to unmount if we're dropped without invoking `close`
        let _ = self.impl_close();
    }
}

/// Main function for the loopback cleanup helper process
/// This function does not return - it either exits normally or via signal
pub async fn run_loopback_cleanup_helper(device_path: &str) -> Result<()> {
    // Check if we're running as a cleanup helper
    if std::env::var("BOOTC_LOOPBACK_CLEANUP_HELPER").is_err() {
        anyhow::bail!("This function should only be called as a cleanup helper");
    }

    // Set up death signal notification - we want to be notified when parent dies
    rustix::process::set_parent_process_death_signal(Some(rustix::process::Signal::TERM))
        .context("Failed to set parent death signal")?;

    // Wait for SIGTERM (either from parent death or normal cleanup)
    tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
        .expect("Failed to create signal stream")
        .recv()
        .await;

    // Clean up the loopback device
    let output = std::process::Command::new("losetup")
        .args(["-d", device_path])
        .output();

    match output {
        Ok(output) if output.status.success() => {
            // Log to systemd journal instead of stderr
            tracing::info!("Cleaned up leaked loopback device {}", device_path);
            std::process::exit(0);
        }
        Ok(output) => {
            let stderr = String::from_utf8_lossy(&output.stderr);
            tracing::error!(
                "Failed to clean up loopback device {}: {}. Stderr: {}",
                device_path,
                output.status,
                stderr.trim()
            );
            std::process::exit(1);
        }
        Err(e) => {
            tracing::error!(
                "Error executing losetup to clean up loopback device {}: {}",
                device_path,
                e
            );
            std::process::exit(1);
        }
    }
}

/// Parse a string into mibibytes
pub fn parse_size_mib(mut s: &str) -> Result<u64> {
    let suffixes = [
        ("MiB", 1u64),
        ("M", 1u64),
        ("GiB", 1024),
        ("G", 1024),
        ("TiB", 1024 * 1024),
        ("T", 1024 * 1024),
    ];
    let mut mul = 1u64;
    for (suffix, imul) in suffixes {
        if let Some((sv, rest)) = s.rsplit_once(suffix) {
            if !rest.is_empty() {
                anyhow::bail!("Trailing text after size: {rest}");
            }
            s = sv;
            mul = imul;
        }
    }
    let v = s.parse::<u64>()?;
    Ok(v * mul)
}

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

    #[test]
    fn test_parse_size_mib() {
        let ident_cases = [0, 10, 9, 1024].into_iter().map(|k| (k.to_string(), k));
        let cases = [
            ("0M", 0),
            ("10M", 10),
            ("10MiB", 10),
            ("1G", 1024),
            ("9G", 9216),
            ("11T", 11 * 1024 * 1024),
        ]
        .into_iter()
        .map(|(k, v)| (k.to_string(), v));
        for (s, v) in ident_cases.chain(cases) {
            assert_eq!(parse_size_mib(&s).unwrap(), v as u64, "Parsing {s}");
        }
    }

    #[test]
    fn test_parse_lsblk() {
        let fixture = include_str!("../tests/fixtures/lsblk.json");
        let devs: DevicesOutput = serde_json::from_str(fixture).unwrap();
        let dev = devs.blockdevices.into_iter().next().unwrap();
        // The parent device has no partition number
        assert_eq!(dev.partn, None);
        let children = dev.children.as_deref().unwrap();
        assert_eq!(children.len(), 3);
        let first_child = &children[0];
        assert_eq!(first_child.partn, Some(1));
        assert_eq!(
            first_child.parttype.as_deref().unwrap(),
            "21686148-6449-6e6f-744e-656564454649"
        );
        assert_eq!(
            first_child.partuuid.as_deref().unwrap(),
            "3979e399-262f-4666-aabc-7ab5d3add2f0"
        );
        // Verify find_device_by_partno works
        let part2 = dev.find_device_by_partno(2).unwrap();
        assert_eq!(part2.partn, Some(2));
        assert_eq!(part2.parttype.as_deref().unwrap(), ESP);
        // Verify find_partition_of_esp works
        let esp = dev.find_partition_of_esp().unwrap();
        assert_eq!(esp.partn, Some(2));
        // Verify find_partition_of_bios_boot works (vda1 is BIOS-BOOT)
        let bios = dev.find_partition_of_bios_boot().unwrap();
        assert_eq!(bios.partn, Some(1));
        assert_eq!(bios.parttype.as_deref().unwrap(), BIOS_BOOT);
    }

    #[test]
    fn test_parse_lsblk_mbr() {
        let fixture = include_str!("../tests/fixtures/lsblk-mbr.json");
        let devs: DevicesOutput = serde_json::from_str(fixture).unwrap();
        let dev = devs.blockdevices.into_iter().next().unwrap();
        // The parent device has no partition number and is MBR
        assert_eq!(dev.partn, None);
        assert_eq!(dev.pttype.as_deref().unwrap(), "dos");
        let children = dev.children.as_deref().unwrap();
        assert_eq!(children.len(), 3);
        // First partition: FAT16 boot partition (MBR type 0x06, an ESP type)
        let first_child = &children[0];
        assert_eq!(first_child.partn, Some(1));
        assert_eq!(first_child.parttype.as_deref().unwrap(), "0x06");
        assert_eq!(first_child.partuuid.as_deref().unwrap(), "a1b2c3d4-01");
        assert_eq!(first_child.fstype.as_deref().unwrap(), "vfat");
        // MBR partitions have no partlabel
        assert!(first_child.partlabel.is_none());
        // Second partition: Linux root (MBR type 0x83)
        let second_child = &children[1];
        assert_eq!(second_child.partn, Some(2));
        assert_eq!(second_child.parttype.as_deref().unwrap(), "0x83");
        assert_eq!(second_child.partuuid.as_deref().unwrap(), "a1b2c3d4-02");
        // Third partition: EFI System Partition (MBR type 0xef)
        let third_child = &children[2];
        assert_eq!(third_child.partn, Some(3));
        assert_eq!(third_child.parttype.as_deref().unwrap(), "0xef");
        assert_eq!(third_child.partuuid.as_deref().unwrap(), "a1b2c3d4-03");
        // Verify find_device_by_partno works on MBR
        let part1 = dev.find_device_by_partno(1).unwrap();
        assert_eq!(part1.partn, Some(1));
        // find_partition_of_esp returns the first matching ESP type (0x06 on partition 1)
        let esp = dev.find_partition_of_esp().unwrap();
        assert_eq!(esp.partn, Some(1));
    }

    /// Helper to construct a minimal MBR disk Device with given child partition types.
    fn make_mbr_disk(parttypes: &[&str]) -> Device {
        Device {
            name: "vda".into(),
            serial: None,
            model: None,
            partlabel: None,
            parttype: None,
            partuuid: None,
            partn: None,
            size: 10737418240,
            maj_min: None,
            start: None,
            label: None,
            fstype: None,
            uuid: None,
            path: Some("/dev/vda".into()),
            pttype: Some("dos".into()),
            children: Some(
                parttypes
                    .iter()
                    .enumerate()
                    .map(|(i, pt)| Device {
                        name: format!("vda{}", i + 1),
                        serial: None,
                        model: None,
                        partlabel: None,
                        parttype: Some(pt.to_string()),
                        partuuid: None,
                        partn: Some(i as u32 + 1),
                        size: 1048576,
                        maj_min: None,
                        start: Some(2048),
                        label: None,
                        fstype: None,
                        uuid: None,
                        path: None,
                        pttype: Some("dos".into()),
                        children: None,
                    })
                    .collect(),
            ),
        }
    }

    #[test]
    fn test_parse_lsblk_vroc() {
        let fixture = include_str!("../tests/fixtures/lsblk-vroc.json");
        let devs: DevicesOutput = serde_json::from_str(fixture).unwrap();
        assert_eq!(devs.blockdevices.len(), 2);

        // find_partition_of_esp recurses through the md126 RAID array to
        // locate the ESP (md126p1) even though it is not a direct child of
        // the NVMe disk.
        for nvme in &devs.blockdevices {
            let esp = nvme.find_partition_of_esp().unwrap();
            assert_eq!(esp.name, "md126p1");
            assert_eq!(esp.partn, Some(1));
            assert_eq!(esp.parttype.as_deref().unwrap(), ESP);
            assert_eq!(esp.fstype.as_deref().unwrap(), "vfat");
        }
    }

    #[test]
    fn test_parse_lsblk_swraid() {
        let fixture = include_str!("../tests/fixtures/lsblk-swraid.json");
        let devs: DevicesOutput = serde_json::from_str(fixture).unwrap();
        assert_eq!(devs.blockdevices.len(), 2);

        // In a software RAID (mdadm) setup each disk is individually
        // partitioned with its own GPT table and ESP.  The root partition
        // (sda3/sdb3) is a linux_raid_member assembled into md0.
        // find_partition_of_esp should locate the ESP as a direct child of
        // each disk — no recursion through an md array is needed here.
        let sda = &devs.blockdevices[0];
        let esp = sda.find_partition_of_esp().unwrap();
        assert_eq!(esp.name, "sda1");
        assert_eq!(esp.partn, Some(1));
        assert_eq!(esp.parttype.as_deref().unwrap(), ESP);
        assert_eq!(esp.fstype.as_deref().unwrap(), "vfat");

        let sdb = &devs.blockdevices[1];
        let esp = sdb.find_partition_of_esp().unwrap();
        assert_eq!(esp.name, "sdb1");
        assert_eq!(esp.partn, Some(1));
        assert_eq!(esp.parttype.as_deref().unwrap(), ESP);
        assert_eq!(esp.fstype.as_deref().unwrap(), "vfat");

        // Verify the md0 RAID array is visible as a child of the root
        // partition on each disk.
        let sda3 = sda
            .children
            .as_ref()
            .unwrap()
            .iter()
            .find(|c| c.name == "sda3")
            .unwrap();
        assert_eq!(sda3.fstype.as_deref().unwrap(), "linux_raid_member");
        let md0 = sda3
            .children
            .as_ref()
            .unwrap()
            .iter()
            .find(|c| c.name == "md0")
            .unwrap();
        assert_eq!(md0.fstype.as_deref().unwrap(), "ext4");
    }

    #[test]
    fn test_mbr_esp_detection() {
        // 0x06 (FAT16) is recognized as ESP
        let dev = make_mbr_disk(&["0x06"]);
        assert_eq!(dev.find_partition_of_esp().unwrap().partn, Some(1));

        // 0xef (EFI System Partition) is recognized as ESP
        let dev = make_mbr_disk(&["0x83", "0xef"]);
        assert_eq!(dev.find_partition_of_esp().unwrap().partn, Some(2));

        // No ESP types present: 0x83 (Linux) and 0x82 (swap)
        let dev = make_mbr_disk(&["0x83", "0x82"]);
        assert!(dev.find_partition_of_esp().is_err());
    }
}