Several soundness issues

[Evian-Zhang]

Describe the bug

Hello, we are security researchers targeting Rust security. By running our tool in this repository, we found several soundness issues. This report is written by 100% human. We promise that all you read will never be generated by LLM. Moreover, we are responsible, so if you confirm these issues do exist, we are happy to file PRs to fix it.

Issue 1: arrow-buffer's MutableBuffer::into_buffer() violates exception safety, leading to double-free

Sources are located here:

arrow-rs/arrow-buffer/src/buffer/mutable.rs

Lines 536 to 546 in d991919

	#[inline]
	pub(super) fn into_buffer(self) -> Buffer {
	let bytes = unsafe { Bytes::new(self.data, self.len, Deallocation::Standard(self.layout)) };
	#[cfg(feature = "pool")]
	{
	let reservation = self.reservation.lock().unwrap().take();
	*bytes.reservation.lock().unwrap() = reservation;
	}
	std::mem::forget(self);
	Buffer::from(bytes)
	}

If the mutex self.reservation was poisoned before, then the unwrap will leading to panic. However, The mem::forget has not yet been invoked, so both bytes and self's drop implementation will be invoked when unwinding, and these two variable owns the same memory region, leading to a double free (both Bytes and MutableBuffer's drop implementation involves std::alloc::dealloc)

To fix it, I think we can just use if let Some instead of unwrap to the mutex.

Issue 2: arrow-buffer's Buffer::shrink_to_fit() violates exception safety, leading to use-after-free

Sources are located here:

arrow-rs/arrow-buffer/src/buffer/immutable.rs

Lines 194 to 227 in d991919

	/// Tries to shrink the capacity of the buffer as much as possible, freeing unused memory.
	///
	/// If the buffer is shared, this is a no-op.
	///
	/// If the memory was allocated with a custom allocator, this is a no-op.
	///
	/// If the capacity is already less than or equal to the desired capacity, this is a no-op.
	///
	/// The memory region will be reallocated using `std::alloc::realloc`.
	pub fn shrink_to_fit(&mut self) {
	let offset = self.ptr_offset();
	let is_empty = self.is_empty();
	let desired_capacity = if is_empty {
	0
	} else {
	// For realloc to work, we cannot free the elements before the offset
	offset + self.len()
	};
	if desired_capacity < self.capacity() {
	if 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.
	}
	}
	}
	}

arrow-rs/arrow-buffer/src/bytes.rs

Lines 131 to 178 in d991919

	/// Try to reallocate the underlying memory region to a new size (smaller or larger).
	///
	/// Only works for bytes allocated with the standard allocator.
	/// 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<(), ()> {
	if let Deallocation::Standard(old_layout) = self.deallocation {
	if old_layout.size() == new_len {
	return Ok(()); // Nothing to do
	}
	if let Ok(new_layout) = std::alloc::Layout::from_size_align(new_len, old_layout.align())
	{
	let old_ptr = self.ptr.as_ptr();
	let new_ptr = match new_layout.size() {
	0 => {
	// SAFETY: Verified that old_layout.size != new_len (0)
	unsafe { std::alloc::dealloc(self.ptr.as_ptr(), old_layout) };
	Some(dangling_ptr())
	}
	// SAFETY: the call to `realloc` is safe if all the following hold (from https://doc.rust-lang.org/stable/std/alloc/trait.GlobalAlloc.html#method.realloc):
	// * `old_ptr` must be currently allocated via this allocator (guaranteed by the invariant/contract of `Bytes`)
	// * `old_layout` must be the same layout that was used to allocate that block of memory (same)
	// * `new_len` must be greater than zero
	// * `new_len`, when rounded up to the nearest multiple of `layout.align()`, must not overflow `isize` (guaranteed by the success of `Layout::from_size_align`)
	_ => NonNull::new(unsafe { std::alloc::realloc(old_ptr, old_layout, new_len) }),
	};
	if let Some(ptr) = new_ptr {
	self.ptr = ptr;
	self.len = new_len;
	self.deallocation = Deallocation::Standard(new_layout);
	#[cfg(feature = "pool")]
	{
	// Resize reservation
	self.resize_reservation(new_len);
	}
	return Ok(());
	}
	}
	}
	Err(())
	}

