diskotech 0.0.5

//! Utilities for interacting with devices and paths on unixy systems.

use std::{
    collections::BTreeMap,
    fs, io,
    path::{Path, PathBuf},
};
use tracing::{debug, error, info};

pub type Map = BTreeMap<PathBuf, Record>;

#[derive(Clone, Debug, Default, Eq, Ord, PartialEq, PartialOrd)]
pub struct Record {
    pub name: Option<PathBuf>,
    pub diskseq: Option<PathBuf>,
    pub ids: Vec<PathBuf>,
    pub label: Option<PathBuf>,
    pub partlabel: Option<PathBuf>,
    pub partuuid: Option<PathBuf>,
    pub path: Option<PathBuf>,
    pub uuid: Option<PathBuf>,
}

/// Collect and correlate all the disk information.
///
/// # Errors
///
/// Returns an error during any file i/o errors.
pub fn collect() -> io::Result<Map> {
    // This will get populated with all the devices we find.
    let mut out = BTreeMap::new();

    // Go through and populate the map with records.
    by_id(&mut out)?;
    by_diskseq(&mut out)?;
    by_label(&mut out)?;
    by_partlabel(&mut out)?;
    by_partuuid(&mut out)?;
    by_path(&mut out)?;
    by_uuid(&mut out)?;

    Ok(out)
}

/// Get map of disk uuid -> logical name. Equivalent to an `ls -l
/// /dev/disk/by-uuid` command.
///
/// # Errors
///
/// Returns an error if any of the file I/O fails.
fn by_uuid(data: &mut Map) -> io::Result<()> {
    for (uuid, name) in lookup("/dev/disk/by-uuid")? {
        if let Some(rec) = data.get_mut(&name) {
            rec.uuid = Some(uuid);

            continue;
        }

        if let Some(rec) = data.insert(
            name.clone(),
            Record {
                name: Some(name),
                uuid: Some(uuid),
                ..Default::default()
            },
        ) {
            error!(?rec, "double insert");
        }
    }

    Ok(())
}

/// Get map of disk path -> logical name. Equivalent to an `ls -l
/// /dev/disk/by-path` command.
///
/// # Errors
///
/// Returns an error if any of the file I/O fails.
fn by_path(data: &mut Map) -> io::Result<()> {
    for (path, name) in lookup("/dev/disk/by-path")? {
        if let Some(rec) = data.get_mut(&name) {
            rec.path = Some(path);

            continue;
        }

        if let Some(rec) = data.insert(
            name.clone(),
            Record {
                name: Some(name),
                path: Some(path),
                ..Default::default()
            },
        ) {
            error!(?rec, "double insert");
        }
    }

    Ok(())
}

/// Get map of partition uuid -> logical name. Equivalent to an `ls -l
/// /dev/disk/by-partuuid` command.
///
/// # Errors
///
/// Returns an error if any of the file I/O fails.
fn by_partuuid(data: &mut Map) -> io::Result<()> {
    for (uuid, name) in lookup("/dev/disk/by-partuuid")? {
        if let Some(rec) = data.get_mut(&name) {
            rec.partuuid = Some(uuid);

            continue;
        }

        if let Some(rec) = data.insert(
            name.clone(),
            Record {
                name: Some(name),
                partuuid: Some(uuid),
                ..Default::default()
            },
        ) {
            error!(?rec, "double insert");
        }
    }

    Ok(())
}

/// Get map of partition label -> logical name. Equivalent to an `ls -l
/// /dev/disk/by-partlabel` command.
///
/// # Errors
///
/// Returns an error if any of the file I/O fails.
fn by_partlabel(data: &mut Map) -> io::Result<()> {
    for (label, name) in lookup("/dev/disk/by-partlabel")? {
        if let Some(rec) = data.get_mut(&name) {
            rec.partlabel = Some(label);

            continue;
        }

        if let Some(rec) = data.insert(
            name.clone(),
            Record {
                name: Some(name),
                partlabel: Some(label),
                ..Default::default()
            },
        ) {
            error!(?rec, "double insert");
        }
    }

    Ok(())
}

/// Get map of disk/part id -> logical name. Equivalent to an `ls -l
/// /dev/disk/by-id` command.
///
/// # Errors
///
/// Returns an error if any of the file I/O fails.
pub fn by_id(data: &mut Map) -> io::Result<()> {
    for (id, name) in lookup("/dev/disk/by-id")? {
        if let Some(rec) = data.get_mut(&name) {
            rec.ids.push(id);

            continue;
        }

        if let Some(rec) = data.insert(
            name.clone(),
            Record {
                name: Some(name),
                ids: vec![id],
                ..Default::default()
            },
        ) {
            error!(?rec, "double insert");
        }
    }

    Ok(())
}

/// Get map of disk sequence (1, 2, etc.) to logical name (/dev/sda, etc.).
/// Equivalent to running `ls -l /dev/disk/by-diskseq`.
///
/// # Errors
///
/// Returns an error if any of the file I/O fails.
pub fn by_diskseq(data: &mut Map) -> io::Result<()> {
    for (seq, name) in lookup("/dev/disk/by-diskseq")? {
        if let Some(rec) = data.get_mut(&name) {
            rec.diskseq = Some(seq);

            continue;
        }

        if let Some(rec) = data.insert(
            name.clone(),
            Record {
                name: Some(name),
                diskseq: Some(seq),
                ..Default::default()
            },
        ) {
            error!(?rec, "double insert");
        }
    }

    Ok(())
}

/// Get map of disk labels (mydisk, etc.) to logical name (/dev/sda, etc.).
/// Equivalent to running `ls -l /dev/disk/by-label`.
///
/// # Errors
///
/// Returns an error if any of the file I/O fails.
pub fn by_label(data: &mut Map) -> io::Result<()> {
    for (label, name) in lookup("/dev/disk/by-label")? {
        if let Some(rec) = data.get_mut(&name) {
            rec.label = Some(label);

            continue;
        }

        if let Some(rec) = data.insert(
            name.clone(),
            Record {
                name: Some(name),
                label: Some(label),
                ..Default::default()
            },
        ) {
            error!(?rec, "double insert");
        }
    }

    Ok(())
}

fn lookup<T: AsRef<Path>>(root: T) -> io::Result<Vec<(PathBuf, PathBuf)>> {
    let mut out = vec![];
    let root = root.as_ref();

    debug!(?root, "reading directory");
    for entry in fs::read_dir(root)? {
        // Make sure we can read the entry.
        let entry = entry?;
        debug!(?entry);

        // Get the path.
        let path = entry.path();
        debug!(?path);

        // If this isn't a symlink, not sure what to do with it. Not sure at
        // this point how prevalent this is, so I won't make it a WARN.
        if !entry.file_type()?.is_symlink() {
            info!(
                ?path,
                "found entry that wasn't a symlink; is this expected?"
            );
            continue;
        }

        // Resolve the symlink.
        let link = fs::read_link(&path)?;
        debug!(?link, "resolved link");
        let link = root.join(link).canonicalize()?;
        debug!(?link, "canonicalized link");

        // Insert the pair into the map.
        debug!(src = ?path, dest = ?link, "inserting src/dest pair into map");
        out.push((path, link));
    }

    Ok(out)
}