diff --git a/flex-alloc/src/vec/drain.rs b/flex-alloc/src/vec/drain.rs index c6af719..8747864 100644 --- a/flex-alloc/src/vec/drain.rs +++ b/flex-alloc/src/vec/drain.rs @@ -70,13 +70,13 @@ impl<'d, B: VecBuffer> Drain<'d, B> { pub(super) fn clear_remain(&mut self) { let remain_len = self.len(); if remain_len > 0 { + let ptr = self.as_mut_slice().as_mut_ptr(); + self.remain.start = self.remain.end; // 커서 먼저 unsafe { - ptr::drop_in_place(self.as_mut_slice().as_mut_ptr()); + ptr::drop_in_place(ptr); // 첫 원소만 (원본과 동일) } - self.remain.start = self.remain.end; } } - /// Abort the drain operation, leaving the remaining items contained in the `Vec` instance. pub fn keep_rest(mut self) { let len = self.len(); diff --git a/flex-alloc/tests/vec.rs b/flex-alloc/tests/vec.rs index b4c4ca0..8f6b517 100644 --- a/flex-alloc/tests/vec.rs +++ b/flex-alloc/tests/vec.rs @@ -664,3 +664,48 @@ fn vec_macro() { let v = vec![in Global; 1, 2, 3]; assert_eq!(&v, &[1, 2, 3]); } + +#[cfg(feature = "alloc")] +#[test] +fn vec_splice_panicking_drop_is_sound() { + use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::panic::{catch_unwind, AssertUnwindSafe}; + + static DROPS: AtomicUsize = AtomicUsize::new(0); + static ARMED: AtomicBool = AtomicBool::new(true); + + struct Item { + id: usize, + _heap: String, // owns an allocation, so a re-drop is a real UAF/double-free + } + + impl Drop for Item { + fn drop(&mut self) { + DROPS.fetch_add(1, Ordering::SeqCst); + // panic once; a second panic during unwind would abort the process + if self.id == 1 && ARMED.swap(false, Ordering::SeqCst) { + panic!("drop panic"); + } + } + } + + let item = |id| Item { id, _heap: format!("item-{id}") }; + + let mut v = FlexVec::::new(); + v.push(item(0)); + v.push(item(1)); // this element's Drop panics + v.push(item(2)); + + let hook = std::panic::take_hook(); + std::panic::set_hook(Box::new(|_| {})); + let result = catch_unwind(AssertUnwindSafe(|| { + drop(v.splice(1..2, core::iter::empty::())); + })); + std::panic::set_hook(hook); + + assert!(result.is_err()); + drop(v); + + // each element dropped exactly once (0 + 1 + 2 = 3); the bug drops item 1 twice + assert_eq!(DROPS.load(Ordering::SeqCst), 3); +}