From 83f5fa2902a8376906362f2ddc18ad2fb867047c Mon Sep 17 00:00:00 2001
From: Oliver Balfour <oliver.leo.balfour@gmail.com>
Date: Tue, 14 Feb 2023 22:38:12 +0800
Subject: [PATCH] update reference pool counts at the end of allow_threads to
 avoid use-after-free

---
 newsfragments/2952.fixed.md |  1 +
 src/gil.rs                  | 32 ++++++++++++++++++++++----------
 src/marker.rs               |  2 ++
 3 files changed, 25 insertions(+), 10 deletions(-)
 create mode 100644 newsfragments/2952.fixed.md

diff --git a/newsfragments/2952.fixed.md b/newsfragments/2952.fixed.md
new file mode 100644
index 00000000000..508631283e3
--- /dev/null
+++ b/newsfragments/2952.fixed.md
@@ -0,0 +1 @@
+Fix a reference counting race condition affecting `PyObject`s cloned in `allow_threads` blocks.
diff --git a/src/gil.rs b/src/gil.rs
index 9ad8ead68d1..678f12eecad 100644
--- a/src/gil.rs
+++ b/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
diff --git a/src/marker.rs b/src/marker.rs
index 714af761020..eb2d1be4339 100644
--- a/src/marker.rs
+++ b/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());
                 }
             }
         }