Support "fastboot flash gpt" for updating GPT

Support updating GPT via flashing a reserved name partition "gpt".

Bug: 374792430

Change-Id: I4930b0889393d59ba740b8fce52123f47b4f9c49
diff --git a/gbl/libgbl/src/fastboot/mod.rs b/gbl/libgbl/src/fastboot/mod.rs
index 1697d25..86af91a 100644
--- a/gbl/libgbl/src/fastboot/mod.rs
+++ b/gbl/libgbl/src/fastboot/mod.rs
@@ -59,6 +59,9 @@
 // Re-exports dependency types
 pub use fastboot::{TcpStream, Transport};
 
+/// Reserved name for indicating flashing GPT.
+const FLASH_GPT_PART: &str = "gpt";
+
 /// Represents a GBL Fastboot async task.
 enum Task<'a, 'b, G: GblOps<'a>, B: BufferPool> {
     /// Image flashing task. (partition io, downloaded data, data size)
@@ -119,6 +122,13 @@
         let blk_id = next_arg_u64(&mut args, Err("".into())).ok();
         let blk_id = blk_id.map(|v| usize::try_from(v)).transpose()?;
         let blk_id = blk_id.or(self.default_block);
+
+        // Reserved "gpt" for flashing GPT partition.
+        if part == Some(FLASH_GPT_PART) {
+            let blk_id = blk_id.ok_or("Block ID is required for flashing GPT")?;
+            return Ok((part, blk_id, 0, 0));
+        }
+
         // Parses sub window offset.
         let window_offset = next_arg_u64(&mut args, Ok(0))?;
         // Parses sub window size.
@@ -134,6 +144,11 @@
         Ok((part, blk_id, window_offset, window_size))
     }
 
