Fix double free when a comparator panics mid-split_off

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() {

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,13 +45,40 @@ impl<K, V> Root<K, V> {
         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
                 Found(kv) => kv.left_edge(),
                 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<K, V> Root<K, V> {
 
         left_root.fix_right_border(alloc.clone());
         right_root.fix_left_border(alloc);
+        mem::forget(guard);
         right_root
     }