tokio/util/
ptr_expose.rs

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
7use std::marker::PhantomData;
8#[cfg(miri)]
9use {crate::loom::sync::Mutex, std::collections::BTreeMap};
10
11pub(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.
18unsafe impl<T> Sync for PtrExposeDomain<T> {}
19
20impl<T> PtrExposeDomain<T> {
21    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]
30    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
48    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]
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}
tokio/util/ptr_expose.rs

tokio/util/
ptr_expose.rs
