Fix GC use-after-free via atomic try_to_owned

youknowone · youknowone · commit f56a0b3b91cb · 2026-03-05T11:32:42.000+09:00
The GC's strong reference creation had a TOCTOU race: between
checking strong_count() &gt; 0 and calling to_owned() (which calls
inc()), another thread could dec() the count to 0 and proceed
with deallocation. For objects without __del__, no resurrection
check occurs, so the memory is freed while GC holds a dangling
reference. The subsequent drop accesses freed memory, corrupting
malloc metadata (malloc(): unaligned tcache chunk detected).

Fix by replacing the check-then-act pattern with CAS-based
try_to_owned()/safe_inc() that atomically verifies count &gt; 0
and increments. Apply the same fix to WeakRefList callback
collection. Also add atomic dict snapshot for list(dict) to
prevent RuntimeError during concurrent dict iteration.

Unskip test_thread_safety which now passes reliably.
diff --git a/Lib/test/_test_multiprocessing.py b/Lib/test/_test_multiprocessing.py
@@ -4815,8 +4815,6 @@ def test_finalize(self):
         result = [obj for obj in iter(conn.recv, 'STOP')]
         self.assertEqual(result, ['a', 'b', 'd10', 'd03', 'd02', 'd01', 'e'])
 
-    # TODO: RUSTPYTHON - gc.get_threshold() and gc.set_threshold() not implemented
-    @unittest.expectedFailure
     @support.requires_resource('cpu')
     def test_thread_safety(self):
         # bpo-24484: _run_finalizers() should be thread-safe
