gc: re-derive shadow-stack slots across visitor calls, and guard a walk in debug

pyre/pyre-object/src/gc_roots.rs

@@ -51,6 +51,7 @@
 #[cfg(debug_assertions)]
 use std::cell::Cell;
 use std::cell::UnsafeCell;
+use std::marker::PhantomData;
 
 use crate::pyobject::PyObjectRef;
 
@@ -60,16 +61,22 @@ thread_local! {
     /// each [`RootScope`]'s lifetime; the matching [`Drop`] truncates
     /// back to the saved length.
     ///
-    /// Access is raw, matching RPython's base/top shadow-stack shape.
-    /// Debug builds retain the old re-entrancy safety net through
-    /// `SHADOW_STACK_ACCESS_DEPTH`; release builds pay no borrow-check cost.
+    /// Access is raw, matching RPython's base/top shadow-stack shape. Debug
+    /// builds check overlapping ordinary accesses through
+    /// `SHADOW_STACK_ACCESS_DEPTH` and separately reject mutation during a
+    /// root walk through `SHADOW_STACK_WALK_IN_PROGRESS`; release builds pay
+    /// no bookkeeping cost.
     static SHADOW_STACK: UnsafeCell<Vec<PyObjectRef>> =
         const { UnsafeCell::new(Vec::new()) };
 
     /// Signed debug borrow depth for [`SHADOW_STACK`]: non-negative values
     /// count shared accesses, while `-1` marks an exclusive access.
     #[cfg(debug_assertions)]
     static SHADOW_STACK_ACCESS_DEPTH: Cell<i32> = const { Cell::new(0) };
+
+    /// Whether this thread is currently driving a shadow-stack root walk.
+    #[cfg(debug_assertions)]
+    static SHADOW_STACK_WALK_IN_PROGRESS: Cell<bool> = const { Cell::new(false) };
 }
 
 #[cfg(debug_assertions)]
@@ -108,6 +115,41 @@ impl Drop for ShadowStackAccess<'_> {
     }
 }
 
