From 8ca225a86012a0241acb6a04d885652f22652864 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alexander=20Mei=C3=9Fner?= Date: Wed, 5 Oct 2022 08:08:00 +0200 Subject: [PATCH] Backport, changes taken from #28079 of the solana monorepo. (#389) --- src/aligned_memory.rs | 125 +++++++++++++++++++++--------------------- src/elf.rs | 2 +- 2 files changed, 64 insertions(+), 63 deletions(-) diff --git a/src/aligned_memory.rs b/src/aligned_memory.rs index bab8ef92..29413802 100644 --- a/src/aligned_memory.rs +++ b/src/aligned_memory.rs @@ -1,8 +1,20 @@ -#![allow(clippy::integer_arithmetic)] //! Aligned memory use std::{mem, ptr}; +/// Scalar types, aka "plain old data" +pub trait Pod {} + +impl Pod for bool {} +impl Pod for u8 {} +impl Pod for u16 {} +impl Pod for u32 {} +impl Pod for u64 {} +impl Pod for i8 {} +impl Pod for i16 {} +impl Pod for i32 {} +impl Pod for i64 {} + /// Provides u8 slices at a specified alignment #[derive(Debug, PartialEq, Eq)] pub struct AlignedMemory { @@ -14,7 +26,7 @@ pub struct AlignedMemory { impl AlignedMemory { fn get_mem(max_len: usize) -> (Vec, usize) { - let mut mem: Vec = Vec::with_capacity(max_len + ALIGN); + let mut mem: Vec = Vec::with_capacity(max_len.saturating_add(ALIGN)); mem.push(0); let align_offset = mem.as_ptr().align_offset(ALIGN); mem.resize(align_offset, 0); @@ -26,7 +38,7 @@ impl AlignedMemory { // https://github.com/rust-lang/rust/issues/54628 let mut mem = vec![0; max_len]; let align_offset = mem.as_ptr().align_offset(ALIGN); - mem.resize(align_offset + max_len, 0); + mem.resize(max_len.saturating_add(align_offset), 0); (mem, align_offset) } /// Returns a filled AlignedMemory by copying the given slice @@ -74,15 +86,15 @@ impl AlignedMemory { } /// Calculate memory size pub fn mem_size(&self) -> usize { - mem::size_of::() + self.mem.capacity() + self.mem.capacity().saturating_add(mem::size_of::()) } /// Get the length of the data pub fn len(&self) -> usize { - self.mem.len() - self.align_offset + self.mem.len().saturating_sub(self.align_offset) } /// Is the memory empty pub fn is_empty(&self) -> bool { - self.mem.len() - self.align_offset == 0 + self.mem.len() == self.align_offset } /// Get the current write index pub fn write_index(&self) -> usize { @@ -102,44 +114,38 @@ impl AlignedMemory { } /// Grows memory with `value` repeated `num` times starting at the `write_index` pub fn fill_write(&mut self, num: usize, value: u8) -> std::io::Result<()> { - if self.mem.len() + num > self.align_offset + self.max_len { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "aligned memory resize failed", - )); - } + let new_len = match ( + self.mem.len().checked_add(num), + self.align_offset.checked_add(self.max_len), + ) { + (Some(new_len), Some(allocation_end)) if new_len <= allocation_end => new_len, + _ => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "aligned memory resize failed", + )) + } + }; if self.zero_up_to_max_len && value == 0 { // Safe because everything up to `max_len` is zeroed and no shrinking is allowed unsafe { - self.mem.set_len(self.mem.len() + num); + self.mem.set_len(new_len); } } else { - self.mem.resize(self.mem.len() + num, value); + self.mem.resize(new_len, value); } Ok(()) } - /// Write one byte into the memory. + /// Write a generic type T into the memory. /// /// # Safety /// /// Unsafe since it assumes that there is enough capacity. - pub unsafe fn write_u8_unchecked(&mut self, value: u8) { + pub unsafe fn write_unchecked(&mut self, value: T) { let pos = self.mem.len(); - debug_assert!(pos < self.align_offset + self.max_len); - self.mem.set_len(pos + 1); - *self.mem.get_unchecked_mut(pos) = value; - } - - /// Write one u64 into the memory. - /// - /// # Safety - /// - /// Unsafe since it assumes that there is enough capacity. - pub unsafe fn write_u64_unchecked(&mut self, value: u64) { - let pos = self.mem.len(); - let new_len = pos + mem::size_of::(); - debug_assert!(new_len <= self.align_offset + self.max_len); + let new_len = pos.saturating_add(mem::size_of::()); + debug_assert!(new_len <= self.align_offset.saturating_add(self.max_len)); self.mem.set_len(new_len); ptr::write_unaligned( self.mem.get_unchecked_mut(pos..new_len).as_mut_ptr().cast(), @@ -154,10 +160,11 @@ impl AlignedMemory { /// Unsafe since it assumes that there is enough capacity. pub unsafe fn write_all_unchecked(&mut self, value: &[u8]) { let pos = self.mem.len(); - debug_assert!(pos + value.len() <= self.align_offset + self.max_len); - self.mem.set_len(pos + value.len()); + let new_len = pos.saturating_add(value.len()); + debug_assert!(new_len <= self.align_offset.saturating_add(self.max_len)); + self.mem.set_len(new_len); self.mem - .get_unchecked_mut(pos..pos + value.len()) + .get_unchecked_mut(pos..new_len) .copy_from_slice(value); } } @@ -173,11 +180,17 @@ impl Clone for AlignedMemory { impl std::io::Write for AlignedMemory { fn write(&mut self, buf: &[u8]) -> std::io::Result { - if self.mem.len() + buf.len() > self.align_offset + self.max_len { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "aligned memory write failed", - )); + match ( + self.mem.len().checked_add(buf.len()), + self.align_offset.checked_add(self.max_len), + ) { + (Some(new_len), Some(allocation_end)) if new_len <= allocation_end => {} + _ => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "aligned memory write failed", + )) + } } self.mem.extend_from_slice(buf); Ok(buf.len()) @@ -193,18 +206,17 @@ impl> From for AlignedMemory { } } -/// Returns true if `data` is aligned to `align`. -pub fn is_memory_aligned(data: &[u8], align: usize) -> bool { - (data.as_ptr() as usize) - .checked_rem(align) +/// Returns true if `ptr` is aligned to `align`. +pub fn is_memory_aligned(ptr: usize, align: usize) -> bool { + ptr.checked_rem(align) .map(|remainder| remainder == 0) .unwrap_or(false) } #[cfg(test)] mod tests { - use super::*; - use std::io::Write; + #![allow(clippy::integer_arithmetic)] + use {super::*, std::io::Write}; fn do_test() { let mut aligned_memory = AlignedMemory::::with_capacity(10); @@ -234,9 +246,9 @@ mod tests { let mut aligned_memory = AlignedMemory::::with_capacity_zeroed(15); unsafe { - aligned_memory.write_u8_unchecked(42); + aligned_memory.write_unchecked::(42); assert_eq!(aligned_memory.len(), 1); - aligned_memory.write_u64_unchecked(0xCAFEBADDDEADCAFE); + aligned_memory.write_unchecked::(0xCAFEBADDDEADCAFE); assert_eq!(aligned_memory.len(), 9); aligned_memory.fill_write(3, 0).unwrap(); aligned_memory.write_all_unchecked(b"foo"); @@ -262,29 +274,18 @@ mod tests { #[cfg(debug_assertions)] #[test] - #[should_panic(expected = "< self.align_offset + self.max_len")] - fn test_write_u8_unchecked_debug_assert() { - let mut aligned_memory = AlignedMemory::<8>::with_capacity(1); - unsafe { - aligned_memory.write_u8_unchecked(42); - aligned_memory.write_u8_unchecked(24); - } - } - - #[cfg(debug_assertions)] - #[test] - #[should_panic(expected = "<= self.align_offset + self.max_len")] - fn test_write_u64_unchecked_debug_assert() { + #[should_panic(expected = "<= self.align_offset.saturating_add(self.max_len)")] + fn test_write_unchecked_debug_assert() { let mut aligned_memory = AlignedMemory::<8>::with_capacity(15); unsafe { - aligned_memory.write_u64_unchecked(42); - aligned_memory.write_u64_unchecked(24); + aligned_memory.write_unchecked::(42); + aligned_memory.write_unchecked::(24); } } #[cfg(debug_assertions)] #[test] - #[should_panic(expected = "<= self.align_offset + self.max_len")] + #[should_panic(expected = "<= self.align_offset.saturating_add(self.max_len)")] fn test_write_all_unchecked_debug_assert() { let mut aligned_memory = AlignedMemory::<8>::with_capacity(5); unsafe { diff --git a/src/elf.rs b/src/elf.rs index abd689a7..182fd650 100644 --- a/src/elf.rs +++ b/src/elf.rs @@ -414,7 +414,7 @@ impl Executable { // it must be properly aligned. We assume that HOST_ALIGN is a // multiple of the ELF "natural" alignment. See test_load_unaligned. let aligned; - let bytes = if is_memory_aligned(bytes, HOST_ALIGN) { + let bytes = if is_memory_aligned(bytes.as_ptr() as usize, HOST_ALIGN) { bytes } else { aligned = AlignedMemory::<{ HOST_ALIGN }>::from_slice(bytes);