From 9257cbbfa26a96e5ba22437eeb0c714c4bff10be Mon Sep 17 00:00:00 2001 From: bjorn3 <17426603+bjorn3@users.noreply.github.com> Date: Thu, 6 Mar 2025 22:25:24 +0100 Subject: [PATCH] storage/partitionlib: Couple of cleanups --- storage/partitionlib/Cargo.toml | 2 +- storage/partitionlib/src/lib.rs | 6 +-- storage/partitionlib/src/mbr.rs | 70 ++++++++++----------------- storage/partitionlib/src/partition.rs | 49 ++++--------------- storage/partitionlib/tests/test.rs | 41 +++++----------- 5 files changed, 49 insertions(+), 119 deletions(-) diff --git a/storage/partitionlib/Cargo.toml b/storage/partitionlib/Cargo.toml index 27df61ce16..8e5e3fad44 100644 --- a/storage/partitionlib/Cargo.toml +++ b/storage/partitionlib/Cargo.toml @@ -2,7 +2,7 @@ name = "partitionlib" version = "0.1.0" authors = ["Deepak Sirone "] -edition = "2018" +edition = "2021" license = "MIT" [dependencies] diff --git a/storage/partitionlib/src/lib.rs b/storage/partitionlib/src/lib.rs index 1b0b71b4e5..be369ce64e 100644 --- a/storage/partitionlib/src/lib.rs +++ b/storage/partitionlib/src/lib.rs @@ -1,7 +1,3 @@ -extern crate gpt; -extern crate uuid; - -pub type Result = std::io::Result; mod mbr; -pub mod partition; +mod partition; pub use self::partition::*; diff --git a/storage/partitionlib/src/mbr.rs b/storage/partitionlib/src/mbr.rs index cf129ce356..67ae5a6dad 100644 --- a/storage/partitionlib/src/mbr.rs +++ b/storage/partitionlib/src/mbr.rs @@ -1,77 +1,57 @@ use scroll::{Pread, Pwrite}; -use std::io; -use std::io::prelude::*; +use std::io::{self, Read, Seek}; #[derive(Clone, Copy, Debug, Pread, Pwrite)] -pub struct Entry { - pub drive_attrs: u8, - pub start_head: u8, - pub start_cs: u16, - pub sys_id: u8, - pub end_head: u8, - pub end_cs: u16, - pub rel_sector: u32, - pub len: u32, +pub(crate) struct Entry { + pub(crate) drive_attrs: u8, + pub(crate) start_head: u8, + pub(crate) start_cs: u16, + pub(crate) sys_id: u8, + pub(crate) end_head: u8, + pub(crate) end_cs: u16, + pub(crate) rel_sector: u32, + pub(crate) len: u32, } #[derive(Pread, Pwrite)] -pub struct Header { - pub bootstrap: [u8; 446], - pub first_entry: Entry, - pub second_entry: Entry, - pub third_entry: Entry, - pub fourth_entry: Entry, - pub last_signature: u16, // 0xAA55 +pub(crate) struct Header { + pub(crate) bootstrap: [u8; 446], + pub(crate) first_entry: Entry, + pub(crate) second_entry: Entry, + pub(crate) third_entry: Entry, + pub(crate) fourth_entry: Entry, + pub(crate) last_signature: u16, // 0xAA55 } -#[derive(Debug)] -pub enum Error { - IoError(io::Error), - ParsingError(scroll::Error), -} - -impl From for Error { - fn from(err: io::Error) -> Self { - Self::IoError(err) - } -} -impl From for Error { - fn from(err: scroll::Error) -> Self { - Self::ParsingError(err) - } -} - -pub fn read_header(device: &mut D) -> Result { +pub(crate) fn read_header(device: &mut D) -> io::Result> { device.seek(io::SeekFrom::Start(0))?; let mut bytes = [0u8; 512]; device.read_exact(&mut bytes)?; - let header: Header = bytes.pread_with(0, scroll::LE)?; + let header: Header = bytes.pread_with(0, scroll::LE).unwrap(); if header.last_signature != 0xAA55 { - return Err(scroll::Error::BadInput { - size: 2, - msg: "no 0xAA55 signature", - } - .into()); + return Ok(None); } - Ok(header) + Ok(Some(header)) } impl Header { - pub fn partitions(&self) -> [Entry; 4] { + pub(crate) fn partitions(&self) -> impl Iterator { [ self.first_entry, self.second_entry, self.third_entry, self.fourth_entry, ] + .into_iter() + .filter(Entry::is_valid) } } impl Entry { - pub fn is_valid(&self) -> bool { + fn is_valid(&self) -> bool { (self.drive_attrs == 0 || self.drive_attrs == 0x80) && self.len != 0 } } diff --git a/storage/partitionlib/src/partition.rs b/storage/partitionlib/src/partition.rs index 030f292852..949171511b 100644 --- a/storage/partitionlib/src/partition.rs +++ b/storage/partitionlib/src/partition.rs @@ -1,8 +1,5 @@ -use super::Result; pub use gpt::disk::LogicalBlockSize; -use std::fs::File; -use std::io::prelude::*; -use std::path::Path; +use std::io::{self, Read, Seek}; use uuid::Uuid; /// A union of the MBR and GPT partition entry @@ -17,46 +14,23 @@ pub struct Partition { pub uuid: Option, } -#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum PartitionTableKind { Mbr, Gpt, } -impl Default for PartitionTableKind { - fn default() -> Self { - Self::Gpt - } -} -impl PartitionTableKind { - pub fn is_mbr(self) -> bool { - self == Self::Mbr - } - pub fn is_gpt(self) -> bool { - self == Self::Gpt - } -} -#[derive(Clone, Debug, Eq, PartialEq, Hash)] +#[derive(Clone, Debug, Eq, PartialEq)] pub struct PartitionTable { pub partitions: Vec, pub kind: PartitionTableKind, } -pub fn get_partitions_from_file>( - path: P, - sector_size: LogicalBlockSize, -) -> Result> { - let mut file = File::open(path)?; - get_partitions(&mut file, sector_size) -} fn get_gpt_partitions( device: &mut D, sector_size: LogicalBlockSize, -) -> Result { - let header = match gpt::header::read_header_from_arbitrary_device(device, sector_size) { - Ok(res) => res, - Err(err) => return Err(err), - }; +) -> io::Result { + let header = gpt::header::read_header_from_arbitrary_device(device, sector_size)?; Ok(PartitionTable { partitions: gpt::partition::file_read_partitions(device, &header, sector_size).map( |btree| { @@ -75,19 +49,14 @@ fn get_gpt_partitions( kind: PartitionTableKind::Gpt, }) } -fn get_mbr_partitions(device: &mut D) -> Result> { - let header = match crate::mbr::read_header(device) { - Ok(h) => h, - Err(crate::mbr::Error::ParsingError(_)) => return Ok(None), - Err(crate::mbr::Error::IoError(ioerr)) => return Err(ioerr), +fn get_mbr_partitions(device: &mut D) -> io::Result> { + let Some(header) = crate::mbr::read_header(device)? else { + return Ok(None); }; Ok(Some(PartitionTable { kind: PartitionTableKind::Mbr, partitions: header .partitions() - .iter() - .copied() - .filter(crate::mbr::Entry::is_valid) .map(|partition: crate::mbr::Entry| Partition { name: None, uuid: None, // TODO: Some kind of one-way conversion should be possible @@ -101,7 +70,7 @@ fn get_mbr_partitions(device: &mut D) -> Result( device: &mut D, sector_size: LogicalBlockSize, -) -> Result> { +) -> io::Result> { get_gpt_partitions(device, sector_size) .map(Some) .or_else(|_| get_mbr_partitions(device)) diff --git a/storage/partitionlib/tests/test.rs b/storage/partitionlib/tests/test.rs index c2a542a1c7..c4ef93220e 100644 --- a/storage/partitionlib/tests/test.rs +++ b/storage/partitionlib/tests/test.rs @@ -1,35 +1,21 @@ -extern crate partitionlib; +use std::fs::File; -use partitionlib::{get_partitions_from_file, LogicalBlockSize, Partition}; +use partitionlib::{ + get_partitions, LogicalBlockSize, Partition, PartitionTable, PartitionTableKind, +}; -#[test] -fn part_is_gpt() { - assert!(dbg!( - get_partitions_from_file("./resources/disk.img", LogicalBlockSize::Lb512) - .unwrap() - .unwrap() - ) - .kind - .is_gpt()); -} - -#[test] -fn part_is_mbr() { - assert!( - get_partitions_from_file("./resources/disk_mbr.img", LogicalBlockSize::Lb512) - .unwrap() - .unwrap() - .kind - .is_mbr() - ); +fn get_partitions_from_file(path: &str) -> PartitionTable { + let mut file = File::open(path).unwrap(); + get_partitions(&mut file, LogicalBlockSize::Lb512) + .unwrap() + .unwrap() } // NOTE: The following tests rely on outside resource files being correct. #[test] fn gpt() { - let table = get_partitions_from_file("./resources/disk.img", LogicalBlockSize::Lb512) - .unwrap() - .unwrap(); + let table = get_partitions_from_file("./resources/disk.img"); + assert_eq!(table.kind, PartitionTableKind::Gpt); assert_eq!( &table.partitions, &[Partition { @@ -44,9 +30,8 @@ fn gpt() { #[test] fn mbr() { - let table = get_partitions_from_file("./resources/disk_mbr.img", LogicalBlockSize::Lb512) - .unwrap() - .unwrap(); + let table = get_partitions_from_file("./resources/disk_mbr.img"); + assert_eq!(table.kind, PartitionTableKind::Mbr); assert_eq!( &table.partitions, &[Partition {