Another bug in `BTreeMap` from allocator panicking on drop; this time in `remove`

[maxdexh]

cc #157428

See also #159334

playground

#![feature(allocator_api, btreemap_alloc)]

use std::alloc::{AllocError, Allocator, Global, Layout};
use std::cell::Cell;
use std::collections::BTreeMap;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::ptr::NonNull;
use std::rc::Rc;

#[derive(Clone)]
struct PanicAllocator {
    countdown: Rc<Cell<u32>>,
}

unsafe impl Allocator for PanicAllocator {
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        Global.allocate(layout)
    }

    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
        unsafe { Global.deallocate(ptr, layout) }
    }
}
impl Drop for PanicAllocator {
    fn drop(&mut self) {
        let cd = self.countdown.get();
        self.countdown.set(cd.saturating_sub(1));

        if cd == 1 {
            panic!("its the final countdown");
        }
    }
}

const PANIC_COUNTDOWN_START: u32 = 1;
const ENTRY_COUNT: usize = 137;

fn main() {
    let allocator = PanicAllocator {
        countdown: Rc::new(Cell::new(0)),
    };
    let mut map = BTreeMap::new_in(allocator.clone());

    for i in 0..ENTRY_COUNT {
        map.insert(i, vec![vec![i]]);
    }

    let insert_count = map.len();

    allocator.countdown.set(PANIC_COUNTDOWN_START);
    catch_unwind(AssertUnwindSafe(|| {
        map.remove(&(ENTRY_COUNT - 1));
    }))
    .expect_err("removal should panic");

    eprintln!("\nSurvived unwind in `remove`");

    let elem_count = map.range(..).count();
    let broken_len = map.len();

    // map did not update the len, so the element should still be there
    assert_eq!(insert_count, broken_len);

    // but it was removed from the tree!
    assert_ne!(elem_count, broken_len);

    eprintln!("\nbtree reported len={broken_len}, but only {elem_count} are accessible!\n");

    let mut iter = map.into_iter();
    for _ in 0..elem_count {
        iter.next_back();
    }
    // make the allocator detect the double free
    let _ = vec![0x41; 10000];
}

@rustbot label A-allocators requires-nightly T-libs I-unsound A-panic A-collections A-destructors

[maxdexh]

The countdown can be replaced with a simple bool, I was too lazy to do that.

[maxdexh]

Bug is in the implementation of remove_entry:

rust/library/alloc/src/collections/btree/map/entry.rs

Lines 606 to 618 in 4b0c9d7

	pub(super) fn remove_kv(self) -> (K, V) {
	let mut emptied_internal_root = false;
	let (old_kv, _) =
	self.handle.remove_kv_tracking(|| emptied_internal_root = true, self.alloc.clone());
	// SAFETY: we consumed the intermediate root borrow, `self.handle`.
	let map = unsafe { self.dormant_map.awaken() };
	map.length -= 1;
	if emptied_internal_root {
	let root = map.root.as_mut().unwrap();
	root.pop_internal_level(self.alloc);
	}
	old_kv
	}

The panic happens between removal from the tree and the length being decremented. This is a problem no matter in which order it happens. Here, it allows us to unwind out of the method after the tree operations, before the length is updated.