Merge "Add GBL interfaces for block device, console, boot" into main

23 files changed, 568 insertions, 221 deletions

diff --git a/gbl/efi/src/error.rs b/gbl/efi/src/error.rs
index 601bf04..0d56949 100644
--- a/gbl/efi/src/error.rs
+++ b/gbl/efi/src/error.rs
@@ -20,6 +20,7 @@ use efi::EfiError;
 use fastboot::TransportError;
 use fdt::FdtError;
 use gbl_storage::StorageError;
+use libgbl::composite_enum;
 use misc::BcbError;
 use smoltcp::socket::tcp::{ListenError, RecvError, SendError};
 use zbi::ZbiError;
@@ -40,74 +41,6 @@ pub enum EfiAppError {
     Unsupported,
 }
 
-/// A convenient macro for declaring a composite enum type that simply wraps other types as
-/// entries. It auto-generate `From<...>` implementation for each entry type. i.e.:
-///
-/// ```rust
-///   composite_enum! {
-///       pub enum MyEnum {
-///           Usize(usize),
-///           I64(i64),
-///       }
-///   }
-/// ```
-///
-/// expands to
-///
-/// ```rust
-///   pub enum MyEnum {
-///       Usize(usize),
-///       I64(i64),
-///   }
-///
-///   impl From<usize> for MyEnum {
-///       fn from(entry: usize) -> MyEnum {
-///           MyEnum::Usize(entry)
-///       }
-///   }
-///
-///   impl From<i64> for MyEnum {
-///       fn from(entry: i64) -> MyEnum {
-///           MyEnum::I64(entry)
-///       }
-///   }
-/// ```
-///
-/// The macro assumes that each entry is a different type.
-macro_rules! composite_enum {
-    // Top level macro entry. Match enum declaration code and call recursively for `From<>`
-    // generation.
-    (
-        $(#[$outer:meta])*
-        $vis:vis enum $EnumName:ident {
-            $($entry:ident($entry_type:ty)),*
-            $(,)*
-        }
-    ) => {
-        // Copy over enum declaration as it is.
-        $(#[$outer])*
-        $vis enum $EnumName {
-            $($entry($entry_type)),*
-        }
-
-        // Generate `From<...>` implementation.
-        composite_enum!{$EnumName,  $($entry($entry_type)),*}
-    };
-    // `From<>` implementation generation. Base case.
-    ($EnumName:ident, $entry:ident($entry_type:ty)) => {
-        impl From<$entry_type> for $EnumName {
-            fn from(entry: $entry_type) -> $EnumName {
-                $EnumName::$entry(entry)
-            }
-        }
-    };
-    // `From<>` implementation generation. Recursive case.
-    ($EnumName:ident, $entry:ident($entry_type:ty), $($entry_next:ident($entry_type_next:ty)),+) => {
-        composite_enum!{$EnumName, $entry($entry_type)}
-        composite_enum!{$EnumName, $($entry_next($entry_type_next)),*}
-    };
-}
-
 composite_enum! {
     /// A top level error type that consolidates errors from different libraries.
     #[derive(Debug)]
diff --git a/gbl/efi/src/utils.rs b/gbl/efi/src/utils.rs
index e98aec4..fca175b 100644
--- a/gbl/efi/src/utils.rs
+++ b/gbl/efi/src/utils.rs
@@ -109,12 +109,14 @@ impl AsBlockDevice for EfiGptDevice<'_> {
 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) {
+    fn for_each(
+        &mut self,
+        f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64),
+    ) -> core::result::Result<(), Option<&'static str>> {
         for (idx, ele) in self.0.iter_mut().enumerate() {
-            if f(ele, u64::try_from(idx).unwrap()) {
-                return;
-            }
+            f(ele, u64::try_from(idx).unwrap());
         }
+        Ok(())
     }
 }
 
diff --git a/gbl/libgbl/BUILD b/gbl/libgbl/BUILD
index 9f86e2b..70da4df 100644
--- a/gbl/libgbl/BUILD
+++ b/gbl/libgbl/BUILD
@@ -52,8 +52,27 @@ rust_test(
     ],
     deps = [
         "@avb//:avb_test",
+        "@gbl//libavb:sysdeps",
         "@gbl//libstorage:libstorage_testlib",
         "@static_assertions",
         "@uuid",
     ],
 )
+
+rust_test(
+    name = "integration_test",
+    srcs = ["tests/integration_tests.rs"],
+    compile_data = [
+        "@gbl//libgbl/testdata:zircon_a.bin",
+        "@gbl//libgbl/testdata:zircon_b.bin",
+        "@gbl//libgbl/testdata:zircon_r.bin",
+        "@gbl//libgbl/testdata:zircon_gpt.bin",
+    ],
+    deps = [
+        ":libgbl",
+        "@avb//:avb_crypto_ops_sha_impl_staticlib",
+        "@gbl//libavb:sysdeps",
+        "@gbl//libstorage",
+        "@gbl//libstorage:libstorage_testlib",
+    ],
+)
diff --git a/gbl/libgbl/src/error.rs b/gbl/libgbl/src/error.rs
index 5855d50..b4aa1f6 100644
--- a/gbl/libgbl/src/error.rs
+++ b/gbl/libgbl/src/error.rs
@@ -14,16 +14,21 @@
 
 //! Error types used in libgbl.
 
+use crate::GblOpsError;
 use avb::{DescriptorError, SlotVerifyError};
 use core::ffi::{FromBytesUntilNulError, FromBytesWithNulError};
 use core::fmt::{Debug, Display, Formatter};
+use gbl_storage::StorageError;
 
 /// Helper type GBL functions will return.
-pub type Result<T> = core::result::Result<T, Error>;
+pub type Result<T> = core::result::Result<T, IntegrationError>;
 
 #[derive(Debug, PartialEq)]
