Fix leak in Vec::extend_from_within

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