Skip to content

Commit fb00ac6

Browse files
committed
Fix VecDeque::shrink_to UB when handle_alloc_error unwinds.
Luckily it's comparatively simple to just restore the `VecDeque` into a valid state on unwinds.
1 parent 62cffee commit fb00ac6

File tree

3 files changed

+120
-2
lines changed

3 files changed

+120
-2
lines changed

library/alloc/src/collections/vec_deque/mod.rs

+65-1
Original file line numberDiff line numberDiff line change
@@ -978,6 +978,8 @@ impl<T, A: Allocator> VecDeque<T, A> {
978978
// `head` and `len` are at most `isize::MAX` and `target_cap < self.capacity()`, so nothing can
979979
// overflow.
980980
let tail_outside = (target_cap + 1..=self.capacity()).contains(&(self.head + self.len));
981+
// Used in the drop guard below.
982+
let old_head = self.head;
981983

982984
if self.len == 0 {
983985
self.head = 0;
@@ -1030,12 +1032,74 @@ impl<T, A: Allocator> VecDeque<T, A> {
10301032
}
10311033
self.head = new_head;
10321034
}
1033-
self.buf.shrink_to_fit(target_cap);
1035+
1036+
struct Guard<'a, T, A: Allocator> {
1037+
deque: &'a mut VecDeque<T, A>,
1038+
old_head: usize,
1039+
target_cap: usize,
1040+
}
1041+
1042+
impl<T, A: Allocator> Drop for Guard<'_, T, A> {
1043+
#[cold]
1044+
fn drop(&mut self) {
1045+
unsafe {
1046+
// SAFETY: This is only called if `buf.shrink_to_fit` unwinds,
1047+
// which is the only time it's safe to call `abort_shrink`.
1048+
self.deque.abort_shrink(self.old_head, self.target_cap)
1049+
}
1050+
}
1051+
}
1052+
1053+
let guard = Guard { deque: self, old_head, target_cap };
1054+
1055+
guard.deque.buf.shrink_to_fit(target_cap);
1056+
1057+
// Don't drop the guard if we didn't unwind.
1058+
mem::forget(guard);
10341059

10351060
debug_assert!(self.head < self.capacity() || self.capacity() == 0);
10361061
debug_assert!(self.len <= self.capacity());
10371062
}
10381063

1064+
/// Reverts the deque back into a consistent state in case `shrink_to` failed.
1065+
/// This is necessary to prevent UB if the backing allocator returns an error
1066+
/// from `shrink` and `handle_alloc_error` subsequently unwinds (see #123369).
1067+
///
1068+
/// `old_head` refers to the head index before `shrink_to` was called. `target_cap`
1069+
/// is the capacity that it was trying to shrink to.
1070+
unsafe fn abort_shrink(&mut self, old_head: usize, target_cap: usize) {
1071+
// Moral equivalent of self.head + self.len <= target_cap. Won't overflow
1072+
// because `self.len <= target_cap`.
1073+
if self.head <= target_cap - self.len {
1074+
// The deque's buffer is contiguous, so no need to copy anything around.
1075+
return;
1076+
}
1077+
1078+
// `shrink_to` already copied the head to fit into the new capacity, so this won't overflow.
1079+
let head_len = target_cap - self.head;
1080+
// `self.head > target_cap - self.len` => `self.len > target_cap - self.head =: head_len` so this must be positive.
1081+
let tail_len = self.len - head_len;
1082+
1083+
if tail_len <= cmp::min(head_len, self.capacity() - target_cap) {
1084+
// There's enough spare capacity to copy the tail to the back (because `tail_len < self.capacity() - target_cap`),
1085+
// and copying the tail should be cheaper than copying the head (because `tail_len <= head_len`).
1086+
1087+
unsafe {
1088+
// The old tail and the new tail can't overlap because the head slice lies between them. The
1089+
// head slice ends at `target_cap`, so that's where we copy to.
1090+
self.copy_nonoverlapping(0, target_cap, tail_len);
1091+
}
1092+
} else {
1093+
// Either there's not enough spare capacity to make the deque contiguous, or the head is shorter than the tail
1094+
// (and therefore hopefully cheaper to copy).
1095+
unsafe {
1096+
// The old and the new head slice can overlap, so we can't use `copy_nonoverlapping` here.
1097+
self.copy(self.head, old_head, head_len);
1098+
self.head = old_head;
1099+
}
1100+
}
1101+
}
1102+
10391103
/// Shortens the deque, keeping the first `len` elements and dropping
10401104
/// the rest.
10411105
///

