Allocators panicking on drop cause yet another panic safety issue in `BTreeMap::insert`

[maxdexh]

Entirely based on #156490.

This issue confirms my speculation about allocators panicking in their drop, as in that they cause the same problems as unwinding handle_alloc_error or clone, like seen in prior issues.

playground link

#![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 INSERT_COUNT: usize = 137;
const PANIC_COUNTDOWN_START: u32 = 2;

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

    for id in 0..INSERT_COUNT {
        map.insert(id, ());
    }

    allocator.countdown.set(PANIC_COUNTDOWN_START);

    eprintln!("Allocator will panic on drop in {PANIC_COUNTDOWN_START} drops");

    catch_unwind(AssertUnwindSafe(|| {
        map.insert(INSERT_COUNT, ());
    }))
    .expect_err("insertion should panic");

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

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

    // map thinks all the elements made it
    assert_eq!(INSERT_COUNT, broken_len);

    // but some were lost!
    assert_ne!(elem_count, broken_len);

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

    // code from #156490 that exploits the mismatch
    let mut iter = map.into_iter();
    for _ in 0..elem_count {
        iter.next_back();
    }
    eprintln!("Dropped all the entries. BTreeMap::IntoIter drop should SIGSEGV now\n");
}

original demo

playground link

#![feature(allocator_api, btreemap_alloc)]

use std::collections::BTreeMap;
use std::panic::{AssertUnwindSafe, catch_unwind};

use panic_allocator::PanicAllocator;

const ENTRY_COUNT: usize = 137;
const INSERT_KEY: usize = 263;
const PANIC_ON_INSERT_ALLOCATION: usize = 2;

fn main() {
    let allocator = PanicAllocator::new();
    let mut map = BTreeMap::new_in(allocator.clone());

    for id in 0..ENTRY_COUNT {
        map.insert(id * 2, Box::new(Payload::new(id)));
    }

    let panic_at = allocator.drop_count() + PANIC_ON_INSERT_ALLOCATION;
    allocator.panic_at_drop(panic_at);

    eprintln!(
        "armed allocator panic at allocation #{panic_at}; \
         inserting key {INSERT_KEY} into {ENTRY_COUNT}-entry map"
    );

    let insertion = catch_unwind(AssertUnwindSafe(|| {
        map.insert(INSERT_KEY, Box::new(Payload::new(ENTRY_COUNT)));
    }));

    assert!(insertion.is_err(), "insertion should have panicked");

    // Count the actual reachable elements. Range is position-based (front/back
    // leaf-edge cursors), so it correctly reports the real count regardless of
    // the corrupted `length` field.
    let actual_count = map.range(..).count();
    let reported_len = map.len();

    assert!(
        actual_count < reported_len,
        "expected length corruption: map.len()={reported_len}, actual={actual_count}"
    );

    let discrepancy = reported_len - actual_count;
    eprintln!(
        "caught allocator panic; map.len()={reported_len}, \
         actual reachable={actual_count}, discrepancy={discrepancy}"
    );

    // Convert to a consuming, double-ended iterator.
    // IntoIter.length = reported_len (inflated).
    let mut iter = map.into_iter();

    // Phase 1: drain every real element from the back.
    // Each next_back() call:
    //   - decrements IntoIter.length
    //   - calls deallocating_next_back which reads the KV via assume_init_read
    //   - the returned Box<Payload> is dropped immediately → heap memory freed
    // After this loop: IntoIter.length == discrepancy, all values freed,
    // but the leftmost leaf is still allocated with its original `len`.
    for _ in 0..actual_count {
        iter.next_back();
    }

    eprintln!(
        "drained {actual_count} elements from back; \
         IntoIter.length={discrepancy}, reading duplicates from front..."
    );

    // Phase 2: read `discrepancy` already-freed KV slots from the front.
    // The front cursor descends to the leftmost leaf (still alive).
    // right_kv() succeeds because the leaf's len was never decremented.
    // into_key_val() does assume_init_read on an already-moved-out slot,
    // producing a duplicate Box<Payload> whose memory was freed in Phase 1.
    // Dropping this duplicate Box is a double-free → undefined behavior.
    for i in 0..discrepancy {
        if let Some((key, _duplicate_box)) = iter.next() {
            eprintln!("  duplicate #{}: key={key} (double-free on drop)", i + 1);
            // _duplicate_box is dropped here → double-free!
        } else {
            eprintln!("  unexpected None at duplicate #{}", i + 1);
            break;
        }
    }

    // Should not reach here—double-free typically triggers SIGABRT.
    eprintln!("unexpected: double-free did not crash the process");
}

struct Payload {
    _id: usize,
    _bytes: [u8; 256],
}

impl Payload {
    fn new(id: usize) -> Self {
        Self {
            _id: id,
            _bytes: [0x41; 256],
        }
    }
}

mod panic_allocator {
    use std::alloc::{AllocError, Allocator, Global, Layout};
    use std::cell::Cell;
    use std::ptr::NonNull;
    use std::rc::Rc;

    #[derive(Clone)]
    pub(super) struct PanicAllocator {
        state: Rc<State>,
    }

    struct State {
        drops: Cell<usize>,
        panic_at: Cell<Option<usize>>,
    }

    impl PanicAllocator {
        pub(super) fn new() -> Self {
            Self {
                state: Rc::new(State {
                    drops: Cell::new(0),
                    panic_at: Cell::new(None),
                }),
            }
        }

        pub(super) fn drop_count(&self) -> usize {
            self.state.drops.get()
        }

        pub(super) fn panic_at_drop(&self, n: usize) {
            self.state.panic_at.set(Some(n));
        }
    }

    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 drops = self.state.drops.get() + 1;
            self.state.drops.set(drops);

            if self.state.panic_at.get() == Some(drops) {
                panic!("its the final countdown");
            }
        }
    }
}

[nia-e]

so, this is fun. i wonder to what degree we can address this with a wrapper type; nounwind provides the functionality for aborting on unwind (which is what we'd want here), so I see the case for providing a NounwindDrop<T> wrapper (name open to bikeshed) with

struct NounwindDrop<T> {
    val: ManuallyDrop<T>,
}
impl Drop for NounwindDrop {
    fn drop(&mut self) {
        // do the AbortGuard nonsense from `nounwind` and drop `val` with that
    }
}

This might be broadly useful outside allocators, so it might deserve a separate ACP tbh

[maxdexh]

I mean arguably BTreeMap shouldn't even be cloning around the allocator here. IIUC, it's only doing so because using references has overhead for ZSTs. However, the amount of cloning it is doing here makes me think that we need a different solution or that at the very least, AllocatorClone for Box is a bad idea.

[maxdexh]

Drop is already an exception to the exceptions to the orphan rules (e.g. you cannot impl it for Box<LocalType>), so maybe requiring the drop not to panic is fine here?

Still a footgun, due to drop glue...