Check overflow in MutableArrayData extend offsets (apache#3123)

tustvold · Nov 22, 2022 · ab6e4da · ab6e4da
1 parent 475e079
commit ab6e4da
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 7 deletions.
diff --git a/arrow-data/src/transform/list.rs b/arrow-data/src/transform/list.rs
@@ -21,9 +21,9 @@ use super::{
 };
 use crate::ArrayData;
 use arrow_buffer::ArrowNativeType;
-use num::Integer;
+use num::{CheckedAdd, Integer};
 
-pub(super) fn build_extend<T: ArrowNativeType + Integer>(array: &ArrayData) -> Extend {
+pub(super) fn build_extend<T: ArrowNativeType + Integer + CheckedAdd>(array: &ArrayData) -> Extend {
     let offsets = array.buffer::<T>(0);
     if array.null_count() == 0 {
         // fast case where we can copy regions without nullability checks

diff --git a/arrow-data/src/transform/utils.rs b/arrow-data/src/transform/utils.rs
@@ -16,7 +16,7 @@
 // under the License.
 
 use arrow_buffer::{bit_util, ArrowNativeType, MutableBuffer};
-use num::Integer;
+use num::{CheckedAdd, Integer};
 
 /// extends the `buffer` to be able to hold `len` bits, setting all bits of the new size to zero.
 #[inline]
@@ -27,7 +27,7 @@ pub(super) fn resize_for_bits(buffer: &mut MutableBuffer, len: usize) {
     }
 }
 
-pub(super) fn extend_offsets<T: ArrowNativeType + Integer>(
+pub(super) fn extend_offsets<T: ArrowNativeType + Integer + CheckedAdd>(
     buffer: &mut MutableBuffer,
     mut last_offset: T,
     offsets: &[T],
@@ -36,7 +36,7 @@ pub(super) fn extend_offsets<T: ArrowNativeType + Integer>(
     offsets.windows(2).for_each(|offsets| {
         // compute the new offset
         let length = offsets[1] - offsets[0];
-        last_offset = last_offset + length;
+        last_offset = last_offset.checked_add(&length).expect("offset overflow");
         buffer.push(last_offset);
     });
 }
@@ -55,3 +55,17 @@ pub(super) unsafe fn get_last_offset<T: ArrowNativeType>(
     debug_assert!(prefix.is_empty() && suffix.is_empty());
     *offsets.get_unchecked(offsets.len() - 1)
 }
+
+
+#[cfg(test)]
+mod tests {
+    use arrow_buffer::MutableBuffer;
+    use crate::transform::utils::extend_offsets;
+
+    #[test]
+    #[should_panic(expected = "offset overflow")]
+    fn test_overflow() {
+        let mut buffer = MutableBuffer::new(10);
+        extend_offsets(&mut buffer, i32::MAX - 4, &[0, 5]);
+    }
+}
diff --git a/arrow-data/src/transform/variable_size.rs b/arrow-data/src/transform/variable_size.rs
@@ -18,7 +18,7 @@
 use crate::ArrayData;
 use arrow_buffer::{ArrowNativeType, MutableBuffer};
 use num::traits::AsPrimitive;
-use num::Integer;
+use num::{CheckedAdd, Integer};
 
 use super::{
     Extend, _MutableArrayData,
@@ -39,7 +39,9 @@ fn extend_offset_values<T: ArrowNativeType + AsPrimitive<usize>>(
     buffer.extend_from_slice(new_values);
 }
 
-pub(super) fn build_extend<T: ArrowNativeType + Integer + AsPrimitive<usize>>(
+pub(super) fn build_extend<
+    T: ArrowNativeType + Integer + CheckedAdd + AsPrimitive<usize>,
+>(
     array: &ArrayData,
 ) -> Extend {
     let offsets = array.buffer::<T>(0);