[arrow-rs] branch master updated: fix(buffer): panic on end index overflow in `MutableBuffer::set_null_bits` (#4621)

[tustvold]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new d5b713ada2 fix(buffer): panic on end index overflow in 
`MutableBuffer::set_null_bits` (#4621)
d5b713ada2 is described below

commit d5b713ada2823443293b5616789e3c6c75bf48bb
Author: Tomoaki Kawada <kaw...@kmckk.co.jp>
AuthorDate: Wed Aug 2 17:46:05 2023 +0900

    fix(buffer): panic on end index overflow in `MutableBuffer::set_null_bits` 
(#4621)
---
 arrow-buffer/src/buffer/mutable.rs | 36 +++++++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/arrow-buffer/src/buffer/mutable.rs 
b/arrow-buffer/src/buffer/mutable.rs
index 3e66e7f23f..0d2d2ed751 100644
--- a/arrow-buffer/src/buffer/mutable.rs
+++ b/arrow-buffer/src/buffer/mutable.rs
@@ -168,7 +168,14 @@ impl MutableBuffer {
     /// `len` of the buffer and so can be used to initialize the memory region 
from
     /// `len` to `capacity`.
     pub fn set_null_bits(&mut self, start: usize, count: usize) {
-        assert!(start + count <= self.layout.size());
+        assert!(
+            start.saturating_add(count) <= self.layout.size(),
+            "range start index {start} and count {count} out of bounds for \
+            buffer of length {}",
+            self.layout.size(),
+        );
+
+        // Safety: `self.data[start..][..count]` is in-bounds and well-aligned 
for `u8`
         unsafe {
             std::ptr::write_bytes(self.data.as_ptr().add(start), 0, count);
         }
@@ -932,4 +939,31 @@ mod tests {
         buffer.shrink_to_fit();
         assert!(buffer.capacity() >= 64 && buffer.capacity() < 128);
     }
+
+    #[test]
+    fn test_mutable_set_null_bits() {
+        let mut buffer = MutableBuffer::new(8).with_bitset(8, true);
+
+        for i in 0..=buffer.capacity() {
+            buffer.set_null_bits(i, 0);
+            assert_eq!(buffer[..8], [255; 8][..]);
+        }
+
+        buffer.set_null_bits(1, 4);
+        assert_eq!(buffer[..8], [255, 0, 0, 0, 0, 255, 255, 255][..]);
+    }
+
+    #[test]
+    #[should_panic = "out of bounds for buffer of length"]
+    fn test_mutable_set_null_bits_oob() {
+        let mut buffer = MutableBuffer::new(64);
+        buffer.set_null_bits(1, buffer.capacity());
+    }
+
+    #[test]
+    #[should_panic = "out of bounds for buffer of length"]
+    fn test_mutable_set_null_bits_oob_by_overflow() {
+        let mut buffer = MutableBuffer::new(0);
+        buffer.set_null_bits(1, usize::MAX);
+    }
 }