In Buffer::shrink_to_fit, it will invokes Bytes::try_realloc to reallocate memories, and update the ptr if the pointer is updated. In Bytes::try_realloc, after the memory hold by Bytes is reallocated, self.resize_reservation() will be invoked. Bytes::resize_reservation will further invokes user-defined custom reservation function. If such reservation function panics, in view of buffer, the memory is in fact reallocated, but self.ptr has not updated yet. If user uses catch_unwind to recover from unwinding, methods such as Buffer::as_slice will let user access freed-memories, leading to use-after-free.

I can't come up with a graceful simple solution to fix it. For me, there are three approaches:

• Declare shrink_to_fit as unsafe, and requires the memory pool should not panic in the SAFETY requirement
• Modify Bytes::try_realloc to accept a callback function, which will be invoked after the reallocating and before calling self.resize_reservation()
• Before calling Bytes::try_realloc, let self.ptr points to a static valid memory buffer. If try_realloc succeeds, self.ptr will still be updated correctly; If it panics, at least there will be no memory-related issues.

Issue 3 (may be): The SAFETY documentation of arrow-buffer's BooleanBufferBuilder::extend_trusted_len is not accurate, potentially violating exception safety

Sources are located here:

arrow-rs/arrow-buffer/src/builder/boolean.rs

Lines 330 to 343 in d991919

	/// Extends the builder from a trusted length iterator of booleans.
	/// # Safety
	/// Callers must ensure that `iter` reports an exact size via `size_hint`.
	///
	#[inline]
	pub unsafe fn extend_trusted_len<I>(&mut self, iterator: I)
	where
	I: Iterator<Item = bool>,
	{
	let len = iterator.size_hint().0;
	unsafe { self.buffer.extend_bool_trusted_len(iterator, self.len) };
	self.len += len;
	}
	}

arrow-rs/arrow-buffer/src/buffer/mutable.rs