library/alloc/src/collections/vec_deque/tests.rs

+54-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,11 @@
1-
use core::iter::TrustedLen;
1+
#![feature(alloc_error_hook)]
2+
3+
use crate::alloc::{AllocError, Layout};
4+
use core::{iter::TrustedLen, ptr::NonNull};
5+
use std::{
6+
alloc::{set_alloc_error_hook, take_alloc_error_hook, System},
7+
panic::{catch_unwind, AssertUnwindSafe},
8+
};
29

310
use super::*;
411

@@ -790,6 +797,52 @@ fn test_shrink_to() {
790797
}
791798
}
792799

800+
#[test]
801+
fn test_shrink_to_unwind() {
802+
// This tests that `shrink_to` leaves the deque in a consistent state when
803+
// the call to `RawVec::shrink_to_fit` unwinds. The code is adapted from #123369
804+
// but changed to hopefully not have any UB even if the test fails.
805+
806+
struct BadAlloc;
807+
808+
unsafe impl Allocator for BadAlloc {
809+
fn allocate(&self, l: Layout) -> Result<NonNull<[u8]>, AllocError> {
810+
// We allocate zeroed here so that the whole buffer of the deque
811+
// is always initialized. That way, even if the deque is left in
812+
// an inconsistent state, no uninitialized memory should be accessed.
813+
System.allocate_zeroed(l)
814+
}
815+
816+
unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
817+
unsafe { System.deallocate(ptr, layout) }
818+
}
819+
820+
unsafe fn shrink(
821+
&self,
822+
_ptr: NonNull<u8>,
823+
_old_layout: Layout,
824+
_new_layout: Layout,
825+
) -> Result<NonNull<[u8]>, AllocError> {
826+
Err(AllocError)
827+
}
828+
}
829+
830+
// preserve the old error hook just in case.
831+
let old_error_hook = take_alloc_error_hook();
832+
set_alloc_error_hook(|_| panic!("alloc error"));
833+
834+
let mut v = VecDeque::with_capacity_in(15, BadAlloc);
835+
v.push_back(1);
836+
v.push_front(2);
837+
// This should unwind because it calls `BadAlloc::shrink` and then `handle_alloc_error` which unwinds.
838+
assert!(catch_unwind(AssertUnwindSafe(|| v.shrink_to_fit())).is_err());
839+
// This should only pass if the deque is left in a consistent state.
840+
assert_eq!(v, [2, 1]);
841+
842+
// restore the old error hook.
843+
set_alloc_error_hook(old_error_hook);
844+
}
845+
793846
#[test]
794847
fn test_shrink_to_fit() {
795848
// This test checks that every single combination of head and tail position,

library/alloc/src/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@
9292
// tidy-alphabetical-start
9393
#![cfg_attr(not(no_global_oom_handling), feature(const_alloc_error))]
9494
#![cfg_attr(not(no_global_oom_handling), feature(const_btree_len))]
95+
#![cfg_attr(test, feature(alloc_error_hook))]
9596
#![cfg_attr(test, feature(is_sorted))]
9697
#![cfg_attr(test, feature(new_uninit))]
9798
#![feature(alloc_layout_extra)]

0 commit comments

Comments
 (0)