Avoid pointer casts in IO driver on miri

spellcheck.dic

@@ -1,4 +1,4 @@
-286
+287
 &
 +
 <
@@ -152,6 +152,7 @@ metadata
 mio
 Mio
 mio's
+miri
 misconfigured
 mock's
 mpmc

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::ptr_expose::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

@@ -86,3 +86,7 @@ pub(crate) mod memchr;
 pub(crate) mod markers;
 
 pub(crate) mod cacheline;
+
+cfg_io_driver_impl! {
+    pub(crate) mod ptr_expose;
+}

tokio/src/util/ptr_expose.rs

@@ -0,0 +1,77 @@
+//! 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 performance cost).
+
+use std::marker::PhantomData;
+#[cfg(miri)]
+use {crate::loom::sync::Mutex, std::collections::BTreeMap};
+
+pub(crate) struct PtrExposeDomain<T> {
+    #[cfg(miri)]
+    map: Mutex<BTreeMap<usize, *const T>>,
+    _phantom: PhantomData<T>,
+}
+
+// SAFETY: Actually using the pointers is unsafe, so it's sound to transfer them across threads.
+unsafe impl<T> Sync for PtrExposeDomain<T> {}
+
+impl<T> PtrExposeDomain<T> {
+    pub(crate) const fn new() -> Self {
+        Self {
+            #[cfg(miri)]
+            map: Mutex::const_new(BTreeMap::new()),
+            _phantom: PhantomData,
+        }
+    }
+
+    #[inline]
+    pub(crate) fn expose_provenance(&self, ptr: *const T) -> usize {
+        #[cfg(miri)]
+        {
+            // FIXME: Use `pointer:addr` when it is stable.
+            // 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]
+    #[allow(clippy::wrong_self_convention)] // mirrors std name
+    pub(crate) fn from_exposed_addr(&self, addr: usize) -> *const T {
+        #[cfg(miri)]
+        {
+            let maybe_ptr = self.map.lock().get(&addr).copied();
+
+            // SAFETY: Intentionally trigger a miri failure if the provenance we want is not
+            // exposed.
+            unsafe { maybe_ptr.unwrap_unchecked() }
+        }
+
+        #[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) };
+            let maybe_ptr = self.map.lock().remove(&addr);
+
+            // SAFETY: Intentionally trigger a miri failure if the provenance we want is not
+            // exposed.
+            unsafe { maybe_ptr.unwrap_unchecked() };
+        }
+    }
+}