storage/partitionlib: Couple of cleanups

storage/partitionlib/Cargo.toml

@@ -2,7 +2,7 @@
 name = "partitionlib"
 version = "0.1.0"
 authors = ["Deepak Sirone <deepaksirone94@gmail.com>"]
-edition = "2018"
+edition = "2021"
 license = "MIT"
 
 [dependencies]

storage/partitionlib/src/lib.rs

@@ -1,7 +1,3 @@
-extern crate gpt;
-extern crate uuid;
-
-pub type Result<T> = std::io::Result<T>;
 mod mbr;
-pub mod partition;
+mod partition;
 pub use self::partition::*;

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<io::Error> for Error {
-    fn from(err: io::Error) -> Self {
-        Self::IoError(err)
-    }
-}
-impl From<scroll::Error> for Error {
-    fn from(err: scroll::Error) -> Self {
-        Self::ParsingError(err)
-    }
-}
-
-pub fn read_header<D: Read + Seek>(device: &mut D) -> Result<Header, Error> {
+pub(crate) fn read_header<D: Read + Seek>(device: &mut D) -> io::Result<Option<Header>> {
     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<Item = Entry> {
         [
             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
     }
 }

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<Uuid>,
 }
 
-#[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<Partition>,
     pub kind: PartitionTableKind,
 }
 
-pub fn get_partitions_from_file<P: AsRef<Path>>(
-    path: P,
-    sector_size: LogicalBlockSize,
-) -> Result<Option<PartitionTable>> {
-    let mut file = File::open(path)?;
-    get_partitions(&mut file, sector_size)
-}
 fn get_gpt_partitions<D: Read + Seek>(
     device: &mut D,
     sector_size: LogicalBlockSize,
-) -> Result<PartitionTable> {
-    let header = match gpt::header::read_header_from_arbitrary_device(device, sector_size) {
-        Ok(res) => res,
-        Err(err) => return Err(err),
-    };
+) -> io::Result<PartitionTable> {
+    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<D: Read + Seek>(
         kind: PartitionTableKind::Gpt,
     })
 }
-fn get_mbr_partitions<D: Read + Seek>(device: &mut D) -> Result<Option<PartitionTable>> {
-    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<D: Read + Seek>(device: &mut D) -> io::Result<Option<PartitionTable>> {
+    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<D: Read + Seek>(device: &mut D) -> Result<Option<Partition
 pub fn get_partitions<D: Read + Seek>(
     device: &mut D,
     sector_size: LogicalBlockSize,
-) -> Result<Option<PartitionTable>> {
+) -> io::Result<Option<PartitionTable>> {
     get_gpt_partitions(device, sector_size)
         .map(Some)
         .or_else(|_| get_mbr_partitions(device))

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 {