From e992922cdaf1ce9259c501589c1bb84e4398e46c Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sat, 25 Apr 2026 11:19:49 -0400 Subject: [PATCH] Prevent repeat slice length overflow (#9819) - None. MutableBuffer repeated slices used unchecked usize arithmetic when calculating the number of bytes to reserve. In optimized builds, very large repeat counts could wrap the capacity calculation before copying repeated bytes. This adds checked arithmetic for repeated slice byte length validation before reserving capacity and copying repeated data. Yes. This adds regression coverage for overflowing repeated slice length calculations. Invalid repeat counts whose requested byte length cannot be represented without overflow now panic consistently. There are no API changes. --- arrow-buffer/src/buffer/mutable.rs | 211 +++++++++++++++++++++++++++++ 1 file changed, 211 insertions(+) diff --git a/arrow-buffer/src/buffer/mutable.rs b/arrow-buffer/src/buffer/mutable.rs index 63fdbf598bdb..f0baf8647929 100644 --- a/arrow-buffer/src/buffer/mutable.rs +++ b/arrow-buffer/src/buffer/mutable.rs @@ -222,6 +222,81 @@ impl MutableBuffer { } } + /// Adding to this mutable buffer `slice_to_repeat` repeated `repeat_count` times. + /// + /// # Example + /// + /// ## Repeat the same string bytes multiple times + /// ``` + /// # use arrow_buffer::buffer::MutableBuffer; + /// let mut buffer = MutableBuffer::new(0); + /// let bytes_to_repeat = b"ab"; + /// buffer.repeat_slice_n_times(bytes_to_repeat, 3); + /// assert_eq!(buffer.as_slice(), b"ababab"); + /// ``` + pub fn repeat_slice_n_times( + &mut self, + slice_to_repeat: &[T], + repeat_count: usize, + ) { + if repeat_count == 0 || slice_to_repeat.is_empty() { + return; + } + + let bytes_to_repeat = size_of_val(slice_to_repeat); + let repeated_bytes = repeat_count + .checked_mul(bytes_to_repeat) + .expect("repeated slice byte length overflow"); + self.len + .checked_add(repeated_bytes) + .expect("mutable buffer length overflow"); + + // Ensure capacity + self.reserve(repeated_bytes); + + // Save the length before we do all the copies to know where to start from + let length_before = self.len; + + // Copy the initial slice once so we can use doubling strategy on it + self.extend_from_slice(slice_to_repeat); + + // This tracks how much bytes we have added by repeating so far + let added_repeats_length = bytes_to_repeat; + assert_eq!( + self.len - length_before, + added_repeats_length, + "should copy exactly the same number of bytes" + ); + + // Number of times the slice was repeated + let mut already_repeated_times = 1; + + // We will use doubling strategy to fill the buffer in log(repeat_count) steps + while already_repeated_times < repeat_count { + // How many slices can we copy in this iteration + // (either double what we have, or just the remaining ones) + let number_of_slices_to_copy = + already_repeated_times.min(repeat_count - already_repeated_times); + let number_of_bytes_to_copy = number_of_slices_to_copy * bytes_to_repeat; + + unsafe { + // Get to the start of the data before we started copying anything + let src = self.data.as_ptr().add(length_before) as *const u8; + + // Go to the current location to copy to (end of current data) + let dst = self.data.as_ptr().add(self.len); + + // SAFETY: the pointers are not overlapping as there is `number_of_bytes_to_copy` or less between them + std::ptr::copy_nonoverlapping(src, dst, number_of_bytes_to_copy) + } + + // Advance the length by the amount of data we just copied (doubled) + self.len += number_of_bytes_to_copy; + + already_repeated_times += number_of_slices_to_copy; + } + } + #[cold] fn reallocate(&mut self, capacity: usize) { let new_layout = Layout::from_size_align(capacity, self.layout.align()).unwrap(); @@ -1182,4 +1257,140 @@ mod tests { assert_eq!(pool.used(), 0); } } + + fn create_expected_repeated_slice( + slice_to_repeat: &[T], + repeat_count: usize, + ) -> Buffer { + let mut expected = MutableBuffer::new(size_of_val(slice_to_repeat) * repeat_count); + for _ in 0..repeat_count { + // Not using push_slice_repeated as this is the function under test + expected.extend_from_slice(slice_to_repeat); + } + expected.into() + } + + // Helper to test a specific repeat count with various slice sizes + fn test_repeat_count( + repeat_count: usize, + test_data: &[T], + ) { + let mut buffer = MutableBuffer::new(0); + buffer.repeat_slice_n_times(test_data, repeat_count); + + let expected = create_expected_repeated_slice(test_data, repeat_count); + let result: Buffer = buffer.into(); + + assert_eq!( + result, + expected, + "Failed for repeat_count={}, slice_len={}", + repeat_count, + test_data.len() + ); + } + + #[test] + fn test_repeat_slice_count_edge_cases() { + // Empty slice + test_repeat_count(100, &[] as &[i32]); + + // Zero repeats + test_repeat_count(0, &[1i32, 2, 3]); + } + + #[test] + #[should_panic(expected = "repeated slice byte length overflow")] + fn test_repeat_slice_count_multiply_overflow() { + let mut buffer = MutableBuffer::new(0); + buffer.repeat_slice_n_times(&[0_u64], usize::MAX / mem::size_of::() + 1); + } + + #[test] + #[should_panic(expected = "mutable buffer length overflow")] + fn test_repeat_slice_count_len_overflow() { + let mut buffer = MutableBuffer::new(0); + buffer.push(0_u8); + buffer.repeat_slice_n_times(&[0_u8], usize::MAX); + } + + #[test] + fn test_small_repeats_counts() { + // test any special implementation for small repeat counts + let data = &[1u8, 2, 3, 4, 5]; + + for _ in 1..=10 { + test_repeat_count(2, data); + } + } + + #[test] + fn test_different_size_of_i32_repeat_slice() { + let data: &[i32] = &[1, 2, 3]; + let data_with_single_item: &[i32] = &[42]; + + for data in &[data, data_with_single_item] { + for item in 1..=9 { + let base_repeat_count = 2_usize.pow(item); + test_repeat_count(base_repeat_count - 1, data); + test_repeat_count(base_repeat_count, data); + test_repeat_count(base_repeat_count + 1, data); + } + } + } + + #[test] + fn test_different_size_of_u8_repeat_slice() { + let data: &[u8] = &[1, 2, 3]; + let data_with_single_item: &[u8] = &[10]; + + for data in &[data, data_with_single_item] { + for item in 1..=9 { + let base_repeat_count = 2_usize.pow(item); + test_repeat_count(base_repeat_count - 1, data); + test_repeat_count(base_repeat_count, data); + test_repeat_count(base_repeat_count + 1, data); + } + } + } + + #[test] + fn test_different_size_of_u16_repeat_slice() { + let data: &[u16] = &[1, 2, 3]; + let data_with_single_item: &[u16] = &[10]; + + for data in &[data, data_with_single_item] { + for item in 1..=9 { + let base_repeat_count = 2_usize.pow(item); + test_repeat_count(base_repeat_count - 1, data); + test_repeat_count(base_repeat_count, data); + test_repeat_count(base_repeat_count + 1, data); + } + } + } + + #[test] + fn test_various_slice_lengths() { + // Test different slice lengths with same repeat pattern + let repeat_count = 37; // Arbitrary non-power-of-2 + + // Single element + test_repeat_count(repeat_count, &[42i32]); + + // Small slices + test_repeat_count(repeat_count, &[1i32, 2]); + test_repeat_count(repeat_count, &[1i32, 2, 3]); + test_repeat_count(repeat_count, &[1i32, 2, 3, 4]); + test_repeat_count(repeat_count, &[1i32, 2, 3, 4, 5]); + + // Larger slices + let data_10: Vec = (0..10).collect(); + test_repeat_count(repeat_count, &data_10); + + let data_100: Vec = (0..100).collect(); + test_repeat_count(repeat_count, &data_100); + + let data_1000: Vec = (0..1000).collect(); + test_repeat_count(repeat_count, &data_1000); + } }