Merge "Generalize MultiGptDevices to AsMultiBlockDevices" into main

10 files changed, 593 insertions, 359 deletions

diff --git a/gbl/efi/src/android_boot.rs b/gbl/efi/src/android_boot.rs
index 40067b3..899eb65 100644
--- a/gbl/efi/src/android_boot.rs
+++ b/gbl/efi/src/android_boot.rs
@@ -20,13 +20,13 @@ use bootconfig::{BootConfigBuilder, BootConfigError};
 use bootimg::{BootImage, VendorImageHeader};
 use efi::{efi_print, efi_println, exit_boot_services, EfiEntry};
 use fdt::Fdt;
-use gbl_storage::MultiGptDevices;
+use gbl_storage::AsMultiBlockDevices;
 use misc::{AndroidBootMode, BootloaderMessage};
 
 use crate::error::{EfiAppError, GblEfiError, Result};
 use crate::utils::{
     aligned_subslice, cstr_bytes_to_str, find_gpt_devices, get_efi_fdt, usize_add, usize_roundup,
-    EfiGptDevice,
+    EfiMultiBlockDevices,
 };
 
 use crate::avb::GblEfiAvbOps;
@@ -46,7 +46,7 @@ macro_rules! cstr_literal {
 
 /// Helper function for performing libavb verification.
 fn avb_verify_slot<'a, 'b, 'c>(
-    gpt_dev: &'b mut [EfiGptDevice<'a>],
+    gpt_dev: &'b mut EfiMultiBlockDevices,
     kernel: &'b [u8],
     vendor_boot: &'b [u8],
     init_boot: &'b [u8],
@@ -97,7 +97,7 @@ pub fn load_android_simple<'a>(
     efi_entry: &EfiEntry,
     load: &'a mut [u8],
 ) -> Result<(&'a mut [u8], &'a mut [u8], &'a mut [u8], &'a mut [u8])> {
-    let mut gpt_devices = &mut find_gpt_devices(efi_entry)?[..];
+    let mut gpt_devices = find_gpt_devices(efi_entry)?;
 
     const PAGE_SIZE: usize = 4096; // V3/V4 image has fixed page size 4096;
 
diff --git a/gbl/efi/src/avb.rs b/gbl/efi/src/avb.rs
index efd7f07..1669104 100644
--- a/gbl/efi/src/avb.rs
+++ b/gbl/efi/src/avb.rs
@@ -19,20 +19,20 @@
 use core::cmp::{min, Ord};
 use core::ffi::CStr;
 
-use crate::utils::EfiGptDevice;
+use crate::utils::EfiMultiBlockDevices;
 use avb::{IoError, IoResult, Ops, PublicKeyForPartitionInfo};
 use efi::{efi_free, efi_malloc};
-use gbl_storage::MultiGptDevices;
+use gbl_storage::AsMultiBlockDevices;
 use uuid::Uuid;
 
 pub struct GblEfiAvbOps<'a, 'b> {
-    gpt_dev: &'b mut [EfiGptDevice<'a>],
+    gpt_dev: &'b mut EfiMultiBlockDevices<'a>,
     preloaded_partitions: Option<&'b [(&'b str, &'b [u8])]>,
 }
 
 impl<'a, 'b> GblEfiAvbOps<'a, 'b> {
     pub fn new(
-        gpt_dev: &'b mut [EfiGptDevice<'a>],
+        gpt_dev: &'b mut EfiMultiBlockDevices<'a>,
         preloaded_partitions: Option<&'b [(&'b str, &'b [u8])]>,
     ) -> Self {
         Self { gpt_dev, preloaded_partitions }
diff --git a/gbl/efi/src/fuchsia_boot.rs b/gbl/efi/src/fuchsia_boot.rs
index 654d5d6..bb0f020 100644
--- a/gbl/efi/src/fuchsia_boot.rs
+++ b/gbl/efi/src/fuchsia_boot.rs
@@ -18,7 +18,7 @@ use core::fmt::Write;
 use core::mem::size_of;
 use efi::{efi_print, efi_println, EfiEntry};
 use fdt::Fdt;
-use gbl_storage::MultiGptDevices;
+use gbl_storage::AsMultiBlockDevices;
 use zbi::{ZbiContainer, ZbiFlags, ZbiHeader, ZbiType, ZBI_ALIGNMENT_USIZE};
 use zerocopy::{AsBytes, FromBytes, FromZeroes, Ref};
 
@@ -93,7 +93,7 @@ fn relocate_to_tail(kernel: &mut [u8]) -> Result<(&mut [u8], &mut [u8], usize)>
 fn load_fuchsia_simple<'a>(efi_entry: &EfiEntry, load: &'a mut [u8]) -> Result<&'a mut [u8]> {
     let load = aligned_subslice(load, ZBI_ALIGNMENT_USIZE)?;
 
-    let mut gpt_devices = &mut find_gpt_devices(&efi_entry)?[..];
+    let mut gpt_devices = find_gpt_devices(&efi_entry)?;
 
     // Gets FDT from EFI configuration table.
     let (_, fdt_bytes) = get_efi_fdt(&efi_entry).ok_or_else(|| EfiAppError::NoFdt).unwrap();
@@ -148,7 +148,7 @@ fn load_fuchsia_simple<'a>(efi_entry: &EfiEntry, load: &'a mut [u8]) -> Result<&
 
 /// Check if the disk GPT layout is a Fuchsia device layout.
 pub fn is_fuchsia_gpt(efi_entry: &EfiEntry) -> Result<()> {
-    let mut gpt_devices = &mut find_gpt_devices(&efi_entry)?[..];
+    let mut gpt_devices = find_gpt_devices(&efi_entry)?;
     let partitions: [&[&str]; 8] = [
         &["zircon_a"],
         &["zircon_b"],
diff --git a/gbl/efi/src/utils.rs b/gbl/efi/src/utils.rs
index d74989e..186c705 100644
--- a/gbl/efi/src/utils.rs
+++ b/gbl/efi/src/utils.rs
@@ -22,7 +22,7 @@ use efi::{
     EfiEntry, EventType, LoadedImageProtocol, Protocol, SimpleTextInputProtocol,
 };
 use fdt::FdtHeader;
-use gbl_storage::{required_scratch_size, AsBlockDevice, BlockIo};
+use gbl_storage::{required_scratch_size, AsBlockDevice, AsMultiBlockDevices, BlockIo};
 
 pub const EFI_DTB_TABLE_GUID: EfiGuid =
     EfiGuid::new(0xb1b621d5, 0xf19c, 0x41a5, [0x83, 0x0b, 0xd9, 0x15, 0x2c, 0x69, 0xaa, 0xe0]);
@@ -95,13 +95,25 @@ impl<'a> EfiGptDevice<'a> {
 }
 
 impl AsBlockDevice for EfiGptDevice<'_> {
-    fn get(&mut self) -> (&mut dyn BlockIo, &mut [u8], u64) {
-        (&mut self.io, &mut self.scratch[..], MAX_GPT_ENTRIES)
+    fn with(&mut self, f: &mut dyn FnMut(&mut dyn BlockIo, &mut [u8], u64)) {
+        f(&mut self.io, &mut self.scratch[..], MAX_GPT_ENTRIES)
+    }
+}
+
+pub struct EfiMultiBlockDevices<'a>(pub alloc::vec::Vec<EfiGptDevice<'a>>);
+
+impl AsMultiBlockDevices for EfiMultiBlockDevices<'_> {
+    fn for_each_until(&mut self, f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64) -> bool) {
+        for (idx, ele) in self.0.iter_mut().enumerate() {
+            if f(ele, u64::try_from(idx).unwrap()) {
+                return;
+            }
+        }
     }
 }
 
 /// Finds and returns all block devices that have a valid GPT.
