From 1e0a264371cbf85eab50ec3c79ed00a6f579674a Mon Sep 17 00:00:00 2001
From: Jeffrey Vo <jeffrey.vo.australia@gmail.com>
Date: Sat, 24 Feb 2024 11:14:41 +1100
Subject: [PATCH] Ensure addition/multiplications in when allocating buffers
 don't overflow (#5417)

* Ensure addition/multiplications in arrow-buffer saturate

* Assert roundup behaviour and fix tests

* Use checked_add when rounding up to power of 2

* Doc
---
 arrow-buffer/src/buffer/immutable.rs |  9 ++++++++-
 arrow-buffer/src/buffer/mutable.rs   | 17 ++++++++++++++++-
 arrow-buffer/src/util/bit_util.rs    | 10 +++++++++-
 3 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs
index f7fc7cffdc73..ca3ee6260ba7 100644
--- a/arrow-buffer/src/buffer/immutable.rs
+++ b/arrow-buffer/src/buffer/immutable.rs
@@ -431,7 +431,7 @@ impl<T: ArrowNativeType> FromIterator<T> for Buffer {
             None => MutableBuffer::new(0),
             Some(element) => {
                 let (lower, _) = iterator.size_hint();
-                let mut buffer = MutableBuffer::new(lower.saturating_add(1) * size);
+                let mut buffer = MutableBuffer::new(lower.saturating_add(1).saturating_mul(size));
                 unsafe {
                     std::ptr::write(buffer.as_mut_ptr() as *mut T, element);
                     buffer.set_len(size);
@@ -820,4 +820,11 @@ mod tests {
         let b = b.into_vec::<u32>().unwrap();
         assert_eq!(b, &[1, 3, 5]);
     }
+
+    #[test]
+    #[should_panic(expected = "failed to round to next highest power of 2")]
+    fn test_from_iter_overflow() {
+        let iter_len = usize::MAX / std::mem::size_of::<u64>() + 1;
+        let _ = Buffer::from_iter(std::iter::repeat(0_u64).take(iter_len));
+    }
 }
diff --git a/arrow-buffer/src/buffer/mutable.rs b/arrow-buffer/src/buffer/mutable.rs
index 361f7369c91d..e08d9c1906b3 100644
--- a/arrow-buffer/src/buffer/mutable.rs
+++ b/arrow-buffer/src/buffer/mutable.rs
@@ -61,16 +61,24 @@ pub struct MutableBuffer {
 
 impl MutableBuffer {
     /// Allocate a new [MutableBuffer] with initial capacity to be at least `capacity`.
+    ///
+    /// See [`MutableBuffer::with_capacity`].
     #[inline]
     pub fn new(capacity: usize) -> Self {
         Self::with_capacity(capacity)
     }
 
     /// Allocate a new [MutableBuffer] with initial capacity to be at least `capacity`.
+    ///
+    /// # Panics
+    ///
+    /// If `capacity`, when rounded up to the nearest multiple of [`ALIGNMENT`], is greater
+    /// then `isize::MAX`, then this function will panic.
     #[inline]
     pub fn with_capacity(capacity: usize) -> Self {
         let capacity = bit_util::round_upto_multiple_of_64(capacity);
-        let layout = Layout::from_size_align(capacity, ALIGNMENT).unwrap();
+        let layout = Layout::from_size_align(capacity, ALIGNMENT)
+            .expect("failed to create layout for MutableBuffer");
         let data = match layout.size() {
             0 => dangling_ptr(),
             _ => {
@@ -1010,4 +1018,11 @@ mod tests {
         assert_eq!(buffer.len(), 4 * mem::size_of::<u16>());
         assert_eq!(buffer.as_slice(), &[1, 0, 2, 0, 3, 0, 4, 0]);
     }
+
+    #[test]
+    #[should_panic(expected = "failed to create layout for MutableBuffer: LayoutError")]
+    fn test_with_capacity_panics_above_max_capacity() {
+        let max_capacity = isize::MAX as usize - (isize::MAX as usize % ALIGNMENT);
+        let _ = MutableBuffer::with_capacity(max_capacity + 1);
+    }
 }
diff --git a/arrow-buffer/src/util/bit_util.rs b/arrow-buffer/src/util/bit_util.rs
index d2dbf3c84882..f4aaf571b5d6 100644
--- a/arrow-buffer/src/util/bit_util.rs
+++ b/arrow-buffer/src/util/bit_util.rs
@@ -39,7 +39,9 @@ pub fn round_upto_multiple_of_64(num: usize) -> usize {
 /// be a power of 2.
 pub fn round_upto_power_of_2(num: usize, factor: usize) -> usize {
     debug_assert!(factor > 0 && (factor & (factor - 1)) == 0);
-    (num + (factor - 1)) & !(factor - 1)
+    num.checked_add(factor - 1)
+        .expect("failed to round to next highest power of 2")
+        & !(factor - 1)
 }
 
 /// Returns whether bit at position `i` in `data` is set or not
@@ -119,6 +121,12 @@ mod tests {
         assert_eq!(192, round_upto_multiple_of_64(129));
     }
 
+    #[test]
+    #[should_panic(expected = "failed to round to next highest power of 2")]
+    fn test_round_upto_panic() {
+        let _ = round_upto_power_of_2(usize::MAX, 2);
+    }
+
     #[test]
     fn test_get_bit() {
         // 00001101