-/// Error values that can be returned by function in this library
+/// Errors originating from GBL native logic.
 pub enum Error {
+    ArithmeticOverflow,
+    /// Fail to hand off to kernel.
+    BootFailed,
     /// Generic error
     Error,
     /// Missing all images required to boot system
@@ -37,73 +42,114 @@ pub enum Error {
     /// AvbOps were already borrowed. This would happen on second `load_and_verify_image()` call
     /// unless `reuse()` is called before.
     AvbOpsBusy,
-    /// Failed to get descriptor from AvbMeta
-    AvbDescriptorError(DescriptorError),
-    /// Avb slot verification failed.
-    /// SlotVerifyError is used without verify data.
-    AvbSlotVerifyError(SlotVerifyError<'static>),
 }
 
-// Unfortunately thiserror is not available in `no_std` world.
-// Thus `Display` implementation is required.
 impl Display for Error {
     fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
         match self {
+            Error::ArithmeticOverflow => write!(f, "Arithmetic Overflow"),
+            Error::BootFailed => write!(f, "Failed to boot"),
             Error::Error => write!(f, "Generic error"),
             Error::MissingImage => write!(f, "Missing image required to boot system"),
             Error::NotImplemented => write!(f, "Functionality is not implemented"),
             Error::OperationProhibited => write!(f, "Operation is prohibited"),
             Error::Internal => write!(f, "Internal error"),
             Error::AvbOpsBusy => write!(f, "AvbOps were already borrowed"),
-            Error::AvbDescriptorError(error) => {
-                write!(f, "Failed to get descriptor from AvbMeta: {:?}", error)
-            }
-            Error::AvbSlotVerifyError(error) => {
-                write!(f, "Avb slot verification failed: {}", error)
-            }
         }
     }
 }
 
-impl From<DescriptorError> for Error {
-    fn from(value: DescriptorError) -> Self {
-        Error::AvbDescriptorError(value)
-    }
-}
-
-impl<'a> From<SlotVerifyError<'a>> for Error {
-    fn from(value: SlotVerifyError<'a>) -> Self {
-        Error::AvbSlotVerifyError(value.without_verify_data())
-    }
-}
+/// A helper macro for declaring a composite enum type that simply wraps other types as entries.
+/// It auto-generate `From<...>` implementation for each entry type. The type for each entry must
+/// be different from each other. i.e.:
+///
+/// ```rust
+///   composite_enum! {
+///       pub enum MyEnum {
+///           Usize(usize),
+///           I64(i64),
+///       }
+///   }
+/// ```
+///
+/// expands to
+///
+/// ```rust
+///   pub enum MyEnum {
+///       Usize(usize),
+///       I64(i64),
+///   }
+///
+///   impl From<usize> for MyEnum {
+///       fn from(ent: usize) -> MyEnum {
+///           MyEnum::Usize(ent)
+///       }
+///   }
+///
+///   impl From<i64> for MyEnum {
+///       fn from(ent: i64) -> MyEnum {
+///           MyEnum::I64(ent)
+///       }
+///   }
+/// ```
+#[macro_export]
+macro_rules! composite_enum {
+    (
+        $(#[$outer:meta])*
+        $vis:vis enum $name:ident {
+            $(
+                $(#[$inner:ident $($args:tt)*])*
+                $ent:ident($ent_t:ty)
+            ),*
+            $(,)*
+        }
+    ) => {
+        // Copy over enum declaration as it is.
+        $(#[$outer])*
+        $vis enum $name {
+            $(
+                $(#[$inner $($args)*])*
+                $ent($ent_t)
+            ),*
+        }
 
-impl From<FromBytesUntilNulError> for Error {
-    fn from(e: FromBytesUntilNulError) -> Self {
-        Error::Internal
-    }
+        // Generate `From<...>` implementation.
+        composite_enum!{$name,  $($ent($ent_t)),*}
+    };
+    // `From<>` implementation generation. Base case.
+    ($name:ident, $ent:ident($ent_t:ty)) => {
+        impl From<$ent_t> for $name {
+            fn from(ent: $ent_t) -> $name {
+                $name::$ent(ent)
+            }
+        }
+    };
+    // `From<>` implementation generation. Recursive case.
+    ($name:ident, $ent:ident($ent_t:ty), $($next:ident($next_t:ty)),+) => {
+        composite_enum!{$name, $ent($ent_t)}
+        composite_enum!{$name, $($next($next_t)),*}
+    };
 }
 
-impl From<FromBytesWithNulError> for Error {
-    fn from(e: FromBytesWithNulError) -> Self {
-        Error::Internal
+composite_enum! {
+    /// Top level error type that integrates errors from various dependency libraries.
+    #[derive(Debug)]
+    pub enum IntegrationError {
+        /// Failed to get descriptor from AvbMeta
+        AvbDescriptorError(DescriptorError),
+        /// Avb slot verification failed.
+        /// SlotVerifyError is used without verify data.
+        AvbSlotVerifyError(SlotVerifyError<'static>),
+        GblNativeError(Error),
+        GblOpsError(GblOpsError),
+        FromBytesUntilNulError(FromBytesUntilNulError),
+        FromBytesWithNulError(FromBytesWithNulError),
+        StorageError(StorageError),
     }
 }
 
-#[cfg(test)]
-mod tests {
-    use crate::*;
-    use avb::{DescriptorError, SlotVerifyError};
-
-    #[test]
-    fn test_error_output_formats() {
-        assert_eq!("Generic error", format!("{}", Error::Error));
-        assert_eq!(
-            format!("Avb slot verification failed: {}", SlotVerifyError::Io),
-            format!("{}", Error::AvbSlotVerifyError(SlotVerifyError::Io))
-        );
-        assert_eq!(
-            format!("Failed to get descriptor from AvbMeta: {:?}", DescriptorError::InvalidValue),
-            format!("{}", Error::AvbDescriptorError(DescriptorError::InvalidValue))
-        );
+impl Display for IntegrationError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
+        write!(f, "{:?}", self)
     }
 }
diff --git a/gbl/libgbl/src/fastboot/vars.rs b/gbl/libgbl/src/fastboot/vars.rs
index 5ff0fe7..39dcb12 100644
--- a/gbl/libgbl/src/fastboot/vars.rs
+++ b/gbl/libgbl/src/fastboot/vars.rs
@@ -113,7 +113,7 @@ impl Variable for Partition {
                 Ok(())
             })();
             res.is_err()
-        });
+        })?;
         res
     }
 }
@@ -165,7 +165,7 @@ impl Variable for BlockDevice {
                 f(BLOCK_DEVICE, &[id, "block-size"], snprintf!(val, "{:#x}", blk.block_size()?))
             })();
             res.is_err()
-        });
+        })?;
         res
     }
 }