diff --git a/crates/common/src/refcount.rs b/crates/common/src/refcount.rs
@@ -123,7 +123,7 @@ impl RefCount {
     pub fn safe_inc(&self) -> bool {
         let mut old = State::from_raw(self.state.load(Ordering::Relaxed));
         loop {
-            if old.destructed() {
+            if old.destructed() || old.strong() == 0 {
                 return false;
             }
             if (old.strong() as usize) >= STRONG {
diff --git a/crates/vm/src/gc_state.rs b/crates/vm/src/gc_state.rs
@@ -299,13 +299,7 @@ impl GcState {
         fn collect_from_list(
             list: &LinkedList<GcLink, PyObject>,
         ) -> impl Iterator<Item = PyObjectRef> + '_ {
-            list.iter().filter_map(|obj| {
-                if obj.strong_count() > 0 {
-                    Some(obj.to_owned())
-                } else {
-                    None
-                }
-            })
+            list.iter().filter_map(|obj| obj.try_to_owned())
         }
 
         match generation {
@@ -468,27 +462,24 @@ impl GcState {
         // After dropping gen_locks, other threads can untrack+free objects,
         // making the raw pointers in `reachable`/`unreachable` dangling.
         // Strong refs keep objects alive for later phases.
+        //
+        // Use try_to_owned() (CAS-based) instead of strong_count()+to_owned()
+        // to prevent a TOCTOU race: another thread can dec() the count to 0
+        // between the check and the increment, causing a use-after-free when
+        // the destroying thread eventually frees the memory.
         let survivor_refs: Vec<PyObjectRef> = reachable
             .iter()
             .filter_map(|ptr| {
                 let obj = unsafe { ptr.0.as_ref() };
-                if obj.strong_count() > 0 {
-                    Some(obj.to_owned())
-                } else {
-                    None
-                }
+                obj.try_to_owned()
             })
             .collect();
 
         let unreachable_refs: Vec<crate::PyObjectRef> = unreachable
             .iter()
             .filter_map(|ptr| {
                 let obj = unsafe { ptr.0.as_ref() };
-                if obj.strong_count() > 0 {
-                    Some(obj.to_owned())
-                } else {
-                    None
-                }
+                obj.try_to_owned()
             })
             .collect();
 
diff --git a/crates/vm/src/object/core.rs b/crates/vm/src/object/core.rs
@@ -576,11 +576,12 @@ impl WeakRefList {
                 ptrs.as_mut().set_next(None);
             }
 
-            // Collect callback if present and weakref is still alive
-            if wr.0.ref_count.get() > 0 {
+            // Collect callback only if we can still acquire a strong ref.
+            if wr.0.ref_count.safe_inc() {
+                let wr_ref = unsafe { PyRef::from_raw(wr as *const Py<PyWeak>) };
                 let cb = unsafe { wr.0.payload.callback.get().replace(None) };
                 if let Some(cb) = cb {
-                    callbacks.push((wr.to_owned(), cb));
+                    callbacks.push((wr_ref, cb));
                 }
             }
 
@@ -626,11 +627,12 @@ impl WeakRefList {
                 ptrs.as_mut().set_next(None);
             }
 
-            // Collect callback without invoking
-            if wr.0.ref_count.get() > 0 {
+            // Collect callback without invoking only if we can keep weakref alive.
+            if wr.0.ref_count.safe_inc() {
+                let wr_ref = unsafe { PyRef::from_raw(wr as *const Py<PyWeak>) };
                 let cb = unsafe { wr.0.payload.callback.get().replace(None) };
                 if let Some(cb) = cb {
-                    callbacks.push((wr.to_owned(), cb));
+                    callbacks.push((wr_ref, cb));
                 }
             }
 
@@ -660,8 +662,8 @@ impl WeakRefList {
         let mut current = NonNull::new(self.head.load(Ordering::Relaxed));
         while let Some(node) = current {
             let wr = unsafe { node.as_ref() };
-            if wr.0.ref_count.get() > 0 {
-                v.push(wr.to_owned());
+            if wr.0.ref_count.safe_inc() {
+                v.push(unsafe { PyRef::from_raw(wr as *const Py<PyWeak>) });
             }
             current = unsafe { WeakLink::pointers(node).as_ref().get_next() };
         }
@@ -952,6 +954,23 @@ impl ToOwned for PyObject {
     }
 }
 
+impl PyObject {
+    /// Atomically try to create a strong reference.
+    /// Returns `None` if the strong count is already 0 (object being destroyed).
+    /// Uses CAS to prevent the TOCTOU race between checking strong_count and
+    /// incrementing it.
+    #[inline]
+    pub fn try_to_owned(&self) -> Option<PyObjectRef> {
+        if self.0.ref_count.safe_inc() {
+            Some(PyObjectRef {
+                ptr: NonNull::from(self),
+            })
+        } else {
+            None
+        }
+    }
+}
+
 impl PyObjectRef {
     #[inline(always)]
     pub const fn into_raw(self) -> NonNull<PyObject> {
diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs
@@ -1319,6 +1319,10 @@ impl VirtualMachine {
         } else if cls.is(self.ctx.types.list_type) {
             list_borrow = value.downcast_ref::<PyList>().unwrap().borrow_vec();
             &list_borrow
+        } else if cls.is(self.ctx.types.dict_type) {
+            // Atomic snapshot of dict keys for thread-safe iteration.
+            let keys = value.downcast_ref::<PyDict>().unwrap().keys_vec();
+            return keys.into_iter().map(func).collect();
         } else if cls.is(self.ctx.types.dict_keys_type) {
             // Atomic snapshot of dict keys - prevents race condition during iteration
             let keys = value.downcast_ref::<PyDictKeys>().unwrap().dict.keys_vec();

Original file line number	Diff line number	Diff line change
`@@ -123,7 +123,7 @@ impl RefCount {`
`123`	`123`	`pub fn safe_inc(&self) -> bool {`
`124`	`124`	`let mut old = State::from_raw(self.state.load(Ordering::Relaxed));`
`125`	`125`	`loop {`
`126`		`- if old.destructed() {`
	`126`	`+ if old.destructed() \|\| old.strong() == 0 {`
`127`	`127`	`return false;`
`128`	`128`	`}`
`129`	`129`	`if (old.strong() as usize) >= STRONG {`