Skip to content

Commit

Permalink
fix: Drop guard was deallocating with the incorrect size
Browse files Browse the repository at this point in the history
InPlaceDstBufDrop holds onto the allocation before the shrinking happens
which means it must deallocate the destination elements but the source
allocation.

(cherry picked from commit 5796b3c)
  • Loading branch information
the8472 authored and cuviper committed Jan 25, 2024
1 parent e55be1a commit 73b36b1
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 13 deletions.
10 changes: 6 additions & 4 deletions library/alloc/src/vec/in_place_collect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@
//! This is handled by the [`InPlaceDrop`] guard for sink items (`U`) and by
//! [`vec::IntoIter::forget_allocation_drop_remaining()`] for remaining source items (`T`).
//!
//! If dropping any remaining source item (`T`) panics then [`InPlaceDstBufDrop`] will handle dropping
//! If dropping any remaining source item (`T`) panics then [`InPlaceDstDataSrcBufDrop`] will handle dropping
//! the already collected sink items (`U`) and freeing the allocation.
//!
//! [`vec::IntoIter::forget_allocation_drop_remaining()`]: super::IntoIter::forget_allocation_drop_remaining()
Expand Down Expand Up @@ -158,11 +158,12 @@ use crate::alloc::{handle_alloc_error, Global};
use core::alloc::Allocator;
use core::alloc::Layout;
use core::iter::{InPlaceIterable, SourceIter, TrustedRandomAccessNoCoerce};
use core::marker::PhantomData;
use core::mem::{self, ManuallyDrop, SizedTypeProperties};
use core::num::NonZeroUsize;
use core::ptr::{self, NonNull};

use super::{InPlaceDrop, InPlaceDstBufDrop, SpecFromIter, SpecFromIterNested, Vec};
use super::{InPlaceDrop, InPlaceDstDataSrcBufDrop, SpecFromIter, SpecFromIterNested, Vec};

const fn in_place_collectible<DEST, SRC>(
step_merge: Option<NonZeroUsize>,
Expand Down Expand Up @@ -265,7 +266,7 @@ where
);
}

// The ownership of the allocation and the new `T` values is temporarily moved into `dst_guard`.
// The ownership of the source allocation and the new `T` values is temporarily moved into `dst_guard`.
// This is safe because
// * `forget_allocation_drop_remaining` immediately forgets the allocation
// before any panic can occur in order to avoid any double free, and then proceeds to drop
Expand All @@ -276,7 +277,8 @@ where
// Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce
// contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the
// module documentation why this is ok anyway.
let dst_guard = InPlaceDstBufDrop { ptr: dst_buf, len, cap: dst_cap };
let dst_guard =
InPlaceDstDataSrcBufDrop { ptr: dst_buf, len, src_cap, src: PhantomData::<I::Src> };
src.forget_allocation_drop_remaining();

// Adjust the allocation if the source had a capacity in bytes that wasn't a multiple
Expand Down
26 changes: 18 additions & 8 deletions library/alloc/src/vec/in_place_drop.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,10 @@
use core::ptr::{self};
use core::marker::PhantomData;
use core::ptr::{self, drop_in_place};
use core::slice::{self};

use crate::alloc::Global;
use crate::raw_vec::RawVec;

// A helper struct for in-place iteration that drops the destination slice of iteration,
// i.e. the head. The source slice (the tail) is dropped by IntoIter.
pub(super) struct InPlaceDrop<T> {
Expand All @@ -23,17 +27,23 @@ impl<T> Drop for InPlaceDrop<T> {
}
}

// A helper struct for in-place collection that drops the destination allocation and elements,
// to avoid leaking them if some other destructor panics.
pub(super) struct InPlaceDstBufDrop<T> {
pub(super) ptr: *mut T,
// A helper struct for in-place collection that drops the destination items together with
// the source allocation - i.e. before the reallocation happened - to avoid leaking them
// if some other destructor panics.
pub(super) struct InPlaceDstDataSrcBufDrop<Src, Dest> {
pub(super) ptr: *mut Dest,
pub(super) len: usize,
pub(super) cap: usize,
pub(super) src_cap: usize,
pub(super) src: PhantomData<Src>,
}

impl<T> Drop for InPlaceDstBufDrop<T> {
impl<Src, Dest> Drop for InPlaceDstDataSrcBufDrop<Src, Dest> {
#[inline]
fn drop(&mut self) {
unsafe { super::Vec::from_raw_parts(self.ptr, self.len, self.cap) };
unsafe {
let _drop_allocation =
RawVec::<Src>::from_raw_parts_in(self.ptr.cast::<Src>(), self.src_cap, Global);
drop_in_place(core::ptr::slice_from_raw_parts_mut::<Dest>(self.ptr, self.len));
};
}
}
2 changes: 1 addition & 1 deletion library/alloc/src/vec/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ use self::set_len_on_drop::SetLenOnDrop;
mod set_len_on_drop;

#[cfg(not(no_global_oom_handling))]
use self::in_place_drop::{InPlaceDrop, InPlaceDstBufDrop};
use self::in_place_drop::{InPlaceDrop, InPlaceDstDataSrcBufDrop};

#[cfg(not(no_global_oom_handling))]
mod in_place_drop;
Expand Down

0 comments on commit 73b36b1

Please sign in to comment.