+#[cfg(debug_assertions)]
+struct ShadowStackWalk<'a> {
+    in_progress: &'a Cell<bool>,
+}
+
+#[cfg(debug_assertions)]
+impl<'a> ShadowStackWalk<'a> {
+    fn new(in_progress: &'a Cell<bool>) -> Self {
+        // Arm the flag outside the assertion: the walk must be marked whether
+        // or not assertions are compiled in, and an assertion that doubles as
+        // the only writer stops arming the moment it is disabled.
+        let was_walking = in_progress.replace(true);
+        debug_assert!(!was_walking, "a shadow-stack root walk was re-entered");
+        Self { in_progress }
+    }
+}
+
+#[cfg(debug_assertions)]
+impl Drop for ShadowStackWalk<'_> {
+    fn drop(&mut self) {
+        self.in_progress.set(false);
+    }
+}
+
+#[cfg(debug_assertions)]
+#[inline(always)]
+fn assert_shadow_stack_not_walking() {
+    SHADOW_STACK_WALK_IN_PROGRESS.with(|in_progress| {
+        debug_assert!(
+            !in_progress.get(),
+            "the shadow stack was mutated while a root walk was in progress"
+        );
+    });
+}
+
 #[inline(always)]
 fn with_shadow_stack<R>(f: impl FnOnce(&Vec<PyObjectRef>) -> R) -> R {
     SHADOW_STACK.with(|cell| {
@@ -166,6 +208,9 @@ pub struct RootScope {
     /// called. Will be replaced with a backend
     /// `RootSet::Handle` once the active GC consumes the stack.
     save_point: usize,
+    /// A root bracket belongs to the thread whose shadow stack supplied this
+    /// save point. Moving it could truncate an unrelated thread's stack.
+    _not_send: PhantomData<*const ()>,
 }
 
 impl RootScope {
@@ -175,6 +220,7 @@ impl RootScope {
     fn new() -> Self {
         Self {
             save_point: shadow_stack_len(),
+            _not_send: PhantomData,
         }
     }
 }
@@ -185,6 +231,8 @@ impl Drop for RootScope {
     /// bracketed `livevars` slots.
     #[inline]
     fn drop(&mut self) {
+        #[cfg(debug_assertions)]
+        assert_shadow_stack_not_walking();
         with_shadow_stack_mut(|stack| {
             // `truncate` is a no-op if `save_point >= len()`, which is
             // the steady-state case for an empty bracket.
@@ -215,6 +263,8 @@ pub fn push_roots() -> RootScope {
 /// it (`@dont_look_inside`, `rlib/jit.py:139`), the `shadow_stack_len` twin.
 #[majit_macros::dont_look_inside]
 pub fn pin_root(root: PyObjectRef) {
+    #[cfg(debug_assertions)]
+    assert_shadow_stack_not_walking();
     // Publish the raw value first.  `try_gc_current_object_address` is itself
     // a GC operation and may wait behind another thread's collection; the
     // fresh allocation must already be visible to that collector before we
@@ -316,6 +366,8 @@ pub fn shadow_stack_copy_range(base: usize, dst: &mut [PyObjectRef]) {
 /// `rlib/jit.py:139`), the [`shadow_stack_get`] twin.
 #[majit_macros::dont_look_inside]
 pub fn shadow_stack_set(index: usize, root: PyObjectRef) {
+    #[cfg(debug_assertions)]
+    assert_shadow_stack_not_walking();
     // Publish the raw value first, for the reason [`pin_root`] gives: the
     // `try_gc_current_object_address` query is itself a GC operation and may
     // wait behind another thread's collection, so the value must already be
@@ -335,16 +387,26 @@ pub fn shadow_stack_set(index: usize, root: PyObjectRef) {
 /// `#[repr(transparent)]` over `usize` and `PyObjectRef =
 /// *mut PyObject` is also pointer-sized.
 ///
-/// The visitor receives mutable references so a moving collector can
-/// rewrite slot contents post-relocation. In debug builds the signed
-/// `SHADOW_STACK_ACCESS_DEPTH` guard catches re-entrant access from an
-/// accidentally allocating walker; release builds match upstream's raw
-/// root-stack access without that safety net.
+/// The visitor receives mutable references so a moving collector can rewrite
+/// slot contents post-relocation. Each slot is re-derived from the stack after
+/// the previous visitor call, so a re-entrant push cannot leave the walk using
+/// a pointer into an old allocation. Debug builds additionally reject shadow-
+/// stack mutation by an accidentally allocating walker; release builds match
+/// upstream's raw root-stack access without that diagnostic.
 #[inline]
 pub fn walk_shadow_stack(mut visitor: impl FnMut(&mut PyObjectRef)) {
-    with_shadow_stack_mut(|stack| {
-        for slot in stack.iter_mut() {
-            visitor(slot);
+    SHADOW_STACK.with(|cell| {
+        let cell = cell as *const UnsafeCell<Vec<PyObjectRef>>;
+        #[cfg(debug_assertions)]
+        SHADOW_STACK_WALK_IN_PROGRESS.with(|in_progress| {
+            let _walk = ShadowStackWalk::new(in_progress);
+            // SAFETY: `cell` is this thread's live shadow-stack cell.
+            unsafe { walk_shadow_stack_cell(cell, &mut visitor) };
+        });
+        #[cfg(not(debug_assertions))]
+        // SAFETY: `cell` is this thread's live shadow-stack cell.
+        unsafe {
+            walk_shadow_stack_cell(cell, &mut visitor);
         }
     });
 }
@@ -355,14 +417,74 @@ pub fn capture_shadow_stack_area() -> *const () {
 }
 
 /// Visit a captured thread's pinned-root stack without consulting caller TLS.
+/// Like [`walk_shadow_stack`], this re-derives each slot after the previous
+/// visitor call so a push cannot strand the walk in a reallocated buffer.
+/// Debug builds reject shadow-stack mutation for the duration of the walk.
 ///
 /// # Safety
 /// `data` must come from [`capture_shadow_stack_area`], and the owning thread
 /// must be quiesced for the duration of the walk.
 pub unsafe fn walk_shadow_stack_area(data: *const (), mut visitor: impl FnMut(&mut PyObjectRef)) {
-    let stack = unsafe { &mut *(*(data as *const UnsafeCell<Vec<PyObjectRef>>)).get() };
-    for slot in stack.iter_mut() {
-        visitor(slot);
+    let cell = data as *const UnsafeCell<Vec<PyObjectRef>>;
+    #[cfg(debug_assertions)]
+    SHADOW_STACK_WALK_IN_PROGRESS.with(|in_progress| {
+        // For a self-collection this protects the stack being walked. A
+        // foreign walk marks the collector thread instead, which is harmless:
+        // the owning mutator is quiesced and performs no stack accesses.
+        let _walk = ShadowStackWalk::new(in_progress);
+        // SAFETY: the caller guarantees that `cell` is a captured shadow stack
+        // whose owning thread is quiesced for this walk.
+        unsafe { walk_shadow_stack_cell(cell, &mut visitor) };
+    });
+    #[cfg(not(debug_assertions))]
+    // SAFETY: the caller guarantees that `cell` is a captured shadow stack
+    // whose owning thread is quiesced for this walk.
+    unsafe {
+        walk_shadow_stack_cell(cell, &mut visitor);
+    }
+}
+
+/// Walk one captured shadow-stack cell without retaining a reference to its
+/// `Vec` across a visitor call.
+///
+/// # Safety
+/// `cell` must remain a valid shadow-stack cell for the duration of the walk,
+/// and no thread other than a re-entrant visitor may mutate it.
+unsafe fn walk_shadow_stack_cell(
+    cell: *const UnsafeCell<Vec<PyObjectRef>>,
+    visitor: &mut impl FnMut(&mut PyObjectRef),
+) {
+    // The interval is fixed at entry: `walk_stack_root` (`shadowstack.py:43-46`)
+    // receives `start` and `addr` as arguments and runs `while addr != start`,
+    // so a root pushed mid-walk is not part of that walk. Re-reading the length
+    // each iteration instead would extend the walk over roots pinned by the
+    // visitor — and pinning during a walk is precisely what the debug guard
+    // above declares illegal, so there is nothing to serve by diverging here.
+    let end = {
+        // SAFETY: no reference from this block outlives it.
+        let stack = unsafe { &*(*cell).get() };
+        stack.len()
+    };
+    let mut index = 0;
+    while index < end {
+        // Re-derive the slot address every iteration rather than holding one
+        // `iter_mut()` cursor: a re-entrant push can still reallocate the
+        // buffer, which would strand a cached pointer in the old allocation.
+        // Only pushes can occur, so `end` stays within bounds.
+        let slot = {
+            // SAFETY: the caller guarantees exclusive access except for a
+            // possible re-entrant mutation by `visitor`. No reference from a
+            // previous iteration is live here, and this temporary reference
+            // ends before `visitor` is called.
+            let stack = unsafe { &mut *(*cell).get() };
+            // Return only the slot address so the `Vec` reference is not held
+            // across the visitor, which may push and reallocate its buffer.
+            unsafe { stack.as_mut_ptr().add(index) }
+        };
+        // SAFETY: `slot` names the live element at `index`. The reference is
+        // handed directly to the visitor and is not used after it returns.
+        visitor(unsafe { &mut *slot });
+        index += 1;
     }
 }
 
@@ -532,4 +654,51 @@ mod tests {
         });
         assert_eq!(shadow_stack_get(shadow_stack_len() - 2) as usize, 0xAA);
     }
+
+    #[cfg(debug_assertions)]
+    #[test]
+    #[should_panic(expected = "the shadow stack was mutated while a root walk was in progress")]
+    fn walk_shadow_stack_area_rejects_reentrant_push_in_debug_builds() {
+        let _roots = push_roots();
+        pin_root(dummy(0x1));
+        let area = capture_shadow_stack_area();
+        // SAFETY: `area` is this thread's shadow stack, and this synchronous
+        // self-walk keeps it live for the duration of the call.
+        unsafe { walk_shadow_stack_area(area, |_| pin_root(dummy(0x2))) };
+    }
+
+    #[cfg(not(debug_assertions))]
+    #[test]
+    fn walk_shadow_stack_survives_push_that_reallocates() {
+        let _roots = push_roots();
+        pin_root(dummy(0x11));
+        pin_root(dummy(0x22));
+
+        let (capacity_before, additional) =
+            with_shadow_stack(|stack| (stack.capacity(), stack.capacity() - stack.len() + 1));
+        let mut pushed = false;
+        let mut seen = Vec::new();
+        walk_shadow_stack(|slot| {
+            seen.push(*slot as usize);
+            if !pushed {
+                pushed = true;
+                for offset in 0..additional {
+                    pin_root(dummy(0x1000 + offset));
+                }
+            }
+        });
+
+        assert!(pushed);
+        // The visitor's pins really did move the buffer.
+        assert!(with_shadow_stack(Vec::capacity) > capacity_before);
+        // `0x22` is the discriminator: it is read on the iteration AFTER the
+        // reallocation, so getting it right means the walk followed the buffer
+        // to its new address instead of reading the freed one. The walk stops
+        // at the two roots that were live at entry — the interval is fixed
+        // there, as in `walk_stack_root` (`shadowstack.py:43-46`), so the
+        // roots the visitor pinned are not part of this walk.
+        assert_eq!(seen, vec![0x11, 0x22]);
+        assert_eq!(shadow_stack_get(0) as usize, 0x11);
+        assert_eq!(shadow_stack_get(1) as usize, 0x22);
+    }
 }