Handle overflow in a couple places

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

@@ -20,7 +20,7 @@ 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.0[*offset..end].as_mut_ptr() as *mut u8;
         *offset = end;

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)
         })
     }