diff --git a/arrow-buffer/src/util/bit_chunk_iterator.rs b/arrow-buffer/src/util/bit_chunk_iterator.rs index 8c7ec5e9a8f6..6655c397ebbd 100644 --- a/arrow-buffer/src/util/bit_chunk_iterator.rs +++ b/arrow-buffer/src/util/bit_chunk_iterator.rs @@ -220,10 +220,8 @@ pub struct BitChunks<'a> { impl<'a> BitChunks<'a> { /// Create a new [`BitChunks`] from a byte array, and an offset and length in bits pub fn new(buffer: &'a [u8], offset: usize, len: usize) -> Self { - assert!( - ceil(offset + len, 8) <= buffer.len(), - "offset + len out of bounds" - ); + let end = offset.checked_add(len).expect("offset + len out of bounds"); + assert!(ceil(end, 8) <= buffer.len(), "offset + len out of bounds"); let byte_offset = offset / 8; let bit_offset = offset % 8; @@ -550,6 +548,13 @@ mod tests { buffer.bit_chunks(1, ALLOC_SIZE * 8); } + #[test] + #[should_panic(expected = "offset + len out of bounds")] + fn test_out_of_bound_should_panic_when_offset_and_length_overflow() { + let buffer = Buffer::from(vec![0xFF_u8; 8]); + buffer.bit_chunks(1, usize::MAX); + } + #[test] #[allow(clippy::assertions_on_constants)] fn test_unaligned_bit_chunk_iterator() {