fix(arrow-buffer): prevent use-after-free in Buffer::shrink_to_fit when MemoryReservation::resize panics

arrow-buffer/src/buffer/immutable.rs

@@ -224,17 +224,16 @@ impl Buffer {
         if desired_capacity < self.capacity()
             && let Some(bytes) = Arc::get_mut(&mut self.data)
         {
-            if bytes.try_realloc(desired_capacity).is_ok() {
-                // Realloc complete - update our pointer into `bytes`:
-                self.ptr = if is_empty {
-                    bytes.as_ptr()
-                } else {
-                    // SAFETY: we kept all elements leading up to the offset
-                    unsafe { bytes.as_ptr().add(offset) }
-                }
-            } else {
-                // Failure to reallocate is fine; we just failed to free up memory.
-            }
+            bytes
+                .try_realloc(desired_capacity, |base| {
+                    self.ptr = if is_empty {
+                        base.as_ptr()
+                    } else {
+                        // SAFETY: we kept all elements leading up to the offset
+                        unsafe { base.as_ptr().add(offset) }
+                    };
+                })
+                .ok(); // Failure to reallocate is fine; we just failed to free up memory.
         }
     }
 
@@ -1158,4 +1157,58 @@ mod tests {
             assert_eq!(buffer_back.as_slice(), expected.as_slice());
         }
     }
+
+    #[test]
+    #[cfg(feature = "pool")]
+    fn test_shrink_to_fit_panicking_reservation() {
+        use std::panic::{AssertUnwindSafe, catch_unwind};
+
+        use crate::pool::{MemoryPool, MemoryReservation};
+
+        #[derive(Debug)]
+        struct PanicPool;
+
+        #[derive(Debug)]
+        struct PanicReservation {
+            panicked: bool,
+        }
+
+        impl MemoryReservation for PanicReservation {
+            fn size(&self) -> usize {
+                0
+            }
+            fn resize(&mut self, _: usize) {
+                if !self.panicked {
+                    self.panicked = true;
+                    panic!("intentional panic in resize");
+                }
+            }
+        }
+
+        impl MemoryPool for PanicPool {
+            fn reserve(&self, _: usize) -> Box<dyn MemoryReservation> {
+                Box::new(PanicReservation { panicked: false })
+            }
+            fn available(&self) -> isize {
+                isize::MAX
+            }
+            fn used(&self) -> usize {
+                0
+            }
+            fn capacity(&self) -> usize {
+                usize::MAX
+            }
+        }
+
+        let pool = PanicPool;
+        let data: Vec<u8> = (0..8).collect();
+        let mut buf = Buffer::from_slice_ref(data.as_slice());
+        buf.claim(&pool);
+
+        // shrink_to_fit panics because PanicReservation::resize panics, but
+        // Buffer::ptr must stay consistent with Bytes::ptr (no use-after-free).
+        let _ = catch_unwind(AssertUnwindSafe(|| buf.shrink_to_fit()));
+
+        assert_eq!(buf.as_slice(), data.as_slice());
+    }
 }

arrow-buffer/src/bytes.rs

@@ -134,7 +134,15 @@ impl Bytes {
     /// Returns `Err` if the memory was allocated with a custom allocator,
     /// or the call to `realloc` failed, for whatever reason.
     /// In case of `Err`, the [`Bytes`] will remain as it was (i.e. have the old size).
-    pub(crate) fn try_realloc(&mut self, new_len: usize) -> Result<(), ()> {
+    ///
+    /// `on_reallocated` is called after [`Bytes`] has updated its internal
+    /// pointer, but before resizing the memory reservation, which may call user
+    /// code.
+    pub(crate) fn try_realloc(
+        &mut self,
+        new_len: usize,
+        on_reallocated: impl FnOnce(NonNull<u8>),
+    ) -> Result<(), ()> {
         if let Deallocation::Standard(old_layout) = self.deallocation {
             if old_layout.size() == new_len {
                 return Ok(()); // Nothing to do
@@ -162,6 +170,7 @@ impl Bytes {
                     self.ptr = ptr;
                     self.len = new_len;
                     self.deallocation = Deallocation::Standard(new_layout);
+                    on_reallocated(ptr);
 
                     #[cfg(feature = "pool")]
                     {