diff --git a/gbl/libgbl/src/lib.rs b/gbl/libgbl/src/lib.rs
index e4912f6..a244607 100644
--- a/gbl/libgbl/src/lib.rs
+++ b/gbl/libgbl/src/lib.rs
@@ -40,6 +40,7 @@ use avb::{HashtreeErrorMode, SlotVerifyData, SlotVerifyError, SlotVerifyFlags, S
 use core::ffi::CStr;
 use core::fmt::Debug;
 use cstr::cstr;
+use gbl_storage::AsMultiBlockDevices;
 use spin::Mutex;
 
 pub mod boot_mode;
@@ -62,11 +63,15 @@ pub use avb::Descriptor;
 pub use boot_mode::BootMode;
 pub use boot_reason::KnownBootReason;
 pub use digest::{Context, Digest};
-pub use error::{Error, Result};
-pub use ops::{DefaultGblOps, GblOps};
+pub use error::{Error, IntegrationError, Result};
+pub use ops::{
+    AndroidBootImages, BootImages, DefaultGblOps, FuchsiaBootImages, GblOps, GblOpsError,
+};
 #[cfg(feature = "sw_digest")]
 pub use sw_digest::{SwContext, SwDigest};
 
+use ops::GblUtils;
+
 // TODO: b/312607649 - Replace placeholders with actual structures: https://r.android.com/2721974, etc
 /// TODO: b/312607649 - placeholder type
 pub struct Partition {}
@@ -184,7 +189,6 @@ type AvbVerifySlot = for<'b> fn(
     hashtree_error_mode: HashtreeErrorMode,
 ) -> SlotVerifyResult<'b, SlotVerifyData<'b>>;
 
-#[derive(Debug)]
 /// GBL object that provides implementation of helpers for boot process.
 ///
 /// To create this object use [GblBuilder].
@@ -230,13 +234,16 @@ where
         let requested_partitions = [cstr!("")];
         let avb_suffix = CStr::from_bytes_until_nul(&bytes)?;
 
-        let verified_data = VerifiedData((self.verify_slot)(
-            avb_ops,
-            &requested_partitions,
-            Some(avb_suffix),
-            SlotVerifyFlags::AVB_SLOT_VERIFY_FLAGS_NONE,
-            HashtreeErrorMode::AVB_HASHTREE_ERROR_MODE_EIO,
-        )?);
+        let verified_data = VerifiedData(
+            (self.verify_slot)(
+                avb_ops,
+                &requested_partitions,
+                Some(avb_suffix),
+                SlotVerifyFlags::AVB_SLOT_VERIFY_FLAGS_NONE,
+                HashtreeErrorMode::AVB_HASHTREE_ERROR_MODE_EIO,
+            )
+            .map_err(|v| v.without_verify_data())?,
+        );
 
         Ok(verified_data)
     }
@@ -250,14 +257,14 @@ where
     ///
     /// * `Ok(Cursor)` - Cursor object that manages a Manager
     /// * `Err(Error)` - on failure
-    pub fn load_slot_interface<B: gbl_storage::AsBlockDevice, M: Manager>(
+    pub fn load_slot_interface<'b, B: gbl_storage::AsBlockDevice, M: Manager>(
         &mut self,
-        block_device: B,
-    ) -> Result<Cursor<B, M>> {
+        block_device: &'b mut B,
+    ) -> Result<Cursor<'b, B, M>> {
         let boot_token = BOOT_TOKEN.lock().take().ok_or(Error::OperationProhibited)?;
         self.ops
             .load_slot_interface::<B, M>(block_device, boot_token)
-            .map_err(|_| Error::OperationProhibited)
+            .map_err(|_| Error::OperationProhibited.into())
     }
 
     /// Info Load
@@ -431,10 +438,10 @@ where
         self.kernel_jump(kernel_image, ramdisk, dtb, token)
     }
 
