Auto merge of #82760 - WaffleLapkin:unleak_extend_from_within, r=kennytm

Fix leak in Vec::extend_from_within

Fixes #82533

Author: bors

library/alloc/src/vec/mod.rs

@@ -1942,6 +1942,18 @@ impl<T, A: Allocator> Vec<T, A> {
     #[unstable(feature = "vec_split_at_spare", issue = "81944")]
     #[inline]
     pub fn split_at_spare_mut(&mut self) -> (&mut [T], &mut [MaybeUninit<T>]) {
+        // SAFETY:
+        // - len is ignored and so never changed
+        let (init, spare, _) = unsafe { self.split_at_spare_mut_with_len() };
+        (init, spare)
+    }
+
+    /// Safety: changing returned .2 (&mut usize) is considered the same as calling `.set_len(_)`.
+    ///
+    /// This method is used to have unique access to all vec parts at once in `extend_from_within`.
+    unsafe fn split_at_spare_mut_with_len(
+        &mut self,
+    ) -> (&mut [T], &mut [MaybeUninit<T>], &mut usize) {
         let Range { start: ptr, end: spare_ptr } = self.as_mut_ptr_range();
         let spare_ptr = spare_ptr.cast::<MaybeUninit<T>>();
         let spare_len = self.buf.capacity() - self.len;
@@ -1953,7 +1965,7 @@ impl<T, A: Allocator> Vec<T, A> {
             let initialized = slice::from_raw_parts_mut(ptr, self.len);
             let spare = slice::from_raw_parts_mut(spare_ptr, spare_len);
 
-            (initialized, spare)
+            (initialized, spare, &mut self.len)
         }
     }
 }
@@ -2165,22 +2177,23 @@ trait ExtendFromWithinSpec {
 
 impl<T: Clone, A: Allocator> ExtendFromWithinSpec for Vec<T, A> {
     default unsafe fn spec_extend_from_within(&mut self, src: Range<usize>) {
-        let initialized = {
-            let (this, spare) = self.split_at_spare_mut();
-
-            // SAFETY:
-            // - caller guaratees that src is a valid index
-            let to_clone = unsafe { this.get_unchecked(src) };
-
-            to_clone.iter().cloned().zip(spare.iter_mut()).map(|(e, s)| s.write(e)).count()
-        };
+        // SAFETY:
+        // - len is increased only after initializing elements
+        let (this, spare, len) = unsafe { self.split_at_spare_mut_with_len() };
 
         // SAFETY:
-        // - elements were just initialized
-        unsafe {
-            let new_len = self.len() + initialized;
-            self.set_len(new_len);
-        }
+        // - caller guaratees that src is a valid index
+        let to_clone = unsafe { this.get_unchecked(src) };
+
+        to_clone
+            .iter()
+            .cloned()
+            .zip(spare.iter_mut())
+            .map(|(src, dst)| dst.write(src))
+            // Note:
+            // - Element was just initialized with `MaybeUninit::write`, so it's ok to increace len
+            // - len is increased after each element to prevent leaks (see issue #82533)
+            .for_each(|_| *len += 1);
     }
 }
 

library/alloc/tests/vec.rs

@@ -7,6 +7,7 @@ use std::mem::{size_of, swap};
 use std::ops::Bound::*;
 use std::panic::{catch_unwind, AssertUnwindSafe};
 use std::rc::Rc;
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::vec::{Drain, IntoIter};
 
 struct DropCounter<'a> {
@@ -2100,3 +2101,44 @@ fn test_extend_from_within() {
 
     assert_eq!(v, ["a", "b", "c", "b", "c", "a", "b"]);
 }
+
+// Regression test for issue #82533
+#[test]
+fn test_extend_from_within_panicing_clone() {
+    struct Panic<'dc> {
+        drop_count: &'dc AtomicU32,
+        aaaaa: bool,
+    }
+
+    impl Clone for Panic<'_> {
+        fn clone(&self) -> Self {
+            if self.aaaaa {
+                panic!("panic! at the clone");
+            }
+
+            Self { ..*self }
+        }
+    }
+
+    impl Drop for Panic<'_> {
+        fn drop(&mut self) {
+            self.drop_count.fetch_add(1, Ordering::SeqCst);
+        }
+    }
+
+    let count = core::sync::atomic::AtomicU32::new(0);
+    let mut vec = vec![
+        Panic { drop_count: &count, aaaaa: false },
+        Panic { drop_count: &count, aaaaa: true },
+        Panic { drop_count: &count, aaaaa: false },
+    ];
+
+    // This should clone&append one Panic{..} at the end, and then panic while
+    // cloning second Panic{..}. This means that `Panic::drop` should be called
+    // 4 times (3 for items already in vector, 1 for just appended).
+    //
+    // Previously just appended item was leaked, making drop_count = 3, instead of 4.
+    std::panic::catch_unwind(move || vec.extend_from_within(..)).unwrap_err();
+
+    assert_eq!(count.load(Ordering::SeqCst), 4);
+}