Ensure addition/multiplications in when allocating buffers don't over…

…flow (#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

@@ -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));
+    }
 }

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

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