diff options
| -rw-r--r-- | library/alloc/Cargo.toml | 4 | ||||
| -rw-r--r-- | library/alloc/src/collections/vec_deque/mod.rs | 66 | ||||
| -rw-r--r-- | library/alloc/tests/vec_deque_alloc_error.rs | 49 |
3 files changed, 118 insertions, 1 deletions
diff --git a/library/alloc/Cargo.toml b/library/alloc/Cargo.toml index 4646eaced01..024b92790e9 100644 --- a/library/alloc/Cargo.toml +++ b/library/alloc/Cargo.toml @@ -20,6 +20,10 @@ rand_xorshift = "0.3.0" name = "alloctests" path = "tests/lib.rs" +[[test]] +name = "vec_deque_alloc_error" +path = "tests/vec_deque_alloc_error.rs" + [[bench]] name = "allocbenches" path = "benches/lib.rs" diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs index 4643a6bbe2e..61d05afccfd 100644 --- a/library/alloc/src/collections/vec_deque/mod.rs +++ b/library/alloc/src/collections/vec_deque/mod.rs @@ -982,6 +982,8 @@ impl<T, A: Allocator> VecDeque<T, A> { // `head` and `len` are at most `isize::MAX` and `target_cap < self.capacity()`, so nothing can // overflow. let tail_outside = (target_cap + 1..=self.capacity()).contains(&(self.head + self.len)); + // Used in the drop guard below. + let old_head = self.head; if self.len == 0 { self.head = 0; @@ -1034,12 +1036,74 @@ impl<T, A: Allocator> VecDeque<T, A> { } self.head = new_head; } - self.buf.shrink_to_fit(target_cap); + + struct Guard<'a, T, A: Allocator> { + deque: &'a mut VecDeque<T, A>, + old_head: usize, + target_cap: usize, + } + + impl<T, A: Allocator> Drop for Guard<'_, T, A> { + #[cold] + fn drop(&mut self) { + unsafe { + // SAFETY: This is only called if `buf.shrink_to_fit` unwinds, + // which is the only time it's safe to call `abort_shrink`. + self.deque.abort_shrink(self.old_head, self.target_cap) + } + } + } + + let guard = Guard { deque: self, old_head, target_cap }; + + guard.deque.buf.shrink_to_fit(target_cap); + + // Don't drop the guard if we didn't unwind. + mem::forget(guard); debug_assert!(self.head < self.capacity() || self.capacity() == 0); debug_assert!(self.len <= self.capacity()); } + /// Reverts the deque back into a consistent state in case `shrink_to` failed. + /// This is necessary to prevent UB if the backing allocator returns an error + /// from `shrink` and `handle_alloc_error` subsequently unwinds (see #123369). + /// + /// `old_head` refers to the head index before `shrink_to` was called. `target_cap` + /// is the capacity that it was trying to shrink to. + unsafe fn abort_shrink(&mut self, old_head: usize, target_cap: usize) { + // Moral equivalent of self.head + self.len <= target_cap. Won't overflow + // because `self.len <= target_cap`. + if self.head <= target_cap - self.len { + // The deque's buffer is contiguous, so no need to copy anything around. + return; + } + + // `shrink_to` already copied the head to fit into the new capacity, so this won't overflow. + let head_len = target_cap - self.head; + // `self.head > target_cap - self.len` => `self.len > target_cap - self.head =: head_len` so this must be positive. + let tail_len = self.len - head_len; + + if tail_len <= cmp::min(head_len, self.capacity() - target_cap) { + // There's enough spare capacity to copy the tail to the back (because `tail_len < self.capacity() - target_cap`), + // and copying the tail should be cheaper than copying the head (because `tail_len <= head_len`). + + unsafe { + // The old tail and the new tail can't overlap because the head slice lies between them. The + // head slice ends at `target_cap`, so that's where we copy to. + self.copy_nonoverlapping(0, target_cap, tail_len); + } + } else { + // Either there's not enough spare capacity to make the deque contiguous, or the head is shorter than the tail + // (and therefore hopefully cheaper to copy). + unsafe { + // The old and the new head slice can overlap, so we can't use `copy_nonoverlapping` here. + self.copy(self.head, old_head, head_len); + self.head = old_head; + } + } + } + /// Shortens the deque, keeping the first `len` elements and dropping /// the rest. /// diff --git a/library/alloc/tests/vec_deque_alloc_error.rs b/library/alloc/tests/vec_deque_alloc_error.rs new file mode 100644 index 00000000000..c11f4556da9 --- /dev/null +++ b/library/alloc/tests/vec_deque_alloc_error.rs @@ -0,0 +1,49 @@ +#![feature(alloc_error_hook, allocator_api)] + +use std::{ + alloc::{set_alloc_error_hook, AllocError, Allocator, Layout, System}, + collections::VecDeque, + panic::{catch_unwind, AssertUnwindSafe}, + ptr::NonNull, +}; + +#[test] +fn test_shrink_to_unwind() { + // This tests that `shrink_to` leaves the deque in a consistent state when + // the call to `RawVec::shrink_to_fit` unwinds. The code is adapted from #123369 + // but changed to hopefully not have any UB even if the test fails. + + struct BadAlloc; + + unsafe impl Allocator for BadAlloc { + fn allocate(&self, l: Layout) -> Result<NonNull<[u8]>, AllocError> { + // We allocate zeroed here so that the whole buffer of the deque + // is always initialized. That way, even if the deque is left in + // an inconsistent state, no uninitialized memory should be accessed. + System.allocate_zeroed(l) + } + + unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) { + unsafe { System.deallocate(ptr, layout) } + } + + unsafe fn shrink( + &self, + _ptr: NonNull<u8>, + _old_layout: Layout, + _new_layout: Layout, + ) -> Result<NonNull<[u8]>, AllocError> { + Err(AllocError) + } + } + + set_alloc_error_hook(|_| panic!("alloc error")); + + let mut v = VecDeque::with_capacity_in(15, BadAlloc); + v.push_back(1); + v.push_front(2); + // This should unwind because it calls `BadAlloc::shrink` and then `handle_alloc_error` which unwinds. + assert!(catch_unwind(AssertUnwindSafe(|| v.shrink_to_fit())).is_err()); + // This should only pass if the deque is left in a consistent state. + assert_eq!(v, [2, 1]); +} |