-    fn is_unrecoverable_error(error: &Error) -> bool {
+    fn is_unrecoverable_error(error: &IntegrationError) -> bool {
         // Note: these ifs are nested instead of chained because multiple
         //       expressions in an if-let is an unstable features
-        if let Error::AvbSlotVerifyError(ref avb_error) = error {
+        if let IntegrationError::AvbSlotVerifyError(ref avb_error) = error {
             // These are the AVB errors that are not recoverable on a subsequent attempt.
             // If necessary in the future, this helper function can be moved to the GblOps trait
             // and customized for platform specific behavior.
@@ -465,7 +472,7 @@ where
         if oneshot_status == Some(OneShot::Bootloader) {
             match self.ops.do_fastboot(&mut slot_cursor) {
                 Ok(_) => oneshot_status = slot_cursor.ctx.get_oneshot_status(),
-                Err(Error::NotImplemented) => (),
+                Err(IntegrationError::GblNativeError(Error::NotImplemented)) => (),
                 Err(e) => return Err(e),
             }
         }
@@ -482,7 +489,7 @@ where
                 AvbVerificationFlags(0),
                 Some(boot_target),
             )
-            .map_err(|e: Error| {
+            .map_err(|e: IntegrationError| {
                 if let BootTarget::NormalBoot(slot) = boot_target {
                     if Self::is_unrecoverable_error(&e) {
                         let _ = slot_cursor.ctx.set_slot_unbootable(
@@ -528,6 +535,23 @@ where
 
         Ok((kernel_image, token))
     }
+
+    /// Loads and boots a Zircon kernel according to ABR + AVB.
+    pub fn zircon_load_and_boot(&mut self, load_buffer: &mut [u8]) -> Result<()> {
+        let (mut block_devices, load_buffer) = GblUtils::new(self.ops, load_buffer)?;
+        block_devices.sync_gpt_all(&mut |_, _, _| {});
+        // TODO(b/334962583): Implement zircon ABR + AVB.
+        // The following are place holder for test of invocation in the integration test only.
+        let ptn_size = block_devices.find_partition("zircon_a")?.size()?;
+        let (kernel, remains) =
+            load_buffer.split_at_mut(ptn_size.try_into().map_err(|_| Error::ArithmeticOverflow)?);
+        block_devices.read_gpt_partition("zircon_a", 0, kernel)?;
+        self.ops.boot(BootImages::Fuchsia(FuchsiaBootImages {
+            zbi_kernel: kernel,
+            zbi_items: &mut [],
+        }))?;
+        Err(Error::BootFailed.into())
+    }
 }
 
 #[cfg(feature = "sw_digest")]
diff --git a/gbl/libgbl/src/ops.rs b/gbl/libgbl/src/ops.rs
index 2054c3a..901f63b 100644
--- a/gbl/libgbl/src/ops.rs
+++ b/gbl/libgbl/src/ops.rs
@@ -20,15 +20,45 @@ extern crate alloc;
 extern crate static_assertions;
 
 use crate::digest::{Algorithm, Context};
-use crate::error::{Error, Result};
+use crate::error::{Error, Result as GblResult};
 #[cfg(feature = "sw_digest")]
 use crate::sw_digest::SwContext;
 #[cfg(feature = "alloc")]
 use alloc::ffi::CString;
-use core::{fmt::Debug, ptr::NonNull};
+use core::{
+    fmt::{Debug, Write},
+    ptr::NonNull,
+    result::Result,
+};
+use gbl_storage::{
+    required_scratch_size, AsBlockDevice, AsMultiBlockDevices, BlockDevice, BlockIo,
+};
 
 use super::slots;
 
+/// `AndroidBootImages` contains references to loaded images for booting Android.
+pub struct AndroidBootImages<'a> {
+    pub kernel: &'a mut [u8],
+    pub ramdisk: &'a mut [u8],
+    pub fdt: &'a mut [u8],
+}
+
+/// `FuchsiaBootImages` contains references to loaded images for booting Zircon.
+pub struct FuchsiaBootImages<'a> {
+    pub zbi_kernel: &'a mut [u8],
+    pub zbi_items: &'a mut [u8],
+}
+
+/// `BootImages` contains images for booting Android/Zircon kernel.
+pub enum BootImages<'a> {
+    Android(AndroidBootImages<'a>),
+    Fuchsia(FuchsiaBootImages<'a>),
+}
+
+/// `GblOpsError` is the error type returned by required methods in `GblOps`.
+#[derive(Default, Debug)]
+pub struct GblOpsError(Option<&'static str>);
+
 // https://stackoverflow.com/questions/41081240/idiomatic-callbacks-in-rust
 // should we use traits for this? or optional/box FnMut?
 //
@@ -38,10 +68,36 @@ missing:
 - key management => atx extension in callback =>  atx_ops: ptr::null_mut(), // support optional ATX.
 */
 /// Trait that defines callbacks that can be provided to Gbl.
-pub trait GblOps: Debug {
+pub trait GblOps {
     /// Digest context type
     type Context: Context;
 
+    /// Iterates block devices on the platform.
+    ///
+    /// For each block device, implementation should call `f` with its 1) `BlockIo` trait
+    /// implementation, 2) a unique u64 ID and 3) maximum number of gpt entries. If the maximum
+    /// entries is 0, it is considered that the block should not use GPT.
+    ///
+    /// The list of block devices and visit order should remain the same for the life time of the
+    /// object that implements this trait. If this can not be met due to media change, error should
+    /// be returned. Dynamic media change is not supported for now.
+    fn visit_block_devices(
+        &mut self,
+        f: &mut dyn FnMut(&mut dyn BlockIo, u64, u64),
+    ) -> Result<(), GblOpsError>;
+
+    /// Prints a ASCII character to the platform console.
+    fn console_put_char(&mut self, ch: u8) -> Result<(), GblOpsError>;
+
+    /// This method can be used to implement platform specific mechanism for deciding whether boot
+    /// should abort and enter Fastboot mode.
+    fn should_stop_in_fastboot(&mut self) -> Result<bool, GblOpsError>;
+
+    /// Platform specific kernel boot implementation.
+    ///
+    /// Implementation is not expected to return on success.
+    fn boot(&mut self, boot_images: BootImages) -> Result<(), GblOpsError>;
+
     /// Create digest object to use for hash computations.
     ///
     /// Context interface allows to update value adding more data to process.
@@ -65,24 +121,89 @@ pub trait GblOps: Debug {
     fn do_fastboot<B: gbl_storage::AsBlockDevice>(
         &self,
         cursor: &mut slots::Cursor<B, impl slots::Manager>,
-    ) -> Result<()> {
-        Err(Error::NotImplemented)
+    ) -> GblResult<()> {
+        Err(Error::NotImplemented.into())
     }
 
     /// TODO: b/312607649 - placeholder interface for Gbl specific callbacks that uses alloc.
     #[cfg(feature = "alloc")]
-    fn gbl_alloc_extra_action(&mut self, s: &str) -> Result<()> {
+    fn gbl_alloc_extra_action(&mut self, s: &str) -> GblResult<()> {
         let _c_string = CString::new(s);
-        Err(Error::Error)
+        Err(Error::NotImplemented.into())
     }
 
     /// Load and initialize a slot manager and return a cursor over the manager on success.
-    fn load_slot_interface<B: gbl_storage::AsBlockDevice, M: slots::Manager>(
+    fn load_slot_interface<'b, B: gbl_storage::AsBlockDevice, M: slots::Manager>(
         &mut self,
-        block_device: B,
+        block_device: &'b mut B,
         boot_token: slots::BootToken,
-    ) -> Result<slots::Cursor<B, M>> {
-        Err(Error::OperationProhibited)
+    ) -> GblResult<slots::Cursor<'b, B, M>> {
+        Err(Error::OperationProhibited.into())
+    }
+
+    /// Computes the sum of required scratch size for all block devices.
+    fn required_scratch_size(&mut self) -> GblResult<usize> {
+        let mut total = 0usize;
+        let mut res = Ok(());
+        self.visit_block_devices(&mut |io, id, max_gpt_entries| {
+            res = (|| {
+                let scratch_size = required_scratch_size(io, max_gpt_entries)?;
+                total = total.checked_add(scratch_size).ok_or(Error::ArithmeticOverflow)?;
+                Ok(())
+            })();
+        })?;
+        res.map(|_| total)
+    }
+}
+
+/// `GblUtils` takes a reference to `GblOps` and implements various traits.
+pub(crate) struct GblUtils<'a, 'b, T: GblOps> {
+    ops: &'a mut T,
+    scratch: &'b mut [u8],
+}
+
+impl<'a, 'b, T: GblOps> GblUtils<'a, 'b, T> {
+    /// Create a new instance with user provided scratch buffer.
+    ///
+    /// # Args
+    ///
+    /// * `ops`: A reference to a `GblOps`,
+    /// * `scratch`: A scratch buffer.
+    ///
+    /// # Returns
+    ///
+    /// Returns a new instance and the trailing unused part of the input scratch buffer.
+    pub fn new(ops: &'a mut T, scratch: &'b mut [u8]) -> GblResult<(Self, &'b mut [u8])> {
+        let total_scratch_size = ops.required_scratch_size()?;
+        let (scratch, remaining) = scratch.split_at_mut(total_scratch_size);
+        Ok((Self { ops: ops, scratch: scratch }, remaining))
+    }
+}
+
+impl<T: GblOps> AsMultiBlockDevices for GblUtils<'_, '_, T> {
+    fn for_each(
+        &mut self,
+        f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64),
+    ) -> core::result::Result<(), Option<&'static str>> {
+        let mut scratch_offset = 0;
+        self.ops
+            .visit_block_devices(&mut |io, id, max_gpt_entries| {
+                // Not expected to fail as `Self::new()` should have checked any overflow.
+                let scratch_size = required_scratch_size(io, max_gpt_entries).unwrap();
+                let scratch = &mut self.scratch[scratch_offset..][..scratch_size];
+                scratch_offset = scratch_offset.checked_add(scratch_size).unwrap();
+                f(&mut BlockDevice::new(io, scratch, max_gpt_entries), id);
+            })
+            .map_err(|v| v.0)
+    }
+}
+
+impl<T: GblOps> Write for GblUtils<'_, '_, T> {
+    fn write_str(&mut self, s: &str) -> core::fmt::Result {
+        for ch in s.as_bytes() {
+            self.ops.console_put_char(*ch).map_err(|_| core::fmt::Error {})?;
+        }
+        Ok(())
     }
 }
 
diff --git a/gbl/libgbl/testdata/gen_sparse_test_bin.py b/gbl/libgbl/testdata/gen_test_data.py
index 8283a94..a968917 100644
--- a/gbl/libgbl/testdata/gen_sparse_test_bin.py
+++ b/gbl/libgbl/testdata/gen_test_data.py
@@ -13,21 +13,25 @@
 # 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.
-"""Generate test files for sparse image flash test"""
+"""Generate test data files for libgbl tests"""
 
 import os
 import pathlib
 import subprocess
 
 SCRIPT_DIR = pathlib.Path(os.path.dirname(os.path.realpath(__file__)))
+GPT_TOOL = pathlib.Path(SCRIPT_DIR.parents[1]) / "tools" / "gen_gpt_disk.py"
+SZ_KB = 1024
 
-# Writes bytes to a file at a given offset.
+
+# A helper for writing bytes to a file at a given offset.
 def write_file(file, offset, data):
     file.seek(offset, 0)
     file.write(data)
 
 
-if __name__ == '__main__':
+# Generates sparse image for flashing test
+def gen_sparse_test_file():
     sz_kb = 1024
     out_file_raw = SCRIPT_DIR / "sparse_test_raw.bin"
     with open(out_file_raw, "wb") as f:
@@ -54,3 +58,32 @@ if __name__ == '__main__':
         SCRIPT_DIR / "sparse_test_blk1024.bin",
         "1024",
     ])
+
+
+# Generates GPT disk, kernel data for Zircon tests
+def gen_zircon_gpt():
+    gen_gpt_args = []
+    for suffix in ["a", "b", "r"]:
+        zircon = os.urandom(16 * SZ_KB)
+        out_file = SCRIPT_DIR / f"zircon_{suffix}.bin"
+        out_file.write_bytes(zircon)
+        gen_gpt_args.append(f"--partition=zircon_{suffix},16K,{str(out_file)}")
+
+    subprocess.run([GPT_TOOL, SCRIPT_DIR / "zircon_gpt.bin", "128K"] +
+                   gen_gpt_args,
+                   check=True)
+
+
+# Generates test data for A/B slot Manager writeback test
+def gen_writeback_test_bin():
+    subprocess.run([
+        GPT_TOOL, SCRIPT_DIR / "writeback_test_disk.bin", "64K",
+        "--partition=test_partition,4k,/dev/zero"
+    ],
+                   check=True)
+
+
+if __name__ == '__main__':
+    gen_writeback_test_bin()
+    gen_sparse_test_file()
+    gen_zircon_gpt()
diff --git a/gbl/libgbl/testdata/gen_writeback_test_bin.sh b/gbl/libgbl/testdata/gen_writeback_test_bin.sh
deleted file mode 100755
index 7eba93e..0000000
--- a/gbl/libgbl/testdata/gen_writeback_test_bin.sh
+++ /dev/null
@@ -1,24 +0,0 @@
-#! /usr/bin/env bash
-# Copyright (C) 2024  Google LLC
-#
-# 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.
-
-set -e
-
-readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
-
-PARTITION_NAME="test_partition"
-PARTITION_SIZE="4k"
-
-python3 ${SCRIPT_DIR}/../../tools/gen_gpt_disk.py ${SCRIPT_DIR}/writeback_test_disk.bin 64K \
-        --partition "${PARTITION_NAME},${PARTITION_SIZE},/dev/zero"
diff --git a/gbl/libgbl/testdata/sparse_test.bin b/gbl/libgbl/testdata/sparse_test.bin
index 1721841..009689e 100644
--- a/gbl/libgbl/testdata/sparse_test.bin
+++ b/gbl/libgbl/testdata/sparse_test.bin
diff --git a/gbl/libgbl/testdata/sparse_test_blk1024.bin b/gbl/libgbl/testdata/sparse_test_blk1024.bin
index bf79c68..e273137 100644
--- a/gbl/libgbl/testdata/sparse_test_blk1024.bin
+++ b/gbl/libgbl/testdata/sparse_test_blk1024.bin
diff --git a/gbl/libgbl/testdata/sparse_test_raw.bin b/gbl/libgbl/testdata/sparse_test_raw.bin
index 1272534..ab13637 100644
--- a/gbl/libgbl/testdata/sparse_test_raw.bin
+++ b/gbl/libgbl/testdata/sparse_test_raw.bin
diff --git a/gbl/libgbl/testdata/writeback_test_disk.bin b/gbl/libgbl/testdata/writeback_test_disk.bin
index 12380d5..ebb428d 100644
--- a/gbl/libgbl/testdata/writeback_test_disk.bin
+++ b/gbl/libgbl/testdata/writeback_test_disk.bin
diff --git a/gbl/libgbl/testdata/zircon_a.bin b/gbl/libgbl/testdata/zircon_a.bin
new file mode 100644
index 0000000..bd30d49
--- /dev/null
+++ b/gbl/libgbl/testdata/zircon_a.bin
diff --git a/gbl/libgbl/testdata/zircon_b.bin b/gbl/libgbl/testdata/zircon_b.bin
new file mode 100644
index 0000000..aa9faa9
--- /dev/null
+++ b/gbl/libgbl/testdata/zircon_b.bin
diff --git a/gbl/libgbl/testdata/zircon_gpt.bin b/gbl/libgbl/testdata/zircon_gpt.bin
new file mode 100644
index 0000000..54624ff
--- /dev/null
+++ b/gbl/libgbl/testdata/zircon_gpt.bin
diff --git a/gbl/libgbl/testdata/zircon_r.bin b/gbl/libgbl/testdata/zircon_r.bin
new file mode 100644
index 0000000..653837e
--- /dev/null
+++ b/gbl/libgbl/testdata/zircon_r.bin
diff --git a/gbl/libgbl/tests/integration_tests.rs b/gbl/libgbl/tests/integration_tests.rs
new file mode 100644
index 0000000..082ac44
--- /dev/null
+++ b/gbl/libgbl/tests/integration_tests.rs
@@ -0,0 +1,170 @@
+// 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 gbl_storage::BlockIo;
+use gbl_storage_testlib::TestBlockIo;
+use libgbl::{
+    digest::Algorithm, BootImages, Context, Digest, FuchsiaBootImages, GblBuilder, GblOps,
+    GblOpsError,
+};
+use std::{collections::VecDeque, vec::Vec};
+
+extern crate avb_sysdeps;
+
+struct GblTestBlockIo {
+    io: TestBlockIo,
+    max_gpt_entries: u64,
+}
+
+/// `TestGblOps` provides mock implementation of GblOps for integration test.
+#[derive(Default)]
+struct TestGblOps<'a> {
+    block_io: Vec<GblTestBlockIo>,
+    console_out: VecDeque<u8>,
+    boot_cb: Option<MustUse<&'a mut dyn FnMut(BootImages)>>,
+}
+
+impl TestGblOps<'_> {
+    /// Adds a new block device.
+    pub(crate) fn add_block_device<T: AsRef<[u8]>>(
+        &mut self,
+        alignment: u64,
+        block_size: u64,
+        max_gpt_entries: u64,
+        data: T,
+    ) {
+        self.block_io.push(GblTestBlockIo {
+            io: TestBlockIo::new(alignment, block_size, data),
+            max_gpt_entries: max_gpt_entries,
+        })
+    }
+}
+
+impl GblOps for TestGblOps<'_> {
+    type Context = TestDigestContext;
+
+    fn visit_block_devices(
+        &mut self,
+        f: &mut dyn FnMut(&mut dyn BlockIo, u64, u64),
+    ) -> Result<(), GblOpsError> {
+        for (idx, ele) in self.block_io.iter_mut().enumerate() {
+            f(&mut ele.io, idx.try_into().unwrap(), ele.max_gpt_entries);
+        }
+        Ok(())
+    }
+
+    fn console_put_char(&mut self, ch: u8) -> Result<(), GblOpsError> {
+        Ok(self.console_out.push_back(ch))
+    }
+
+    fn should_stop_in_fastboot(&mut self) -> Result<bool, GblOpsError> {
+        Ok(false)
+    }
+
+    fn boot(&mut self, boot_images: BootImages) -> Result<(), GblOpsError> {
+        Ok((self.boot_cb.as_mut().unwrap().get())(boot_images))
+    }
+}
+
+/// Placeholder.
+struct DigestBytes(Vec<u8>);
+
+impl AsRef<[u8]> for DigestBytes {
+    fn as_ref(&self) -> &[u8] {
+        self.0.as_ref()
+    }
+}
+
+impl Digest for DigestBytes {
+    fn algorithm(&self) -> &Algorithm {
+        unimplemented!();
+    }
+}
+
+/// Placeholder.
+struct TestDigestContext {}
+
+impl Context for TestDigestContext {
+    type Digest = DigestBytes;
+
+    fn new(_: Algorithm) -> Self {
+        unimplemented!();
+    }
+
+    fn update(&mut self, _: &[u8]) {
+        unimplemented!();
+    }
+
+    fn finish(self) -> Self::Digest {
+        unimplemented!();
+    }
+
+    fn algorithm(&self) -> &Algorithm {
+        unimplemented!();
+    }
+}
+
+/// `MustUse` wraps an object and checks that it is accessed at least once before it's dropped.
+/// In this integration test, it is mainly used to check that test provided ops callbacks are run.
+struct MustUse<T: ?Sized> {
+    used: bool,
+    val: T,
+}
+
+impl<T: ?Sized> MustUse<T> {
+    /// Create a new instance.
+    fn new(val: T) -> Self
+    where
+        T: Sized,
+    {
+        Self { used: false, val: val }
+    }
+
+    /// Returns a mutable reference to the object.
+    fn get(&mut self) -> &mut T {
+        self.used = true;
+        &mut self.val
+    }
+}
+
+impl<T: ?Sized> Drop for MustUse<T> {
+    fn drop(&mut self) {
+        assert!(self.used)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_zircon_load_and_boot() {
+        // TODO(b/334962583): Invocation test only. Update this test once
+        // `Gbl::zircon_load_and_boot()` is implemented.
+        let mut boot_cb = |boot_images: BootImages| {
+            let BootImages::Fuchsia(FuchsiaBootImages { zbi_kernel, zbi_items }) = boot_images
+            else {
+                panic!("Wrong image type");
+            };
+            assert_eq!(zbi_kernel, include_bytes!("../testdata/zircon_a.bin"));
+            assert_eq!(zbi_items, []);
+        };
+        let mut ops: TestGblOps = Default::default();
+        ops.add_block_device(512, 512, 128, include_bytes!("../testdata/zircon_gpt.bin"));
+        ops.boot_cb = Some(MustUse::new(&mut boot_cb));
+        let mut gbl = GblBuilder::new(&mut ops).build();
+        let mut load_buffer = vec![0u8; 64 * 1024];
+        let _ = gbl.zircon_load_and_boot(&mut load_buffer[..]);
+    }
+}
diff --git a/gbl/libstorage/src/lib.rs b/gbl/libstorage/src/lib.rs
index 46276d7..5a10aa2 100644
--- a/gbl/libstorage/src/lib.rs
+++ b/gbl/libstorage/src/lib.rs
@@ -114,7 +114,7 @@ use gpt::Gpt;
 pub use gpt::{GptEntry, GPT_NAME_LEN_U16};
 
 mod multi_blocks;
-pub use multi_blocks::{with_id, AsMultiBlockDevices};
+pub use multi_blocks::AsMultiBlockDevices;
 
 /// The type of Result used in this library.
 pub type Result<T> = core::result::Result<T, StorageError>;
@@ -128,6 +128,7 @@ pub enum StorageError {
     BlockDeviceNotFound,
     BlockIoError,
     BlockIoNotProvided,
+    FailedGettingBlockDevices(Option<&'static str>),
     InvalidInput,
     NoValidGpt,
     NotExist,
@@ -135,27 +136,26 @@ pub enum StorageError {
     U64toUSizeOverflow,
 }
 
-impl Into<&'static str> for StorageError {
-    fn into(self) -> &'static str {
-        match self {
-            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",
-        }
-    }
-}
-
 impl core::fmt::Display for StorageError {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
-        write!(f, "{}", Into::<&'static str>::into(*self))
+        match self {
+            StorageError::ArithmeticOverflow => write!(f, "Arithmetic overflow"),
+            StorageError::OutOfRange => write!(f, "Out of range"),
+            StorageError::ScratchTooSmall => write!(f, "Not enough scratch buffer"),
+            StorageError::BlockDeviceNotFound => write!(f, "Block device not found"),
+            StorageError::BlockIoError => write!(f, "Block IO error"),
+            StorageError::BlockIoNotProvided => write!(f, "Block IO is not provided"),
+            StorageError::FailedGettingBlockDevices(v) => {
+                write!(f, "Failed to iterate all block devices {:?}", v)
+            }
+            StorageError::InvalidInput => write!(f, "Invalid input"),
+            StorageError::NoValidGpt => write!(f, "GPT not found"),
+            StorageError::NotExist => write!(f, "The specified partition could not be found"),
+            StorageError::PartitionNotUnique => {
+                write!(f, "Partition is found on multiple block devices")
+            }
+            StorageError::U64toUSizeOverflow => write!(f, "u64 to usize fails"),
+        }
     }
 }
 
diff --git a/gbl/libstorage/src/multi_blocks.rs b/gbl/libstorage/src/multi_blocks.rs
index 3f92472..0d63eb3 100644
--- a/gbl/libstorage/src/multi_blocks.rs
+++ b/gbl/libstorage/src/multi_blocks.rs
@@ -17,9 +17,24 @@ use crate::{AsBlockDevice, BlockIo, Partition, 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);
+    /// Calls closure `f` for each `AsBlockDevice` object and its unique `id` until reaching end.
+    fn for_each(
+        &mut self,
+        f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64),
+    ) -> core::result::Result<(), Option<&'static str>>;
+
+    /// Calls closure `f` for each `AsBlockDevice` object and its unique `id` until reaching end of
+    /// returnning true.
+    fn for_each_until(
+        &mut self,
+        f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64) -> bool,
+    ) -> Result<()> {
+        let mut stop = false;
+        self.for_each(&mut |io, id| {
+            stop = stop || f(io, id);
+        })
+        .map_err(|v| StorageError::FailedGettingBlockDevices(v))
+    }
 
     /// Gets the block device with the given id.
     fn get(&mut self, id: u64) -> Result<SelectedBlockDevice>
@@ -33,7 +48,7 @@ pub trait AsMultiBlockDevices {
     /// 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| {
+        let _ = self.for_each_until(&mut |v, id| {
             match v.sync_gpt() {
                 Err(e) => f(v, id, e),
                 _ => {}
@@ -53,7 +68,7 @@ pub trait AsMultiBlockDevices {
                 v => v.or(res),
             };
             res.err() == Some(StorageError::PartitionNotUnique)
-        });
+        })?;
         res
     }
 
@@ -97,8 +112,11 @@ pub trait AsMultiBlockDevices {
 }
 
 impl<T: ?Sized + AsMultiBlockDevices> AsMultiBlockDevices for &mut T {
-    fn for_each_until(&mut self, f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64) -> bool) {
-        (*self).for_each_until(&mut |io, id| f(io, id))
+    fn for_each(
+        &mut self,
+        f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64),
+    ) -> core::result::Result<(), Option<&'static str>> {
+        (*self).for_each(&mut |io, id| f(io, id))
     }
 }
 