-pub fn find_gpt_devices(efi_entry: &EfiEntry) -> Result<Vec<EfiGptDevice>> {
+pub fn find_gpt_devices(efi_entry: &EfiEntry) -> Result<EfiMultiBlockDevices> {
     let bs = efi_entry.system_table().boot_services();
     let block_dev_handles = bs.locate_handle_buffer_by_protocol::<BlockIoProtocol>()?;
     let mut gpt_devices = Vec::<EfiGptDevice>::new();
@@ -114,7 +126,7 @@ pub fn find_gpt_devices(efi_entry: &EfiEntry) -> Result<Vec<EfiGptDevice>> {
             _ => {}
         };
     }
-    Ok(gpt_devices)
+    Ok(EfiMultiBlockDevices(gpt_devices))
 }
 
 /// Helper function to get the `DevicePathText` from a `DeviceHandle`.
diff --git a/gbl/libgbl/src/lib.rs b/gbl/libgbl/src/lib.rs
index d0ff9cf..93c7a97 100644
--- a/gbl/libgbl/src/lib.rs
+++ b/gbl/libgbl/src/lib.rs
@@ -589,7 +589,9 @@ mod tests {
     use avb::IoError;
     use avb::IoResult as AvbIoResult;
     use avb::PublicKeyForPartitionInfo;
+    #[cfg(feature = "sw_digest")]
     use avb_test::TestOps;
+    #[cfg(feature = "sw_digest")]
     use std::fs;
 
     struct AvbOpsUnimplemented {}
diff --git a/gbl/libgbl/src/slots/fuchsia.rs b/gbl/libgbl/src/slots/fuchsia.rs
index 52bb8d5..cfb65e1 100644
--- a/gbl/libgbl/src/slots/fuchsia.rs
+++ b/gbl/libgbl/src/slots/fuchsia.rs
@@ -664,8 +664,8 @@ mod test {
     }
 
     impl gbl_storage::AsBlockDevice for TestBlockDevice {
-        fn get(&mut self) -> (&mut dyn gbl_storage::BlockIo, &mut [u8], u64) {
-            (&mut self.io, &mut self.scratch[..], Self::GPT_ENTRIES)
+        fn with(&mut self, f: &mut dyn FnMut(&mut dyn gbl_storage::BlockIo, &mut [u8], u64)) {
+            f(&mut self.io, &mut self.scratch[..], Self::GPT_ENTRIES)
         }
     }
 
diff --git a/gbl/libstorage/BUILD b/gbl/libstorage/BUILD
index caf2e25..3ea904d 100644
--- a/gbl/libstorage/BUILD
+++ b/gbl/libstorage/BUILD
@@ -23,6 +23,7 @@ rust_library(
     srcs = [
         "src/gpt.rs",
         "src/lib.rs",
+        "src/multi_blocks.rs",
     ],
     crate_name = "gbl_storage",
     edition = "2021",
diff --git a/gbl/libstorage/src/gpt.rs b/gbl/libstorage/src/gpt.rs
index 2d09435..d1c0f44 100644
--- a/gbl/libstorage/src/gpt.rs
+++ b/gbl/libstorage/src/gpt.rs
@@ -13,8 +13,8 @@
 // limitations under the License.
 
 use crate::{
-    add, aligned_subslice, div, mul, read, sub, to_usize, write_bytes, write_bytes_mut,
-    AsBlockDevice, BlockIo, Result, StorageError,
+    add, aligned_subslice, div, mul, read, sub, to_usize, write_bytes, write_bytes_mut, BlockIo,
+    Result, StorageError,
 };
 use core::default::Default;
 use core::mem::{align_of, size_of};
@@ -23,7 +23,7 @@ use crc32fast::Hasher;
 use zerocopy::{AsBytes, FromBytes, FromZeroes, Ref};
 
 const GPT_GUID_LEN: usize = 16;
-const GPT_NAME_LEN: usize = 36;
+pub(crate) const GPT_NAME_LEN: usize = 36;
 
 #[repr(C, packed)]
 #[derive(Debug, Default, Copy, Clone, AsBytes, FromBytes, FromZeroes)]
@@ -145,7 +145,7 @@ impl GptInfo {
 }
 
 /// An object that contains the GPT header/entries information.
-pub struct Gpt<'a> {
+pub(crate) struct Gpt<'a> {
     info: &'a mut GptInfo,
     /// Raw bytes of primary GPT header.
     primary_header: &'a mut [u8],
@@ -217,7 +217,7 @@ impl<'a> Gpt<'a> {
     /// Return the list of GPT entries.
     ///
     /// If the object does not contain a valid GPT, the method returns Error.
-    pub fn entries(&self) -> Result<&[GptEntry]> {
+    pub(crate) fn entries(&self) -> Result<&[GptEntry]> {
         self.check_valid()?;
         Ok(&Ref::<_, [GptEntry]>::new_slice(&self.primary_entries[..]).unwrap().into_slice()
             [..to_usize(self.info.num_valid_entries()?)?])
@@ -228,15 +228,15 @@ impl<'a> Gpt<'a> {
     /// If partition doesn't exist, the method returns `Ok(None)`.
     ///
     /// If the object does not contain a valid GPT, the method returns Error.
-    pub fn find_partition(&self, part: &str) -> Result<Option<&GptEntry>> {
+    pub(crate) fn find_partition(&self, part: &str) -> Result<&GptEntry> {
         for entry in self.entries()? {
             let mut name_conversion_buffer = [0u8; GPT_NAME_LEN * 2];
             if entry.name_to_str(&mut name_conversion_buffer)? != part {
                 continue;
             }
-            return Ok(Some(entry));
+            return Ok(entry);
         }
-        Ok(None)
+        Err(StorageError::NotExist)
     }
 
     /// Check whether the Gpt has been initialized.
@@ -393,13 +393,9 @@ pub(crate) fn read_gpt_partition(
     out: &mut [u8],
     scratch: &mut [u8],
 ) -> Result<()> {
-    match gpt.find_partition(part_name) {
-        Ok(Some(e)) => {
-            let abs_offset = check_offset(blk_dev, e, offset, out.len() as u64)?;
-            read(blk_dev, abs_offset, out, scratch)
-        }
-        _ => Err(StorageError::NotExist),
-    }
+    let e = gpt.find_partition(part_name)?;
+    let abs_offset = check_offset(blk_dev, e, offset, out.len() as u64)?;
+    read(blk_dev, abs_offset, out, scratch)
 }
 
 /// Write GPT partition. Library internal helper for AsBlockDevice::write_gpt_partition().
@@ -411,13 +407,9 @@ pub(crate) fn write_gpt_partition(
     data: &[u8],
     scratch: &mut [u8],
 ) -> Result<()> {
-    match gpt.find_partition(part_name) {
-        Ok(Some(e)) => {
-            let abs_offset = check_offset(blk_dev, e, offset, data.len() as u64)?;
-            write_bytes(blk_dev, abs_offset, data, scratch)
-        }
-        _ => Err(StorageError::NotExist),
-    }
+    let e = gpt.find_partition(part_name)?;
+    let abs_offset = check_offset(blk_dev, e, offset, data.len() as u64)?;
+    write_bytes(blk_dev, abs_offset, data, scratch)
 }
 
 /// Write GPT partition. Library internal helper for AsBlockDevice::write_gpt_partition().
@@ -430,13 +422,9 @@ pub(crate) fn write_gpt_partition_mut(
     data: &mut [u8],
     scratch: &mut [u8],
 ) -> Result<()> {
-    match gpt.find_partition(part_name) {
-        Ok(Some(e)) => {
-            let abs_offset = check_offset(blk_dev, e, offset, data.as_ref().len() as u64)?;
-            write_bytes_mut(blk_dev, abs_offset, data.as_mut(), scratch)
-        }
-        _ => Err(StorageError::NotExist),
-    }
+    let e = gpt.find_partition(part_name)?;
+    let abs_offset = check_offset(blk_dev, e, offset, data.as_ref().len() as u64)?;
+    write_bytes_mut(blk_dev, abs_offset, data.as_mut(), scratch)
 }
 
 fn crc32(data: &[u8]) -> u32 {
@@ -445,111 +433,13 @@ fn crc32(data: &[u8]) -> u32 {
     hasher.finalize()
 }
 
-/// `MultiGptDevices` provides APIs for finding/reading partitions from multiple block devices.
-///
-/// The APIs use first match when searching for a partition. The caller should ensure that the
-/// partition of interest is unique among all devices. Otherwise it may lead to unexpected behavor.
-/// The intended use of this trait is for cases where a single collection of functional/bootable
-/// partitions are scattered on multiple GPT devices due to design constraint such as storage size,
-/// or GPT A/B etc. It should not be used to handle multiple GPT devices that contains multiple
-/// collections or copipes of functional/bootable partitions.
-pub trait MultiGptDevices {
-    /// Calls closure `f` for each `AsBlockDevice` object until reaching end or `f` returns true.
-    fn for_each_until(&mut self, f: &mut dyn FnMut(&mut dyn AsBlockDevice) -> bool);
-
-    /// Calls AsBlockDevice::sync_gpt() for all block devices.
-    fn sync_gpt(&mut self) -> Result<()> {
-        let mut res: Result<()> = Ok(());
-        self.for_each_until(&mut |v| {
-            res = v.sync_gpt().map(|_| ());
-            res.is_err()
-        });
-        res
-    }
-
-    /// Checks if there is exactly one partition with the given name among all devices.
-    fn check_unique(&mut self, part: &str) -> bool {
-        let mut count = 0usize;
-        self.for_each_until(&mut |v| {
-            count += (|| -> Result<bool> { Ok(v.gpt()?.find_partition(part)?.is_some()) })()
-                .unwrap_or(false) as usize;
-            count > 1
-        });
-        count == 1
-    }
-
-    /// Returns the block size and `GptEntry` for a partition on the first match.
-    fn find_partition(&mut self, part: &str) -> Result<(u64, GptEntry)> {
-        let mut res = Err(StorageError::NotExist);
-        self.for_each_until(&mut |v| {
-            res = (|| match v.gpt()?.find_partition(part)?.map(|v| *v) {
-                Some(p) => Ok((v.block_io().block_size(), p)),
-                _ => res,
-            })();
-            res.is_ok()
-        });
-        res
-    }
-
-    /// Returns the size of a partition on the first match.
-    fn partition_size(&mut self, part: &str) -> Result<u64> {
-        let (block_size, entry) = self.find_partition(part)?;
-        Ok(mul(block_size, entry.blocks()?)?)
-    }
-
-    /// Reads a GPT partition on the first match.
-    fn read_gpt_partition(&mut self, part_name: &str, offset: u64, out: &mut [u8]) -> Result<()> {
-        let mut res = Err(StorageError::NotExist);
-        self.for_each_until(&mut |v| {
-            res = v.read_gpt_partition(part_name, offset, out);
-            res.is_ok()
-        });
-        res
-    }
-
-    /// Writes a GPT partition on the first match.
-    fn write_gpt_partition(&mut self, part_name: &str, offset: u64, data: &[u8]) -> Result<()> {
-        let mut res = Err(StorageError::NotExist);
-        self.for_each_until(&mut |v| {
-            res = v.write_gpt_partition(part_name, offset, data);
-            res.is_ok()
-        });
-        res
-    }
-
-    /// Writes a GPT partition on the first match.
-    /// Optimization for mutable buffers.
-    fn write_gpt_partition_mut(
-        &mut self,
-        part_name: &str,
-        offset: u64,
-        data: &mut [u8],
-    ) -> Result<()> {
-        let mut res = Err(StorageError::NotExist);
-        self.for_each_until(&mut |v| {
-            res = v.write_gpt_partition_mut(part_name, offset, data);
-            res.is_ok()
-        });
-        res
-    }
-}
-
-impl<B: AsBlockDevice> MultiGptDevices for &mut [B] {
-    fn for_each_until(&mut self, f: &mut dyn FnMut(&mut dyn AsBlockDevice) -> bool) {
-        for ele in self.iter_mut() {
-            if f(ele) {
-                return;
-            }
-        }
-    }
-}
-
 #[cfg(test)]
-mod test {
+pub(crate) mod test {
     use super::*;
     use crate::test::TestBlockDevice;
+    use crate::AsBlockDevice;
 
-    fn gpt_block_device(max_entries: u64, data: &[u8]) -> TestBlockDevice {
+    pub(crate) fn gpt_block_device(max_entries: u64, data: &[u8]) -> TestBlockDevice {
         TestBlockDevice::new_with_data(512, 512, max_entries, data)
     }
 
@@ -557,12 +447,12 @@ mod test {
     fn test_new_from_buffer() {
         let disk = include_bytes!("../test/gpt_test_1.bin").to_vec();
         let mut dev = gpt_block_device(128, &disk);
-        let gpt = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
 
-        assert_eq!(gpt.entries().unwrap().len(), 2);
-        gpt.find_partition("boot_a").unwrap().unwrap();
-        gpt.find_partition("boot_b").unwrap().unwrap();
-        assert!(gpt.find_partition("boot_c").unwrap().is_none());
+        assert_eq!(dev.partition_iter().count(), 2);
+        dev.find_partition("boot_a").unwrap();
+        dev.find_partition("boot_b").unwrap();
+        assert!(dev.find_partition("boot_c").is_err());
     }
 
     #[test]
@@ -587,12 +477,12 @@ mod test {
 
         // Corrupt secondary.
         dev.io.storage[disk.len() - 512..].fill(0);
-        let gpt = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
 
-        assert_eq!(gpt.entries().unwrap().len(), 2);
-        gpt.find_partition("boot_a").unwrap().unwrap();
-        gpt.find_partition("boot_b").unwrap().unwrap();
-        assert!(gpt.find_partition("boot_c").unwrap().is_none());
+        assert_eq!(dev.partition_iter().count(), 2);
+        dev.find_partition("boot_a").unwrap();
+        dev.find_partition("boot_b").unwrap();
+        assert!(dev.find_partition("boot_c").is_err());
 
         // Check that secondary is restored
         assert_eq!(dev.io.storage, disk);
@@ -605,11 +495,11 @@ mod test {
 
         // Corrupt primary.
         dev.io.storage[512..1024].fill(0);
-        let gpt = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
 
-        assert_eq!(gpt.entries().unwrap().len(), 2);
-        gpt.find_partition("boot_a").unwrap().unwrap();
-        gpt.find_partition("boot_b").unwrap().unwrap();
+        assert_eq!(dev.partition_iter().count(), 2);
+        dev.find_partition("boot_a").unwrap();
+        dev.find_partition("boot_b").unwrap();
 
         // Check that primary is restored
         assert_eq!(dev.io.storage, disk);
@@ -619,7 +509,7 @@ mod test {
     fn test_good_gpt_no_repair_write() {
         let disk = include_bytes!("../test/gpt_test_1.bin").to_vec();
         let mut dev = gpt_block_device(128, &disk);
-        let _ = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
 
         assert_eq!(dev.io.num_writes, 0);
     }
@@ -628,8 +518,9 @@ mod test {
     fn test_load_gpt_incorrect_magic() {
         let disk = include_bytes!("../test/gpt_test_1.bin").to_vec();
         let mut dev = gpt_block_device(128, &disk);
-        let gpt = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
 
+        let gpt = dev.gpt();
         let primary_header = &mut gpt.primary_header[..to_usize(GPT_HEADER_SIZE).unwrap()];
         let gpt_header = GptHeader::from_bytes(primary_header);
         gpt_header.magic = 0x123456;
@@ -637,7 +528,7 @@ mod test {
         let primary_header = Vec::from(primary_header);
         dev.io.storage[512..512 + primary_header.len()].clone_from_slice(&primary_header);
 
-        let _ = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
 
         // Check that incorrect magic header is restored
         assert_eq!(dev.io.storage, disk);
@@ -654,8 +545,9 @@ mod test {
         // Create a header with non-max entries_count
         let disk = include_bytes!("../test/gpt_test_1.bin").to_vec();
         let mut dev = gpt_block_device(128, &disk);
-        let gpt = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
 
+        let gpt = dev.gpt();
         let primary_header = &mut gpt.primary_header[..to_usize(GPT_HEADER_SIZE).unwrap()];
         let gpt_header = GptHeader::from_bytes(primary_header);
         gpt_header.entries_count = 2;
@@ -669,11 +561,11 @@ mod test {
 
         // Corrupt secondary. Sync ok
         dev.io.storage[disk.len() - 512..].fill(0);
-        let _ = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
 
         // Corrup primary. Sync ok
         dev.io.storage[512..1024].fill(0);
-        let _ = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
     }
 
     #[test]
@@ -681,18 +573,17 @@ mod test {
         // Load a good GPT first.
         let disk = include_bytes!("../test/gpt_test_1.bin").to_vec();
         let mut dev = gpt_block_device(128, &disk);
-        let _ = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
         dev.io.storage[..64 * 1024].fill(0);
         // Load a bad GPT. Validate that the valid state is reset.
         assert!(dev.sync_gpt().is_err());
-        assert!(dev.gpt().unwrap().entries().is_err());
-        assert!(dev.gpt().unwrap().find_partition("").is_err());
+        assert!(dev.find_partition("").is_err());
     }
 
     #[test]
     fn test_gpt_read() {
         let mut dev = gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin"));
-        let _ = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
 
         let expect_boot_a = include_bytes!("../test/boot_a.bin");
         let expect_boot_b = include_bytes!("../test/boot_b.bin");
@@ -718,7 +609,7 @@ mod test {
     #[test]
     fn test_gpt_write() {
         let mut dev = gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin"));
-        let _ = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
 
         let mut expect_boot_a = include_bytes!("../test/boot_a.bin").to_vec();
         expect_boot_a.reverse();
@@ -768,7 +659,7 @@ mod test {
     #[test]
     fn test_gpt_rw_overflow() {
         let mut dev = gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin"));
-        let _ = dev.sync_gpt().unwrap();
+        dev.sync_gpt().unwrap();
 
         let mut boot_a = [0u8; include_bytes!("../test/boot_a.bin").len()];
         let mut boot_b = [0u8; include_bytes!("../test/boot_b.bin").len()];
@@ -781,90 +672,4 @@ mod test {
         assert!(dev.write_gpt_partition_mut("boot_b", 1, boot_b.as_mut_slice()).is_err());
         assert!(dev.write_gpt_partition("boot_b", 1, &boot_b).is_err());
     }
-
-    #[test]
-    fn test_multi_gpt_check_unique() {
-        let mut devs = &mut vec![
-            gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin")),
-            gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin")),
-            gpt_block_device(128, include_bytes!("../test/gpt_test_2.bin")),
-        ][..];
-
-        devs.sync_gpt().unwrap();
-        assert!(!devs.check_unique("boot_a"));
-        assert!(!devs.check_unique("boot_b"));
-        assert!(devs.check_unique("vendor_boot_a"));
-        assert!(devs.check_unique("vendor_boot_b"));
-    }
-
-    #[test]
-    fn test_multi_gpt_read() {
-        let mut devs = &mut vec![
-            gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin")),
-            gpt_block_device(128, include_bytes!("../test/gpt_test_2.bin")),
-        ][..];
-
-        devs.sync_gpt().unwrap();
-        assert_eq!(devs.partition_size("boot_a").unwrap(), 8 * 1024);
-        assert_eq!(devs.partition_size("boot_b").unwrap(), 12 * 1024);
-        assert_eq!(devs.partition_size("vendor_boot_a").unwrap(), 4 * 1024);
-        assert_eq!(devs.partition_size("vendor_boot_b").unwrap(), 6 * 1024);
-
-        let expect_boot_a = include_bytes!("../test/boot_a.bin");
-        let expect_boot_b = include_bytes!("../test/boot_b.bin");
-        let mut actual_boot_a = vec![0u8; expect_boot_a.len()];
-        let mut actual_boot_b = vec![0u8; expect_boot_b.len()];
-
-        devs.read_gpt_partition("boot_a", 0, &mut actual_boot_a).unwrap();
-        assert_eq!(expect_boot_a.to_vec(), actual_boot_a);
-        devs.read_gpt_partition("boot_b", 0, &mut actual_boot_b).unwrap();
-        assert_eq!(expect_boot_b.to_vec(), actual_boot_b);
-
-        let expect_vendor_boot_a = include_bytes!("../test/vendor_boot_a.bin");
-        let expect_vendor_boot_b = include_bytes!("../test/vendor_boot_b.bin");
-        let mut actual_vendor_boot_a = vec![0u8; expect_vendor_boot_a.len()];
-        let mut actual_vendor_boot_b = vec![0u8; expect_vendor_boot_b.len()];
-
-        devs.read_gpt_partition("vendor_boot_a", 0, &mut actual_vendor_boot_a).unwrap();
-        assert_eq!(expect_vendor_boot_a.to_vec(), actual_vendor_boot_a);
-        devs.read_gpt_partition("vendor_boot_b", 0, &mut actual_vendor_boot_b).unwrap();
-        assert_eq!(expect_vendor_boot_b.to_vec(), actual_vendor_boot_b);
-    }
-
-    #[test]
-    fn test_multi_gpt_write() {
-        let mut devs = &mut vec![
-            gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin")),
-            gpt_block_device(128, include_bytes!("../test/gpt_test_2.bin")),
-        ][..];
-        devs.sync_gpt().unwrap();
-
-        let mut expect_boot_a = include_bytes!("../test/boot_a.bin").to_vec();
-        expect_boot_a.reverse();
-        let mut expect_boot_b = include_bytes!("../test/boot_b.bin").to_vec();
-        expect_boot_b.reverse();
-        let mut actual_boot_a = vec![0u8; expect_boot_a.len()];
-        let mut actual_boot_b = vec![0u8; expect_boot_b.len()];
-
-        devs.write_gpt_partition_mut("boot_a", 0, expect_boot_a.as_mut_slice()).unwrap();
-        devs.read_gpt_partition("boot_a", 0, &mut actual_boot_a).unwrap();
-        assert_eq!(expect_boot_a.to_vec(), actual_boot_a);
-        devs.write_gpt_partition_mut("boot_b", 0, expect_boot_b.as_mut_slice()).unwrap();
-        devs.read_gpt_partition("boot_b", 0, &mut actual_boot_b).unwrap();
-        assert_eq!(expect_boot_b.to_vec(), actual_boot_b);
-
-        let mut expect_vendor_boot_a = include_bytes!("../test/vendor_boot_a.bin").to_vec();
-        expect_vendor_boot_a.reverse();
-        let mut expect_vendor_boot_b = include_bytes!("../test/vendor_boot_b.bin").to_vec();
-        expect_vendor_boot_b.reverse();
-        let mut actual_vendor_boot_a = vec![0u8; expect_vendor_boot_a.len()];
-        let mut actual_vendor_boot_b = vec![0u8; expect_vendor_boot_b.len()];
-
-        devs.write_gpt_partition_mut("boot_a", 0, expect_vendor_boot_a.as_mut_slice()).unwrap();
-        devs.read_gpt_partition("boot_a", 0, &mut actual_vendor_boot_a).unwrap();
-        assert_eq!(expect_vendor_boot_a.to_vec(), actual_vendor_boot_a);
-        devs.write_gpt_partition_mut("boot_b", 0, expect_vendor_boot_b.as_mut_slice()).unwrap();
-        devs.read_gpt_partition("boot_b", 0, &mut actual_vendor_boot_b).unwrap();
-        assert_eq!(expect_vendor_boot_b.to_vec(), actual_vendor_boot_b);
-    }
 }
diff --git a/gbl/libstorage/src/lib.rs b/gbl/libstorage/src/lib.rs
index f885962..408b86f 100644
--- a/gbl/libstorage/src/lib.rs
+++ b/gbl/libstorage/src/lib.rs
@@ -81,7 +81,7 @@
 //! // Sync GPT
 //! let _ = ram_block_dev.sync_gpt();
 //! // Access GPT entries
-//! let _ = ram_block_dev.gpt();
+//! let _ = ram_block_dev.find_partition("some-partition");
 //! // Read/Write GPT partitions with arbitrary offset, size, buffer
 //! let _ = ram_block_dev.read_gpt_partition("partition", 4321, &mut out[..]);
 //! let _ = ram_block_dev.write_gpt_partition_mut("partition", 8765, data.as_mut_slice());
@@ -95,8 +95,8 @@
 //! }
 //!
 //! impl AsBlockDevice for OwnedBlockDevice {
-//!     fn get(&mut self) -> (&mut dyn BlockIo, &mut [u8], u64) {
-//!         (&mut self.io, &mut self.scratch[..], MAX_GPT_ENTRIES)
+//!     fn with(&mut self, f: &mut dyn FnMut(&mut dyn BlockIo, &mut [u8], u64)) {
+//!         f(&mut self.io, &mut self.scratch[..], MAX_GPT_ENTRIES)
 //!     }
 //! }
 //!
@@ -110,9 +110,11 @@ use core::cmp::min;
 
 // Selective export of submodule types.
 mod gpt;
-pub use gpt::Gpt;
+use gpt::Gpt;
 pub use gpt::GptEntry;
-pub use gpt::MultiGptDevices;
+
+mod multi_blocks;
+pub use multi_blocks::{with_id, AsMultiBlockDevices};
 
 /// The type of Result used in this library.
 pub type Result<T> = core::result::Result<T, StorageError>;
@@ -123,10 +125,13 @@ pub enum StorageError {
     ArithmeticOverflow,
     OutOfRange,
     ScratchTooSmall,
+    BlockDeviceNotFound,
     BlockIoError,
+    BlockIoNotProvided,
     InvalidInput,
     NoValidGpt,
     NotExist,
+    PartitionNotUnique,
     U64toUSizeOverflow,
 }
 
@@ -136,10 +141,13 @@ impl Into<&'static str> for StorageError {
             StorageError::ArithmeticOverflow => "Arithmetic overflow",
             StorageError::OutOfRange => "Out of range",
             StorageError::ScratchTooSmall => "Not enough scratch buffer",
+            StorageError::BlockDeviceNotFound => "Block device not found",
             StorageError::BlockIoError => "Block IO error",
+            StorageError::BlockIoNotProvided => "Block IO is not provided",
             StorageError::InvalidInput => "Invalid input",
             StorageError::NoValidGpt => "GPT not found",
             StorageError::NotExist => "The specified partition could not be found",
+            StorageError::PartitionNotUnique => "Partition is found on multiple block devices",
             StorageError::U64toUSizeOverflow => "u64 to usize fails",
         }
     }
@@ -195,10 +203,32 @@ pub trait BlockIo {
     fn write_blocks(&mut self, blk_offset: u64, data: &[u8]) -> bool;
 }
 
+/// `GptEntryIterator` is returned by `AsBlockDevice::partition_iter()` and can be used to iterate
+/// all GPT partition entries.
+pub struct GptEntryIterator<'a> {
+    dev: &'a mut dyn AsBlockDevice,
+    idx: usize,
+}
+
+impl Iterator for GptEntryIterator<'_> {
+    type Item = GptEntry;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let res =
+            with_partitioned_scratch(self.dev, |_, _, gpt_buffer, _| -> Result<Option<GptEntry>> {
+                Ok(Gpt::from_existing(gpt_buffer)?.entries()?.get(self.idx).map(|v| *v))
+            })
+            .ok()?
+            .ok()??;
+        self.idx += 1;
+        Some(res)
+    }
+}
+
 /// `AsBlockDevice` provides APIs for reading raw block content and GPT partitions with
 /// arbirary offset, size and input/output buffer.
 pub trait AsBlockDevice {
-    /// Returns a tuple containing the following information:
+    /// Runs the provided closure `f` with the following parameters:
     ///
     ///   1. An implementation of block IO `&mut dyn BlockIo`.
     ///   2. A scratch buffer `&mut [u8]`.
@@ -220,8 +250,7 @@ pub trait AsBlockDevice {
     ///   `Self::write_gpt_partition()`, and `Self::write_gpt_partition_mut()`
     ///   will look up partition entries from the cached GPT header.
     ///   Thus callers should make sure to always return the same scratch buffer and avoid
-    ///   modifying its content. `Self::gpt()` returns a `Gpt` type that provides APIs for
-    ///   accessing the content of the GPT header.
+    ///   modifying its content.
     ///
     /// * A smaller value of maximum allowed GPT entries gives smaller required scratch buffer
     ///   size. However if the `entries_count` field in the GPT header is greater than this value,
@@ -229,69 +258,29 @@ pub trait AsBlockDevice {
     ///   max value 128 regardless of the actual number of partition entries used. Thus unless you
     ///   have full control of GPT generation in your entire system where you can always ensure a
     ///   smaller bound on it, it is recommended to always return 128.
-    fn get(&mut self) -> (&mut dyn BlockIo, &mut [u8], u64);
+    fn with(&mut self, f: &mut dyn FnMut(&mut dyn BlockIo, &mut [u8], u64));
 
-    /// Read data from the block device.
-    ///
-    /// # Args
-    ///
-    /// * `offset`: Offset in number of bytes.
-    ///
-    /// * `out`: Buffer to store the read data.
-    ///
-    /// * Returns success when exactly `out.len()` number of bytes are read.
-    fn read(&mut self, offset: u64, out: &mut [u8]) -> Result<()> {
-        let (io, alignment, _, _) = get_with_partitioned_scratch(self)?;
-        read(io, offset, out, alignment)
+    // Returns the block size of the underlying `BlockIo`
+    fn block_size(&mut self) -> Result<u64> {
+        with_partitioned_scratch(self, |io, _, _, _| io.block_size())
     }
 
-    /// Parse and sync GPT from a block device.
-    ///
-    /// The API validates and restores primary/secondary GPT header.
-    ///
-    /// # Returns
-    ///
-    /// Returns success if GPT is loaded/restored successfully.
-    fn sync_gpt(&mut self) -> Result<Gpt> {
-        let (io, alignment_scratch, gpt_buffer, max_entries) = get_with_partitioned_scratch(self)?;
-        gpt::gpt_sync(io, &mut Gpt::new_from_buffer(max_entries, gpt_buffer)?, alignment_scratch)?;
-        self.gpt()
+    // Returns the number of blocks of the underlying `BlockIo`
+    fn num_blocks(&mut self) -> Result<u64> {
+        with_partitioned_scratch(self, |io, _, _, _| io.num_blocks())
     }
 
-    /// Read a GPT partition on a block device
+    /// Read data from the block device.
     ///
     /// # Args
     ///
-    /// * `part_name`: Name of the partition.
-    ///
-    /// * `offset`: Offset in number of bytes into the partition.
+    /// * `offset`: Offset in number of bytes.
     ///
     /// * `out`: Buffer to store the read data.
     ///
-    /// # Returns
-    ///
-    /// Returns success when exactly `out.len()` of bytes are read successfully.
-    fn read_gpt_partition(&mut self, part_name: &str, offset: u64, out: &mut [u8]) -> Result<()> {
-        let (io, alignment_scratch, gpt_buffer, _) = get_with_partitioned_scratch(self)?;
-        gpt::read_gpt_partition(
-            io,
-            &mut Gpt::from_existing(gpt_buffer)?,
-            part_name,
-            offset,
-            out,
-            alignment_scratch,
-        )
-    }
-
-    /// Returns the GPT header.
-    fn gpt(&mut self) -> Result<Gpt> {
-        let (_, _, gpt_buffer, _) = get_with_partitioned_scratch(self)?;
-        Gpt::from_existing(gpt_buffer)
-    }
-
-    /// Returns the `BlockIo` implementation.
-    fn block_io(&mut self) -> &mut dyn BlockIo {
-        self.get().0
+    /// * Returns success when exactly `out.len()` number of bytes are read.
+    fn read(&mut self, offset: u64, out: &mut [u8]) -> Result<()> {
+        with_partitioned_scratch(self, |io, alignment, _, _| read(io, offset, out, alignment))?
     }
 
     /// Write data to the device.
@@ -308,8 +297,9 @@ pub trait AsBlockDevice {
     ///
     /// * Returns success when exactly `data.len()` number of bytes are written.
     fn write(&mut self, offset: u64, data: &[u8]) -> Result<()> {
-        let (io, alignment_scratch, _, _) = get_with_partitioned_scratch(self)?;
-        write_bytes(io, offset, data, alignment_scratch)
+        with_partitioned_scratch(self, |io, alignment_scratch, _, _| {
+            write_bytes(io, offset, data, alignment_scratch)
+        })?
     }
 
     /// Write data to the device.
@@ -326,8 +316,82 @@ pub trait AsBlockDevice {
     ///
     /// * Returns success when exactly `data.len()` number of bytes are written.
     fn write_mut(&mut self, offset: u64, data: &mut [u8]) -> Result<()> {
-        let (io, alignment_scratch, _, _) = get_with_partitioned_scratch(self)?;
-        write_bytes_mut(io, offset, data, alignment_scratch)
+        with_partitioned_scratch(self, |io, alignment_scratch, _, _| {
+            write_bytes_mut(io, offset, data, alignment_scratch)
+        })?
+    }
+
+    /// Parse and sync GPT from a block device.
+    ///
+    /// The API validates and restores primary/secondary GPT header.
+    ///
+    /// # Returns
+    ///
+    /// Returns success if GPT is loaded/restored successfully.
+    fn sync_gpt(&mut self) -> Result<()> {
+        with_partitioned_scratch(self, |io, alignment_scratch, gpt_buffer, max_entries| {
+            gpt::gpt_sync(
+                io,
+                &mut Gpt::new_from_buffer(max_entries, gpt_buffer)?,
+                alignment_scratch,
+            )
+        })?
+    }
+
+    /// Returns an iterator to GPT partition entries.
+    fn partition_iter(&mut self) -> GptEntryIterator
+    where
+        Self: Sized,
+    {
+        GptEntryIterator { dev: self, idx: 0 }
+    }
+
+    /// Returns the `GptEntry` for a partition.
+    ///
+    /// # Args
+    ///
+    /// * `part`: Name of the partition.
+    fn find_partition(&mut self, part: &str) -> Result<GptEntry> {
+        with_partitioned_scratch(self, |_, _, gpt_buffer, _| {
+            Ok(Gpt::from_existing(gpt_buffer)?.find_partition(part).map(|v| *v)?)
+        })?
+    }
+
+    /// Returns the size of a partition.
+    ///
+    /// # Args
+    ///
+    /// * `part`: Name of the partition.
+    fn partition_size(&mut self, part: &str) -> Result<u64> {
+        let blk_sz = self.block_size()?;
+        let entry = self.find_partition(part)?;
+        mul(blk_sz, entry.blocks()?)
+    }
+
+    /// Read a GPT partition on a block device
+    ///
+    /// # Args
+    ///
+    /// * `part_name`: Name of the partition.
+    ///
+    /// * `offset`: Offset in number of bytes into the partition.
+    ///
+    /// * `out`: Buffer to store the read data.
+    ///
+    /// # Returns
+    ///
+    /// Returns success when exactly `out.len()` of bytes are read successfully.
+    fn read_gpt_partition(&mut self, part_name: &str, offset: u64, out: &mut [u8]) -> Result<()> {
+        with_partitioned_scratch(self, |io, alignment_scratch, gpt_buffer, _| {
+            gpt::read_gpt_partition(
+                io,
+                &mut Gpt::from_existing(gpt_buffer)?,
+                part_name,
+                offset,
+                out,
+                alignment_scratch,
+            )
+        })?
     }
 
     /// Write a GPT partition on a block device
@@ -344,15 +408,16 @@ pub trait AsBlockDevice {
     ///
     /// Returns success when exactly `data.len()` of bytes are written successfully.
     fn write_gpt_partition(&mut self, part_name: &str, offset: u64, data: &[u8]) -> Result<()> {
-        let (io, alignment_scratch, gpt_buffer, _) = get_with_partitioned_scratch(self)?;
-        gpt::write_gpt_partition(
-            io,
-            &mut Gpt::from_existing(gpt_buffer)?,
-            part_name,
-            offset,
-            data,
-            alignment_scratch,
-        )
+        with_partitioned_scratch(self, |io, alignment_scratch, gpt_buffer, _| {
+            gpt::write_gpt_partition(
+                io,
+                &mut Gpt::from_existing(gpt_buffer)?,
+                part_name,
+                offset,
+                data,
+                alignment_scratch,
+            )
+        })?
     }
 
     /// Write a GPT partition on a block device.
@@ -377,21 +442,22 @@ pub trait AsBlockDevice {
         offset: u64,
         data: &mut [u8],
     ) -> Result<()> {
-        let (io, alignment_scratch, gpt_buffer, _) = get_with_partitioned_scratch(self)?;
-        gpt::write_gpt_partition_mut(
-            io,
-            &mut Gpt::from_existing(gpt_buffer)?,
-            part_name,
-            offset,
-            data.into(),
-            alignment_scratch,
-        )
+        with_partitioned_scratch(self, |io, alignment_scratch, gpt_buffer, _| {
+            gpt::write_gpt_partition_mut(
+                io,
+                &mut Gpt::from_existing(gpt_buffer)?,
+                part_name,
+                offset,
+                data.into(),
+                alignment_scratch,
+            )
+        })?
     }
 }
 
 impl<T: ?Sized + AsBlockDevice> AsBlockDevice for &mut T {
-    fn get(&mut self) -> (&mut dyn BlockIo, &mut [u8], u64) {
-        (*self).get()
+    fn with(&mut self, f: &mut dyn FnMut(&mut dyn BlockIo, &mut [u8], u64)) {
+        (*self).with(f)
     }
 }
 
@@ -409,8 +475,8 @@ impl<'a, 'b> BlockDevice<'a, 'b> {
 }
 
 impl AsBlockDevice for BlockDevice<'_, '_> {
-    fn get(&mut self) -> (&mut dyn BlockIo, &mut [u8], u64) {
-        (self.io, self.scratch, self.max_gpt_entries)
+    fn with(&mut self, f: &mut dyn FnMut(&mut dyn BlockIo, &mut [u8], u64)) {
+        f(self.io, self.scratch, self.max_gpt_entries)
     }
 }
 
@@ -427,17 +493,26 @@ pub fn required_scratch_size(
     alignment_size.checked_add(gpt_buffer_size).ok_or(StorageError::ArithmeticOverflow)
 }
 
-/// The helper calls `AsBlockDevice:get()`, partitions the scratch buffer into alignment scratch and
-/// GPT buffers, and returns them.
-fn get_with_partitioned_scratch(
+/// A helper that wraps `AsBlockDevice::with` and additionally partitions the scratch buffer into
+/// alignment scratch and GPT buffers.
+pub(crate) fn with_partitioned_scratch<F, R>(
     dev: &mut (impl AsBlockDevice + ?Sized),
-) -> Result<(&mut dyn BlockIo, &mut [u8], &mut [u8], u64)> {
-    let (io, scratch, max_entries) = dev.get();
-    if scratch.len() < required_scratch_size(io, max_entries)? {
-        return Err(StorageError::ScratchTooSmall);
-    }
-    let (alignment, gpt) = scratch.split_at_mut(alignment_scratch_size(io)?);
-    Ok((io, alignment, gpt, max_entries))
+    mut f: F,
+) -> Result<R>
+where
+    F: FnMut(&mut dyn BlockIo, &mut [u8], &mut [u8], u64) -> R,
+{
+    let mut res: Result<R> = Err(StorageError::BlockIoNotProvided);
+    dev.with(&mut |io, scratch, max_entries| {
+        res = (|| {
+            if scratch.len() < required_scratch_size(io, max_entries)? {
+                return Err(StorageError::ScratchTooSmall);
+            }
+            let (alignment, gpt) = scratch.split_at_mut(alignment_scratch_size(io)?);
+            Ok(f(io, alignment, gpt, max_entries))
+        })();
+    });
+    res
 }
 
 /// Add two u64 integers and check overflow
@@ -962,11 +1037,17 @@ mod test {
             let storage: Vec<u8> = (0..storage_size).map(|x| x as u8).collect();
             Self::new_with_data(alignment, block_size, max_gpt_entries, &storage)
         }
+
+        /// Extract the internal Gpt structure for examination by test.
+        pub fn gpt(&mut self) -> Gpt {
+            let (_, gpt) = self.scratch.split_at_mut(alignment_scratch_size(&mut self.io).unwrap());
+            Gpt::from_existing(gpt).unwrap()
+        }
     }
 
     impl AsBlockDevice for TestBlockDevice {
-        fn get(&mut self) -> (&mut dyn BlockIo, &mut [u8], u64) {
-            (&mut self.io, &mut self.scratch[..], self.max_gpt_entries)
+        fn with(&mut self, f: &mut dyn FnMut(&mut dyn BlockIo, &mut [u8], u64)) {
+            f(&mut self.io, &mut self.scratch[..], self.max_gpt_entries)
         }
     }
 
@@ -1365,7 +1446,7 @@ mod test {
     #[test]
     fn test_no_alignment_require_zero_size_scratch() {
         let mut blk = TestBlockDevice::new(1, 1, 0, 1);
-        assert_eq!(required_scratch_size(blk.block_io(), 0).unwrap(), 0);
+        assert_eq!(required_scratch_size(&mut blk.io, 0).unwrap(), 0);
     }
 
     #[test]
diff --git a/gbl/libstorage/src/multi_blocks.rs b/gbl/libstorage/src/multi_blocks.rs
new file mode 100644
index 0000000..7df42a4
--- /dev/null
+++ b/gbl/libstorage/src/multi_blocks.rs
@@ -0,0 +1,333 @@
+// Copyright 2024, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::{mul, AsBlockDevice, BlockIo, GptEntry, Result, StorageError};
+
+/// `AsMultiBlockDevices` provides APIs for finding/reading/writing raw data or GPT partitions from
+/// multiple block devices.
+pub trait AsMultiBlockDevices {
+    /// Calls closure `f` for each `AsBlockDevice` object and its unique `id` until reaching end or
+    /// `f` returns true.
+    fn for_each_until(&mut self, f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64) -> bool);
+
+    /// Gets the block device with the given id.
+    fn get(&mut self, id: u64) -> Result<SelectedBlockDevice>
+    where
+        Self: Sized,
+    {
+        with_id(self, id, |_| {})?;
+        Ok(SelectedBlockDevice { devs: self, id: id })
+    }
+
+    /// Syncs gpt for all block devices. Caller provides a callback for handling sync error for
+    /// each block device.
+    fn sync_gpt_all(&mut self, f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64, StorageError)) {
+        self.for_each_until(&mut |v, id| {
+            match v.sync_gpt() {
+                Err(e) => f(v, id, e),
+                _ => {}
+            }
+            false
+        });
+    }
+
+    /// Checks that a partition exists and is unique among all block devices with GPT.
+    fn check_part(&mut self, part: &str) -> Result<()> {
+        let mut count = 0usize;
+        self.for_each_until(&mut |v, _| {
+            count += (|| -> Result<bool> { Ok(v.find_partition(part).is_ok()) })().unwrap_or(false)
+                as usize;
+            count > 1
+        });
+        match count {
+            1 => Ok(()),
+            0 => Err(StorageError::NotExist),
+            _ => Err(StorageError::PartitionNotUnique),
+        }
+    }
+
+    /// Returns the block size and `GptEntry` for a partition.
+    ///
+    /// Returns Ok(()) if the partition is found and unique among all block devices.
+    fn find_partition(&mut self, part: &str) -> Result<(u64, GptEntry)> {
+        self.check_part(part)?;
+        until_ok(self, |dev, _| {
+            let blk_sz = dev.block_size()?;
+            dev.find_partition(part).map(|v| (blk_sz, v))
+        })
+    }
+
+    /// Returns the size of a partition.
+    ///
+    /// Returns Ok(()) if the partition is found and unique among all block devices.
+    fn partition_size(&mut self, part: &str) -> Result<u64> {
+        let (block_size, entry) = self.find_partition(part)?;
+        Ok(mul(block_size, entry.blocks()?)?)
+    }
+
+    /// Reads a GPT partition.
+    ///
+    /// Returns Ok(()) if the partition is unique among all block devices and read is successful.
+    fn read_gpt_partition(&mut self, part_name: &str, offset: u64, out: &mut [u8]) -> Result<()> {
+        self.check_part(part_name)?;
+        until_ok(self, |dev, _| dev.read_gpt_partition(part_name, offset, out))
+    }
+
+    /// Writes a GPT partition with mutable input buffer.
+    ///
+    /// Returns Ok(()) if the partition is unique among all block devices and write is successful.
+    fn write_gpt_partition_mut(
+        &mut self,
+
+        part_name: &str,
+        offset: u64,
+        data: &mut [u8],
+    ) -> Result<()> {
+        self.check_part(part_name)?;
+        until_ok(self, |dev, _| dev.write_gpt_partition_mut(part_name, offset, &mut data[..]))
+    }
+
+    /// Writes a GPT partition with const input buffer.
+    ///
+    /// Returns Ok(()) if the partition is unique among all block devices and write is successful.
+    fn write_gpt_partition(&mut self, part_name: &str, offset: u64, data: &mut [u8]) -> Result<()> {
+        self.check_part(part_name)?;
+        until_ok(self, |dev, _| dev.write_gpt_partition(part_name, offset, &mut data[..]))
+    }
+}
+
+/// Iterates and runs a closure on each block device until `Ok(R)` is returned.
+fn until_ok<F, R>(devs: &mut (impl AsMultiBlockDevices + ?Sized), mut f: F) -> Result<R>
+where
+    F: FnMut(&mut dyn AsBlockDevice, u64) -> Result<R>,
+{
+    let mut res: Result<R> = Err(StorageError::BlockDeviceNotFound);
+    devs.for_each_until(&mut |v, id| {
+        res = f(v, id);
+        res.is_ok()
+    });
+    res
+}
+
+/// Finds the first block device with the given ID and runs a closure with it.
+pub fn with_id<F, R>(
+    devs: &mut (impl AsMultiBlockDevices + ?Sized),
+    dev_id: u64,
+    mut f: F,
+) -> Result<R>
+where
+    F: FnMut(&mut dyn AsBlockDevice) -> R,
+{
+    until_ok(devs, |dev, id| match dev_id == id {
+        true => Ok(f(dev)),
+        _ => Err(StorageError::BlockDeviceNotFound),
+    })
+}
+
+/// `SelectedBlockDevice` is returned by `AsMultiBlockDevices::get()` and provides access to
+/// the `AsBlockDevice` object of the given id.
+pub struct SelectedBlockDevice<'a> {
+    devs: &'a mut dyn AsMultiBlockDevices,
+    id: u64,
+}
+
+impl AsBlockDevice for SelectedBlockDevice<'_> {
+    fn with(&mut self, f: &mut dyn FnMut(&mut dyn BlockIo, &mut [u8], u64)) {
+        let _ = with_id(self.devs, self.id, |dev| dev.with(f));
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::gpt::test::gpt_block_device;
+    use crate::test::TestBlockDevice;
+    use crate::AsMultiBlockDevices;
+
+    impl<B: AsBlockDevice> AsMultiBlockDevices for Vec<B> {
+        fn for_each_until(&mut self, f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64) -> bool) {
+            for (idx, ele) in self.iter_mut().enumerate() {
+                if f(ele, u64::try_from(idx).unwrap()) {
+                    return;
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn test_get() {
+        let devs = &mut vec![
+            gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin")),
+            gpt_block_device(128, include_bytes!("../test/gpt_test_2.bin")),
+        ];
+        devs.sync_gpt_all(&mut |_, _, _| panic!("GPT sync failed"));
+        devs.get(0).unwrap();
+        devs.get(1).unwrap();
+        assert!(devs.get(2).is_err());
+    }
+
+    #[test]
+    fn test_multi_block_read() {
+        let off = 512; // Randomly selected offset.
+        let blk_0 = include_bytes!("../test/gpt_test_1.bin");
+        let blk_1 = include_bytes!("../test/gpt_test_2.bin");
+        let mut devs = vec![gpt_block_device(128, blk_0), gpt_block_device(128, blk_1)];
+
+        let mut out = vec![0u8; blk_0[off..].len()];
+        devs.get(0).unwrap().read(u64::try_from(off).unwrap(), &mut out[..]).unwrap();
+        assert_eq!(out, blk_0[off..]);
+
+        let mut out = vec![0u8; blk_1[off..].len()];
+        devs.get(1).unwrap().read(u64::try_from(off).unwrap(), &mut out[..]).unwrap();
+        assert_eq!(out, blk_1[off..]);
+    }
+
+    #[test]
+    fn test_multi_block_write() {
+        let off = 512; // Randomly selected offset.
+        let mut blk_0 = Vec::from(include_bytes!("../test/gpt_test_1.bin"));
+        let mut blk_1 = Vec::from(include_bytes!("../test/gpt_test_2.bin"));
+        let mut devs = vec![
+            gpt_block_device(128, &vec![0u8; blk_0.len()][..]),
+            gpt_block_device(128, &vec![0u8; blk_1.len()][..]),
+        ];
+
+        devs.get(0).unwrap().write(u64::try_from(off).unwrap(), &mut blk_0[off..]).unwrap();
+        assert_eq!(blk_0[off..], devs[0].io.storage[off..]);
+
+        devs.get(1).unwrap().write(u64::try_from(off).unwrap(), &mut blk_1[off..]).unwrap();
+        assert_eq!(blk_1[off..], devs[1].io.storage[off..]);
+    }
+
+    #[test]
+    fn test_multi_block_gpt_partition_size() {
+        let devs = &mut vec![
+            gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin")),
+            gpt_block_device(128, include_bytes!("../test/gpt_test_2.bin")),
+        ];
+        devs.sync_gpt_all(&mut |_, _, _| panic!("GPT sync failed"));
+
+        assert_eq!(devs.partition_size("boot_a").unwrap(), 8 * 1024);
+        assert_eq!(devs.get(0).unwrap().partition_size("boot_a").unwrap(), 8 * 1024);
+
+        assert_eq!(devs.partition_size("boot_b").unwrap(), 12 * 1024);
+        assert_eq!(devs.get(0).unwrap().partition_size("boot_b").unwrap(), 12 * 1024);
+
+        assert_eq!(devs.partition_size("vendor_boot_a").unwrap(), 4 * 1024);
+        assert_eq!(devs.get(1).unwrap().partition_size("vendor_boot_a").unwrap(), 4 * 1024);
+
+        assert_eq!(devs.partition_size("vendor_boot_b").unwrap(), 6 * 1024);
+        assert_eq!(devs.get(1).unwrap().partition_size("vendor_boot_b").unwrap(), 6 * 1024);
+    }
+
+    /// A test helper for `AsMultiBlockDevices::read_gpt_partition`
+    /// It verifies that data read partition `part` at offset `off` is the same as
+    /// `expected[off..]`.
+    fn check_read_partition(
+        devs: &mut Vec<TestBlockDevice>,
+        part: &str,
+        off: u64,
+        part_data: &[u8],
+    ) {
+        let expected = &part_data[off.try_into().unwrap()..];
+        let mut out = vec![0u8; expected.len()];
+        devs.read_gpt_partition(part, off, &mut out[..]).unwrap();
+        assert_eq!(out, expected.to_vec());
+    }
+
+    #[test]
+    fn test_multi_block_gpt_read() {
+        let off = 512u64; // Randomly selected offset.
+
+        let mut devs = vec![
+            gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin")),
+            gpt_block_device(128, include_bytes!("../test/gpt_test_2.bin")),
+        ];
+        devs.sync_gpt_all(&mut |_, _, _| panic!("GPT sync failed"));
+
+        let expect_boot_a = include_bytes!("../test/boot_a.bin");
+        let expect_boot_b = include_bytes!("../test/boot_b.bin");
+
+        check_read_partition(&mut devs, "boot_a", off, expect_boot_a);
+        check_read_partition(&mut devs, "boot_b", off, expect_boot_b);
+
+        let expect_vendor_boot_a = include_bytes!("../test/vendor_boot_a.bin");
+        let expect_vendor_boot_b = include_bytes!("../test/vendor_boot_b.bin");
+
+        check_read_partition(&mut devs, "vendor_boot_a", off, expect_vendor_boot_a);
+        check_read_partition(&mut devs, "vendor_boot_b", off, expect_vendor_boot_b);
+    }
+
+    /// A test helper for `AsMultiBlockDevices::write_gpt_partition`
+    /// It verifies that `data[off..]` is correctly written to partition `part` at offset `off`.
+    fn check_write_partition(
+        devs: &mut Vec<TestBlockDevice>,
+        part: &str,
+        off: u64,
+        data: &mut [u8],
+    ) {
+        let to_write = &mut data[off.try_into().unwrap()..];
+
+        let mut out = vec![0u8; to_write.len()];
+        devs.write_gpt_partition_mut(part, off, to_write).unwrap();
+        devs.read_gpt_partition(part, off, &mut out[..]).unwrap();
+        assert_eq!(out, to_write.to_vec());
+
+        to_write.reverse();
+        devs.write_gpt_partition(part, off, to_write).unwrap();
+        devs.read_gpt_partition(part, off, &mut out[..]).unwrap();
+        assert_eq!(out, to_write.to_vec());
+    }
+
+    #[test]
+    fn test_multi_block_gpt_write() {
+        let off = 512u64; // Randomly selected offset.
+
+        let mut devs = vec![
+            gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin")),
+            gpt_block_device(128, include_bytes!("../test/gpt_test_2.bin")),
+        ];
+        devs.sync_gpt_all(&mut |_, _, _| panic!("GPT sync failed"));
+
+        let expect_boot_a = &mut include_bytes!("../test/boot_a.bin").to_vec();
+        let expect_boot_b = &mut include_bytes!("../test/boot_b.bin").to_vec();
+
+        expect_boot_a.reverse();
+        expect_boot_b.reverse();
+        check_write_partition(&mut devs, "boot_a", off, expect_boot_a);
+        check_write_partition(&mut devs, "boot_b", off, expect_boot_b);
+
+        let expect_vendor_boot_a = &mut include_bytes!("../test/vendor_boot_a.bin").to_vec();
+        let expect_vendor_boot_b = &mut include_bytes!("../test/vendor_boot_b.bin").to_vec();
+
+        expect_boot_a.reverse();
+        expect_boot_b.reverse();
+        check_write_partition(&mut devs, "vendor_boot_a", off, expect_vendor_boot_a);
+        check_write_partition(&mut devs, "vendor_boot_b", off, expect_vendor_boot_b);
+    }
+
+    #[test]
+    fn test_none_block_id_fail_with_non_unique_partition() {
+        let mut devs = vec![
+            gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin")),
+            gpt_block_device(128, include_bytes!("../test/gpt_test_1.bin")),
+        ];
+        devs.sync_gpt_all(&mut |_, _, _| panic!("GPT sync failed"));
+        assert!(devs.read_gpt_partition("boot_a", 0, &mut []).is_err());
+        assert!(devs.write_gpt_partition_mut("boot_a", 0, &mut []).is_err());
+        assert!(devs.write_gpt_partition("boot_a", 0, &mut []).is_err());
+        assert!(devs.find_partition("boot_a").is_err());
+        assert!(devs.partition_size("boot_a").is_err());
+    }
+}