flexor/src/lsblk.rs

// Copyright 2023 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::process;

use anyhow::bail;
use anyhow::Result;
use log::info;
use serde::Deserialize;
use std::process::Command;

/// Struct for deserializing the JSON output of `lsblk`.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq)]
pub struct LsBlkDevice {
    /// Device name.
    ///
    /// This is a full path because lsblk is run with "--paths".
    pub name: String,

    /// Device type.
    #[serde(rename = "type")]
    pub device_type: String,
}

#[derive(Clone, Debug, Deserialize, PartialEq)]
struct LsBlkDeviceWithChildren {
    #[serde(flatten)]
    details: LsBlkDevice,

    /// Child devices.
    #[serde(default)]
    children: Vec<LsBlkDeviceWithChildren>,
}

#[derive(Debug, Deserialize, PartialEq)]
struct LsBlkOutput {
    #[serde(rename = "blockdevices")]
    block_devices: Vec<LsBlkDeviceWithChildren>,
}

impl LsBlkOutput {
    fn parse(input: &[u8]) -> Result<LsBlkOutput> {
        Ok(serde_json::from_slice(input)?)
    }

    fn flattened(self) -> Vec<LsBlkDevice> {
        let mut output = Vec::new();
        let mut stack = self.block_devices;
        while let Some(device) = stack.pop() {
            output.push(device.details);
            stack.extend(device.children);
        }
        output
    }
}

/// Capture information about block devices from lsblk.
///
/// lsblk is a convenient tool that already exists on CrOS base builds
/// and in most other linux distributions. Using the "--json" flag
/// makes the output easily parsible.
///
/// target: Block device to show information about. It will limit
/// lsblk to only return information about partitions on the target
/// device. If target is None lsblk will return information about most
/// block devices, excluding the zram device and slow devices such as
/// floppy drives.
///
/// Returns the raw output of lsblk.
fn get_lsblk_output() -> Result<Vec<u8>> {
    let mut command = process::Command::new("/bin/lsblk");
    command.args([
        // Select the fields to output
        "--output",
        "NAME,TYPE",
        // Format output as JSON
        "--json",
        // Print full device paths
        "--paths",
        // Exclude some devices by major number. See
        // https://www.kernel.org/doc/Documentation/admin-guide/devices.txt
        // for a list of major numbers.
        //
        // - Exclude floppy drives (2), as they are slow.
        // - Exclude scsi cdrom drives (11), as they are slow.
        // - Exclude zram (253), not a valid install target.
        "--exclude",
        "2,11,253",
    ]);

    get_command_output(command)
}

/// Capture information about block devices from lsblk.
///
/// Returns a flattened vector of devices.
pub fn get_lsblk_devices() -> Result<Vec<LsBlkDevice>> {
    let output = get_lsblk_output()?;
    let parsed = LsBlkOutput::parse(&output)?;
    Ok(parsed.flattened())
}

/// Run a command and get its stdout as raw bytes. An error is
/// returned if the process fails to launch, or if it exits non-zero.
fn get_command_output(mut command: Command) -> Result<Vec<u8>> {
    info!("running command: {:?}", command);
    let output = match command.output() {
        Ok(output) => output,
        Err(err) => {
            bail!("Failed to execute command: {err}");
        }
    };

    if !output.status.success() {
        bail!("Failed to execute command");
    }
    Ok(output.stdout)
}

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

    fn mkdev(name: &str, dtype: &str) -> LsBlkDevice {
        LsBlkDevice {
            name: name.into(),
            device_type: dtype.into(),
        }
    }

    #[test]
    fn test_lsblk_deserialization() {
        // This test input was generated by running this command in a VM:
        //
        //     lsblk --bytes --output NAME,TYPE \
        //         --json --paths --exclude 2,11,253
        let input = include_bytes!("test_lsblk_output.json");

        #[rustfmt::skip]
        let expected = vec![
            mkdev("/dev/sda", "disk"),
            mkdev("/dev/sda12", "part"),
            mkdev("/dev/sda11", "part"),
            mkdev("/dev/sda10", "part"),
            mkdev("/dev/sda9", "part"),
            mkdev("/dev/sda8", "part"),
            mkdev("/dev/sda7", "part"),
            mkdev("/dev/sda6", "part"),
            mkdev("/dev/sda5", "part"),
            mkdev("/dev/sda4", "part"),
            mkdev("/dev/sda3", "part"),
            mkdev("/dev/sda2", "part"),
            mkdev("/dev/sda1", "part"),
            mkdev("/dev/loop4", "loop"),
            mkdev("/dev/loop3", "loop"),
            mkdev("/dev/loop2", "loop"),
            mkdev("/dev/loop1", "loop"),
            mkdev("/dev/mapper/encstateful", "dm"),
            mkdev("/dev/loop0", "loop"),
        ];

        let output = LsBlkOutput::parse(input).unwrap();
        assert_eq!(output.flattened(), expected);
    }
}