Skip to content

Commit

Permalink
btree: don't leak value if destructor of key panics
Browse files Browse the repository at this point in the history
  • Loading branch information
Lukas Markeffsky committed Nov 4, 2024
1 parent 84fae7e commit 9cdbf39
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 2 deletions.
48 changes: 48 additions & 0 deletions alloc/src/collections/btree/map/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2270,6 +2270,54 @@ fn test_into_iter_drop_leak_height_0() {
assert_eq!(e.dropped(), 1);
}

#[test]
#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
fn test_into_iter_drop_leak_kv_panic_in_key() {
let a_k = CrashTestDummy::new(0);
let a_v = CrashTestDummy::new(1);
let b_k = CrashTestDummy::new(2);
let b_v = CrashTestDummy::new(3);
let c_k = CrashTestDummy::new(4);
let c_v = CrashTestDummy::new(5);
let mut map = BTreeMap::new();
map.insert(a_k.spawn(Panic::Never), a_v.spawn(Panic::Never));
map.insert(b_k.spawn(Panic::InDrop), b_v.spawn(Panic::Never));
map.insert(c_k.spawn(Panic::Never), c_v.spawn(Panic::Never));

catch_unwind(move || drop(map.into_iter())).unwrap_err();

assert_eq!(a_k.dropped(), 1);
assert_eq!(a_v.dropped(), 1);
assert_eq!(b_k.dropped(), 1);
assert_eq!(b_v.dropped(), 1);
assert_eq!(c_k.dropped(), 1);
assert_eq!(c_v.dropped(), 1);
}

#[test]
#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
fn test_into_iter_drop_leak_kv_panic_in_val() {
let a_k = CrashTestDummy::new(0);
let a_v = CrashTestDummy::new(1);
let b_k = CrashTestDummy::new(2);
let b_v = CrashTestDummy::new(3);
let c_k = CrashTestDummy::new(4);
let c_v = CrashTestDummy::new(5);
let mut map = BTreeMap::new();
map.insert(a_k.spawn(Panic::Never), a_v.spawn(Panic::Never));
map.insert(b_k.spawn(Panic::Never), b_v.spawn(Panic::InDrop));
map.insert(c_k.spawn(Panic::Never), c_v.spawn(Panic::Never));

catch_unwind(move || drop(map.into_iter())).unwrap_err();

assert_eq!(a_k.dropped(), 1);
assert_eq!(a_v.dropped(), 1);
assert_eq!(b_k.dropped(), 1);
assert_eq!(b_v.dropped(), 1);
assert_eq!(c_k.dropped(), 1);
assert_eq!(c_v.dropped(), 1);
}

#[test]
#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
fn test_into_iter_drop_leak_height_1() {
Expand Down
18 changes: 16 additions & 2 deletions alloc/src/collections/btree/node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1173,11 +1173,25 @@ impl<K, V, NodeType> Handle<NodeRef<marker::Dying, K, V, NodeType>, marker::KV>
/// The node that the handle refers to must not yet have been deallocated.
#[inline]
pub unsafe fn drop_key_val(mut self) {
// Run the destructor of the value even if the destructor of the key panics.
struct Dropper<'a, T>(&'a mut MaybeUninit<T>);
impl<T> Drop for Dropper<'_, T> {
#[inline]
fn drop(&mut self) {
unsafe {
self.0.assume_init_drop();
}
}
}

debug_assert!(self.idx < self.node.len());
let leaf = self.node.as_leaf_dying();
unsafe {
leaf.keys.get_unchecked_mut(self.idx).assume_init_drop();
leaf.vals.get_unchecked_mut(self.idx).assume_init_drop();
let key = leaf.keys.get_unchecked_mut(self.idx);
let val = leaf.vals.get_unchecked_mut(self.idx);
let _guard = Dropper(val);
key.assume_init_drop();
// dropping the guard will drop the value
}
}
}
Expand Down

0 comments on commit 9cdbf39

Please sign in to comment.