update reference pool counts at the end of allow_threads to avoid use-after-free

newsfragments/2952.fixed.md

@@ -0,0 +1 @@
+Fix a reference counting race condition affecting `PyObject`s cloned in `allow_threads` blocks.

src/gil.rs

@@ -277,7 +277,7 @@ impl Drop for GILGuard {
 type PyObjVec = Vec<NonNull<ffi::PyObject>>;
 
 /// Thread-safe storage for objects which were inc_ref / dec_ref while the GIL was not held.
-struct ReferencePool {
+pub(crate) struct ReferencePool {
     dirty: atomic::AtomicBool,
     // .0 is INCREFs, .1 is DECREFs
     pointer_ops: Mutex<(PyObjVec, PyObjVec)>,
@@ -301,7 +301,7 @@ impl ReferencePool {
         self.dirty.store(true, atomic::Ordering::Release);
     }
 
-    fn update_counts(&self, _py: Python<'_>) {
+    pub(crate) fn update_counts(&self, _py: Python<'_>) {
         let prev = self.dirty.swap(false, atomic::Ordering::Acquire);
         if !prev {
             return;
@@ -325,7 +325,7 @@ impl ReferencePool {
 
 unsafe impl Sync for ReferencePool {}
 
-static POOL: ReferencePool = ReferencePool::new();
+pub(crate) static POOL: ReferencePool = ReferencePool::new();
 
 /// A RAII pool which PyO3 uses to store owned Python references.
 ///
@@ -632,7 +632,9 @@ mod tests {
         // Next time the GIL is acquired, the reference is released
         Python::with_gil(|py| {
             assert_eq!(obj.get_refcnt(py), 1);
-            assert!(pool_not_dirty());
+            let non_null = unsafe { NonNull::new_unchecked(obj.as_ptr()) };
+            assert!(!POOL.pointer_ops.lock().0.contains(&non_null));
+            assert!(!POOL.pointer_ops.lock().1.contains(&non_null));
         });
     }
 
@@ -693,6 +695,22 @@ mod tests {
         assert!(gil_is_acquired());
     }
 
+    #[test]
+    fn test_allow_threads_updates_refcounts() {
+        Python::with_gil(|py| {
+            // Make a simple object with 1 reference
+            let obj = get_object(py);
+            assert!(obj.get_refcnt(py) == 1);
+            // Clone the object without the GIL to use internal tracking
+            let escaped_ref = py.allow_threads(|| obj.clone());
+            // But after the block the refcounts are updated
+            assert!(obj.get_refcnt(py) == 2);
+            drop(escaped_ref);
+            assert!(obj.get_refcnt(py) == 1);
+            drop(obj);
+        });
+    }
+
     #[test]
     fn dropping_gil_does_not_invalidate_references() {
         // Acquiring GIL for the second time should be safe - see #864
@@ -800,12 +818,6 @@ mod tests {
 
             REFCNT_CHECKED.wait();
 
-            // Returning from allow_threads doesn't clear the pool
-            py.allow_threads(|| {
-                // Acquiring GIL will clear the pending change
-                Python::with_gil(|_| {});
-            });
-
             println!("6. The main thread has acquired the GIL again and processed the pool.");
 
             // Total reference count should be one higher

src/marker.rs

@@ -478,6 +478,8 @@ impl<'py> Python<'py> {
                 gil::GIL_COUNT.with(|c| c.set(self.count));
                 unsafe {
                     ffi::PyEval_RestoreThread(self.tstate);
+                    // Update counts of PyObjects / Py that were cloned or dropped during `f`.
+                    gil::POOL.update_counts(Python::assume_gil_acquired());
                 }
             }
         }