@@ -111,16 +129,12 @@ where
     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>
+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,
 {
diff --git a/gbl/libstorage/src/testlib.rs b/gbl/libstorage/src/testlib.rs
index d06a410..85a8225 100644
--- a/gbl/libstorage/src/testlib.rs
+++ b/gbl/libstorage/src/testlib.rs
@@ -31,6 +31,16 @@ pub struct TestBlockIo {
 }
 
 impl TestBlockIo {
+    pub fn new<T: AsRef<[u8]>>(block_size: u64, alignment: u64, data: T) -> Self {
+        Self {
+            block_size: block_size,
+            alignment: alignment,
+            storage: Vec::from(data.as_ref()),
+            num_writes: 0,
+            num_reads: 0,
+        }
+    }
+
     fn check_alignment(&mut self, buffer: &[u8]) -> bool {
         matches!(is_buffer_aligned(buffer, self.alignment()), Ok(true))
             && matches!(is_aligned(buffer.len() as u64, self.block_size()), Ok(true))
@@ -217,13 +227,7 @@ impl<'a> TestBlockDeviceBuilder<'a> {
             BackingStore::Size(size) => vec![0u8; size],
         };
         assert!(storage.len() % (self.block_size as usize) == 0);
-        let mut io = TestBlockIo {
-            block_size: self.block_size,
-            alignment: self.alignment,
-            storage,
-            num_reads: 0,
-            num_writes: 0,
-        };
+        let mut io = TestBlockIo::new(self.block_size, self.alignment, storage);
         let scratch_size = match self.scratch_size {
             Some(s) => s,
             None => required_scratch_size(&mut io, self.max_gpt_entries).unwrap(),
@@ -240,11 +244,15 @@ impl<'a> TestBlockDeviceBuilder<'a> {
 pub struct TestMultiBlockDevices(pub Vec<TestBlockDevice>);
 
 impl AsMultiBlockDevices for TestMultiBlockDevices {
-    fn for_each_until(&mut self, f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64) -> bool) {
+    fn for_each(
+        &mut self,
+        f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64),
+    ) -> core::result::Result<(), Option<&'static str>> {
         let _ = self
             .0
             .iter_mut()
             .enumerate()
-            .find_map(|(idx, ele)| f(ele, u64::try_from(idx).unwrap()).then_some(()));
+            .for_each(|(idx, ele)| f(ele, u64::try_from(idx).unwrap()));
+        Ok(())
     }
 }
diff --git a/gbl/tests/BUILD b/gbl/tests/BUILD
index a072b39..27b6952 100644
--- a/gbl/tests/BUILD
+++ b/gbl/tests/BUILD
@@ -20,6 +20,7 @@ test_suite(
         "@gbl//libefi:libefi_test",
         "@gbl//libfastboot:libfastboot_test",
         "@gbl//libfdt:libfdt_test",
+        "@gbl//libgbl:integration_test",
         "@gbl//libgbl:libgbl_test",
         "@gbl//libsafemath:libsafemath_test",
         "@gbl//libstorage:libstorage_doc_test",
diff --git a/gbl/tools/gen_gpt_disk.py b/gbl/tools/gen_gpt_disk.py
index 09428e7..09428e7 100644..100755
--- a/gbl/tools/gen_gpt_disk.py
+++ b/gbl/tools/gen_gpt_disk.py