Avoid pointer casts in IO driver on miri

tokio/src/runtime/io/driver.rs

@@ -168,8 +168,7 @@ impl Driver {
                 self.signal_ready = true;
             } else {
                 let ready = Ready::from_mio(event);
-                // Use std::ptr::from_exposed_addr when stable
-                let ptr: *const ScheduledIo = token.0 as *const _;
+                let ptr = super::EXPOSE_IO.from_exposed_addr(token.0);
 
                 // Safety: we ensure that the pointers used as tokens are not freed
                 // until they are both deregistered from mio **and** we know the I/O

tokio/src/runtime/io/mod.rs

@@ -14,3 +14,6 @@ use scheduled_io::ScheduledIo;
 
 mod metrics;
 use metrics::IoDriverMetrics;
+
+use crate::util::PtrExposeDomain;
+static EXPOSE_IO: PtrExposeDomain<ScheduledIo> = PtrExposeDomain::new();

tokio/src/runtime/io/registration_set.rs

@@ -115,6 +115,7 @@ impl RegistrationSet {
     // This function is marked as unsafe, because the caller must make sure that
     // `io` is part of the registration set.
     pub(super) unsafe fn remove(&self, synced: &mut Synced, io: &ScheduledIo) {
+        super::EXPOSE_IO.unexpose_provenance(io);
         let _ = synced.registrations.remove(io.into());
     }
 }

tokio/src/runtime/io/scheduled_io.rs

@@ -187,8 +187,7 @@ impl Default for ScheduledIo {
 
 impl ScheduledIo {
     pub(crate) fn token(&self) -> mio::Token {
-        // use `expose_addr` when stable
-        mio::Token(self as *const _ as usize)
+        mio::Token(super::EXPOSE_IO.expose_provenance(self))
     }
 
     /// Invoked when the IO driver is shut down; forces this `ScheduledIo` into a

tokio/src/util/mod.rs

@@ -69,6 +69,9 @@ cfg_rt! {
 
     mod rc_cell;
     pub(crate) use rc_cell::RcCell;
+
+    mod ptr_expose;
+    pub(crate) use ptr_expose::PtrExposeDomain;
 }
 
 cfg_rt_multi_thread! {

tokio/src/util/ptr_expose.rs

@@ -0,0 +1,64 @@
+//! Utility for helping miri understand our exposed pointers.
+//!
+//! During normal execution, this module is equivalent to pointer casts. However, when running
+//! under miri, pointer casts are replaced with lookups in a hash map. This makes Tokio compatible
+//! with strict provenance when running under miri (which comes with a perf cost).
+
+use std::marker::PhantomData;
+#[cfg(miri)]
+use {std::collections::HashMap, crate::loom::sync::Mutex};
+
+pub(crate) struct PtrExposeDomain<T> {
+    #[cfg(miri)]
+    map: Mutex<HashMap<usize, *mut T>>,
+    _phantom: PhantomData<T>,
+}
+
+impl<T> PtrExposeDomain<T> {
+    pub(crate) const fn new() -> Self {
+        Self {
+            #[cfg(miri)]
+            map: Mutex::new(HashMap::new()),
+            _phantom: PhantomData,
+        }
+    }
+
+    #[inline]
+    pub(crate) fn expose_provenance(&self, ptr: *const T) -> usize {
+        #[cfg(miri)]
+        {
+            // SAFETY: Equivalent to `pointer::addr` which is safe.
+            let addr: usize = unsafe { std::mem::transmute(ptr) };
+            self.map.lock().insert(addr, ptr);
+            addr
+        }
+
+        #[cfg(not(miri))]
+        {
+            ptr as usize
+        }
+    }
+
+    #[inline]
+    pub(crate) fn from_exposed_addr(&self, addr: usize) -> *const T {
+        #[cfg(miri)]
+        {
+            self.map.lock().get(&addr).expect("Provided address is not exposed.")
+        }
+
+        #[cfg(not(miri))]
+        {
+            addr as *const T
+        }
+    }
+
+    #[inline]
+    pub(crate) fn unexpose_provenance(&self, _ptr: *const T) {
+        #[cfg(miri)]
+        {
+            // SAFETY: Equivalent to `pointer::addr` which is safe.
+            let addr: usize = unsafe { std::mem::transmute(_ptr) };
+            self.map.lock().remove(addr).expect("Provided address is not exposed.");
+        }
+    }
+}