Fix double free when retain unwinds

Cargo.toml

@@ -5,9 +5,7 @@ edition = "2024"
 authors = ["Kyle Thapar <kylethapar2007@outlook.com>"]
 description = "Stack-allocated collections for Rust"
 license = "MIT OR Apache-2.0"
-readme = "README.md"
 repository = "https://github.com/OverwrittenCode/stack_collections"
-homepage = "https://github.com/OverwrittenCode/stack_collections"
 documentation = "https://docs.rs/stack_collections"
 keywords = ["stack", "collections", "no-heap"]
 categories = ["data-structures", "no-std", "no-std::no-alloc"]

src/stack_vec.rs

@@ -461,34 +461,80 @@ impl<T, const CAP: usize> StackVec<T, CAP> {
     where
         F: FnMut(&mut T) -> bool,
     {
-        let base_ptr = self.as_mut_ptr();
-        let mut kept = 0;
+        struct BackshiftOnDrop<'vec, T, const CAP: usize> {
+            vec: &'vec mut StackVec<T, CAP>,
+            processed: usize,
+            deleted: usize,
+            original_len: usize,
+        }
+
+        impl<T, const CAP: usize> Drop for BackshiftOnDrop<'_, T, CAP> {
+            fn drop(&mut self) {
+                if self.deleted > 0 {
+                    let base = self.vec.as_mut_ptr();
+                    // SAFETY: processed and processed - deleted are both
+                    //         within bounds of the buffer.
+                    let src = unsafe { base.add(self.processed) };
+                    // SAFETY: see above.
+                    let dst = unsafe { base.add(self.processed - self.deleted) };
+                    // SAFETY: the unprocessed tail is still initialized and
+                    //         the destination range is inside the buffer.
+                    //         src/dst derive from a single reborrow, so
+                    //         neither invalidates the other.
+                    unsafe {
+                        ptr::copy(src, dst, self.original_len - self.processed);
+                    }
+                }
+                self.vec.len = self.original_len - self.deleted;
+            }
+        }
+
+        // Both the predicate and a removed element's `Drop` are user code and
+        // may unwind. Retire the whole buffer up front and let a guard restore
+        // a correct length on every exit path, so a duplicated or
+        // already-dropped slot can never stay inside `0..len`.
+        let original_len = self.len;
+        self.len = 0;
+
+        let mut g = BackshiftOnDrop {
+            vec: self,
+            processed: 0,
+            deleted: 0,
+            original_len,
+        };
+
+        while g.processed < original_len {
+            let base = g.vec.as_mut_ptr();
 
-        for i in 0..self.len {
-            // SAFETY: i < self.len, so this is a valid and initialized element
-            let ptr = unsafe { base_ptr.add(i) };
-            // SAFETY: ptr points to valid initialized element
+            // SAFETY: processed < original_len, so this element is initialized.
+            let ptr = unsafe { base.add(g.processed) };
+            // SAFETY: ptr points to a valid initialized element.
             let elem = unsafe { &mut *ptr };
 
             if f(elem) {
-                if kept != i {
-                    // SAFETY: kept < i < self.len, ranges don't overlap
-                    let dst = unsafe { base_ptr.add(kept) };
-                    // SAFETY: Copying single element from valid source to valid dest
+                if g.deleted > 0 {
+                    // SAFETY: dst is an earlier, already-vacated slot; ptr and dst
+                    //         both derive from the same reborrow, so neither
+                    //         invalidates the other.
+                    let dst = unsafe { base.add(g.processed - g.deleted) };
+                    // SAFETY: copying a single initialized element to a vacated slot.
                     unsafe {
                         ptr::copy_nonoverlapping(ptr, dst, 1);
                     }
                 }
-                kept += 1;
+                g.processed += 1;
             } else {
-                // SAFETY: Dropping initialized element that won't be kept
+                // Count the slot out before destroying it, so an unwinding
+                // `Drop` leaves the guard with a consistent view.
+                g.processed += 1;
+                g.deleted += 1;
+                // SAFETY: dropping an initialized element that won't be kept.
                 unsafe {
-                    ptr::drop_in_place(elem);
+                    ptr::drop_in_place(ptr);
                 }
             }
         }
-
-        self.len = kept;
+        // `g`'s `Drop` commits the final length.
     }
 
     // index
@@ -1837,4 +1883,84 @@ mod tests {
         drop(iter);
         assert_eq!(counter.load(Ordering::SeqCst), 3);
     }
+
+    /// A panic in `retain` must not leave a duplicated or already-destroyed
+    /// slot inside `0..len`.
+    #[cfg(feature = "std")]
+    #[test]
+    fn retain_panicking_predicate() {
+        extern crate std;
+        use core::panic::AssertUnwindSafe;
+        use std::panic::catch_unwind;
+
+        struct DropCounter(i32, Arc<AtomicUsize>);
+        impl Drop for DropCounter {
+            fn drop(&mut self) {
+                self.1.fetch_add(1, Ordering::SeqCst);
+            }
+        }
+
+        let counter = Arc::new(AtomicUsize::new(0));
+        let mut vec = StackVec::<DropCounter, 8>::new();
+        for i in 0..4_i32 {
+            vec.push(DropCounter(i, Arc::clone(&counter)));
+        }
+
+        // Reject element 0 so the next kept one is compacted forward, then
+        // unwind on element 2 before the length is committed.
+        let r = catch_unwind(AssertUnwindSafe(|| {
+            vec.retain(|e| {
+                assert_ne!(e.0, 2_i32, "predicate panics");
+                e.0 != 0_i32
+            });
+        }));
+        assert!(r.is_err(), "the predicate should have panicked");
+
+        drop(vec);
+
+        // Four elements exist. Fewer drops mean a leak, which is sound; more
+        // mean a value was destroyed twice.
+        let drops = counter.load(Ordering::SeqCst);
+        assert!(drops <= 4, "{drops} drops for 4 elements");
+    }
+
+    /// The same, with the panic coming from a removed element's `Drop`.
+    #[cfg(feature = "std")]
+    #[test]
+    fn retain_panicking_element_drop() {
+        extern crate std;
+        use core::cell::Cell;
+        use core::panic::AssertUnwindSafe;
+        use std::panic::catch_unwind;
+
+        std::thread_local! {
+            static ARMED: Cell<bool> = const { Cell::new(false) };
+        }
+
+        struct Boom(i32, Arc<AtomicUsize>);
+        impl Drop for Boom {
+            fn drop(&mut self) {
+                self.1.fetch_add(1, Ordering::SeqCst);
+                assert!(!ARMED.replace(false), "element Drop panics");
+            }
+        }
+
+        let counter = Arc::new(AtomicUsize::new(0));
+        let mut vec = StackVec::<Boom, 8>::new();
+        for i in 0..4_i32 {
+            vec.push(Boom(i, Arc::clone(&counter)));
+        }
+
+        ARMED.set(true);
+        let r = catch_unwind(AssertUnwindSafe(|| {
+            vec.retain(|e| e.0 != 1_i32);
+        }));
+        ARMED.set(false);
+        assert!(r.is_err(), "the armed Drop should have panicked");
+
+        drop(vec);
+
+        let drops = counter.load(Ordering::SeqCst);
+        assert!(drops <= 4, "{drops} drops for 4 elements");
+    }
 }