Skip to content

Commit

Permalink
Rollup merge of #120145 - the8472:fix-inplace-dest-drop, r=cuviper
Browse files Browse the repository at this point in the history
fix: Drop guard was deallocating with the incorrect size

InPlaceDstBufDrop holds onto the allocation before the shrinking happens which means it must deallocate the destination elements but the source allocation.

Thanks `@cuviper` for spotting this.
  • Loading branch information
GuillaumeGomez authored Jan 20, 2024
2 parents 5729ecc + 5796b3c commit 4317705
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 4317705

Please sign in to comment.