Skip to content

Commit

Permalink
Fix integer overflow in BlobVec::push for ZST (#10799)
Browse files Browse the repository at this point in the history
  • Loading branch information
stepancheg committed Jan 4, 2024
1 parent cc2a77b commit cf70f53
Showing 1 changed file with 42 additions and 7 deletions.
49 changes: 42 additions & 7 deletions crates/bevy_ecs/src/storage/blob_vec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,8 @@ impl BlobVec {
if item_layout.size() == 0 {
BlobVec {
data,
// ZST `BlobVec` max size is `usize::MAX`, and `reserve_exact` for ZST assumes
// the capacity is always `usize::MAX` and panics if it overflows.
capacity: usize::MAX,
len: 0,
item_layout,
Expand Down Expand Up @@ -109,20 +111,29 @@ impl BlobVec {
///
/// Note that the allocator may give the collection more space than it requests. Therefore, capacity can not be relied upon
/// to be precisely minimal.
///
/// # Panics
///
/// Panics if new capacity overflows `usize`.
pub fn reserve_exact(&mut self, additional: usize) {
let available_space = self.capacity - self.len;
if available_space < additional && self.item_layout.size() > 0 {
if available_space < additional {
// SAFETY: `available_space < additional`, so `additional - available_space > 0`
let increment = unsafe { NonZeroUsize::new_unchecked(additional - available_space) };
// SAFETY: not called for ZSTs
unsafe { self.grow_exact(increment) };
self.grow_exact(increment);
}
}

// SAFETY: must not be called for a ZST item layout
#[warn(unsafe_op_in_unsafe_fn)] // to allow unsafe blocks in unsafe fn
unsafe fn grow_exact(&mut self, increment: NonZeroUsize) {
debug_assert!(self.item_layout.size() != 0);
/// Grows the capacity by `increment` elements.
///
/// # Panics
///
/// Panics if the new capacity overflows `usize`.
/// For ZST it panics unconditionally because ZST `BlobVec` capacity
/// is initialized to `usize::MAX` and always stays that way.
fn grow_exact(&mut self, increment: NonZeroUsize) {
// If we got here with ZST, we requested total capacity > usize::MAX.
assert!(self.item_layout.size() != 0, "ZST capacity overflow");

let new_capacity = self.capacity + increment.get();
let new_layout =
Expand Down Expand Up @@ -609,6 +620,30 @@ mod tests {
let _ = unsafe { BlobVec::new(item_layout, Some(drop), 0) };
}

#[test]
#[should_panic(expected = "ZST capacity overflow")]
fn blob_vec_zst_size_overflow() {
// SAFETY: no drop is correct drop for `()`.
let mut blob_vec = unsafe { BlobVec::new(Layout::new::<()>(), None, 0) };

assert_eq!(usize::MAX, blob_vec.capacity(), "Self-check");

// SAFETY: Because `()` is a ZST trivial drop type, and because `BlobVec` capacity
// is always `usize::MAX` for ZSTs, we can arbitrarily set the length
// and still be sound.
unsafe {
blob_vec.set_len(usize::MAX);
}

// SAFETY: `BlobVec` was initialized for `()`, so it is safe to push `()` to it.
unsafe {
OwningPtr::make((), |ptr| {
// This should panic because len is usize::MAX, remaining capacity is 0.
blob_vec.push(ptr);
});
}
}

#[test]
fn aligned_zst() {
// NOTE: This test is explicitly for uncovering potential UB with miri.
Expand Down

0 comments on commit cf70f53

Please sign in to comment.