diff --git a/library/core/src/hint.rs b/library/core/src/hint.rs index c72eeb9a9c976..89c64f6d28fef 100644 --- a/library/core/src/hint.rs +++ b/library/core/src/hint.rs @@ -786,12 +786,45 @@ pub fn select_unpredictable(condition: bool, true_val: T, false_val: T) -> T // Change this to use ManuallyDrop instead. let mut true_val = MaybeUninit::new(true_val); let mut false_val = MaybeUninit::new(false_val); + + struct DropOnPanic { + // Invariant: valid pointer and points to an initialized value that is not further used, + // i.e. it can be dropped by this guard. + inner: *mut T, + } + + impl Drop for DropOnPanic { + fn drop(&mut self) { + // SAFETY: Must be guaranteed on construction of local type `DropOnPanic`. + unsafe { self.inner.drop_in_place() } + } + } + + let true_ptr = true_val.as_mut_ptr(); + let false_ptr = false_val.as_mut_ptr(); + // SAFETY: The value that is not selected is dropped, and the selected one // is returned. This is necessary because the intrinsic doesn't drop the // value that is not selected. unsafe { - crate::intrinsics::select_unpredictable(!condition, &mut true_val, &mut false_val) - .assume_init_drop(); + // Extract the selected value first, ensure it is dropped as well if dropping the unselected + // value panics. We construct a temporary by-pointer guard around the selected value while + // dropping the unselected value. Arguments overlap here, so we can not use mutable + // reference for these arguments. + let guard = crate::intrinsics::select_unpredictable(condition, true_ptr, false_ptr); + let drop = crate::intrinsics::select_unpredictable(condition, false_ptr, true_ptr); + + // SAFETY: both pointers are well-aligned and point to initialized values inside a + // `MaybeUninit` each. In both possible values for `condition` the pointer `guard` and + // `drop` do not alias (even though the two argument pairs we have selected from did alias + // each other). + let guard = DropOnPanic { inner: guard }; + drop.drop_in_place(); + crate::mem::forget(guard); + + // Note that it is important to use the values here. Reading from the pointer we got makes + // LLVM forget the !unpredictable annotation sometimes (in tests, integer sized values in + // particular seemed to confuse it, also observed in llvm/llvm-project #82340). crate::intrinsics::select_unpredictable(condition, true_val, false_val).assume_init() } } diff --git a/library/coretests/tests/hint.rs b/library/coretests/tests/hint.rs index 032bbc1dcc80f..24de27b24b802 100644 --- a/library/coretests/tests/hint.rs +++ b/library/coretests/tests/hint.rs @@ -21,3 +21,39 @@ fn select_unpredictable_drop() { assert!(a_dropped.get()); assert!(b_dropped.get()); } + +#[test] +#[should_panic = "message canary"] +fn select_unpredictable_drop_on_panic() { + use core::cell::Cell; + + struct X<'a> { + cell: &'a Cell, + expect: u16, + write: u16, + } + + impl Drop for X<'_> { + fn drop(&mut self) { + let value = self.cell.get(); + self.cell.set(self.write); + assert_eq!(value, self.expect, "message canary"); + } + } + + let cell = Cell::new(0); + + // Trigger a double-panic if the selected cell was not dropped during panic. + let _armed = X { cell: &cell, expect: 0xdead, write: 0 }; + let selected = X { cell: &cell, write: 0xdead, expect: 1 }; + let unselected = X { cell: &cell, write: 1, expect: 0xff }; + + // The correct drop order is: + // + // 1. `unselected` drops, writes 1, and panics as 0 != 0xff + // 2. `selected` drops during unwind, writes 0xdead and does not panic as 1 == 1 + // 3. `armed` drops during unwind, writes 0 and does not panic as 0xdead == 0xdead + // + // If `selected` is not dropped, `armed` panics as 1 != 0xdead + let _unreachable = core::hint::select_unpredictable(true, selected, unselected); +}