Detect when overflow would happen and panic

fitzgen · fitzgen · commit e53e77132059 · 2019-02-11T10:48:23.000-08:00
diff --git a/src/lib.rs b/src/lib.rs
@@ -404,6 +404,10 @@ impl Bump {
     /// Allocate an object in this `Bump` and return an exclusive reference to
     /// it.
     ///
+    /// ## Panics
+    ///
+    /// Panics if reserving space for `T` would cause an overflow.
+    ///
     /// ## Example
     ///
     /// ```
@@ -423,8 +427,17 @@ impl Bump {
         }
     }
 
+    /// Allocate space for an object with the given `Layout`.
+    ///
+    /// The returned pointer points at uninitialized memory, and should be
+    /// initialized with
+    /// [`std::ptr::write`](https://doc.rust-lang.org/stable/std/ptr/fn.write.html).
+    ///
+    /// ## Panics
+    ///
+    /// Panics if reserving space for `T` would cause an overflow.
     #[inline(always)]
-    fn alloc_layout(&self, layout: Layout) -> NonNull<u8> {
+    pub fn alloc_layout(&self, layout: Layout) -> NonNull<u8> {
         unsafe {
             let footer = self.current_chunk_footer.get();
             let footer = footer.as_ref();
@@ -433,7 +446,11 @@ impl Bump {
             let end = footer as *const _ as usize;
             debug_assert!(ptr <= end);
 
-            let new_ptr = ptr + layout.size();
+            let new_ptr = match ptr.checked_add(layout.size()) {
+                Some(p) => p,
+                None => self.overflow(),
+            };
+
             if new_ptr <= end {
                 let p = ptr as *mut u8;
                 debug_assert!(new_ptr <= footer as *const _ as usize);
@@ -445,6 +462,12 @@ impl Bump {
         self.alloc_layout_slow(layout)
     }
 
+    #[inline(never)]
+    #[cold]
+    fn overflow(&self) -> ! {
+        panic!("allocation too large, caused overflow")
+    }
+
     // Slow path allocation for when we need to allocate a new chunk from the
     // parent bump set because there isn't enough room in our current chunk.
     #[inline(never)]
diff --git a/tests/tests.rs b/tests/tests.rs
@@ -1,8 +1,10 @@
 extern crate bumpalo;
 
 use bumpalo::Bump;
+use std::alloc::Layout;
 use std::mem;
 use std::slice;
+use std::usize;
 
 #[test]
 fn can_iterate_over_allocated_things() {
@@ -59,3 +61,26 @@ fn can_iterate_over_allocated_things() {
 
     assert!(seen.iter().all(|s| *s));
 }
+
+#[test]
+#[should_panic(expected = "allocation too large, caused overflow")]
+fn alloc_overflow() {
+    let bump = Bump::new();
+    let x = bump.alloc(0_u8);
+    let p = x as *mut u8 as usize;
+
+    // A size guaranteed to overflow.
+    let size = usize::MAX - p + 1;
+    let align = 1;
+    let layout = match Layout::from_size_align(size, align) {
+        Err(e) => {
+            // Return on error so that we don't panic and the test fails.
+            eprintln!("Layout::from_size_align errored: {}", e);
+            return;
+        }
+        Ok(l) => l,
+    };
+
+    // This should panic.
+    bump.alloc_layout(layout);
+}