+    /// Takes the download data and resets download size.
+    fn take_download(&mut self) -> Option<(ScopedBuffer<'b, B>, usize)> {
+        Some((self.current_download_buffer.take()?, take(&mut self.current_download_size)))
+    }
+
     /// Waits for all block devices to be ready.
     async fn sync_all_blocks(&self) -> CommandResult<()> {
         for ele in self.gbl_ops.partitions()? {
@@ -257,10 +272,21 @@
     async fn flash(&mut self, part: &str, mut responder: impl InfoSender) -> CommandResult<()> {
         let (part, blk_idx, start, sz) = self.parse_partition(part)?;
         let partitions = self.gbl_ops.partitions()?;
+
+        if part == Some(FLASH_GPT_PART) {
+            partitions[blk_idx].wait_partition_io(None).await?.last_err()?;
+            let (mut gpt, size) = self.take_download().ok_or("No GPT downloaded")?;
+            responder.send_info("Updating GPT...").await?;
+            return match partitions[blk_idx].update_gpt(&mut gpt[..size]).await {
+                Err(Error::NotReady) => panic!("Should not be busy"),
+                Err(Error::Unsupported) => Err("Block device is not for GPT".into()),
+                v => Ok(v?),
+            };
+        }
+
         let part_io = partitions[blk_idx].wait_partition_io(part).await?.sub(start, sz)?;
         part_io.last_err()?;
-        let download_buffer = self.current_download_buffer.take().ok_or("No download")?;
-        let data_size = take(&mut self.current_download_size);
+        let (download_buffer, data_size) = self.take_download().ok_or("No download")?;
         let write_task = Task::Flash(part_io, download_buffer, data_size);
         match self.enable_async_block_io {
             true => {
@@ -334,7 +360,7 @@
     async fn oem<'s>(
         &mut self,
         cmd: &str,
-        responder: impl InfoSender,
+        mut responder: impl InfoSender,
         res: &'s mut [u8],
     ) -> CommandResult<&'s [u8]> {
         let mut args = cmd.split(' ');
@@ -356,6 +382,9 @@
             "gbl-set-default-block" => {
                 let id = next_arg_u64(&mut args, Err("Missing block device ID".into()))?;
                 self.default_block = Some(id.try_into()?);
+                responder
+                    .send_formatted_info(|f| write!(f, "Default block device: {id:#x}").unwrap())
+                    .await?;
                 Ok(b"")
             }
             "add-staged-bootloader-file" => {
@@ -1371,7 +1400,7 @@
             let mut res = String::from("");
             for v in self.lock().usb_out_queue.iter() {
                 let v = String::from_utf8(v.clone()).unwrap_or(format!("{:?}", v));
-                res += format!("b\"{}\",\n", v).as_str();
+                res += format!("b{:?},\n", v).as_str();
             }
             res
         }
@@ -1387,7 +1416,7 @@
                 let (len, rest) = remains.split_first_chunk::<{ size_of::<u64>() }>().unwrap();
                 (v, remains) = rest.split_at(u64::from_be_bytes(*len).try_into().unwrap());
                 let s = String::from_utf8(v.to_vec()).unwrap_or(format!("{:?}", v));
-                res += format!("b\"{}\",\n", s).as_str();
+                res += format!("b{:?},\n", s).as_str();
             }
             res
         }
@@ -1649,4 +1678,195 @@
         assert!(fuchsia_fastboot_mdns_packet("fuchsia-5254-0012-345", ip6_addr).is_err());
         assert!(fuchsia_fastboot_mdns_packet("fuchsia-5254-0012-34567", ip6_addr).is_err());
     }
+
+    #[test]
+    fn test_oem_update_gpt() {
+        let disk_orig = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
+        // Erase the primary and secondary header.
+        let mut disk = disk_orig.to_vec();
+        disk[512..][..512].fill(0);
+        disk.last_chunk_mut::<512>().unwrap().fill(0);
+
+        let mut storage = FakeGblOpsStorage::default();
+        storage.add_gpt_device(&disk);
+        storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_2.bin"));
+        let partitions = storage.as_partition_block_devices();
+        let buffers = vec![vec![0u8; 128 * 1024]; 2];
+        let mut gbl_ops = FakeGblOps::new(&partitions);
+        let listener: SharedTestListener = Default::default();
+        let (usb, tcp) = (&listener, &listener);
+
+        // Checks that there is no valid partitions for block #0.
+        listener.add_usb_input(b"getvar:partition-size:boot_a");
+        listener.add_usb_input(b"getvar:partition-size:boot_b");
+        // No partitions on block #0 should show up in `getvar:all` despite being a GPT device,
+        // since the GPTs are corrupted.
+        listener.add_usb_input(b"getvar:all");
+        // Download a GPT
+        let gpt = &disk_orig[..34 * 512];
+        listener.add_usb_input(format!("download:{:#x}", gpt.len()).as_bytes());
+        listener.add_usb_input(gpt);
+        listener.add_usb_input(b"flash:gpt/0");
+        // Checks that we can get partition info now.
+        listener.add_usb_input(b"getvar:partition-size:boot_a");
+        listener.add_usb_input(b"getvar:partition-size:boot_b");
+        listener.add_usb_input(b"getvar:all");
+
+        listener.add_usb_input(b"continue");
+
+        block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp)));
+
+        assert_eq!(
+            listener.usb_out_queue(),
+            make_expected_usb_out(&[
+                b"FAILNotFound",
+                b"FAILNotFound",
+                b"INFOmax-download-size: 0x20000",
+                b"INFOversion-bootloader: 1.0",
+                b"INFOmax-fetch-size: 0xffffffffffffffff",
+                b"INFOblock-device:0:total-blocks: 0x80",
+                b"INFOblock-device:0:block-size: 0x200",
+                b"INFOblock-device:0:status: idle",
+                b"INFOblock-device:1:total-blocks: 0x100",
+                b"INFOblock-device:1:block-size: 0x200",
+                b"INFOblock-device:1:status: idle",
+                b"INFOgbl-default-block: None",
+                b"INFOpartition-size:vendor_boot_a/1: 0x1000",
+                b"INFOpartition-type:vendor_boot_a/1: raw",
+                b"INFOpartition-size:vendor_boot_b/1: 0x1800",
+                b"INFOpartition-type:vendor_boot_b/1: raw",
+                b"OKAY",
+                b"DATA00004400",
+                b"OKAY",
+                b"INFOUpdating GPT...",
+                b"OKAY",
+                b"OKAY0x2000",
+                b"OKAY0x3000",
+                b"INFOmax-download-size: 0x20000",
+                b"INFOversion-bootloader: 1.0",
+                b"INFOmax-fetch-size: 0xffffffffffffffff",
+                b"INFOblock-device:0:total-blocks: 0x80",
+                b"INFOblock-device:0:block-size: 0x200",
+                b"INFOblock-device:0:status: idle",
+                b"INFOblock-device:1:total-blocks: 0x100",
+                b"INFOblock-device:1:block-size: 0x200",
+                b"INFOblock-device:1:status: idle",
+                b"INFOgbl-default-block: None",
+                b"INFOpartition-size:boot_a/0: 0x2000",
+                b"INFOpartition-type:boot_a/0: raw",
+                b"INFOpartition-size:boot_b/0: 0x3000",
+                b"INFOpartition-type:boot_b/0: raw",
+                b"INFOpartition-size:vendor_boot_a/1: 0x1000",
+                b"INFOpartition-type:vendor_boot_a/1: raw",
+                b"INFOpartition-size:vendor_boot_b/1: 0x1800",
+                b"INFOpartition-type:vendor_boot_b/1: raw",
+                b"OKAY",
+                b"INFOSyncing storage...",
+                b"OKAY",
+            ]),
+            "\nActual USB output:\n{}",
+            listener.dump_usb_out_queue()
+        );
+    }
+
+    #[test]
+    fn test_oem_update_gpt_bad_gpt() {
+        let disk = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
+        let mut storage = FakeGblOpsStorage::default();
+        storage.add_gpt_device(&disk);
+        let partitions = storage.as_partition_block_devices();
+        let buffers = vec![vec![0u8; 128 * 1024]; 2];
+        let mut gbl_ops = FakeGblOps::new(&partitions);
+        let listener: SharedTestListener = Default::default();
+        let (usb, tcp) = (&listener, &listener);
+        // Download a bad GPT.
+        let mut gpt = disk[..34 * 512].to_vec();
+        gpt[512] = !gpt[512];
+        listener.add_usb_input(format!("download:{:#x}", gpt.len()).as_bytes());
+        listener.add_usb_input(&gpt);
+        listener.add_usb_input(b"flash:gpt/0");
+        listener.add_usb_input(b"continue");
+
+        block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp)));
+
+        assert_eq!(
+            listener.usb_out_queue(),
+            make_expected_usb_out(&[
+                b"DATA00004400",
+                b"OKAY",
+                b"INFOUpdating GPT...",
+                b"FAILGptError(\n    IncorrectMagic(\n        6075990659671082682,\n    ),\n)",
+                b"INFOSyncing storage...",
+                b"OKAY",
+            ]),
+            "\nActual USB output:\n{}",
+            listener.dump_usb_out_queue()
+        );
+    }
+
+    #[test]
+    fn test_oem_update_gpt_invalid_input() {
+        let disk_orig = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
+        let mut storage = FakeGblOpsStorage::default();
+        storage.add_gpt_device(&disk_orig);
+        let partitions = storage.as_partition_block_devices();
+        let buffers = vec![vec![0u8; 128 * 1024]; 2];
+        let mut gbl_ops = FakeGblOps::new(&partitions);
+        let listener: SharedTestListener = Default::default();
+        let (usb, tcp) = (&listener, &listener);
+
+        let gpt = &disk_orig[..34 * 512];
+        listener.add_usb_input(format!("download:{:#x}", gpt.len()).as_bytes());
+        listener.add_usb_input(gpt);
+        // Missing block device ID.
+        listener.add_usb_input(b"flash:gpt");
+        listener.add_usb_input(b"continue");
+        block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp)));
+
+        assert_eq!(
+            listener.usb_out_queue(),
+            make_expected_usb_out(&[
+                b"DATA00004400",
+                b"OKAY",
+                b"FAILBlock ID is required for flashing GPT",
+                b"INFOSyncing storage...",
+                b"OKAY",
+            ]),
+            "\nActual USB output:\n{}",
+            listener.dump_usb_out_queue()
+        );
+    }
+
+    #[test]
+    fn test_oem_update_gpt_fail_on_raw_blk() {
+        let disk_orig = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
+        let mut storage = FakeGblOpsStorage::default();
+        storage.add_raw_device("raw_0", [0u8; 1024 * 1024]);
+        let partitions = storage.as_partition_block_devices();
+        let buffers = vec![vec![0u8; 128 * 1024]; 2];
+        let mut gbl_ops = FakeGblOps::new(&partitions);
+        let listener: SharedTestListener = Default::default();
+        let (usb, tcp) = (&listener, &listener);
+
+        let gpt = &disk_orig[..34 * 512];
+        listener.add_usb_input(format!("download:{:#x}", gpt.len()).as_bytes());
+        listener.add_usb_input(gpt);
+        listener.add_usb_input(b"flash:gpt/0");
+        listener.add_usb_input(b"continue");
+        block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp)));
+
+        assert_eq!(
+            listener.usb_out_queue(),
+            make_expected_usb_out(&[
+                b"DATA00004400",
+                b"OKAY",
+                b"INFOUpdating GPT...",
+                b"FAILBlock device is not for GPT",
+                b"INFOSyncing storage...",
+                b"OKAY",
+            ]),
+            "\nActual USB output:\n{}",
+            listener.dump_usb_out_queue()
+        );
+    }
 }
