From 9c4489c2f700d19ae74553bccf1e585bf3b452f8 Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Mon, 28 Oct 2019 13:19:49 -0700 Subject: [PATCH 1/3] Ensure that the static array backend's scratch heap is page-aligned --- wee_alloc/src/imp_static_array.rs | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/wee_alloc/src/imp_static_array.rs b/wee_alloc/src/imp_static_array.rs index 9a85366..d30ad1c 100644 --- a/wee_alloc/src/imp_static_array.rs +++ b/wee_alloc/src/imp_static_array.rs @@ -1,13 +1,20 @@ -use const_init::ConstInit; use super::AllocErr; +use const_init::ConstInit; #[cfg(feature = "extra_assertions")] use core::cell::Cell; use core::ptr::NonNull; use memory_units::{Bytes, Pages}; use spin::Mutex; -const SCRATCH_LEN_BYTES: usize = include!(concat!(env!("OUT_DIR"), "/wee_alloc_static_array_backend_size_bytes.txt")); -static mut SCRATCH_HEAP: [u8; SCRATCH_LEN_BYTES] = [0; SCRATCH_LEN_BYTES]; +const SCRATCH_LEN_BYTES: usize = include!(concat!( + env!("OUT_DIR"), + "/wee_alloc_static_array_backend_size_bytes.txt" +)); + +#[repr(align(4096))] +struct ScratchHeap([u8; SCRATCH_LEN_BYTES]); + +static mut SCRATCH_HEAP: ScratchHeap = ScratchHeap([0; SCRATCH_LEN_BYTES]); static mut OFFSET: Mutex = Mutex::new(0); pub(crate) unsafe fn alloc_pages(pages: Pages) -> Result, AllocErr> { @@ -15,7 +22,7 @@ pub(crate) unsafe fn alloc_pages(pages: Pages) -> Result, AllocErr> let mut offset = OFFSET.lock(); let end = bytes.0 + *offset; if end < SCRATCH_LEN_BYTES { - let ptr = SCRATCH_HEAP[*offset..end].as_mut_ptr() as *mut u8; + let ptr = SCRATCH_HEAP.0[*offset..end].as_mut_ptr() as *mut u8; *offset = end; NonNull::new(ptr).ok_or_else(|| AllocErr) } else { From d855e42090e396716e5ba423c93c51c940f46b8a Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Mon, 28 Oct 2019 13:20:13 -0700 Subject: [PATCH 2/3] Use a smaller heap size when testing the static array backend --- test.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test.sh b/test.sh index 2e216b1..f5a84d5 100755 --- a/test.sh +++ b/test.sh @@ -20,7 +20,7 @@ time cargo test --release --features "extra_assertions" time cargo test --release --features "size_classes" time cargo test --release -export WEE_ALLOC_STATIC_ARRAY_BACKEND_BYTES=$((1024 * 1024 * 1024)) +export WEE_ALLOC_STATIC_ARRAY_BACKEND_BYTES=$((512 * 1024 * 1024)) time cargo test --release --features "static_array_backend extra_assertions size_classes" time cargo test --release --features "static_array_backend extra_assertions" From 816eca62e3317dbc431bd7a27bc3d8582af6c28e Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Mon, 28 Oct 2019 13:50:40 -0700 Subject: [PATCH 3/3] Handle overflow in a couple places --- test/src/lib.rs | 11 +++++++++++ wee_alloc/src/imp_static_array.rs | 2 +- wee_alloc/src/lib.rs | 21 +++++++++++++++++---- 3 files changed, 29 insertions(+), 5 deletions(-) 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) }) }