Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

btree: don't leak value if destructor of key panics #132597

Merged
merged 1 commit into from
Nov 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 48 additions & 0 deletions library/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 library/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
Loading