gbl/libgbl/src/image_buffer.rs - platform/bootable/libbootloader - Git at Google

 // 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.

 //! ImageBuffer is helper to store uninitialized memory buffer. And work with it allowing to read
 //! into the buffer and retrieve initialized part.
 //!
 //! Similar to [ReadBuf](https://docs.rs/tokio/latest/tokio/io/struct.ReadBuf.html) but works in
 //! `no_std`.

 use core::mem::MaybeUninit;
 use liberror::{Error, Result};

 /// Wrapper class for buffer received with [get_buffer] function.
 ///
 /// Helps to keep track of allocated/init memory and avoid getting same buffer more than once.
 #[derive(Debug)]
 pub struct ImageBuffer<'a> {
     buffer: Option<&'a mut [MaybeUninit<u8>]>,
     // number of initialized and filled bytes.
     used_bytes: usize,
 }

 // Unstable MaybeUninit API
 // feature = "maybe_uninit_slice", issue = "63569"

 // Assuming all the elements are initialized, get a mutable slice to them.
 //
 // # Safety
 //
 // It is up to the caller to guarantee that the `MaybeUninit<T>` elements
 // really are in an initialized state.
 // Calling this when the content is not yet fully initialized causes undefined behavior.
 #[inline(always)]
 unsafe fn slice_assume_init_mut<T>(slice: &mut [MaybeUninit<T>]) -> &mut [T] {
     // SAFETY: Caller must make sure provided data is initialized.
     unsafe { &mut *(slice as *mut [MaybeUninit<T>] as *mut [T]) }
 }

 // Assuming all the elements are initialized, get a slice of them.
 //
 // # Safety
 //
 // It is up to the caller to guarantee that the `MaybeUninit<T>` elements
 // really are in an initialized state.
 // Calling this when the content is not yet fully initialized causes undefined behavior.
 #[inline(always)]
 unsafe fn slice_assume_init_ref<T>(slice: &[MaybeUninit<T>]) -> &[T] {
     // SAFETY: Caller must make sure provided data is initialized.
     unsafe { &*(slice as *const [MaybeUninit<T>] as *const [T]) }
 }

 impl ImageBuffer<'_> {
     /// Create new ImageBuffer from buffer and used_bytes
     pub fn new(buffer: &mut [MaybeUninit<u8>]) -> ImageBuffer {
         ImageBuffer { buffer: Some(buffer), used_bytes: 0 }
     }

     /// Total buffer capacity.
     pub fn capacity(&self) -> usize {
         self.buffer.as_ref().unwrap().len()
     }

     /// Increase used part of the buffer by `len`
     ///
     /// Return:
     /// Error() - if current used_bytes + len > capacity, or causes arithmetic overflow.
     /// Ok(()) - on success
     ///
     /// SAFETY:
     /// It is up to the user to guarantee that `len` bytes for tail was initialized and filled.
     pub unsafe fn advance_used(&mut self, len: usize) -> Result<()> {
         let Some(new_len) = self.used_bytes.checked_add(len) else {
             return Err(Error::Other(Some("Used bytes overflow")));
         };
         if new_len > self.capacity() {
             return Err(Error::BufferTooSmall(Some(new_len)));
         }
         self.used_bytes = new_len;
         Ok(())
     }

     /// Return used and tail parts of the buffer
     pub fn get_split(&self) -> (&[u8], &[MaybeUninit<u8>]) {
         let (used, tail) = self.buffer.as_ref().unwrap().split_at(self.used_bytes);
         // SAFETY:
         //
         // ImageBuffer user guarantees that changing used elements means they were initialized.
         // And object assumes initialized only for slice [..used_bytes]
         let initialized = unsafe { slice_assume_init_ref(used) };
         (initialized, tail)
     }

     /// Return used and tail parts of the buffer
     pub fn get_split_mut(&mut self) -> (&mut [u8], &mut [MaybeUninit<u8>]) {
         let (used, tail) = self.buffer.as_mut().unwrap().split_at_mut(self.used_bytes);
         // SAFETY:
         //
         // ImageBuffer user guaranties that changing used elements means they were initialized.
         // And object assumes initialized only for slice [..used_bytes]
         let initialized = unsafe { slice_assume_init_mut(used) };
         (initialized, tail)
     }

     /// Slice of the buffer that is used
     pub fn used(&self) -> &[u8] {
         let (used, _) = self.get_split();
         used
     }

     /// Slice of the buffer that is used
     pub fn used_mut(&mut self) -> &mut [u8] {
         let (used, _) = self.get_split_mut();
         used
     }

     /// Return part of the buffer that is not used
     pub fn tail(&mut self) -> &mut [MaybeUninit<u8>] {
         let (_, tail) = self.get_split_mut();
         tail
     }
 }

 impl AsRef<[MaybeUninit<u8>]> for ImageBuffer<'_> {
     fn as_ref(&self) -> &[MaybeUninit<u8>] {
         self.buffer.as_ref().unwrap()
     }
 }

 impl AsMut<[MaybeUninit<u8>]> for ImageBuffer<'_> {
     fn as_mut(&mut self) -> &mut [MaybeUninit<u8>] {
         self.buffer.as_mut().unwrap()
     }
 }

 impl PartialEq for ImageBuffer<'_> {
     fn eq(&self, other: &Self) -> bool {
         self.used() == other.used()
     }
 }

 #[cfg(test)]
 /// Helper to create ImageBuffers from Vec<u8>
 pub struct ImageBufferVec {
     buf: Vec<u8>,
 }

 #[cfg(test)]
 impl ImageBufferVec {
     pub fn new(buf: Vec<u8>) -> Self {
         Self { buf }
     }

     pub fn get(&mut self) -> ImageBuffer {
         ImageBuffer::new(Self::slice_assume_not_init_mut(self.buf.as_mut_slice()))
     }

     fn slice_assume_not_init_mut<T>(slice: &mut [T]) -> &mut [MaybeUninit<T>] {
         // SAFETY: similar to safety notes for `slice_get_ref`, but we have a
         // mutable reference which is also guaranteed to be valid for writes.
         unsafe { &mut *(slice as *mut [T] as *mut [MaybeUninit<T>]) }
     }
 }

 #[cfg(test)]
 mod test {
     use super::*;

     #[test]
     fn test_image_buffer_capacity() {
         assert_eq!(ImageBufferVec::new(vec![0u8; 0]).get().capacity(), 0);
         assert_eq!(ImageBufferVec::new(vec![0u8; 0]).get().capacity(), 0);
         assert_eq!(ImageBufferVec::new(vec![0u8; 1]).get().capacity(), 1);
         assert_eq!(ImageBufferVec::new(vec![0u8; 100]).get().capacity(), 100);
         assert_eq!(
             ImageBufferVec::new(vec![0u8; 128 * 1024 * 1024]).get().capacity(),
             128 * 1024 * 1024
         );
     }

     #[test]
     fn test_image_buffer_used() {
         let mut img_buf_vec = ImageBufferVec::new(vec![0u8; 100]);
         let mut img_buf = img_buf_vec.get();
         assert_eq!(img_buf.used().len(), 0);
         assert_eq!(img_buf.used_mut().len(), 0);
         // SAFETY:
         // All data in img_buf is initialized since it was created from vec
         unsafe { img_buf.advance_used(1).unwrap() };
         assert_eq!(img_buf.used().len(), 1);
         assert_eq!(img_buf.used_mut().len(), 1);
         // SAFETY:
         // All data in img_buf is initialized since it was created from vec
         unsafe { img_buf.advance_used(3).unwrap() };
         assert_eq!(img_buf.used().len(), 4);
         assert_eq!(img_buf.used_mut().len(), 4);
         assert_eq!(
             // SAFETY:
             // All data in img_buf is initialized since it was created from vec
             unsafe { img_buf.advance_used(1024) },
             Err(Error::BufferTooSmall(Some(1028)))
         );
         assert_eq!(img_buf.used().len(), 4);
         assert_eq!(img_buf.used_mut().len(), 4);
     }

     #[test]
     fn test_image_buffer_get_split() {
         let mut img_buf_vec = ImageBufferVec::new(vec![0u8, 1, 2, 3]);
         let mut img_buf = img_buf_vec.get();

         assert_eq!(img_buf.used(), [].as_slice());
         assert_eq!(img_buf.used_mut(), [].as_mut_slice());
         assert_eq!(img_buf.tail().len(), 4);
         let (used, tail) = img_buf.get_split();
         assert_eq!(used, [].as_slice());
         assert_eq!(tail.len(), 4);

         let (used, tail) = img_buf.get_split_mut();
         assert_eq!(used, [].as_mut_slice());
         assert_eq!(tail.len(), 4);

         // SAFETY:
         // All data in img_buf is initialized since it was created from vec
         unsafe { img_buf.advance_used(2).unwrap() };
         assert_eq!(img_buf.used(), [0, 1].as_slice());
         assert_eq!(img_buf.used_mut(), [0, 1].as_mut_slice());
         assert_eq!(img_buf.tail().len(), 2);
         let (used, tail) = img_buf.get_split();
         assert_eq!(used, [0, 1].as_slice());
         assert_eq!(tail.len(), 2);
         let (used, tail) = img_buf.get_split_mut();
         assert_eq!(used, [0, 1].as_mut_slice());
         assert_eq!(tail.len(), 2);

         // SAFETY:
         // All data in img_buf is initialized since it was created from vec
         unsafe { img_buf.advance_used(2).unwrap() };
         assert_eq!(img_buf.used(), [0, 1, 2, 3].as_slice());
         assert_eq!(img_buf.used_mut(), [0, 1, 2, 3].as_mut_slice());
         assert_eq!(img_buf.tail().len(), 0);
         let (used, tail) = img_buf.get_split();
         assert_eq!(used, [0, 1, 2, 3].as_slice());
         assert_eq!(tail.len(), 0);
         let (used, tail) = img_buf.get_split_mut();
         assert_eq!(used, [0, 1, 2, 3].as_mut_slice());
         assert_eq!(tail.len(), 0);
     }

     #[test]
     fn test_image_buffer_eq_not_init() {
         assert_eq!(
             ImageBufferVec::new(vec![0u8, 1, 2]).get(),
             ImageBufferVec::new(vec![0u8, 1, 2]).get()
         );
     }

     #[test]
     fn test_image_buffer_eq_init_same() {
         let mut v1 = ImageBufferVec::new(vec![0u8, 1, 2]);
         let mut v2 = ImageBufferVec::new(vec![0u8, 1, 2]);
         let mut image_buffer_1 = v1.get();
         let mut image_buffer_2 = v2.get();

         // SAFETY:
         // Buffers initialised on creation.
         unsafe {
             image_buffer_1.advance_used(3).unwrap();
             image_buffer_2.advance_used(3).unwrap();
         }

         assert_eq!(image_buffer_1, image_buffer_2);
     }

     #[test]
     fn test_image_buffer_eq_diff_capacity() {
         let mut v1 = ImageBufferVec::new(vec![0u8, 1, 2]);
         let mut v2 = ImageBufferVec::new(vec![0u8, 1, 2, 3]);
         let mut image_buffer_1 = v1.get();
         let mut image_buffer_2 = v2.get();

         // SAFETY:
         // Buffers initialised on creation.
         unsafe {
             image_buffer_1.advance_used(2).unwrap();
             image_buffer_2.advance_used(2).unwrap();
         }

         assert_eq!(image_buffer_1, image_buffer_2);
     }
 }
	// 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.

	//! ImageBuffer is helper to store uninitialized memory buffer. And work with it allowing to read
	//! into the buffer and retrieve initialized part.
	//!
	//! Similar to [ReadBuf](https://docs.rs/tokio/latest/tokio/io/struct.ReadBuf.html) but works in
	//! `no_std`.

	use core::mem::MaybeUninit;
	use liberror::{Error, Result};

	/// Wrapper class for buffer received with [get_buffer] function.
	///
	/// Helps to keep track of allocated/init memory and avoid getting same buffer more than once.
	#[derive(Debug)]
	pub struct ImageBuffer<'a> {
	buffer: Option<&'a mut [MaybeUninit<u8>]>,
	// number of initialized and filled bytes.
	used_bytes: usize,
	}

	// Unstable MaybeUninit API
	// feature = "maybe_uninit_slice", issue = "63569"

	// Assuming all the elements are initialized, get a mutable slice to them.
	//
	// # Safety
	//
	// It is up to the caller to guarantee that the `MaybeUninit<T>` elements
	// really are in an initialized state.
	// Calling this when the content is not yet fully initialized causes undefined behavior.
	#[inline(always)]
	unsafe fn slice_assume_init_mut<T>(slice: &mut [MaybeUninit<T>]) -> &mut [T] {
	// SAFETY: Caller must make sure provided data is initialized.
	unsafe { &mut (slice as mut [MaybeUninit<T>] as *mut [T]) }
	}

	// Assuming all the elements are initialized, get a slice of them.
	//
	// # Safety
	//
	// It is up to the caller to guarantee that the `MaybeUninit<T>` elements
	// really are in an initialized state.
	// Calling this when the content is not yet fully initialized causes undefined behavior.
	#[inline(always)]
	unsafe fn slice_assume_init_ref<T>(slice: &[MaybeUninit<T>]) -> &[T] {
	// SAFETY: Caller must make sure provided data is initialized.
	unsafe { &(slice as const [MaybeUninit<T>] as *const [T]) }
	}

	impl ImageBuffer<'_> {
	/// Create new ImageBuffer from buffer and used_bytes
	pub fn new(buffer: &mut [MaybeUninit<u8>]) -> ImageBuffer {
	ImageBuffer { buffer: Some(buffer), used_bytes: 0 }
	}

	/// Total buffer capacity.
	pub fn capacity(&self) -> usize {
	self.buffer.as_ref().unwrap().len()
	}

	/// Increase used part of the buffer by `len`
	///
	/// Return:
	/// Error() - if current used_bytes + len > capacity, or causes arithmetic overflow.
	/// Ok(()) - on success
	///
	/// SAFETY:
	/// It is up to the user to guarantee that `len` bytes for tail was initialized and filled.
	pub unsafe fn advance_used(&mut self, len: usize) -> Result<()> {
	let Some(new_len) = self.used_bytes.checked_add(len) else {
	return Err(Error::Other(Some("Used bytes overflow")));
	};
	if new_len > self.capacity() {
	return Err(Error::BufferTooSmall(Some(new_len)));
	}
	self.used_bytes = new_len;
	Ok(())
	}

	/// Return used and tail parts of the buffer
	pub fn get_split(&self) -> (&[u8], &[MaybeUninit<u8>]) {
	let (used, tail) = self.buffer.as_ref().unwrap().split_at(self.used_bytes);
	// SAFETY:
	//
	// ImageBuffer user guarantees that changing used elements means they were initialized.
	// And object assumes initialized only for slice [..used_bytes]
	let initialized = unsafe { slice_assume_init_ref(used) };
	(initialized, tail)
	}

	/// Return used and tail parts of the buffer
	pub fn get_split_mut(&mut self) -> (&mut [u8], &mut [MaybeUninit<u8>]) {
	let (used, tail) = self.buffer.as_mut().unwrap().split_at_mut(self.used_bytes);
	// SAFETY:
	//
	// ImageBuffer user guaranties that changing used elements means they were initialized.
	// And object assumes initialized only for slice [..used_bytes]
	let initialized = unsafe { slice_assume_init_mut(used) };
	(initialized, tail)
	}

	/// Slice of the buffer that is used
	pub fn used(&self) -> &[u8] {
	let (used, _) = self.get_split();
	used
	}

	/// Slice of the buffer that is used
	pub fn used_mut(&mut self) -> &mut [u8] {
	let (used, _) = self.get_split_mut();
	used
	}

	/// Return part of the buffer that is not used
	pub fn tail(&mut self) -> &mut [MaybeUninit<u8>] {
	let (_, tail) = self.get_split_mut();
	tail
	}
	}

	impl AsRef<[MaybeUninit<u8>]> for ImageBuffer<'_> {
	fn as_ref(&self) -> &[MaybeUninit<u8>] {
	self.buffer.as_ref().unwrap()
	}
	}

	impl AsMut<[MaybeUninit<u8>]> for ImageBuffer<'_> {
	fn as_mut(&mut self) -> &mut [MaybeUninit<u8>] {
	self.buffer.as_mut().unwrap()
	}
	}

	impl PartialEq for ImageBuffer<'_> {
	fn eq(&self, other: &Self) -> bool {
	self.used() == other.used()
	}
	}

	#[cfg(test)]
	/// Helper to create ImageBuffers from Vec<u8>
	pub struct ImageBufferVec {
	buf: Vec<u8>,
	}

	#[cfg(test)]
	impl ImageBufferVec {
	pub fn new(buf: Vec<u8>) -> Self {
	Self { buf }
	}

	pub fn get(&mut self) -> ImageBuffer {
	ImageBuffer::new(Self::slice_assume_not_init_mut(self.buf.as_mut_slice()))
	}

	fn slice_assume_not_init_mut<T>(slice: &mut [T]) -> &mut [MaybeUninit<T>] {
	// SAFETY: similar to safety notes for `slice_get_ref`, but we have a
	// mutable reference which is also guaranteed to be valid for writes.
	unsafe { &mut (slice as mut [T] as *mut [MaybeUninit<T>]) }
	}
	}

	#[cfg(test)]
	mod test {
	use super::*;

	#[test]
	fn test_image_buffer_capacity() {
	assert_eq!(ImageBufferVec::new(vec![0u8; 0]).get().capacity(), 0);
	assert_eq!(ImageBufferVec::new(vec![0u8; 0]).get().capacity(), 0);
	assert_eq!(ImageBufferVec::new(vec![0u8; 1]).get().capacity(), 1);
	assert_eq!(ImageBufferVec::new(vec![0u8; 100]).get().capacity(), 100);
	assert_eq!(
	ImageBufferVec::new(vec![0u8; 128 * 1024 * 1024]).get().capacity(),
	128 * 1024 * 1024
	);
	}

	#[test]
	fn test_image_buffer_used() {
	let mut img_buf_vec = ImageBufferVec::new(vec![0u8; 100]);
	let mut img_buf = img_buf_vec.get();
	assert_eq!(img_buf.used().len(), 0);
	assert_eq!(img_buf.used_mut().len(), 0);
	// SAFETY:
	// All data in img_buf is initialized since it was created from vec
	unsafe { img_buf.advance_used(1).unwrap() };
	assert_eq!(img_buf.used().len(), 1);
	assert_eq!(img_buf.used_mut().len(), 1);
	// SAFETY:
	// All data in img_buf is initialized since it was created from vec
	unsafe { img_buf.advance_used(3).unwrap() };
	assert_eq!(img_buf.used().len(), 4);
	assert_eq!(img_buf.used_mut().len(), 4);
	assert_eq!(
	// SAFETY:
	// All data in img_buf is initialized since it was created from vec
	unsafe { img_buf.advance_used(1024) },
	Err(Error::BufferTooSmall(Some(1028)))
	);
	assert_eq!(img_buf.used().len(), 4);
	assert_eq!(img_buf.used_mut().len(), 4);
	}

	#[test]
	fn test_image_buffer_get_split() {
	let mut img_buf_vec = ImageBufferVec::new(vec![0u8, 1, 2, 3]);
	let mut img_buf = img_buf_vec.get();

	assert_eq!(img_buf.used(), [].as_slice());
	assert_eq!(img_buf.used_mut(), [].as_mut_slice());
	assert_eq!(img_buf.tail().len(), 4);
	let (used, tail) = img_buf.get_split();
	assert_eq!(used, [].as_slice());
	assert_eq!(tail.len(), 4);

	let (used, tail) = img_buf.get_split_mut();
	assert_eq!(used, [].as_mut_slice());
	assert_eq!(tail.len(), 4);

	// SAFETY:
	// All data in img_buf is initialized since it was created from vec
	unsafe { img_buf.advance_used(2).unwrap() };
	assert_eq!(img_buf.used(), [0, 1].as_slice());
	assert_eq!(img_buf.used_mut(), [0, 1].as_mut_slice());
	assert_eq!(img_buf.tail().len(), 2);
	let (used, tail) = img_buf.get_split();
	assert_eq!(used, [0, 1].as_slice());
	assert_eq!(tail.len(), 2);
	let (used, tail) = img_buf.get_split_mut();
	assert_eq!(used, [0, 1].as_mut_slice());
	assert_eq!(tail.len(), 2);

	// SAFETY:
	// All data in img_buf is initialized since it was created from vec
	unsafe { img_buf.advance_used(2).unwrap() };
	assert_eq!(img_buf.used(), [0, 1, 2, 3].as_slice());
	assert_eq!(img_buf.used_mut(), [0, 1, 2, 3].as_mut_slice());
	assert_eq!(img_buf.tail().len(), 0);
	let (used, tail) = img_buf.get_split();
	assert_eq!(used, [0, 1, 2, 3].as_slice());
	assert_eq!(tail.len(), 0);
	let (used, tail) = img_buf.get_split_mut();
	assert_eq!(used, [0, 1, 2, 3].as_mut_slice());
	assert_eq!(tail.len(), 0);
	}

	#[test]
	fn test_image_buffer_eq_not_init() {
	assert_eq!(
	ImageBufferVec::new(vec![0u8, 1, 2]).get(),
	ImageBufferVec::new(vec![0u8, 1, 2]).get()
	);
	}

	#[test]
	fn test_image_buffer_eq_init_same() {
	let mut v1 = ImageBufferVec::new(vec![0u8, 1, 2]);
	let mut v2 = ImageBufferVec::new(vec![0u8, 1, 2]);
	let mut image_buffer_1 = v1.get();
	let mut image_buffer_2 = v2.get();

	// SAFETY:
	// Buffers initialised on creation.
	unsafe {
	image_buffer_1.advance_used(3).unwrap();
	image_buffer_2.advance_used(3).unwrap();
	}

	assert_eq!(image_buffer_1, image_buffer_2);
	}

	#[test]
	fn test_image_buffer_eq_diff_capacity() {
	let mut v1 = ImageBufferVec::new(vec![0u8, 1, 2]);
	let mut v2 = ImageBufferVec::new(vec![0u8, 1, 2, 3]);
	let mut image_buffer_1 = v1.get();
	let mut image_buffer_2 = v2.get();

	// SAFETY:
	// Buffers initialised on creation.
	unsafe {
	image_buffer_1.advance_used(2).unwrap();
	image_buffer_2.advance_used(2).unwrap();
	}

	assert_eq!(image_buffer_1, image_buffer_2);
	}
	}