Fix leaking in inplace collection when destructor panics

3 files changed, 22 insertions, 3 deletions

diff --git a/library/alloc/src/vec/in_place_collect.rs b/library/alloc/src/vec/in_place_collect.rs
index 55dcb84ad16..359f1a9cec2 100644
--- a/library/alloc/src/vec/in_place_collect.rs
+++ b/library/alloc/src/vec/in_place_collect.rs
@@ -138,7 +138,7 @@ use core::iter::{InPlaceIterable, SourceIter, TrustedRandomAccessNoCoerce};
 use core::mem::{self, ManuallyDrop};
 use core::ptr::{self};
 
-use super::{InPlaceDrop, SpecFromIter, SpecFromIterNested, Vec};
+use super::{InPlaceDrop, InPlaceDstBufDrop, SpecFromIter, SpecFromIterNested, Vec};
 
 /// Specialization marker for collecting an iterator pipeline into a Vec while reusing the
 /// source allocation, i.e. executing the pipeline in place.
@@ -193,12 +193,16 @@ where
 
         // Drop any remaining values at the tail of the source but prevent drop of the allocation
         // itself once IntoIter goes out of scope.
-        // If the drop panics then we also leak any elements collected into dst_buf.
+        // If the drop panics then we also try to drop the destination buffer and its elements.
+        // This is safe because `forget_allocation_drop_remaining` forgets the allocation *before*
+        // trying to drop the remaining elements.
         //
         // 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 documenttation why this is ok anyway.
+        let dst_guard = InPlaceDstBufDrop { ptr: dst_buf, len, cap };
         src.forget_allocation_drop_remaining();
+        mem::forget(dst_guard);
 
         let vec = unsafe { Vec::from_raw_parts(dst_buf, len, cap) };
 
diff --git a/library/alloc/src/vec/in_place_drop.rs b/library/alloc/src/vec/in_place_drop.rs
index 1b1ef9130fa..25ca33c6a7b 100644
--- a/library/alloc/src/vec/in_place_drop.rs
+++ b/library/alloc/src/vec/in_place_drop.rs
@@ -22,3 +22,18 @@ 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,
+    pub(super) len: usize,
+    pub(super) cap: usize,
+}
+
+impl<T> Drop for InPlaceDstBufDrop<T> {
+    #[inline]
+    fn drop(&mut self) {
+        unsafe { super::Vec::from_raw_parts(self.ptr, self.len, self.cap) };
+    }
+}
diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
index fa9f2131c0c..acfbb98272d 100644
--- a/library/alloc/src/vec/mod.rs
+++ b/library/alloc/src/vec/mod.rs
@@ -125,7 +125,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;
+use self::in_place_drop::{InPlaceDrop, InPlaceDstBufDrop};
 
 #[cfg(not(no_global_oom_handling))]
 mod in_place_drop;