Lines 705 to 838 in d991919

	/// Extends this buffer with boolean values.
	///
	/// This requires `iter` to report an exact size via `size_hint`.
	/// `offset` indicates the starting offset in bits in this buffer to begin writing to
	/// and must be less than or equal to the current length of this buffer.
	/// All bits not written to (but readable due to byte alignment) will be zeroed out.
	///
	/// # Panics
	///
	/// Panics if `iter` does not report an exact size via `size_hint`, or if it yields fewer
	/// items than reported, or if extending the buffer requires reserving a capacity that fails
	/// for the same reasons as [`MutableBuffer::reserve`].
	///
	/// # Safety
	/// Callers must ensure that `iter` reports an exact size via `size_hint`.
	#[inline]
	pub unsafe fn extend_bool_trusted_len<I: Iterator<Item = bool>>(
	&mut self,
	mut iter: I,
	offset: usize,
	) {
	let (lower, upper) = iter.size_hint();
	let len = upper.expect("Iterator must have exact size_hint");
	assert_eq!(lower, len, "Iterator must have exact size_hint");
	debug_assert!(
	offset <= self.len * 8,
	"offset must be <= buffer length in bits"
	);
	if len == 0 {
	return;
	}
	let start_len = offset;
	let end_bit = start_len + len;
	// SAFETY: we will initialize all newly exposed bytes before they are read
	let new_len_bytes = bit_util::ceil(end_bit, 8);
	if new_len_bytes > self.len {
	self.reserve(new_len_bytes - self.len);
	// SAFETY: caller will initialize all newly exposed bytes before they are read
	unsafe { self.set_len(new_len_bytes) };
	}
	let slice = self.as_slice_mut();
	let mut bit_idx = start_len;
	// ---- Unaligned prefix: advance to the next 64-bit boundary ----
	let misalignment = bit_idx & 63;
	let prefix_bits = if misalignment == 0 {
	0
	} else {
	(64 - misalignment).min(end_bit - bit_idx)
	};
	if prefix_bits != 0 {
	let byte_start = bit_idx / 8;
	let byte_end = bit_util::ceil(bit_idx + prefix_bits, 8);
	let bit_offset = bit_idx % 8;
	// Clear any newly-visible bits in the existing partial byte
	if bit_offset != 0 {
	let keep_mask = (1u8 << bit_offset).wrapping_sub(1);
	slice[byte_start] &= keep_mask;
	}
	// Zero any new bytes we will partially fill in this prefix
	let zero_from = if bit_offset == 0 {
	byte_start
	} else {
	byte_start + 1
	};
	if byte_end > zero_from {
	slice[zero_from..byte_end].fill(0);
	}
	for _ in 0..prefix_bits {
	let v = iter.next().unwrap();
	if v {
	let byte_idx = bit_idx / 8;
	let bit = bit_idx % 8;
	slice[byte_idx] |= 1 << bit;
	}
	bit_idx += 1;
	}
	}
	if bit_idx < end_bit {
	// ---- Aligned middle: write u64 chunks ----
	debug_assert_eq!(bit_idx & 63, 0);
	let remaining_bits = end_bit - bit_idx;
	let chunks = remaining_bits / 64;
	let words_start = bit_idx / 8;
	let words_end = words_start + chunks * 8;
	for dst in slice[words_start..words_end].chunks_exact_mut(8) {
	let mut packed: u64 = 0;
	for i in 0..64 {
	packed |= (iter.next().unwrap() as u64) << i;
	}
	dst.copy_from_slice(&packed.to_le_bytes());
	bit_idx += 64;
	}
	// ---- Unaligned suffix: remaining < 64 bits ----
	let suffix_bits = end_bit - bit_idx;
	if suffix_bits != 0 {
	debug_assert_eq!(bit_idx % 8, 0);
	let byte_start = bit_idx / 8;
	let byte_end = bit_util::ceil(end_bit, 8);
	slice[byte_start..byte_end].fill(0);
	for _ in 0..suffix_bits {
	let v = iter.next().unwrap();
	if v {
	let byte_idx = bit_idx / 8;
	let bit = bit_idx % 8;
	slice[byte_idx] |= 1 << bit;
	}
	bit_idx += 1;
	}
	}
	}
	// Clear any unused bits in the last byte
	let remainder = end_bit % 8;
	if remainder != 0 {
	let mask = (1u8 << remainder).wrapping_sub(1);
	slice[bit_util::ceil(end_bit, 8) - 1] &= mask;
	}
	debug_assert_eq!(bit_idx, end_bit);
	}

BooleanBufferBuilder::extend_trusted_len will call MutableBuffer::extend_bool_trusted_len. In this function, there will be firstly a set_len (line 746), and then invokes iter.next to fill those slots. If I::next panics, then there will be a typical exception safety violation. The SAFETY requirement of these two functions are "Callers must ensure that iter reports an exact size via size_hint." I think maybe this is a little bit inaccurate, although this can implicitly indicate that I::next must never panic. Maybe we should explicitly write this requirement in the SAFETY documentation?

To Reproduce

Issue 1

1. User create a MaliciousMemoryPool that implements arrow_buffer::pool::MemoryPool, which panics at resize function.
2. User register this pool with MutableBuffer::claim(), in which the MutableBuffer will hold this memory pool
3. User invokes MutableBuffer::resize() under std::panics::catch_unwind(). Although the resize will immediately panic, while the catch_unwind helps recover from panic, and user can access the MutableBuffer instance after panicking. However, the mutex self.reservation is poisoned here.
4. User invokes MutableBuffer::into_buffer(). Since the mutex is poisoned, the code will panic by unwrap. And during unwinding, the destructors of MutableBuffer and Bytes are automatically invoked, leading to UAF.

Issue 2

1. User create a MaliciousMemoryPool that implements arrow_buffer::pool::MemoryPool. The reserve function will create a malicious MemoryReservation implementor, which will panic at resize.
2. User register this pool with Buffer::claim(), in which the self.data will hold this memory reservation
3. User invokes Buffer::shrink_to_fit() under std::panics::catch_unwind()
4. User invokes Buffer::as_slice() after the panic is recovered

Expected behavior

Be safe and sound, even in context of exception safety.

Additional context

No response