From e32a118e6cec5cfbd3552559787c98f6371d5ac1 Mon Sep 17 00:00:00 2001 From: Lukas Markeffsky <@> Date: Mon, 4 Nov 2024 11:21:22 +0000 Subject: [PATCH] btree: don't leak value if destructor of key panics --- .../alloc/src/collections/btree/map/tests.rs | 48 +++++++++++++++++++ library/alloc/src/collections/btree/node.rs | 18 ++++++- 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index db16d82be7dcc..5975134382e7d 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -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() { diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index 2a853ef421629..64a13bb6a0b3a 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -1173,11 +1173,25 @@ impl Handle, 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); + impl 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 } } }