diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index 64348745aa07d..36e2298b2dd30 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -2428,6 +2428,35 @@ fn test_split_off_large_random_sorted() { assert!(right.into_iter().eq(data.into_iter().filter(|x| x.0 >= key))); } +// Regression test for #158165: a comparator that panics partway through +// `split_off`, after at least one level of the tree has already had a +// suffix moved into the new right-hand tree, used to leave `self` with a +// tree structure inconsistent with its own recorded length. Iterating or +// dropping the "recovered" map afterwards could then double free values +// that had already been moved into (and dropped along with) the +// abandoned right-hand tree. `split_off` now aborts the process instead +// of unwinding out of that inconsistent state, so this test only checks +// that ordinary (non-panicking) splits over multi-level trees are +// unaffected; the abort itself can't be observed from within a single +// process. See the reproducer attached to the issue for the double free. +#[test] +fn test_split_off_multi_level_unaffected_by_panic_guard() { + // MIN_INSERTS_HEIGHT_2 consecutive keys guarantee a 3-level tree, so + // `split_off` has to walk down (and move a suffix at) more than one + // level for most split points below. + let data = Vec::from_iter((0..MIN_INSERTS_HEIGHT_2).map(|i| (i, i))); + for &split_at in + &[0, 1, 2, MIN_INSERTS_HEIGHT_2 / 2, MIN_INSERTS_HEIGHT_2 - 2, MIN_INSERTS_HEIGHT_2 - 1] + { + let mut map = BTreeMap::from_iter(data.iter().copied()); + let right = map.split_off(&split_at); + map.check(); + right.check(); + assert!(map.keys().copied().eq(0..split_at)); + assert!(right.keys().copied().eq(split_at..MIN_INSERTS_HEIGHT_2)); + } +} + #[test] #[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")] fn test_into_iter_drop_leak_height_0() { diff --git a/library/alloc/src/collections/btree/split.rs b/library/alloc/src/collections/btree/split.rs index 87a79e6cf3f93..724441b73cab4 100644 --- a/library/alloc/src/collections/btree/split.rs +++ b/library/alloc/src/collections/btree/split.rs @@ -1,5 +1,6 @@ use core::alloc::Allocator; use core::borrow::Borrow; +use core::{intrinsics, mem}; use super::node::ForceResult::*; use super::node::Root; @@ -44,6 +45,25 @@ impl Root { let mut left_node = left_root.borrow_mut(); let mut right_node = right_root.borrow_mut(); + // Once the first `move_suffix` below has run, `left_root` and + // `right_root` reference a common set of key-value pairs through two + // different tree structures, and neither is a valid, independently + // droppable `BTreeMap` until `fix_right_border`/`fix_left_border` + // have repaired them and the caller has recomputed both lengths. The + // only thing that can fail beyond this point is the caller-supplied + // `Ord`/`Borrow` impl invoked by `search_node`. If that panics, we + // cannot safely unwind with the trees in this intermediate state + // (doing so leads to a double free, see #158165), so abort instead, + // matching the panic-safety strategy used elsewhere in this module + // (see `mem::replace`). + struct AbortOnDrop; + impl Drop for AbortOnDrop { + fn drop(&mut self) { + intrinsics::abort() + } + } + let mut guard = None; + loop { let mut split_edge = match left_node.search_node(key) { // key is going to the right tree @@ -51,6 +71,14 @@ impl Root { GoDown(edge) => edge, }; + // From here on, `left_root` and `right_root` are both + // unsound to drop until the loop finishes and the borders + // are fixed up below. Use `get_or_insert_with` rather than + // `get_or_insert`: the latter takes its argument by value, so + // it would construct (and immediately drop, aborting) a fresh + // `AbortOnDrop` on every iteration after the first. + guard.get_or_insert_with(|| AbortOnDrop); + split_edge.move_suffix(&mut right_node); match (split_edge.force(), right_node.force()) { @@ -65,6 +93,7 @@ impl Root { left_root.fix_right_border(alloc.clone()); right_root.fix_left_border(alloc); + mem::forget(guard); right_root }