1 //! Utility for helping miri understand our exposed pointers.
2 //!
3 //! During normal execution, this module is equivalent to pointer casts. However, when running
4 //! under miri, pointer casts are replaced with lookups in a hash map. This makes Tokio compatible
5 //! with strict provenance when running under miri (which comes with a performance cost).
6 
7 use std::marker::PhantomData;
8 #[cfg(miri)]
9 use {crate::loom::sync::Mutex, std::collections::BTreeMap};
10 
11 pub(crate) struct PtrExposeDomain<T> {
12     #[cfg(miri)]
13     map: Mutex<BTreeMap<usize, *const T>>,
14     _phantom: PhantomData<T>,
15 }
16 
17 // SAFETY: Actually using the pointers is unsafe, so it's sound to transfer them across threads.
18 unsafe impl<T> Sync for PtrExposeDomain<T> {}
19 
20 impl<T> PtrExposeDomain<T> {
new() -> Self21     pub(crate) const fn new() -> Self {
22         Self {
23             #[cfg(miri)]
24             map: Mutex::const_new(BTreeMap::new()),
25             _phantom: PhantomData,
26         }
27     }
28 
29     #[inline]
expose_provenance(&self, ptr: *const T) -> usize30     pub(crate) fn expose_provenance(&self, ptr: *const T) -> usize {
31         #[cfg(miri)]
32         {
33             // FIXME: Use `pointer:addr` when it is stable.
34             // SAFETY: Equivalent to `pointer::addr` which is safe.
35             let addr: usize = unsafe { std::mem::transmute(ptr) };
36             self.map.lock().insert(addr, ptr);
37             addr
38         }
39 
40         #[cfg(not(miri))]
41         {
42             ptr as usize
43         }
44     }
45 
46     #[inline]
47     #[allow(clippy::wrong_self_convention)] // mirrors std name
from_exposed_addr(&self, addr: usize) -> *const T48     pub(crate) fn from_exposed_addr(&self, addr: usize) -> *const T {
49         #[cfg(miri)]
50         {
51             let maybe_ptr = self.map.lock().get(&addr).copied();
52 
53             // SAFETY: Intentionally trigger a miri failure if the provenance we want is not
54             // exposed.
55             unsafe { maybe_ptr.unwrap_unchecked() }
56         }
57 
58         #[cfg(not(miri))]
59         {
60             addr as *const T
61         }
62     }
63 
64     #[inline]
unexpose_provenance(&self, _ptr: *const T)65     pub(crate) fn unexpose_provenance(&self, _ptr: *const T) {
66         #[cfg(miri)]
67         {
68             // SAFETY: Equivalent to `pointer::addr` which is safe.
69             let addr: usize = unsafe { std::mem::transmute(_ptr) };
70             let maybe_ptr = self.map.lock().remove(&addr);
71 
72             // SAFETY: Intentionally trigger a miri failure if the provenance we want is not
73             // exposed.
74             unsafe { maybe_ptr.unwrap_unchecked() };
75         }
76     }
77 }
78