| // Copyright 2021 The Chromium OS Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| use std::path::Path; |
| use std::str::FromStr; |
| use std::{fmt, process}; |
| |
| use anyhow::Error; |
| use serde::de::{self, Visitor}; |
| use serde::{Deserialize, Deserializer}; |
| |
| use crate::util::get_command_output; |
| |
| /// Struct for deserializing the JSON output of `lsblk`. |
| #[derive(Clone, Debug, Deserialize, PartialEq)] |
| pub struct LsBlkDevice { |
| /// Device name. |
| /// |
| /// This is a full path because lsblk is run with "--paths". |
| pub name: String, |
| |
| /// Internal kernel device name. |
| #[serde(rename = "kname")] |
| pub kernel_name: String, |
| |
| /// Whether the device is removable. This is not completely |
| /// reliable, for example USB SSDs may not show as removable. |
| #[serde(rename = "rm", deserialize_with = "deserialize_rm")] |
| pub is_removable: bool, |
| |
| /// Size in bytes. |
| #[serde(rename = "size", deserialize_with = "deserialize_size")] |
| pub size_in_bytes: u64, |
| |
| /// Device type. |
| #[serde(rename = "type")] |
| pub device_type: String, |
| } |
| |
| impl LsBlkDevice { |
| /// Get the partition number, e.g. for `/dev/sda12` this returns 12. |
| pub fn partition_number(&self) -> Option<u64> { |
| if self.device_type != "part" { |
| return None; |
| } |
| |
| // Find the index of the last non-numeric character. |
| let index = self.name.rfind(|c: char| !c.is_ascii_digit())?; |
| |
| // Parse the rest of the string past that index into a number. |
| let num_part = &self.name[index + 1..]; |
| num_part.parse().ok() |
| } |
| } |
| |
| /// Deserialize the "rm" field of an lsblk device. |
| /// |
| /// This is a boolean field, but the representation depends on the |
| /// version of lsblk. In 2.32 it uses "0" for false and "1" for |
| /// true, in 2.36.1 it uses proper true and false values. |
| fn deserialize_rm<'de, D>(deserializer: D) -> Result<bool, D::Error> |
| where |
| D: Deserializer<'de>, |
| { |
| struct StringOrBool; |
| |
| impl<'de> Visitor<'de> for StringOrBool { |
| type Value = bool; |
| |
| fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { |
| formatter.write_str("string or bool") |
| } |
| |
| fn visit_str<E>(self, value: &str) -> Result<bool, E> |
| where |
| E: de::Error, |
| { |
| match value { |
| "0" => Ok(false), |
| "1" => Ok(true), |
| _ => Err(de::Error::custom("expected either '0' or '1'")), |
| } |
| } |
| |
| fn visit_bool<E>(self, value: bool) -> Result<bool, E> |
| where |
| E: de::Error, |
| { |
| Ok(value) |
| } |
| } |
| |
| deserializer.deserialize_any(StringOrBool) |
| } |
| |
| /// Deserialize the "size" field of an lsblk device. |
| /// |
| /// This is a numeric field, but the representation depends on the |
| /// version of lsblk. In 2.32 it wraps the number in a string, in |
| /// 2.36.1 it uses an actual number. |
| fn deserialize_size<'de, D>(deserializer: D) -> Result<u64, D::Error> |
| where |
| D: Deserializer<'de>, |
| { |
| struct StringOrU64; |
| |
| impl<'de> Visitor<'de> for StringOrU64 { |
| type Value = u64; |
| |
| fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { |
| formatter.write_str("string or u64") |
| } |
| |
| fn visit_str<E>(self, value: &str) -> Result<u64, E> |
| where |
| E: de::Error, |
| { |
| FromStr::from_str(value).map_err(|_| de::Error::custom("cannot parse string as u64")) |
| } |
| |
| fn visit_u64<E>(self, value: u64) -> Result<u64, E> |
| where |
| E: de::Error, |
| { |
| Ok(value) |
| } |
| } |
| |
| deserializer.deserialize_any(StringOrU64) |
| } |
| |
| #[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, serde_json::Error> { |
| 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(target_drive: Option<&Path>) -> Result<Vec<u8>, Error> { |
| let mut command = process::Command::new("lsblk"); |
| command.args(&[ |
| // Print size in bytes |
| "--bytes", |
| // Select the fields to output |
| "--output", |
| "KNAME,NAME,RM,SIZE,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", |
| ]); |
| if let Some(target_drive) = target_drive { |
| command.arg(target_drive); |
| } |
| Ok(get_command_output(command)?) |
| } |
| |
| /// Capture information about block devices from lsblk. |
| /// |
| /// 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 a flattened vector of devices. |
| pub fn get_lsblk_devices(target_drive: Option<&Path>) -> Result<Vec<LsBlkDevice>, Error> { |
| let output = get_lsblk_output(target_drive)?; |
| let parsed = LsBlkOutput::parse(&output)?; |
| Ok(parsed.flattened()) |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| |
| #[test] |
| fn test_partition_number() { |
| let mkdev = |path: &str| LsBlkDevice { |
| kernel_name: path.into(), |
| name: path.into(), |
| is_removable: false, |
| size_in_bytes: 0, |
| device_type: "part".into(), |
| }; |
| |
| // Valid partition devices. |
| assert_eq!(mkdev("/dev/sda1").partition_number(), Some(1)); |
| assert_eq!(mkdev("/dev/sda12").partition_number(), Some(12)); |
| assert_eq!(mkdev("/dev/nvme0n1p3").partition_number(), Some(3)); |
| |
| // Doesn't end in a number. |
| assert_eq!(mkdev("/dev/dev").partition_number(), None); |
| |
| // Not a partition-type device. |
| let mut dev = mkdev("/dev/sda1"); |
| dev.device_type = "disk".into(); |
| assert_eq!(dev.partition_number(), None); |
| } |
| |
| fn mkdev( |
| kname: &str, |
| name: &str, |
| is_removable: bool, |
| size_in_bytes: u64, |
| dtype: &str, |
| ) -> LsBlkDevice { |
| LsBlkDevice { |
| kernel_name: kname.into(), |
| name: name.into(), |
| is_removable, |
| size_in_bytes, |
| device_type: dtype.into(), |
| } |
| } |
| |
| #[test] |
| fn test_lsblk_deserialization() { |
| // This test input was generated by running this command in a VM: |
| // |
| // lsblk --bytes --output KNAME,NAME,RM,SIZE,TYPE \ |
| // --json --paths --exclude 2,11,253 |
| let input = include_bytes!("test_lsblk_output.json"); |
| |
| #[rustfmt::skip] |
| let expected = vec![ |
| mkdev("/dev/sda", "/dev/sda", false, 6807435776, "disk"), |
| mkdev("/dev/sda12", "/dev/sda12", false, 134217728, "part"), |
| mkdev("/dev/sda11", "/dev/sda11", false, 8388608, "part"), |
| mkdev("/dev/sda10", "/dev/sda10", false, 512, "part"), |
| mkdev("/dev/sda9", "/dev/sda9", false, 512, "part"), |
| mkdev("/dev/sda8", "/dev/sda8", false, 16777216, "part"), |
| mkdev("/dev/sda7", "/dev/sda7", false, 512, "part"), |
| mkdev("/dev/sda6", "/dev/sda6", false, 512, "part"), |
| mkdev("/dev/sda5", "/dev/sda5", false, 2097152, "part"), |
| mkdev("/dev/sda4", "/dev/sda4", false, 67108864, "part"), |
| mkdev("/dev/sda3", "/dev/sda3", false, 2147483648, "part"), |
| mkdev("/dev/sda2", "/dev/sda2", false, 67108864, "part"), |
| mkdev("/dev/sda1", "/dev/sda1", false, 4295023104, "part"), |
| mkdev("/dev/loop1", "/dev/loop1", false, 113876992, "loop"), |
| mkdev("/dev/loop0", "/dev/loop0", false, 1248124928, "loop"), |
| mkdev("/dev/dm-0", "/dev/mapper/encstateful", false, 1248124928, "dm"), |
| ]; |
| |
| let output = LsBlkOutput::parse(input).unwrap(); |
| assert_eq!(output.flattened(), expected); |
| } |
| |
| fn mkdev2( |
| kname: &str, |
| name: &str, |
| is_removable: bool, |
| size_in_bytes: u64, |
| dtype: &str, |
| ) -> LsBlkDevice { |
| LsBlkDevice { |
| kernel_name: kname.into(), |
| name: name.into(), |
| is_removable, |
| size_in_bytes, |
| device_type: dtype.into(), |
| } |
| } |
| |
| #[test] |
| fn test_lsblk_deserialization_v2_36_1() { |
| // This test input was generated by running this command on a |
| // machine with lsblk v2.36.1 (newer than what CrOS currently |
| // has). |
| // |
| // lsblk --bytes --output KNAME,NAME,RM,SIZE,TYPE \ |
| // --json --paths --exclude 2,11,253 |
| let input = include_bytes!("test_lsblk_output_v2_36_1.json"); |
| |
| #[rustfmt::skip] |
| let expected = vec![ |
| mkdev2("/dev/nvme0n1", "/dev/nvme0n1", false, 1024209543168, "disk"), |
| mkdev2("/dev/nvme0n1p3", "/dev/nvme0n1p3", false, 1020208477696, "part"), |
| mkdev2("/dev/nvme0n1p2", "/dev/nvme0n1p2", false, 2000000000, "part"), |
| mkdev2("/dev/nvme0n1p1", "/dev/nvme0n1p1", false, 2000000000, "part"), |
| mkdev2("/dev/nvme1n1", "/dev/nvme1n1", false, 1000204886016, "disk"), |
| ]; |
| |
| let output = LsBlkOutput::parse(input).unwrap(); |
| assert_eq!(output.flattened(), expected); |
| } |
| |
| #[test] |
| fn test_lsblk_removable_deserialization() { |
| // The "rm" field can be either a string or a bool depending |
| // on the version of lsblk, verify that both work. |
| let device: Vec<LsBlkDevice> = serde_json::from_value(serde_json::json!( |
| [ |
| {"name":"a", "kname": "b", "rm":false, "size":"1", "type":"part"}, |
| {"name":"a", "kname": "b", "rm":true, "size":"1", "type":"part"}, |
| |
| {"name":"a", "kname": "b", "rm":"0", "size":"1", "type":"part"}, |
| {"name":"a", "kname": "b", "rm":"1", "size":"1", "type":"part"} |
| ] |
| )) |
| .unwrap(); |
| assert_eq!( |
| device.iter().map(|d| d.is_removable).collect::<Vec<_>>(), |
| vec![false, true, false, true] |
| ); |
| } |
| } |