diff --git a/gbl/libgbl/src/fastboot/vars.rs b/gbl/libgbl/src/fastboot/vars.rs
index 167370d..7bdf842 100644
--- a/gbl/libgbl/src/fastboot/vars.rs
+++ b/gbl/libgbl/src/fastboot/vars.rs
@@ -98,7 +98,8 @@
         let partitions = gbl_fb.gbl_ops.partitions()?;
         let mut size_str = [0u8; 32];
         for (idx, blk) in partitions.iter().enumerate() {
-            for ptn in blk.partition_iter() {
+            for ptn_idx in 0..blk.num_partitions().unwrap_or(0) {
+                let ptn = blk.get_partition_by_idx(ptn_idx)?;
                 let sz: u64 = ptn.size()?;
                 let part = ptn.name()?;
                 // Assumes max partition name length of 72 plus max u64 hex string length 18.
diff --git a/gbl/libgbl/src/partition.rs b/gbl/libgbl/src/partition.rs
index b576a1c..29f9791 100644
--- a/gbl/libgbl/src/partition.rs
+++ b/gbl/libgbl/src/partition.rs
@@ -15,7 +15,10 @@
 //! This file implements storage and partition logic for libgbl.
 
 use crate::fastboot::sparse::{is_sparse_image, write_sparse_image, SparseRawWriter};
-use core::mem::swap;
+use core::{
+    mem::swap,
+    ops::{Deref, DerefMut},
+};
 use fastboot::CommandError;
 use gbl_async::yield_now;
 use gbl_storage::{
@@ -65,25 +68,6 @@
     Gpt(GptCache<'a>),
 }
 
-/// Internal partition entry iterator type.
-enum PartitionIter<'a, G> {
-    /// Raw partition block device is simply a 1-partition device.
-    Raw(Option<Partition<'a>>),
-    /// Gpt partition block device holds a `GptPartition` iterator.
-    Gpt(G),
-}
-
-impl<'a, G: Iterator<Item = GptPartition>> Iterator for PartitionIter<'a, G> {
-    type Item = Partition<'a>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        match self {
-            PartitionIter::Raw(part) => part.take(),
-            PartitionIter::Gpt(gpt) => gpt.next().map(|v| Partition::Gpt(v)),
-        }
-    }
-}
-
 /// The status of block device
 pub enum BlockStatus {
     /// Idle,
@@ -110,8 +94,17 @@
 pub struct PartitionBlockDevice<'a, B: BlockIoAsync> {
     // Contains an `AsyncBlockDevice` for block IO and `Result` to track the most recent error.
     // Wraps in `Mutex` as it will be used in parallel fastboot task.
+    //
+    // `blk` and `partitions` are separately guarded because we need to get partition info for
+    // `fastboot getvar` even when block IO itself is busy. Also we need interior mutability on
+    // `partitions` for updating and syncing GPT.
+    //
+    // To prevent deadlock, locking of `partitions` is only managed internally and does not block.
+    // Failure to lock returns Error. No method returns a locked `partitions` to the caller.
+    // Thus there won't be two callers locking one of the resource and blocking each other on
+    // acquiring the other one.
     blk: Mutex<(AsyncBlockDevice<'a, B>, Result<(), Error>)>,
-    partitions: PartitionTable<'a>,
+    partitions: Mutex<PartitionTable<'a>>,
     info_cache: BlockInfo,
 }
 
@@ -119,7 +112,7 @@
     /// Creates a new instance as a GPT device.
     pub fn new_gpt(mut blk: AsyncBlockDevice<'a, B>, gpt: GptCache<'a>) -> Self {
         let info_cache = blk.io().info();
-        Self { blk: (blk, Ok(())).into(), info_cache, partitions: PartitionTable::Gpt(gpt) }
+        Self { blk: (blk, Ok(())).into(), info_cache, partitions: PartitionTable::Gpt(gpt).into() }
     }
 
     /// Creates a new instance as a raw storage partition.
@@ -128,7 +121,7 @@
         Ok(Self {
             blk: (blk, Ok(())).into(),
             info_cache,
-            partitions: PartitionTable::Raw(name, info_cache.total_size()?),
+            partitions: PartitionTable::Raw(name, info_cache.total_size()?).into(),
         })
     }
 
@@ -177,20 +170,27 @@
             return Ok(Partition::Raw("", self.info_cache.total_size()?));
         };
 
-        match &self.partitions {
+        match self.partitions.try_lock().ok_or(Error::NotReady)?.deref() {
             PartitionTable::Gpt(gpt) => Ok(Partition::Gpt(gpt.find_partition(part)?)),
             PartitionTable::Raw(name, size) if *name == part => Ok(Partition::Raw(name, *size)),
             _ => Err(Error::NotFound),
         }
     }
 
-    /// Gets an iterator to partition entries.
-    pub fn partition_iter(&self) -> impl Iterator<Item = Partition<'a>> + '_ {
-        match &self.partitions {
-            PartitionTable::Gpt(gpt) => PartitionIter::Gpt(gpt.partition_iter()),
-            PartitionTable::Raw(name, size) => {
-                PartitionIter::Raw(Some(Partition::Raw(name, *size)))
-            }
+    /// Get total number of partitions.
+    pub fn num_partitions(&self) -> Result<usize, Error> {
+        match self.partitions.try_lock().ok_or(Error::NotReady)?.deref() {
+            PartitionTable::Raw(name, _) => Ok(1),
+            PartitionTable::Gpt(gpt) => gpt.num_partitions(),
+        }
+    }
+
+    /// Gets a partition by index.
+    pub fn get_partition_by_idx(&self, idx: usize) -> Result<Partition<'a>, Error> {
+        match self.partitions.try_lock().ok_or(Error::NotReady)?.deref() {
+            PartitionTable::Raw(name, v) if idx == 0 => Ok(Partition::Raw(name, *v)),
+            PartitionTable::Gpt(gpt) => Ok(Partition::Gpt(gpt.get_partition(idx)?)),
+            _ => Err(Error::InvalidInput),
         }
     }
 
@@ -202,8 +202,8 @@
     ///  `sync_res` contains the GPT verification and restoration result.
     /// * Returns `Ok(None)` if partition type is not GPT.
     /// * Returns `Err` in other cases.
-    pub async fn sync_gpt(&mut self) -> Result<Option<GptSyncResult>, Error> {
-        match &mut self.partitions {
+    pub async fn sync_gpt(&self) -> Result<Option<GptSyncResult>, Error> {
+        match self.partitions.try_lock().ok_or(Error::NotReady)?.deref_mut() {
             PartitionTable::Raw(name, _) => Ok(None),
             PartitionTable::Gpt(ref mut gpt) => {
                 let mut blk = self.blk.try_lock().ok_or(Error::NotReady)?;
@@ -211,6 +211,27 @@
             }
         }
     }
+
+    /// Updates GPT to the block device and sync primary and secondary GPT.
+    ///
+    /// # Args
+    ///
+    /// * `mbr_primary`: A buffer containing the MBR block, primary GPT header and entries.
+    ///
+    /// # Returns
+    ///
+    /// * Return `Err(Error::NotReady)` if device is busy.
+    /// * Return `Err(Error::Unsupported)` if partition type is not GPT.
+    /// * Return `Ok(())` new GPT is valid and device is updated and synced successfully.
+    pub async fn update_gpt(&self, mbr_primary: &mut [u8]) -> Result<(), Error> {
+        match self.partitions.try_lock().ok_or(Error::NotReady)?.deref_mut() {
+            PartitionTable::Raw(name, _) => Err(Error::Unsupported),
+            PartitionTable::Gpt(ref mut gpt) => {
+                let mut blk = self.blk.try_lock().ok_or(Error::NotReady)?;
+                blk.0.update_gpt(mbr_primary, gpt).await
+            }
+        }
+    }
 }
 
 /// `PartitionIo` provides read/write APIs to a partition.
@@ -403,7 +424,7 @@
     #[test]
     fn test_find_partition_gpt() {
         let mut gpt = (&include_bytes!("../../libstorage/test/gpt_test_1.bin")[..]).into();
-        let mut gpt = as_gpt_part(&mut gpt);
+        let gpt = as_gpt_part(&mut gpt);
         assert_eq!(block_on(gpt.sync_gpt()).unwrap(), Some(GptSyncResult::BothValid));
 
         let boot_a = gpt.find_partition(Some("boot_a")).unwrap();
@@ -468,7 +489,7 @@
     fn test_read_partition_gpt() {
         let disk = include_bytes!("../../libstorage/test/gpt_test_1.bin");
         let mut gpt = (&disk[..]).into();
-        let mut gpt = as_gpt_part(&mut gpt);
+        let gpt = as_gpt_part(&mut gpt);
         assert_eq!(block_on(gpt.sync_gpt()).unwrap(), Some(GptSyncResult::BothValid));
 
         let expect_boot_a = include_bytes!("../../libstorage/test/boot_a.bin");
@@ -516,7 +537,7 @@
     #[test]
     fn test_write_partition_gpt() {
         let mut gpt = (&include_bytes!("../../libstorage/test/gpt_test_1.bin")[..]).into();
-        let mut gpt = as_gpt_part(&mut gpt);
+        let gpt = as_gpt_part(&mut gpt);
         assert_eq!(block_on(gpt.sync_gpt()).unwrap(), Some(GptSyncResult::BothValid));
         test_part_write(&gpt, Some("boot_a"), 1, 1024);
         test_part_write(&gpt, Some("boot_b"), 1, 1024);
@@ -535,7 +556,7 @@
     fn test_read_write_partition_overflow() {
         let disk = include_bytes!("../../libstorage/test/gpt_test_1.bin");
         let mut gpt = (&disk[..]).into();
-        let mut gpt = as_gpt_part(&mut gpt);
+        let gpt = as_gpt_part(&mut gpt);
         assert_eq!(block_on(gpt.sync_gpt()).unwrap(), Some(GptSyncResult::BothValid));
 
         let mut part_io = gpt.partition_io(Some("boot_a")).unwrap();
@@ -561,7 +582,7 @@
     fn test_sub_overflow() {
         let disk = include_bytes!("../../libstorage/test/gpt_test_1.bin");
         let mut gpt = (&disk[..]).into();
-        let mut gpt = as_gpt_part(&mut gpt);
+        let gpt = as_gpt_part(&mut gpt);
         assert_eq!(block_on(gpt.sync_gpt()).unwrap(), Some(GptSyncResult::BothValid));
         assert!(gpt.partition_io(Some("boot_a")).unwrap().sub(0, BOOT_A_SZ + 1).is_err());
         assert!(gpt.partition_io(Some("boot_a")).unwrap().sub(1, BOOT_A_SZ).is_err());
@@ -661,17 +682,17 @@
     fn test_partition_iter() {
         let mut raw = (&vec![0u8; 1024][..]).into();
         let raw = as_raw_part(&mut raw, "raw");
-        assert_eq!(raw.partition_iter().collect::<Vec<_>>(), [Partition::Raw("raw", 1024)]);
+        assert_eq!(raw.num_partitions().unwrap(), 1);
+        assert_eq!(raw.get_partition_by_idx(0).unwrap(), Partition::Raw("raw", 1024));
 
         let mut gpt = (&include_bytes!("../../libstorage/test/gpt_test_1.bin")[..]).into();
-        let mut gpt = as_gpt_part(&mut gpt);
+        let gpt = as_gpt_part(&mut gpt);
         block_on(gpt.sync_gpt()).unwrap();
-        let actual = gpt.partition_iter().collect::<Vec<_>>();
-        assert_eq!(actual.len(), 2);
-        assert_eq!(actual[0].name().unwrap(), "boot_a");
-        assert_eq!(actual[0].size().unwrap(), 0x2000);
-        assert_eq!(actual[1].name().unwrap(), "boot_b");
-        assert_eq!(actual[1].size().unwrap(), 0x3000);
+        assert_eq!(gpt.num_partitions().unwrap(), 2);
+        assert_eq!(gpt.get_partition_by_idx(0).unwrap().name().unwrap(), "boot_a");
+        assert_eq!(gpt.get_partition_by_idx(0).unwrap().size().unwrap(), 0x2000);
+        assert_eq!(gpt.get_partition_by_idx(1).unwrap().name().unwrap(), "boot_b");
+        assert_eq!(gpt.get_partition_by_idx(1).unwrap().size().unwrap(), 0x3000);
     }
 
     /// A test helper for `read_unique_partition`