diff --git a/test/src/lib.rs b/test/src/lib.rs index 554d918..7f391f4 100755 --- a/test/src/lib.rs +++ b/test/src/lib.rs @@ -502,6 +502,17 @@ fn smoke() { } } +#[test] +fn cannot_alloc_max_usize() { + let mut a = &wee_alloc::WeeAlloc::INIT; + unsafe { + let layout = Layout::from_size_align(std::usize::MAX, 1) + .expect("should be able to create a `Layout` with size = std::usize::MAX"); + let result = a.alloc(layout); + assert!(result.is_err()); + } +} + // This takes too long with our extra assertion checks enabled, // and the fixed-sized static array backend is too small. #[test] diff --git a/wee_alloc/src/imp_static_array.rs b/wee_alloc/src/imp_static_array.rs index d30ad1c..0b98784 100644 --- a/wee_alloc/src/imp_static_array.rs +++ b/wee_alloc/src/imp_static_array.rs @@ -20,7 +20,7 @@ static mut OFFSET: Mutex = Mutex::new(0); pub(crate) unsafe fn alloc_pages(pages: Pages) -> Result, AllocErr> { let bytes: Bytes = pages.into(); let mut offset = OFFSET.lock(); - let end = bytes.0 + *offset; + let end = bytes.0.checked_add(*offset).ok_or(AllocErr)?; if end < SCRATCH_LEN_BYTES { let ptr = SCRATCH_HEAP.0[*offset..end].as_mut_ptr() as *mut u8; *offset = end; diff --git a/wee_alloc/src/lib.rs b/wee_alloc/src/lib.rs index a5650f2..f4aff45 100644 --- a/wee_alloc/src/lib.rs +++ b/wee_alloc/src/lib.rs @@ -231,7 +231,7 @@ use core::cmp; use core::marker::Sync; use core::mem; use core::ptr::{self, NonNull}; -use memory_units::{size_of, Bytes, Pages, RoundUpTo, Words}; +use memory_units::{size_of, ByteSize, Bytes, Pages, RoundUpTo, Words}; use neighbors::Neighbors; /// The WebAssembly page size, in bytes. @@ -255,6 +255,19 @@ extra_only! { } } +#[inline] +fn checked_round_up_to(b: Bytes) -> Option +where + T: ByteSize, + Bytes: RoundUpTo, +{ + if b.0.checked_add(T::BYTE_SIZE.0).is_none() { + return None; + } else { + Some(b.round_up_to()) + } +} + #[repr(C)] #[derive(Default, Debug)] struct CellHeader<'a> { @@ -1029,11 +1042,11 @@ impl<'a> WeeAlloc<'a> { return Ok(NonNull::new_unchecked(align.0 as *mut u8)); } - let size: Words = size.round_up_to(); + let word_size: Words = checked_round_up_to(size).ok_or(AllocErr)?; - self.with_free_list_and_policy_for_size(size, align, |head, policy| { + self.with_free_list_and_policy_for_size(word_size, align, |head, policy| { assert_is_valid_free_list(head.get(), policy); - alloc_with_refill(size, align, head, policy) + alloc_with_refill(word_size, align, head, policy) }) }