Snap for 11579193 from 7b7fe1bfd57f5c508b042e2e868ee4ac3cce7208 to sdk-releaseplatform-tools-35.0.1

Change-Id: I7001cf671b3bb9e5be9423b236db9a4774519137

10 files changed, 418 insertions, 246 deletions

diff --git a/gbl/libgbl/BUILD b/gbl/libgbl/BUILD
index cdc974d..0ef63fb 100644
--- a/gbl/libgbl/BUILD
+++ b/gbl/libgbl/BUILD
@@ -27,6 +27,7 @@ rust_library(
         "@bitflags",
         "@crc32fast",
         "@cstr",
+        "@gbl//libstorage",
         "@gbl//third_party/libzbi",
         "@spin",
         "@zerocopy",
@@ -41,6 +42,7 @@ rust_test(
         "@gbl//libgbl/testdata:test_image.img",
         "@gbl//libgbl/testdata:testkey_rsa4096.pem",
         "@gbl//libgbl/testdata:testkey_rsa4096.pub.pem",
+        "@gbl//libgbl/testdata:writeback_test_disk.bin",
     ],
     deps = [
         "@avb//:avb_test",
diff --git a/gbl/libgbl/Cargo.toml b/gbl/libgbl/Cargo.toml
index 69917f4..f383fb3 100644
--- a/gbl/libgbl/Cargo.toml
+++ b/gbl/libgbl/Cargo.toml
@@ -27,13 +27,13 @@ sw_digest = []
 
 [dependencies]
 zbi = {version = "0.1", path = "../third_party/libzbi"}
+gbl_storage = {version = "0.1", path = "../libstorage"}
 ring = "0.17"
 spin = "0.9"
 static_assertions = "0"
 lazy_static = "1"
-zerocopy = {version = "=0.7.0-alpha.5"}
+zerocopy = {version = "=0.7.32"}
 crc32fast = "1.3"
 
 [dev-dependencies]
 itertools = "0"
-
diff --git a/gbl/libgbl/src/lib.rs b/gbl/libgbl/src/lib.rs
index 29292b9..adf675b 100644
--- a/gbl/libgbl/src/lib.rs
+++ b/gbl/libgbl/src/lib.rs
@@ -32,6 +32,7 @@
 extern crate avb;
 extern crate core;
 extern crate cstr;
+extern crate gbl_storage;
 extern crate spin;
 extern crate zbi;
 
@@ -218,7 +219,6 @@ where
         partitions_ram_map: &mut [PartitionRamMap],
         avb_verification_flags: AvbVerificationFlags,
         boot_target: Option<BootTarget>,
-        slot_cursor: &mut Cursor,
     ) -> Result<VerifiedData<'b>>
     where
         'a: 'b,
@@ -249,9 +249,14 @@ where
     ///
     /// * `Ok(Cursor)` - Cursor object that manages a Manager
     /// * `Err(Error)` - on failure
-    pub fn load_slot_interface(&mut self) -> Result<Cursor> {
+    pub fn load_slot_interface<B: gbl_storage::AsBlockDevice>(
+        &mut self,
+        block_device: B,
+    ) -> Result<Cursor<B>> {
         let boot_token = BOOT_TOKEN.lock().take().ok_or(Error::OperationProhibited)?;
-        self.ops.load_slot_interface(boot_token).map_err(|_| Error::OperationProhibited)
+        self.ops
+            .load_slot_interface(block_device, boot_token)
+            .map_err(|_| Error::OperationProhibited)
     }
 
     /// Info Load
@@ -397,19 +402,16 @@ where
     /// * `Err(Error)` - on failure
     // Nevertype could be used here when it is stable https://github.com/serde-rs/serde/issues/812
     #[allow(clippy::too_many_arguments)]
-    pub fn load_verify_boot<'b: 'c, 'c, 'd: 'b>(
+    pub fn load_verify_boot<'b: 'c, 'c, 'd: 'b, B: gbl_storage::AsBlockDevice>(
         &mut self,
         avb_ops: &'b mut impl avb::Ops,
         partitions_ram_map: &'d mut [PartitionRamMap<'b, 'c>],
         avb_verification_flags: AvbVerificationFlags,
-        slot_cursor: Cursor,
+        slot_cursor: Cursor<B>,
         kernel_load_buffer: &mut [u8],
         ramdisk_load_buffer: &mut [u8],
         fdt: &mut [u8],
     ) -> Result<()> {
-        // TODO(dovs):
-        // * Change the receiver of ops.load_slot_interface to be &mut self
-        // * Add partition write capabilites to slot manager
         let dtb = Dtb(&mut fdt[..]);
         let mut ramdisk = Ramdisk(ramdisk_load_buffer);
 
@@ -447,20 +449,20 @@ where
         false
     }
 
-    fn lvb_inner<'b: 'c, 'c, 'd: 'b, 'e>(
+    fn lvb_inner<'b: 'c, 'c, 'd: 'b, 'e, B: gbl_storage::AsBlockDevice>(
         &mut self,
         avb_ops: &'b mut impl avb::Ops,
         ramdisk: &mut Ramdisk,
         kernel_load_buffer: &'e mut [u8],
         partitions_ram_map: &'d mut [PartitionRamMap<'b, 'c>],
         avb_verification_flags: AvbVerificationFlags,
-        mut slot_cursor: Cursor,
+        mut slot_cursor: Cursor<B>,
     ) -> Result<(KernelImage<'e>, BootToken)> {
         let mut oneshot_status = slot_cursor.ctx.get_oneshot_status();
         slot_cursor.ctx.clear_oneshot_status();
 
         if oneshot_status == Some(OneShot::Bootloader) {
-            match self.ops.do_fastboot(&slot_cursor) {
+            match self.ops.do_fastboot(&mut slot_cursor) {
                 Ok(_) => oneshot_status = slot_cursor.ctx.get_oneshot_status(),
                 Err(Error::NotImplemented) => (),
                 Err(e) => return Err(e),
@@ -478,7 +480,6 @@ where
                 partitions_ram_map,
                 AvbVerificationFlags(0),
                 Some(boot_target),
-                &mut slot_cursor,
             )
             .map_err(|e: Error| {
                 if let BootTarget::NormalBoot(slot) = boot_target {
@@ -578,13 +579,11 @@ where
 #[cfg(test)]
 mod tests {
     extern crate avb_test;
-
     use super::*;
     use avb::IoError;
     use avb::IoResult as AvbIoResult;
     use avb::PublicKeyForPartitionInfo;
     use avb_test::TestOps;
-    use slots::fuchsia::SlotBlock;
     use std::fs;
 
     struct AvbOpsUnimplemented {}
@@ -637,13 +636,11 @@ mod tests {
         let mut avb_ops = AvbOpsUnimplemented {};
         let mut partitions_ram_map: [PartitionRamMap; 0] = [];
         let avb_verification_flags = AvbVerificationFlags(0);
-        let mut slot_cursor = Cursor { ctx: &mut SlotBlock::default() };
         let res = gbl.load_and_verify_image(
             &mut avb_ops,
             &mut partitions_ram_map,
             avb_verification_flags,
             None,
-            &mut slot_cursor,
         );
         assert_eq!(res.unwrap_err(), Error::AvbSlotVerifyError(SlotVerifyError::Io));
     }
@@ -659,7 +656,6 @@ mod tests {
     fn test_load_and_verify_image_stub() {
         let mut gbl = GblBuilder::new(DefaultGblOps::new()).build();
         let mut avb_ops = TestOps::default();
-        let mut slot_cursor = Cursor { ctx: &mut SlotBlock::default() };
 
         avb_ops.add_partition(TEST_PARTITION_NAME, fs::read(TEST_IMAGE_PATH).unwrap());
         avb_ops.add_partition("vbmeta", fs::read(TEST_VBMETA_PATH).unwrap());
@@ -674,7 +670,6 @@ mod tests {
             &mut partitions_ram_map,
             avb_verification_flags,
             None,
-            &mut slot_cursor,
         );
         assert!(res.is_ok());
     }
@@ -690,13 +685,11 @@ mod tests {
         let mut avb_ops = AvbOpsUnimplemented {};
         let mut partitions_ram_map: [PartitionRamMap; 0] = [];
         let avb_verification_flags = AvbVerificationFlags(0);
-        let mut slot_cursor = Cursor { ctx: &mut SlotBlock::default() };
         let res = gbl.load_and_verify_image(
             &mut avb_ops,
             &mut partitions_ram_map,
             avb_verification_flags,
             None,
-            &mut slot_cursor,
         );
         assert_eq!(res.unwrap_err(), Error::AvbSlotVerifyError(TEST_ERROR));
     }
diff --git a/gbl/libgbl/src/ops.rs b/gbl/libgbl/src/ops.rs
index e3b77eb..1a56a08 100644
--- a/gbl/libgbl/src/ops.rs
+++ b/gbl/libgbl/src/ops.rs
@@ -62,7 +62,10 @@ pub trait GblOps: Debug {
 
     /// Callback for when fastboot mode is requested.
     // Nevertype could be used here when it is stable https://github.com/serde-rs/serde/issues/812
-    fn do_fastboot(&self, cursor: &slots::Cursor) -> Result<()> {
+    fn do_fastboot<B: gbl_storage::AsBlockDevice>(
+        &self,
+        cursor: &mut slots::Cursor<B>,
+    ) -> Result<()> {
         Err(Error::NotImplemented)
     }
 
@@ -74,7 +77,11 @@ pub trait GblOps: Debug {
     }
 
     /// Load and initialize a slot manager and return a cursor over the manager on success.
-    fn load_slot_interface(&mut self, boot_token: slots::BootToken) -> Result<slots::Cursor> {
+    fn load_slot_interface<B: gbl_storage::AsBlockDevice>(
+        &mut self,
+        block_device: B,
+        boot_token: slots::BootToken,
+    ) -> Result<slots::Cursor<B>> {
         Err(Error::OperationProhibited)
     }
 }
diff --git a/gbl/libgbl/src/slots.rs b/gbl/libgbl/src/slots.rs
index 5269e51..127f198 100644
--- a/gbl/libgbl/src/slots.rs
+++ b/gbl/libgbl/src/slots.rs
@@ -386,18 +386,21 @@ pub trait Manager: private::SlotGet {
     /// This is useful for partition based slot setups,
     /// where we do not write back every interaction in order to coalesce writes
     /// and preserve disk lifetime.
-    fn write_back(&mut self) {}
+    fn write_back(&mut self, block_dev: &mut dyn gbl_storage::AsBlockDevice) {}
 }
 
 /// RAII helper object for coalescing changes.
-pub struct Cursor<'a> {
+pub struct Cursor<'a, B: gbl_storage::AsBlockDevice> {
     /// The backing manager for slot metadata.
     pub ctx: &'a mut dyn Manager,
+    /// The backing disk. Used for partition-backed metadata implementations
+    /// and for fastboot.
+    pub block_dev: B,
 }
 
-impl<'a> Drop for Cursor<'a> {
+impl<'a, Block: gbl_storage::AsBlockDevice> Drop for Cursor<'a, Block> {
     fn drop(&mut self) {
-        self.ctx.write_back();
+        self.ctx.write_back(&mut self.block_dev);
     }
 }
 
diff --git a/gbl/libgbl/src/slots/fuchsia.rs b/gbl/libgbl/src/slots/fuchsia.rs
index 7aead36..0fd4b67 100644
--- a/gbl/libgbl/src/slots/fuchsia.rs
+++ b/gbl/libgbl/src/slots/fuchsia.rs
@@ -132,6 +132,23 @@ impl AbrData {
         hasher.update(&self.as_bytes()[..(size_of::<Self>() - size_of::<BigEndianU32>())]);
         hasher.finalize()
     }
+
+    fn validate<B: ByteSlice>(buffer: B) -> Result<Ref<B, AbrData>, AbrSlotError> {
+        let abr_data =
+            Ref::<B, AbrData>::new_from_prefix(buffer).ok_or(AbrSlotError::BufferTooSmall)?.0;
+
+        if abr_data.magic != *ABR_MAGIC {
+            return Err(AbrSlotError::BadMagic);
+        }
+        if abr_data.version_major > ABR_VERSION_MAJOR {
+            return Err(AbrSlotError::BadVersion);
+        }
+        if abr_data.crc32.get() != abr_data.calculate_crc32() {
+            return Err(AbrSlotError::BadCrc);
+        }
+
+        Ok(abr_data)
+    }
 }
 
 impl Default for AbrData {
@@ -157,13 +174,15 @@ impl Default for AbrData {
 /// a special recovery partition, R, and support for oneshot boot.
 /// Includes an Option<BootToken> to support `mark_boot_attempt`
 /// and a cache status to support lazy write-back on destruction.
-pub struct SlotBlock {
+pub struct SlotBlock<'a> {
     cache_status: CacheStatus,
     boot_token: Option<BootToken>,
     abr_data: AbrData,
+    partition: &'a str,
+    partition_offset: u64,
 }
 
-impl SlotBlock {
+impl<'a> SlotBlock<'a> {
     fn get_mut_data(&mut self) -> &mut AbrData {
         self.cache_status = CacheStatus::Dirty;
         &mut self.abr_data
@@ -176,51 +195,37 @@ impl SlotBlock {
     /// Attempt to deserialize a slot control block
     ///
     /// # Returns
-    /// * `Ok(SlotBlock)` - on success returns a SlotBlock that wraps a copy of the serialized data
-    /// * `Err(AbrSlotError)` - on failure
+    /// * `SlotBlock` - returns either the deserialized representation of the slot control block
+    ///                 OR a fresh, default valued slot control block if there was an internal error.
+    ///
+    ///                 TODO(b/329116902): errors are logged
     pub fn deserialize<B: ByteSlice>(
         buffer: B,
+        partition: &'a str,
+        partition_offset: u64,
         boot_token: BootToken,
-    ) -> Result<Self, AbrSlotError> {
-        let abr_data =
-            Ref::<B, AbrData>::new_from_prefix(buffer).ok_or(AbrSlotError::BufferTooSmall)?.0;
-
-        if abr_data.magic != *ABR_MAGIC {
-            return Err(AbrSlotError::BadMagic);
-        }
-        if abr_data.version_major > ABR_VERSION_MAJOR {
-            return Err(AbrSlotError::BadVersion);
-        }
-
-        if abr_data.crc32.get() != abr_data.calculate_crc32() {
-            return Err(AbrSlotError::BadCrc);
-        }
+    ) -> Self {
+        // TODO(b/329116902): log failures
+        // validate(buffer)
+        // .inspect_err(|e| {
+        //     eprintln!("ABR metadata failed verification, using metadata defaults: {e}")
+        // })
+        let (abr_data, cache_status) = match AbrData::validate(buffer) {
+            Ok(data) => (*data, CacheStatus::Clean),
+            Err(_) => (Default::default(), CacheStatus::Dirty),
+        };
 
-        Ok(SlotBlock {
-            cache_status: CacheStatus::Clean,
+        SlotBlock {
+            cache_status,
             boot_token: Some(boot_token),
-            abr_data: *abr_data,
-        })
-    }
-
-    pub(crate) fn new(boot_token: BootToken) -> Self {
-        Self { boot_token: Some(boot_token), ..Default::default() }
-    }
-}
-
-impl Default for SlotBlock {
-    /// Returns a default valued SlotBlock.
-    /// Only used in tests because BootToken cannot be constructed out of crate.
-    fn default() -> Self {
-        Self {
-            cache_status: CacheStatus::Clean,
-            boot_token: Some(BootToken(())),
-            abr_data: Default::default(),
+            abr_data,
+            partition,
+            partition_offset,
         }
     }
 }
 
-impl super::private::SlotGet for SlotBlock {
+impl super::private::SlotGet for SlotBlock<'_> {
     fn get_slot_by_number(&self, number: usize) -> Result<Slot, Error> {
         let lower_ascii = 'a'..='z';
         let (suffix, &abr_slot) = core::iter::zip(lower_ascii, self.get_data().slot_data.iter())
@@ -249,7 +254,7 @@ fn suffix_rank(s: Suffix) -> i64 {
     -i64::from(u32::from(s.0))
 }
 
-impl Manager for SlotBlock {
+impl Manager for SlotBlock<'_> {
     fn get_boot_target(&self) -> BootTarget {
         self.slots_iter()
             .filter(Slot::is_bootable)
@@ -357,20 +362,26 @@ impl Manager for SlotBlock {
         }
     }
 
-    fn write_back(&mut self) {
+    fn write_back(&mut self, block_dev: &mut dyn gbl_storage::AsBlockDevice) {
         if self.cache_status == CacheStatus::Clean {
             return;
         }
 
+        let part = self.partition;
+        let offset = self.partition_offset;
+
         let data = self.get_mut_data();
         data.version_minor = ABR_VERSION_MINOR;
         data.crc32 = data.calculate_crc32().into();
 
-        // TODO(dovs): write data back to partition
+        match block_dev.write_gpt_partition_mut(part, offset, data.as_bytes_mut()) {
+            Ok(_) => self.cache_status = CacheStatus::Clean,
+            Err(e) => panic!("{}", e),
+        };
     }
 }
 
-impl SlotBlock {
+impl<'a> SlotBlock<'a> {
     fn get_index_and_slot_with_suffix(&self, slot_suffix: Suffix) -> Result<(usize, Slot), Error> {
         self.slots_iter()
             .enumerate()
@@ -381,7 +392,23 @@ impl SlotBlock {
 
 #[cfg(test)]
 mod test {
+    use super::super::Cursor;
     use super::*;
+    use gbl_storage::AsBlockDevice;
+
+    impl Default for SlotBlock<'_> {
+        /// Returns a default valued SlotBlock.
+        /// Only used in tests because BootToken cannot be constructed out of crate.
+        fn default() -> Self {
+            Self {
+                cache_status: CacheStatus::Clean,
+                boot_token: Some(BootToken(())),
+                abr_data: Default::default(),
+                partition: "",
+                partition_offset: 0,
+            }
+        }
+    }
 
     #[test]
     fn test_slot_block_defaults() {
@@ -404,13 +431,6 @@ mod test {
     }
 
     #[test]
-    fn test_slot_block_new() {
-        let new = SlotBlock::new(BootToken(()));
-        let default: SlotBlock = Default::default();
-        assert_eq!(new, default);
-    }
-
-    #[test]
     fn test_suffix() {
         let slot = Slot { suffix: 'a'.into(), ..Default::default() };
         assert_eq!(BootTarget::Recovery(RecoveryTarget::Dedicated).suffix(), 'r'.into());
@@ -420,37 +440,28 @@ mod test {
 
     #[test]
     fn test_slot_block_parse() {
-        let sb: SlotBlock = Default::default();
-        assert_eq!(SlotBlock::deserialize(sb.get_data().as_bytes(), BootToken(())), Ok(sb));
+        let abr: AbrData = Default::default();
+        assert_eq!(AbrData::validate(abr.as_bytes()), Ok(Ref::new(abr.as_bytes()).unwrap()));
     }
 
     #[test]
     fn test_slot_block_parse_buffer_too_small() {
         let buffer: [u8; 0] = Default::default();
-        assert_eq!(
-            SlotBlock::deserialize::<&[u8]>(&buffer, BootToken(())),
-            Err(AbrSlotError::BufferTooSmall),
-        );
+        assert_eq!(AbrData::validate(&buffer[..]), Err(AbrSlotError::BufferTooSmall),);
     }
 
     #[test]
     fn test_slot_block_parse_bad_magic() {
         let mut sb: SlotBlock = Default::default();
         sb.get_mut_data().magic[0] += 1;
-        assert_eq!(
-            SlotBlock::deserialize(sb.get_data().as_bytes(), BootToken(())),
-            Err(AbrSlotError::BadMagic)
-        );
+        assert_eq!(AbrData::validate(sb.get_data().as_bytes()), Err(AbrSlotError::BadMagic));
     }
 
     #[test]
     fn test_slot_block_parse_bad_version_major() {
         let mut sb: SlotBlock = Default::default();
         sb.get_mut_data().version_major = 15;
-        assert_eq!(
-            SlotBlock::deserialize(sb.get_data().as_bytes(), BootToken(())),
-            Err(AbrSlotError::BadVersion)
-        );
+        assert_eq!(AbrData::validate(sb.get_data().as_bytes()), Err(AbrSlotError::BadVersion));
     }
 
     #[test]
@@ -458,10 +469,7 @@ mod test {
         let mut sb: SlotBlock = Default::default();
         let bad_crc = sb.get_data().crc32.get() ^ BigEndianU32::MAX_VALUE.get();
         sb.get_mut_data().crc32 = bad_crc.into();
-        assert_eq!(
-            SlotBlock::deserialize(sb.get_data().as_bytes(), BootToken(())),
-            Err(AbrSlotError::BadCrc)
-        );
+        assert_eq!(AbrData::validate(sb.get_data().as_bytes()), Err(AbrSlotError::BadCrc));
     }
 
     #[test]
@@ -672,4 +680,145 @@ mod test {
             Err(Error::OperationProhibited)
         );
     }
+
+    #[test]
+    fn test_deserialize_default_to_dirty_cache() {
+        let mut abr_data: AbrData = Default::default();
+        // Changing the success both invalidates the crc
+        // and lets us verify that the deserialized slot block
+        // uses defaulted backing bytes instead of the provided bytes.
+        abr_data.slot_data[0].successful = 1;
+        let sb = SlotBlock::deserialize(abr_data.as_bytes(), "partition_moniker", 0, BootToken(()));
+        assert_eq!(sb.cache_status, CacheStatus::Dirty);
+        assert_eq!(
+            sb.slots_iter().next().unwrap().bootability,
+            Bootability::Retriable(DEFAULT_RETRIES.into())
+        );
+    }
+
+    #[test]
+    fn test_deserialize_modified_to_clean_cache() {
+        let mut abr_data: AbrData = Default::default();
+        abr_data.slot_data[0].successful = 1;
+        // If we recalculate the crc,
+        // that just means we have a metadata block that stores
+        // relevant, non-default information.
+        abr_data.crc32.set(abr_data.calculate_crc32());
+        let sb = SlotBlock::deserialize(abr_data.as_bytes(), "partition_moniker", 0, BootToken(()));
+        assert_eq!(sb.cache_status, CacheStatus::Clean);
+        assert_eq!(sb.slots_iter().next().unwrap().bootability, Bootability::Successful);
+    }
+
+    /// Mocks a block device with content from buffer.
+    /// Copied from libstorage test
+    pub struct TestBlockIo {
+        pub storage: Vec<u8>,
+    }
+
+    const BLOCK_SIZE: u64 = 512;
+    const ALIGNMENT: u64 = 512;
+
+    impl gbl_storage::BlockIo for TestBlockIo {
+        fn block_size(&mut self) -> u64 {
+            BLOCK_SIZE
+        }
+
+        fn num_blocks(&mut self) -> u64 {
+            self.storage.len() as u64 / self.block_size()
+        }
+
+        fn alignment(&mut self) -> u64 {
+            ALIGNMENT
+        }
+
+        fn read_blocks(&mut self, blk_offset: u64, out: &mut [u8]) -> bool {
+            let start = blk_offset * self.block_size();
+            let end = start + out.len() as u64;
+            out.clone_from_slice(&self.storage[start as usize..end as usize]);
+            true
+        }
+
+        fn write_blocks(&mut self, blk_offset: u64, data: &[u8]) -> bool {
+            let start = blk_offset * self.block_size();
+            let end = start + data.len() as u64;
+            self.storage[start as usize..end as usize].clone_from_slice(&data);
+            true
+        }
+    }
+
+    pub struct TestBlockDevice {
+        pub io: TestBlockIo,
+        pub scratch: Vec<u8>,
+    }
+
+    impl TestBlockDevice {
+        const GPT_ENTRIES: u64 = 128;
+
+        /// Creates an instance with provided storage content.
+        pub fn new_with_data(data: &[u8]) -> Self {
+            let mut io = TestBlockIo { storage: Vec::from(data) };
+            let scratch_size =
+                gbl_storage::required_scratch_size(&mut io, Self::GPT_ENTRIES).unwrap();
+            Self { io, scratch: vec![0u8; scratch_size] }
+        }
+    }
+
+    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)
+        }
+    }
+
+    #[test]
+    fn test_writeback() {
+        const PARTITION: &str = "test_partition";
+        const OFFSET: u64 = 2112; // Deliberately wrong to test propagation of parameter.
+        let disk = include_bytes!("../../testdata/writeback_test_disk.bin");
+        let mut block_dev = TestBlockDevice::new_with_data(disk);
+        assert!(block_dev.sync_gpt().is_ok());
+        let mut sb =
+            SlotBlock { partition: PARTITION, partition_offset: OFFSET, ..Default::default() };
+
+        let mut read_buffer: [u8; size_of::<AbrData>()] = Default::default();
+
+        // Clean cache, write_back is a no-op
+        sb.write_back(&mut block_dev);
+        let res = block_dev.read_gpt_partition(PARTITION, OFFSET, &mut read_buffer);
+        assert!(res.is_ok());
+        assert_eq!(read_buffer, [0; std::mem::size_of::<AbrData>()]);
+
+        // Make a change, write_back writes back to the defined partition
+        // at the defined offset.
+        assert_eq!(sb.set_oneshot_status(OneShot::Bootloader), Ok(()));
+        assert_eq!(sb.cache_status, CacheStatus::Dirty);
+
+        sb.write_back(&mut block_dev);
+        let res = block_dev.read_gpt_partition(PARTITION, OFFSET, &mut read_buffer);
+        assert!(res.is_ok());
+        assert_eq!(read_buffer, sb.abr_data.as_bytes());
+        assert_eq!(sb.cache_status, CacheStatus::Clean);
+    }
+
+    #[test]
+    fn test_writeback_with_cursor() {
+        const PARTITION: &str = "test_partition";
+        const OFFSET: u64 = 2112; // Deliberately wrong to test propagation of parameter.
+        let disk = include_bytes!("../../testdata/writeback_test_disk.bin");
+
+        let mut block_dev = TestBlockDevice::new_with_data(disk);
+        assert!(block_dev.sync_gpt().is_ok());
+        let mut sb =
+            SlotBlock { partition: PARTITION, partition_offset: OFFSET, ..Default::default() };
+
+        let mut read_buffer: [u8; size_of::<AbrData>()] = Default::default();
+
+        {
+            let cursor = Cursor { ctx: &mut sb, block_dev: &mut block_dev };
+            assert!(cursor.ctx.set_active_slot('b'.into()).is_ok());
+        }
+
+        let res = block_dev.read_gpt_partition(PARTITION, OFFSET, &mut read_buffer);
+        assert!(res.is_ok());
+        assert_eq!(read_buffer, sb.abr_data.as_bytes());
+    }
 }
diff --git a/gbl/libgbl/testdata/gen_writeback_test_bin.sh b/gbl/libgbl/testdata/gen_writeback_test_bin.sh
new file mode 100755
index 0000000..7eba93e
--- /dev/null
+++ b/gbl/libgbl/testdata/gen_writeback_test_bin.sh
@@ -0,0 +1,24 @@
+#! /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/writeback_test_disk.bin b/gbl/libgbl/testdata/writeback_test_disk.bin
new file mode 100644
index 0000000..12380d5
--- /dev/null
+++ b/gbl/libgbl/testdata/writeback_test_disk.bin
diff --git a/gbl/libstorage/src/gpt.rs b/gbl/libstorage/src/gpt.rs
index ac26aad..2d09435 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, AsBlockDevice, BlockDeviceWrite,
-    BlockIo, Result, StorageError, WriteBuffer,
+    add, aligned_subslice, div, mul, read, sub, to_usize, write_bytes, write_bytes_mut,
+    AsBlockDevice, BlockIo, Result, StorageError,
 };
 use core::default::Default;
 use core::mem::{align_of, size_of};
@@ -318,7 +318,7 @@ impl<'a> Gpt<'a> {
         let primary_valid = self.validate_gpt(blk_dev, scratch, HeaderType::Primary)?;
         let secondary_valid = self.validate_gpt(blk_dev, scratch, HeaderType::Secondary)?;
 
-        let primary_header = GptHeader::from_bytes(&mut self.primary_header[..]);
+        let primary_header = GptHeader::from_bytes(self.primary_header);
         let secondary_header = GptHeader::from_bytes(&mut self.secondary_header[..]);
         if !primary_valid {
             if !secondary_valid {
@@ -331,8 +331,8 @@ impl<'a> Gpt<'a> {
             primary_header.backup = secondary_header_blk;
             primary_header.entries = primary_entries_blk;
             primary_header.update_crc();
-            write(blk_dev, primary_header_pos, primary_header.as_bytes_mut(), scratch)?;
-            write(blk_dev, primary_entries_pos, &mut self.primary_entries[..], scratch)?
+            write_bytes_mut(blk_dev, primary_header_pos, primary_header.as_bytes_mut(), scratch)?;
+            write_bytes_mut(blk_dev, primary_entries_pos, self.primary_entries, scratch)?
         } else if !secondary_valid {
             // Restore to secondary
             secondary_header.as_bytes_mut().clone_from_slice(primary_header.as_bytes());
@@ -341,8 +341,13 @@ impl<'a> Gpt<'a> {
             secondary_header.backup = primary_header_blk;
             secondary_header.entries = secondary_entries_blk;
             secondary_header.update_crc();
-            write(blk_dev, secondary_header_pos, secondary_header.as_bytes_mut(), scratch)?;
-            write(blk_dev, secondary_entries_pos, &mut self.secondary_entries[..], scratch)?;
+            write_bytes_mut(
+                blk_dev,
+                secondary_header_pos,
+                secondary_header.as_bytes_mut(),
+                scratch,
+            )?;
+            write_bytes_mut(blk_dev, secondary_entries_pos, self.secondary_entries, scratch)?;
         }
 
         // Calculate actual number of GPT entries by finding the first invalid entry.
@@ -398,21 +403,37 @@ pub(crate) fn read_gpt_partition(
 }
 
 /// Write GPT partition. Library internal helper for AsBlockDevice::write_gpt_partition().
-pub(crate) fn write_gpt_partition<B>(
+pub(crate) fn write_gpt_partition(
     blk_dev: &mut (impl BlockIo + ?Sized),
     gpt: &Gpt,
     part_name: &str,
     offset: u64,
-    mut data: B,
+    data: &[u8],
     scratch: &mut [u8],
-) -> Result<()>
-where
-    B: WriteBuffer,
-{
+) -> 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),
+    }
+}
+
+/// Write GPT partition. Library internal helper for AsBlockDevice::write_gpt_partition().
+/// Optimized version for mutable buffers.
+pub(crate) fn write_gpt_partition_mut(
+    blk_dev: &mut (impl BlockIo + ?Sized),
+    gpt: &Gpt,
+    part_name: &str,
+    offset: u64,
+    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.data_len()? as u64)?;
-            write(blk_dev, abs_offset, data, scratch)
+            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),
     }
@@ -485,22 +506,33 @@ pub trait MultiGptDevices {
         });
         res
     }
-}
 
-/// `MultiGptDevicesWrite` provides partition write APIs for a `MultiGptDevices`
-pub trait MultiGptDevicesWrite<B: WriteBuffer>: MultiGptDevices {
     /// Writes a GPT partition on the first match.
-    fn write_gpt_partition(&mut self, part_name: &str, offset: u64, mut data: B) -> Result<()> {
+    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, &mut data);
+            res = v.write_gpt_partition(part_name, offset, data);
             res.is_ok()
         });
         res
     }
-}
 
-impl<T: ?Sized + MultiGptDevices, B: WriteBuffer> MultiGptDevicesWrite<B> for T {}
+    /// 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) {
@@ -516,7 +548,6 @@ impl<B: AsBlockDevice> MultiGptDevices for &mut [B] {
 mod test {
     use super::*;
     use crate::test::TestBlockDevice;
-    use crate::BlockDeviceWrite;
 
     fn gpt_block_device(max_entries: u64, data: &[u8]) -> TestBlockDevice {
         TestBlockDevice::new_with_data(512, 512, max_entries, data)
@@ -699,37 +730,37 @@ mod test {
 
         // "boot_a" partition
         // Mutable version
-        dev.write_gpt_partition("boot_a", 0, &mut expect_boot_a[..]).unwrap();
+        dev.write_gpt_partition_mut("boot_a", 0, expect_boot_a.as_mut_slice()).unwrap();
         dev.read_gpt_partition("boot_a", 0, &mut actual_boot_a).unwrap();
         assert_eq!(expect_boot_a.to_vec(), actual_boot_a);
         // Mutable version, partial write.
-        dev.write_gpt_partition("boot_a", 1, &expect_boot_a[1..]).unwrap();
+        dev.write_gpt_partition_mut("boot_a", 1, expect_boot_a[1..].as_mut()).unwrap();
         dev.read_gpt_partition("boot_a", 1, &mut actual_boot_a[1..]).unwrap();
         assert_eq!(expect_boot_a[1..], actual_boot_a[1..]);
         // Immutable version
-        dev.write_gpt_partition("boot_a", 0, &expect_boot_a[..] as &[u8]).unwrap();
+        dev.write_gpt_partition("boot_a", 0, &expect_boot_a).unwrap();
         dev.read_gpt_partition("boot_a", 0, &mut actual_boot_a).unwrap();
         assert_eq!(expect_boot_a.to_vec(), actual_boot_a);
         // Immutable version, partial write.
-        dev.write_gpt_partition("boot_a", 1, &expect_boot_a[1..] as &[u8]).unwrap();
+        dev.write_gpt_partition("boot_a", 1, &expect_boot_a[1..]).unwrap();
         dev.read_gpt_partition("boot_a", 1, &mut actual_boot_a[1..]).unwrap();
         assert_eq!(expect_boot_a[1..], actual_boot_a[1..]);
 
         // "boot_b" partition
         // Mutable version
-        dev.write_gpt_partition("boot_b", 0, &mut expect_boot_b[..]).unwrap();
+        dev.write_gpt_partition_mut("boot_b", 0, expect_boot_b.as_mut_slice()).unwrap();
         dev.read_gpt_partition("boot_b", 0, &mut actual_boot_b).unwrap();
         assert_eq!(expect_boot_b.to_vec(), actual_boot_b);
         // Mutable version, partial write.
-        dev.write_gpt_partition("boot_b", 1, &expect_boot_b[1..]).unwrap();
+        dev.write_gpt_partition_mut("boot_b", 1, expect_boot_b[1..].as_mut()).unwrap();
         dev.read_gpt_partition("boot_b", 1, &mut actual_boot_b[1..]).unwrap();
         assert_eq!(expect_boot_b[1..], actual_boot_b[1..]);
         // Immutable version
-        dev.write_gpt_partition("boot_b", 0, &expect_boot_b[..] as &[u8]).unwrap();
+        dev.write_gpt_partition("boot_b", 0, &expect_boot_b).unwrap();
         dev.read_gpt_partition("boot_b", 0, &mut actual_boot_b).unwrap();
         assert_eq!(expect_boot_b.to_vec(), actual_boot_b);
         // Immutable version, partial write.
-        dev.write_gpt_partition("boot_b", 1, &expect_boot_b[1..] as &[u8]).unwrap();
+        dev.write_gpt_partition("boot_b", 1, &expect_boot_b[1..]).unwrap();
         dev.read_gpt_partition("boot_b", 1, &mut actual_boot_b[1..]).unwrap();
         assert_eq!(expect_boot_b[1..], actual_boot_b[1..]);
     }
@@ -743,12 +774,12 @@ mod test {
         let mut boot_b = [0u8; include_bytes!("../test/boot_b.bin").len()];
 
         assert!(dev.read_gpt_partition("boot_a", 1, &mut boot_a).is_err());
-        assert!(dev.write_gpt_partition("boot_a", 1, &mut boot_a[..]).is_err());
-        assert!(dev.write_gpt_partition("boot_a", 1, &boot_a[..]).is_err());
+        assert!(dev.write_gpt_partition_mut("boot_a", 1, boot_a.as_mut_slice()).is_err());
+        assert!(dev.write_gpt_partition("boot_a", 1, &boot_a).is_err());
 
         assert!(dev.read_gpt_partition("boot_b", 1, &mut boot_b).is_err());
-        assert!(dev.write_gpt_partition("boot_b", 1, &mut boot_b[..]).is_err());
-        assert!(dev.write_gpt_partition("boot_b", 1, &boot_b[..]).is_err());
+        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]
@@ -815,10 +846,10 @@ mod test {
         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("boot_a", 0, &mut expect_boot_a[..]).unwrap();
+        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("boot_b", 0, &mut expect_boot_b[..]).unwrap();
+        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);
 
@@ -829,10 +860,10 @@ mod test {
         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("boot_a", 0, &mut expect_vendor_boot_a[..]).unwrap();
+        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("boot_b", 0, &mut expect_vendor_boot_b[..]).unwrap();
+        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 fb79d85..f885962 100644
--- a/gbl/libstorage/src/lib.rs
+++ b/gbl/libstorage/src/lib.rs
@@ -23,7 +23,7 @@
 //!
 //! ```rust
 //! use gbl_storage::{
-//!     AsBlockDevice, BlockIo, BlockDevice, required_scratch_size, BlockDeviceWrite,
+//!     AsBlockDevice, BlockIo, BlockDevice, required_scratch_size,
 //! };
 //!
 //! /// Mocks a block device using a buffer.
@@ -74,9 +74,9 @@
 //! ram_block_dev.read(4321, &mut out[..]).unwrap();
 //! let mut data = vec![0u8; 5678];
 //! // Mutable input. More efficient
-//! ram_block_dev.write(8765, &mut data[..]).unwrap();
+//! ram_block_dev.write_mut(8765, data.as_mut_slice()).unwrap();
 //! // Immutable input. Works too but not as efficient.
-//! ram_block_dev.write(8765, &data[..]).unwrap();
+//! ram_block_dev.write(8765, &data).unwrap();
 //!
 //! // Sync GPT
 //! let _ = ram_block_dev.sync_gpt();
@@ -84,7 +84,7 @@
 //! let _ = ram_block_dev.gpt();
 //! // 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("partition", 8765, &mut data[..]);
+//! let _ = ram_block_dev.write_gpt_partition_mut("partition", 8765, data.as_mut_slice());
 //!
 //! // Alterantively, you can also define a custom type that internally owns and binds the
 //! // implementation of `BlockIo` and scratch buffer together, and then implement the
@@ -112,7 +112,7 @@ use core::cmp::min;
 mod gpt;
 pub use gpt::Gpt;
 pub use gpt::GptEntry;
-pub use gpt::{MultiGptDevices, MultiGptDevicesWrite};
+pub use gpt::MultiGptDevices;
 
 /// The type of Result used in this library.
 pub type Result<T> = core::result::Result<T, StorageError>;
@@ -139,7 +139,7 @@ impl Into<&'static str> for StorageError {
             StorageError::BlockIoError => "Block IO error",
             StorageError::InvalidInput => "Invalid input",
             StorageError::NoValidGpt => "GPT not found",
-            StorageError::NotExist => "Not exists",
+            StorageError::NotExist => "The specified partition could not be found",
             StorageError::U64toUSizeOverflow => "u64 to usize fails",
         }
     }
@@ -196,8 +196,7 @@ pub trait BlockIo {
 }
 
 /// `AsBlockDevice` provides APIs for reading raw block content and GPT partitions with
-/// arbirary offset, size and input/output buffer. Write APIs are available in trait
-/// `BlockDeviceWrite`, which is automatically implemented for all `AsBlockDevice` types.
+/// arbirary offset, size and input/output buffer.
 pub trait AsBlockDevice {
     /// Returns a tuple containing the following information:
     ///
@@ -217,8 +216,9 @@ pub trait AsBlockDevice {
     ///   total required scratch size is 0.
     ///
     /// * GPT headers will be cached in the scratch buffer after calling `Self::sync_gpt()` and
-    ///   returning success. Subsequent call of `Self:read_gpt_partiton()` and
-    ///   `Self::write_gpt_partition()` will look up partition entries from the cached GPT header.
+    ///   returning success. Subsequent call of `Self:read_gpt_partiton()`,
+    ///   `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.
@@ -293,37 +293,41 @@ pub trait AsBlockDevice {
     fn block_io(&mut self) -> &mut dyn BlockIo {
         self.get().0
     }
-}
 
-/// `BlockDeviceWrite` provides raw/GPT partition write APIs for a `AsBlockDevice`.
-/// The input data needs to implement `WriteBuffer` which tells the API whether input data is
-/// mutable. Mutable input data allows additional internal optimization which makes write more
-/// efficient. In most cases, B should either be a `&mut [u8]` or `&[u8]`. See `write()` for more
-/// details.
-pub trait BlockDeviceWrite<B: WriteBuffer>: AsBlockDevice {
     /// Write data to the device.
     ///
     /// # Args
     ///
     /// * `offset`: Offset in number of bytes.
     ///
-    /// * `data`: Data to write. It can be `&[u8]`, `&mut [u8]` or other types that implement the
-    ///   `WriteBuffer` trait.
+    /// * `data`: Data to write.
     ///
-    /// * If `data` is an immutable bytes slice `&[u8]` or doesn't support
-    ///   `WriteBuffer::as_mut_slice()`, when offset`/`data.len()` is not aligned to
-    ///   `Self::block_size()` or data.as_ptr() is not aligned to `Self::alignment()`, the API may
+    /// * If offset`/`data.len()` is not aligned to `Self::block_size()`
+    ///   or data.as_ptr() is not aligned to `Self::alignment()`, the API may
     ///   reduce to a block by block read-modify-write in the worst case, which can be inefficient.
     ///
-    /// * If `data` is a mutable bytes slice `&mut [u8]`, or support `WriteBuffer::as_mut_slice()`.
-    ///   The API enables optimization which temporarily changes `data` layout internally and
-    ///   reduces the number of calls to `Self::write_blocks()` down to O(1) regardless of inputs
+    /// * 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)
+    }
+
+    /// Write data to the device.
+    ///
+    /// # Args
+    ///
+    /// * `offset`: Offset in number of bytes.
+    ///
+    /// * `data`: Data to write.
+    ///
+    /// * The API enables an optimization which temporarily changes `data` layout internally and
+    ///   reduces the number of calls to `Self::write_blocks()` down to O(1) regardless of input's
     ///   alignment. This is the recommended usage.
     ///
     /// * Returns success when exactly `data.len()` number of bytes are written.
-    fn write(&mut self, offset: u64, data: B) -> Result<()> {
+    fn write_mut(&mut self, offset: u64, data: &mut [u8]) -> Result<()> {
         let (io, alignment_scratch, _, _) = get_with_partitioned_scratch(self)?;
-        write(io, offset, data, alignment_scratch)
+        write_bytes_mut(io, offset, data, alignment_scratch)
     }
 
     /// Write a GPT partition on a block device
@@ -339,7 +343,7 @@ pub trait BlockDeviceWrite<B: WriteBuffer>: AsBlockDevice {
     /// # Returns
     ///
     /// Returns success when exactly `data.len()` of bytes are written successfully.
-    fn write_gpt_partition(&mut self, part_name: &str, offset: u64, data: B) -> Result<()> {
+    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,
@@ -350,6 +354,39 @@ pub trait BlockDeviceWrite<B: WriteBuffer>: AsBlockDevice {
             alignment_scratch,
         )
     }
+
+    /// Write a GPT partition on a block device.
+    /// Optimization for mutable buffers.
+    /// See `AsBlockDevice::write_mut` for details on alignment requirements
+    /// for optimized performance.
+    ///
+    /// # Args
+    ///
+    /// * `part_name`: Name of the partition.
+    ///
+    /// * `offset`: Offset in number of bytes into the partition.
+    ///
+    /// * `data`: Data to write. See `data` passed to `BlockIo::write()` for details.
+    ///
+    /// # Returns
+    ///
+    /// Returns success when exactly `data.len()` of bytes are written successfully.
+    fn write_gpt_partition_mut(
+        &mut self,
+        part_name: &str,
+        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,
+        )
+    }
 }
 
 impl<T: ?Sized + AsBlockDevice> AsBlockDevice for &mut T {
@@ -358,8 +395,6 @@ impl<T: ?Sized + AsBlockDevice> AsBlockDevice for &mut T {
     }
 }
 
-impl<T: ?Sized + AsBlockDevice, B: WriteBuffer> BlockDeviceWrite<B> for T {}
-
 /// `BlockDevice` borrows a `BlockIo`, scratch buffer and implements `AsBlockDevice`.
 pub struct BlockDevice<'a, 'b> {
     io: &'a mut dyn BlockIo,
@@ -843,82 +878,10 @@ fn write_bytes_mut(
     )
 }
 
-/// Top level write API that accept mutable/immutable bytes slice
-fn write<B>(
-    blk_io: &mut (impl BlockIo + ?Sized),
-    offset: u64,
-    mut data: B,
-    scratch: &mut [u8],
-) -> Result<()>
-where
-    B: WriteBuffer,
-{
-    if let Some(slice) = data.as_mut_slice() {
-        return write_bytes_mut(blk_io, offset, slice, scratch);
-    } else if let Some(slice) = data.as_slice() {
-        return write_bytes(blk_io, offset, slice, scratch);
-    }
-
-    Err(StorageError::InvalidInput)
-}
-
-/// The trait is for abstracting the representation of bytes data. The `AsBlockDevice::write()` API
-/// will choose different algorithm based on the return result of the methods. The library
-/// provides native implementation for `&[u8]` and `&mut [u8]`. It's technically possilbe for users
-/// to implement other types such as `Vec<u8>` etc.
-pub trait WriteBuffer {
-    /// Return a mutable bytes slice for the data. If the data cannot be mutable, returns None.
-    fn as_mut_slice(&mut self) -> Option<&mut [u8]>;
-
-    /// Return a bytes slice for the data. Return None if not supported.
-    fn as_slice(&mut self) -> Option<&[u8]>;
-
-    // Provided methods
-
-    /// Return the length of the data.
-    fn data_len(&mut self) -> Result<usize> {
-        self.as_slice()
-            .map(|v| v.len())
-            .or(self.as_mut_slice().map(|v| v.len()))
-            .ok_or(StorageError::InvalidInput)
-    }
-}
-
-impl WriteBuffer for &[u8] {
-    fn as_mut_slice(&mut self) -> Option<&mut [u8]> {
-        None
-    }
-
-    fn as_slice(&mut self) -> Option<&[u8]> {
-        Some(*self)
-    }
-}
-
-impl WriteBuffer for &mut [u8] {
-    fn as_mut_slice(&mut self) -> Option<&mut [u8]> {
-        Some(*self)
-    }
-
-    fn as_slice(&mut self) -> Option<&[u8]> {
-        None
-    }
-}
-
-impl<T: WriteBuffer> WriteBuffer for &mut T {
-    fn as_mut_slice(&mut self) -> Option<&mut [u8]> {
-        (*self).as_mut_slice()
-    }
-
-    fn as_slice(&mut self) -> Option<&[u8]> {
-        (*self).as_slice()
-    }
-}
-
 #[cfg(test)]
 mod test {
     use super::*;
     use core::mem::size_of;
-    use Vec;
 
     /// Mocks a block device with content from buffer.
     pub struct TestBlockIo {
@@ -1133,7 +1096,7 @@ mod test {
                 #[test]
                 fn write_test() {
                     let func = |blk: &mut TestBlockDevice, offset: u64, data: &mut [u8]| {
-                        blk.write(offset, data as &[u8]).unwrap();
+                        blk.write(offset, data).unwrap();
                     };
                     write_test_helper(&TestCase::new($x0, $x1, $x2, $x3, $x4, $x5), func);
                 }
@@ -1142,7 +1105,7 @@ mod test {
                 fn write_scaled_test() {
                     // Scaled all parameters by double and test again.
                     let func = |blk: &mut TestBlockDevice, offset: u64, data: &mut [u8]| {
-                        blk.write(offset, data as &[u8]).unwrap();
+                        blk.write(offset, data).unwrap();
                     };
                     let (x0, x1, x2, x3, x4, x5) =
                         (2 * $x0, 2 * $x1, 2 * $x2, 2 * $x3, 2 * $x4, 2 * $x5);
@@ -1153,7 +1116,7 @@ mod test {
                 #[test]
                 fn write_mut_test() {
                     let func = |blk: &mut TestBlockDevice, offset: u64, data: &mut [u8]| {
-                        blk.write(offset, data).unwrap();
+                        blk.write_mut(offset, data).unwrap();
                         assert!(blk.io.num_reads <= READ_WRITE_BLOCKS_UPPER_BOUND);
                         assert!(blk.io.num_writes <= READ_WRITE_BLOCKS_UPPER_BOUND);
                     };
@@ -1166,7 +1129,7 @@ mod test {
                     let (x0, x1, x2, x3, x4, x5) =
                         (2 * $x0, 2 * $x1, 2 * $x2, 2 * $x3, 2 * $x4, 2 * $x5);
                     let func = |blk: &mut TestBlockDevice, offset: u64, data: &mut [u8]| {
-                        blk.write(offset, data).unwrap();
+                        blk.write_mut(offset, data).unwrap();
                         assert!(blk.io.num_reads <= READ_WRITE_BLOCKS_UPPER_BOUND);
                         assert!(blk.io.num_writes <= READ_WRITE_BLOCKS_UPPER_BOUND);
                     };
@@ -1430,18 +1393,18 @@ mod test {
     #[test]
     fn test_write_overflow() {
         let mut blk = TestBlockDevice::new(1, 1, 0, 512);
-        assert!(blk.write(512, vec![0u8; 1].as_mut_slice()).is_err());
-        assert!(blk.write(0, vec![0u8; 513].as_mut_slice()).is_err());
+        assert!(blk.write_mut(512, vec![0u8; 1].as_mut_slice()).is_err());
+        assert!(blk.write_mut(0, vec![0u8; 513].as_mut_slice()).is_err());
 
-        assert!(blk.write(512, vec![0u8; 1].as_slice()).is_err());
-        assert!(blk.write(0, vec![0u8; 513].as_slice()).is_err());
+        assert!(blk.write(512, &vec![0u8; 1]).is_err());
+        assert!(blk.write(0, &vec![0u8; 513]).is_err());
     }
 
     #[test]
     fn test_write_arithmetic_overflow() {
         let mut blk = TestBlockDevice::new(1, 1, 0, 512);
-        assert!(blk.write(u64::MAX, vec![0u8; 1].as_mut_slice()).is_err());
-        assert!(blk.write(u64::MAX, vec![0u8; 1].as_slice()).is_err());
+        assert!(blk.write_mut(u64::MAX, vec![0u8; 1].as_mut_slice()).is_err());
+        assert!(blk.write(u64::MAX, &vec![0u8; 1]).is_err());
     }
 
     #[test]