From 73b36b14042351b55fcce7afc196669a5bc84c09 Mon Sep 17 00:00:00 2001
From: The 8472 <git@infinite-source.de>
Date: Fri, 19 Jan 2024 22:42:52 +0100
Subject: [PATCH] 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.

(cherry picked from commit 5796b3c16733f3d9cb2f3476ed1dc0fe05b14e25)
---
 library/alloc/src/vec/in_place_collect.rs | 10 +++++----
 library/alloc/src/vec/in_place_drop.rs    | 26 ++++++++++++++++-------
 library/alloc/src/vec/mod.rs              |  2 +-
 3 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/library/alloc/src/vec/in_place_collect.rs b/library/alloc/src/vec/in_place_collect.rs
index d2cc2b23ea6f8..e68cce04ce4aa 100644
--- a/library/alloc/src/vec/in_place_collect.rs
+++ b/library/alloc/src/vec/in_place_collect.rs
@@ -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()
@@ -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>,
@@ -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
@@ -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
diff --git a/library/alloc/src/vec/in_place_drop.rs b/library/alloc/src/vec/in_place_drop.rs
index 25ca33c6a7bf0..40a540b57fc22 100644
--- a/library/alloc/src/vec/in_place_drop.rs
+++ b/library/alloc/src/vec/in_place_drop.rs
@@ -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> {
@@ -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));
+        };
     }
 }
diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
index fca85c6123b3f..b2e9203976244 100644
--- a/library/alloc/src/vec/mod.rs
+++ b/library/alloc/src/vec/mod.rs
@@ -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;