Documentation of `arrow-buffer`'s `BooleanBufferBuilder::extend_trusted_len` does not include exception safety statement, potentially allowing use-of-uninitialized-memory

[Evian-Zhang]

Describe the bug

Split from #10281

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

use arrow_buffer::builder::BooleanBufferBuilder;
use std::panic::{AssertUnwindSafe, catch_unwind};

struct PanicIter {
    remaining: usize,
    yielded: usize,
    panic_after: usize,
}

impl Iterator for PanicIter {
    type Item = bool;

    fn next(&mut self) -> Option<Self::Item> {
        if self.remaining == 0 {
            return None;
        }
        self.remaining -= 1;
        self.yielded += 1;
        if self.yielded > self.panic_after {
            panic!("Issue 3: panic during iter.next() after set_len already ran");
        }
        Some(true)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.remaining, Some(self.remaining))
    }
}

fn main() {
    let mut builder = BooleanBufferBuilder::new(0);

    // Iterator: reports 64 items, yields 8, then panics on the 9th
    let iter = PanicIter {
        remaining: 64,
        yielded: 0,
        panic_after: 8,
    };

    let result = catch_unwind(AssertUnwindSafe(|| unsafe {
        builder.extend_trusted_len(iter);
    }));

    println!(
        "builder.len() = {} (logical len — NOT updated by extend_trusted_len)",
        builder.len()
    );
    println!(
        "builder.as_slice().len() = {} (internal MutableBuffer len — WAS expanded by set_len)",
        builder.as_slice().len()
    );

    let slice = builder.as_slice();

    println!(
        "raw bytes: {:02x?}",
        &slice[..slice.len().min(16)]
    );
}

Expected behavior

No soundness issue happens

Additional context

No response