Ensure that the static array backend's scratch heap is page-aligned

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"

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]

wee_alloc/src/imp_static_array.rs

@@ -1,21 +1,28 @@
-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<usize> = Mutex::new(0);
 
 pub(crate) unsafe fn alloc_pages(pages: Pages) -> Result<NonNull<u8>, 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[*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 {

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<T>(b: Bytes) -> Option<T>
+where
+    T: ByteSize,
+    Bytes: RoundUpTo<T>,
+{
+    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)
         })
     }