update reference pool counts at the end of allow_threads to avoid use-after-free
This commit is contained in:
parent
03ecb805ad
commit
83f5fa2902
|
@ -0,0 +1 @@
|
||||||
|
Fix a reference counting race condition affecting `PyObject`s cloned in `allow_threads` blocks.
|
32
src/gil.rs
32
src/gil.rs
|
@ -277,7 +277,7 @@ impl Drop for GILGuard {
|
||||||
type PyObjVec = Vec<NonNull<ffi::PyObject>>;
|
type PyObjVec = Vec<NonNull<ffi::PyObject>>;
|
||||||
|
|
||||||
/// Thread-safe storage for objects which were inc_ref / dec_ref while the GIL was not held.
|
/// 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,
|
dirty: atomic::AtomicBool,
|
||||||
// .0 is INCREFs, .1 is DECREFs
|
// .0 is INCREFs, .1 is DECREFs
|
||||||
pointer_ops: Mutex<(PyObjVec, PyObjVec)>,
|
pointer_ops: Mutex<(PyObjVec, PyObjVec)>,
|
||||||
|
@ -301,7 +301,7 @@ impl ReferencePool {
|
||||||
self.dirty.store(true, atomic::Ordering::Release);
|
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);
|
let prev = self.dirty.swap(false, atomic::Ordering::Acquire);
|
||||||
if !prev {
|
if !prev {
|
||||||
return;
|
return;
|
||||||
|
@ -325,7 +325,7 @@ impl ReferencePool {
|
||||||
|
|
||||||
unsafe impl Sync for 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.
|
/// 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
|
// Next time the GIL is acquired, the reference is released
|
||||||
Python::with_gil(|py| {
|
Python::with_gil(|py| {
|
||||||
assert_eq!(obj.get_refcnt(py), 1);
|
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());
|
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]
|
#[test]
|
||||||
fn dropping_gil_does_not_invalidate_references() {
|
fn dropping_gil_does_not_invalidate_references() {
|
||||||
// Acquiring GIL for the second time should be safe - see #864
|
// Acquiring GIL for the second time should be safe - see #864
|
||||||
|
@ -800,12 +818,6 @@ mod tests {
|
||||||
|
|
||||||
REFCNT_CHECKED.wait();
|
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.");
|
println!("6. The main thread has acquired the GIL again and processed the pool.");
|
||||||
|
|
||||||
// Total reference count should be one higher
|
// Total reference count should be one higher
|
||||||
|
|
|
@ -478,6 +478,8 @@ impl<'py> Python<'py> {
|
||||||
gil::GIL_COUNT.with(|c| c.set(self.count));
|
gil::GIL_COUNT.with(|c| c.set(self.count));
|
||||||
unsafe {
|
unsafe {
|
||||||
ffi::PyEval_RestoreThread(self.tstate);
|
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());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue