Fix in-place collection leak when remaining element destructor panic

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) };
 

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) };
+    }
+}

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;

library/alloc/tests/vec.rs

@@ -1110,23 +1110,31 @@ fn test_from_iter_specialization_panic_during_iteration_drops() {
 }
 
 #[test]
-fn test_from_iter_specialization_panic_during_drop_leaks() {
-    static mut DROP_COUNTER: usize = 0;
+fn test_from_iter_specialization_panic_during_drop_doesnt_leak() {
+    static mut DROP_COUNTER_SHOULD_BE_DROPPED: usize = 0;
+    static mut DROP_COUNTER_DROPPED_TWICE: usize = 0;
 
     #[derive(Debug)]
     enum Droppable {
+        ShouldBeDropped,
         DroppedTwice(Box<i32>),
         PanicOnDrop,
     }
 
     impl Drop for Droppable {
         fn drop(&mut self) {
             match self {
+                Droppable::ShouldBeDropped => {
+                    unsafe {
+                        DROP_COUNTER_SHOULD_BE_DROPPED += 1;
+                    }
+                    println!("Dropping ShouldBeDropped!")
+                }
                 Droppable::DroppedTwice(_) => {
                     unsafe {
-                        DROP_COUNTER += 1;
+                        DROP_COUNTER_DROPPED_TWICE += 1;
                     }
-                    println!("Dropping!")
+                    println!("Dropping DroppedTwice!")
                 }
                 Droppable::PanicOnDrop => {
                     if !std::thread::panicking() {
@@ -1137,21 +1145,17 @@ fn test_from_iter_specialization_panic_during_drop_leaks() {
         }
     }
 
-    let mut to_free: *mut Droppable = core::ptr::null_mut();
-    let mut cap = 0;
-
     let _ = std::panic::catch_unwind(AssertUnwindSafe(|| {
-        let mut v = vec![Droppable::DroppedTwice(Box::new(123)), Droppable::PanicOnDrop];
-        to_free = v.as_mut_ptr();
-        cap = v.capacity();
-        let _ = v.into_iter().take(0).collect::<Vec<_>>();
+        let v = vec![
+            Droppable::ShouldBeDropped,
+            Droppable::DroppedTwice(Box::new(123)),
+            Droppable::PanicOnDrop,
+        ];
+        let _ = v.into_iter().take(1).collect::<Vec<_>>();
     }));
 
-    assert_eq!(unsafe { DROP_COUNTER }, 1);
-    // clean up the leak to keep miri happy
-    unsafe {
-        drop(Vec::from_raw_parts(to_free, 0, cap));
-    }
+    assert_eq!(unsafe { DROP_COUNTER_SHOULD_BE_DROPPED }, 1);
+    assert_eq!(unsafe { DROP_COUNTER_DROPPED_TWICE }, 1);
 }
 
 // regression test for issue #85322. Peekable previously implemented InPlaceIterable,