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            let addr: usize = ptr.addr();
34            self.map.lock().insert(addr, ptr);
35            addr
36        }
37
38        #[cfg(not(miri))]
39        {
40            ptr as usize
41        }
42    }
43
44    #[inline]
45    #[allow(clippy::wrong_self_convention)] // mirrors std name
46    pub(crate) fn from_exposed_addr(&self, addr: usize) -> *const T {
47        #[cfg(miri)]
48        {
49            let maybe_ptr = self.map.lock().get(&addr).copied();
50
51            // SAFETY: Intentionally trigger a miri failure if the provenance we want is not
52            // exposed.
53            unsafe { maybe_ptr.unwrap_unchecked() }
54        }
55
56        #[cfg(not(miri))]
57        {
58            addr as *const T
59        }
60    }
61
62    #[inline]
63    pub(crate) fn unexpose_provenance(&self, _ptr: *const T) {
64        #[cfg(miri)]
65        {
66            let addr: usize = _ptr.addr();
67            let maybe_ptr = self.map.lock().remove(&addr);
68
69            // SAFETY: Intentionally trigger a miri failure if the provenance we want is not
70            // exposed.
71            unsafe { maybe_ptr.unwrap_unchecked() };
72        }
73    }
74}
tokio/util/ptr_expose.rs

tokio/util/
ptr_expose.rs
