From 185c0953d15c4358723f60bbacf74599916e6d88 Mon Sep 17 00:00:00 2001
From: Raphael Taylor-Davies <r.taylordavies@googlemail.com>
Date: Wed, 12 Apr 2023 20:52:04 +0100
Subject: [PATCH] Checked arithmetic

---
 arrow-buffer/src/buffer/offset.rs | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/arrow-buffer/src/buffer/offset.rs b/arrow-buffer/src/buffer/offset.rs
index ec181d017e15..bfafe3306aed 100644
--- a/arrow-buffer/src/buffer/offset.rs
+++ b/arrow-buffer/src/buffer/offset.rs
@@ -61,7 +61,11 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
 
     /// Create a new [`OffsetBuffer`] containing `len + 1` `0` values
     pub fn new_zeroed(len: usize) -> Self {
-        let buffer = MutableBuffer::from_len_zeroed((len + 1) * std::mem::size_of::<O>());
+        let len_bytes = len
+            .checked_add(1)
+            .and_then(|o| o.checked_mul(std::mem::size_of::<O>()))
+            .expect("overflow");
+        let buffer = MutableBuffer::from_len_zeroed(len_bytes);
         Self(buffer.into_buffer().into())
     }
 
@@ -116,6 +120,18 @@ mod tests {
     #[test]
     fn offsets() {
         OffsetBuffer::new(vec![0, 1, 2, 3].into());
+
+        let offsets = OffsetBuffer::<i32>::new_zeroed(3);
+        assert_eq!(offsets.as_ref(), &[0; 4]);
+
+        let offsets = OffsetBuffer::<i32>::new_zeroed(0);
+        assert_eq!(offsets.as_ref(), &[0; 1]);
+    }
+
+    #[test]
+    #[should_panic(expected = "overflow")]
+    fn offsets_new_zeroed_overflow() {
+        OffsetBuffer::<i32>::new_zeroed(usize::MAX);
     }
 
     #[test]