wrap PyClassObjectContents to prevent it from leaking outside pyo3

Author: mbway

src/impl_/coroutine.rs

@@ -6,7 +6,7 @@ use std::{
 use crate::{
     coroutine::{cancel::ThrowCallback, Coroutine},
     instance::Bound,
-    pycell::impl_::{InternalPyClassObjectLayout, PyClassBorrowChecker},
+    pycell::impl_::{PyClassBorrowChecker, PyClassObjectLayout},
     pyclass::boolean_struct::False,
     types::{PyAnyMethods, PyString},
     IntoPyObject, Py, PyAny, PyClass, PyErr, PyResult, Python,

src/impl_/pycell.rs

@@ -1,5 +1,5 @@
 //! Externally-accessible implementation of pycell
 pub use crate::pycell::impl_::{
-    GetBorrowChecker, PyClassMutability, PyClassObjectBase, PyClassObjectLayout,
+    GetBorrowChecker, PyClassMutability, PyClassObjectBase, PyClassObjectBaseLayout,
     PyStaticClassObject, PyVariableClassObject, PyVariableClassObjectBase,
 };

src/impl_/pyclass.rs

@@ -6,11 +6,11 @@ use crate::{
     ffi,
     impl_::{
         freelist::FreeList,
-        pycell::{GetBorrowChecker, PyClassMutability, PyClassObjectLayout},
+        pycell::{GetBorrowChecker, PyClassMutability, PyClassObjectBaseLayout},
         pyclass_init::PyObjectInit,
         pymethods::{PyGetterDef, PyMethodDefType},
     },
-    pycell::{impl_::InternalPyClassObjectLayout, PyBorrowError},
+    pycell::{impl_::PyClassObjectLayout, PyBorrowError},
     types::{any::PyAnyMethods, PyBool},
     Borrowed, BoundObject, Py, PyAny, PyClass, PyErr, PyRef, PyResult, PyTypeInfo, Python,
 };
@@ -170,7 +170,7 @@ pub trait PyClassImpl: Sized + 'static {
     const IS_SEQUENCE: bool = false;
 
     /// Description of how this class is laid out in memory
-    type Layout: InternalPyClassObjectLayout<Self>;
+    type Layout: PyClassObjectLayout<Self>;
 
     /// Base class
     type BaseType: PyTypeInfo + PyClassBaseType;
@@ -1137,7 +1137,7 @@ impl<T> PyClassThreadChecker<T> for ThreadCheckerImpl {
     )
 )]
 pub trait PyClassBaseType: Sized {
-    type LayoutAsBase: PyClassObjectLayout<Self>;
+    type LayoutAsBase: PyClassObjectBaseLayout<Self>;
     type BaseNativeType;
     type Initializer: PyObjectInit<Self>;
     type PyClassMutability: PyClassMutability;
@@ -1549,7 +1549,7 @@ where
             let class_ptr = obj.cast::<<ClassT as PyClassImpl>::Layout>();
             // Safety: the object `obj` must have the layout `ClassT::Layout`
             let class_obj = unsafe { &mut *class_ptr };
-            let contents = class_obj.contents_mut() as *mut PyClassObjectContents<ClassT>;
+            let contents = (&mut class_obj.contents_mut().0) as *mut PyClassObjectContents<ClassT>;
             (contents.cast::<u8>(), offset)
         }
         #[cfg(not(Py_3_12))]

src/impl_/pyclass_init.rs

@@ -2,15 +2,19 @@
 use crate::ffi_ptr_ext::FfiPtrExt;
 use crate::impl_::pyclass::PyClassImpl;
 use crate::internal::get_slot::TP_ALLOC;
-use crate::pycell::impl_::{InternalPyClassObjectLayout, PyClassObjectContents};
+use crate::pycell::impl_::{
+    PyClassObjectContents, PyClassObjectLayout, WrappedPyClassObjectContents,
+};
 use crate::types::PyType;
 use crate::{ffi, Borrowed, PyErr, PyResult, Python};
 use crate::{ffi::PyTypeObject, sealed::Sealed, type_object::PyTypeInfo};
 use std::marker::PhantomData;
 
 pub unsafe fn initialize_with_default<T: PyClassImpl + Default>(obj: *mut ffi::PyObject) {
     let contents = T::Layout::contents_uninitialised(obj);
-    (*contents).write(PyClassObjectContents::new(T::default()));
+    (*contents).write(WrappedPyClassObjectContents(PyClassObjectContents::new(
+        T::default(),
+    )));
 }
 
 /// Initializer for Python types.

src/impl_/pymethods.rs

@@ -2,9 +2,9 @@ use crate::exceptions::PyStopAsyncIteration;
 use crate::gil::LockGIL;
 use crate::impl_::callback::IntoPyCallbackOutput;
 use crate::impl_::panic::PanicTrap;
-use crate::impl_::pycell::PyClassObjectLayout;
+use crate::impl_::pycell::PyClassObjectBaseLayout;
 use crate::internal::get_slot::{get_slot, TP_BASE, TP_CLEAR, TP_TRAVERSE};
-use crate::pycell::impl_::{InternalPyClassObjectLayout, PyClassBorrowChecker as _};
+use crate::pycell::impl_::{PyClassBorrowChecker as _, PyClassObjectLayout};
 use crate::pycell::{PyBorrowError, PyBorrowMutError};
 use crate::pyclass::boolean_struct::False;
 use crate::types::any::PyAnyMethods;
@@ -310,8 +310,8 @@ where
     if class_object.check_threadsafe().is_ok()
     // ... and we cannot traverse a type which might be being mutated by a Rust thread
     && class_object.borrow_checker().try_borrow().is_ok() {
-        struct TraverseGuard<'a, U: PyClassImpl, V: InternalPyClassObjectLayout<U>>(&'a V, PhantomData<U>);
-        impl<U: PyClassImpl, V: InternalPyClassObjectLayout<U>> Drop for TraverseGuard<'_, U, V> {
+        struct TraverseGuard<'a, U: PyClassImpl, V: PyClassObjectLayout<U>>(&'a V, PhantomData<U>);
+        impl<U: PyClassImpl, V: PyClassObjectLayout<U>> Drop for TraverseGuard<'_, U, V> {
             fn drop(&mut self) {
                 self.0.borrow_checker().release_borrow()
             }
@@ -320,7 +320,7 @@ where
         // `.try_borrow()` above created a borrow, we need to release it when we're done
         // traversing the object. This allows us to read `instance` safely.
         let _guard = TraverseGuard(class_object, PhantomData);
-        let instance = &*class_object.contents().value.get();
+        let instance = &*class_object.contents().0.value.get();
 
         let visit = PyVisit { visit, arg, _guard: PhantomData };
 

src/instance.rs

@@ -2,7 +2,7 @@ use crate::conversion::IntoPyObject;
 use crate::err::{self, PyErr, PyResult};
 use crate::impl_::pyclass::PyClassImpl;
 use crate::internal_tricks::ptr_from_ref;
-use crate::pycell::impl_::InternalPyClassObjectLayout;
+use crate::pycell::impl_::PyClassObjectLayout;
 use crate::pycell::{PyBorrowError, PyBorrowMutError};
 use crate::pyclass::boolean_struct::{False, True};
 use crate::types::{any::PyAnyMethods, string::PyStringMethods, typeobject::PyTypeMethods};

src/pycell.rs

@@ -208,7 +208,7 @@ use std::mem::ManuallyDrop;
 use std::ops::{Deref, DerefMut};
 
 pub(crate) mod impl_;
-use impl_::{InternalPyClassObjectLayout, PyClassBorrowChecker, PyClassObjectLayout};
+use impl_::{PyClassBorrowChecker, PyClassObjectBaseLayout, PyClassObjectLayout};
 
 /// A wrapper type for an immutably borrowed value from a [`Bound<'py, T>`].
 ///

src/pycell/impl_.rs

@@ -187,13 +187,13 @@ pub trait GetBorrowChecker<T: PyClassImpl> {
 
 impl<T: PyClassImpl<PyClassMutability = Self>> GetBorrowChecker<T> for MutableClass {
     fn borrow_checker(class_object: &T::Layout) -> &BorrowChecker {
-        &class_object.contents().borrow_checker
+        &class_object.contents().0.borrow_checker
     }
 }
 
 impl<T: PyClassImpl<PyClassMutability = Self>> GetBorrowChecker<T> for ImmutableClass {
     fn borrow_checker(class_object: &T::Layout) -> &EmptySlot {
-        &class_object.contents().borrow_checker
+        &class_object.contents().0.borrow_checker
     }
 }
 
@@ -226,7 +226,7 @@ pub struct PyVariableClassObjectBase {
 
 unsafe impl<T> PyLayout<T> for PyVariableClassObjectBase {}
 
-impl<T: PyTypeInfo> PyClassObjectLayout<T> for PyVariableClassObjectBase {
+impl<T: PyTypeInfo> PyClassObjectBaseLayout<T> for PyVariableClassObjectBase {
     fn ensure_threadsafe(&self) {}
     fn check_threadsafe(&self) -> Result<(), PyBorrowError> {
         Ok(())
@@ -237,7 +237,7 @@ impl<T: PyTypeInfo> PyClassObjectLayout<T> for PyVariableClassObjectBase {
 }
 
 #[doc(hidden)]
-pub trait PyClassObjectLayout<T>: PyLayout<T> {
+pub trait PyClassObjectBaseLayout<T>: PyLayout<T> {
     fn ensure_threadsafe(&self);
     fn check_threadsafe(&self) -> Result<(), PyBorrowError>;
     /// Implementation of tp_dealloc.
@@ -247,6 +247,30 @@ pub trait PyClassObjectLayout<T>: PyLayout<T> {
     unsafe fn tp_dealloc(py: Python<'_>, slf: *mut ffi::PyObject);
 }
 
+/// Allow [PyClassObjectLayout] to have public visibility without leaking the structure of [PyClassObjectContents].
+#[doc(hidden)]
+#[repr(transparent)]
+pub struct WrappedPyClassObjectContents<T: PyClassImpl>(pub(crate) PyClassObjectContents<T>);
+
+impl<'a, T: PyClassImpl> From<&'a PyClassObjectContents<T>>
+    for &'a WrappedPyClassObjectContents<T>
+{
+    fn from(value: &'a PyClassObjectContents<T>) -> &'a WrappedPyClassObjectContents<T> {
+        // Safety: Wrapped struct must use repr(transparent)
+        unsafe { std::mem::transmute(value) }
+    }
+}
+
+impl<'a, T: PyClassImpl> From<&'a mut PyClassObjectContents<T>>
+    for &'a mut WrappedPyClassObjectContents<T>
+{
+    fn from(value: &'a mut PyClassObjectContents<T>) -> &'a mut WrappedPyClassObjectContents<T> {
+        // Safety: Wrapped struct must use repr(transparent)
+        unsafe { std::mem::transmute(value) }
+    }
+}
+
+/// Functionality required for creating and managing the memory associated with a pyclass annotated struct.
 #[doc(hidden)]
 #[cfg_attr(
     all(diagnostic_namespace),
@@ -256,18 +280,18 @@ pub trait PyClassObjectLayout<T>: PyLayout<T> {
         note = "the python version being built against influences which layouts are valid",
     )
 )]
-pub trait InternalPyClassObjectLayout<T: PyClassImpl>: PyClassObjectLayout<T> {
+pub trait PyClassObjectLayout<T: PyClassImpl>: PyClassObjectBaseLayout<T> {
     /// Obtain a pointer to the contents of an uninitialized PyObject of this type
     /// Safety: the provided object must have the layout that the implementation is expecting
     unsafe fn contents_uninitialised(
         obj: *mut ffi::PyObject,
-    ) -> *mut MaybeUninit<PyClassObjectContents<T>>;
+    ) -> *mut MaybeUninit<WrappedPyClassObjectContents<T>>;
 
     fn get_ptr(&self) -> *mut T;
 
-    fn contents(&self) -> &PyClassObjectContents<T>;
+    fn contents(&self) -> &WrappedPyClassObjectContents<T>;
 
-    fn contents_mut(&mut self) -> &mut PyClassObjectContents<T>;
+    fn contents_mut(&mut self) -> &mut WrappedPyClassObjectContents<T>;
 
     fn ob_base(&self) -> &<T::BaseType as PyClassBaseType>::LayoutAsBase;
 
@@ -287,7 +311,7 @@ pub trait InternalPyClassObjectLayout<T: PyClassImpl>: PyClassObjectLayout<T> {
     fn borrow_checker(&self) -> &<T::PyClassMutability as PyClassMutability>::Checker;
 }
 
-impl<T, U> PyClassObjectLayout<T> for PyClassObjectBase<U>
+impl<T, U> PyClassObjectBaseLayout<T> for PyClassObjectBase<U>
 where
     U: PySizedLayout<T>,
     T: PyTypeInfo,
@@ -301,6 +325,10 @@ where
     }
 }
 
+/// Implementation of tp_dealloc.
+/// # Safety
+/// - obj must be a valid pointer to an instance of the type at `type_ptr` or a subclass.
+/// - obj must not be used after this call (as it will be freed).
 unsafe fn tp_dealloc(py: Python<'_>, obj: *mut ffi::PyObject, type_ptr: *mut ffi::PyTypeObject) {
     // FIXME: there is potentially subtle issues here if the base is overwritten
     // at runtime? To be investigated.
@@ -376,14 +404,14 @@ pub struct PyStaticClassObject<T: PyClassImpl> {
     contents: PyClassObjectContents<T>,
 }
 
-impl<T: PyClassImpl> InternalPyClassObjectLayout<T> for PyStaticClassObject<T> {
+impl<T: PyClassImpl> PyClassObjectLayout<T> for PyStaticClassObject<T> {
     unsafe fn contents_uninitialised(
         obj: *mut ffi::PyObject,
-    ) -> *mut MaybeUninit<PyClassObjectContents<T>> {
+    ) -> *mut MaybeUninit<WrappedPyClassObjectContents<T>> {
         #[repr(C)]
         struct PartiallyInitializedClassObject<T: PyClassImpl> {
             _ob_base: <T::BaseType as PyClassBaseType>::LayoutAsBase,
-            contents: MaybeUninit<PyClassObjectContents<T>>,
+            contents: MaybeUninit<WrappedPyClassObjectContents<T>>,
         }
         let obj: *mut PartiallyInitializedClassObject<T> = obj.cast();
         addr_of_mut!((*obj).contents)
@@ -397,12 +425,12 @@ impl<T: PyClassImpl> InternalPyClassObjectLayout<T> for PyStaticClassObject<T> {
         &self.ob_base
     }
 
-    fn contents(&self) -> &PyClassObjectContents<T> {
-        &self.contents
+    fn contents(&self) -> &WrappedPyClassObjectContents<T> {
+        (&self.contents).into()
     }
 
-    fn contents_mut(&mut self) -> &mut PyClassObjectContents<T> {
-        &mut self.contents
+    fn contents_mut(&mut self) -> &mut WrappedPyClassObjectContents<T> {
+        (&mut self.contents).into()
     }
 
     /// used to set PyType_Spec::basicsize
@@ -453,9 +481,9 @@ impl<T: PyClassImpl> InternalPyClassObjectLayout<T> for PyStaticClassObject<T> {
 unsafe impl<T: PyClassImpl> PyLayout<T> for PyStaticClassObject<T> {}
 impl<T: PyClass> PySizedLayout<T> for PyStaticClassObject<T> {}
 
-impl<T: PyClassImpl> PyClassObjectLayout<T> for PyStaticClassObject<T>
+impl<T: PyClassImpl> PyClassObjectBaseLayout<T> for PyStaticClassObject<T>
 where
-    <T::BaseType as PyClassBaseType>::LayoutAsBase: PyClassObjectLayout<T::BaseType>,
+    <T::BaseType as PyClassBaseType>::LayoutAsBase: PyClassObjectBaseLayout<T::BaseType>,
 {
     fn ensure_threadsafe(&self) {
         self.contents.thread_checker.ensure();
@@ -470,7 +498,7 @@ where
     unsafe fn tp_dealloc(py: Python<'_>, slf: *mut ffi::PyObject) {
         // Safety: Python only calls tp_dealloc when no references to the object remain.
         let class_object = &mut *(slf.cast::<T::Layout>());
-        class_object.contents_mut().dealloc(py, slf);
+        class_object.contents_mut().0.dealloc(py, slf);
         <T::BaseType as PyClassBaseType>::LayoutAsBase::tp_dealloc(py, slf)
     }
 }
@@ -482,41 +510,41 @@ pub struct PyVariableClassObject<T: PyClassImpl> {
 
 impl<T: PyClassImpl> PyVariableClassObject<T> {
     #[cfg(Py_3_12)]
-    fn get_contents_of_obj(obj: *mut ffi::PyObject) -> *mut PyClassObjectContents<T> {
+    fn get_contents_of_obj(obj: *mut ffi::PyObject) -> *mut WrappedPyClassObjectContents<T> {
         // https://peps.python.org/pep-0697/
         let type_obj = unsafe { ffi::Py_TYPE(obj) };
         let pointer = unsafe { ffi::PyObject_GetTypeData(obj, type_obj) };
-        pointer as *mut PyClassObjectContents<T>
+        pointer as *mut WrappedPyClassObjectContents<T>
     }
 
     #[cfg(Py_3_12)]
-    fn get_contents_ptr(&self) -> *mut PyClassObjectContents<T> {
+    fn get_contents_ptr(&self) -> *mut WrappedPyClassObjectContents<T> {
         Self::get_contents_of_obj(self as *const PyVariableClassObject<T> as *mut ffi::PyObject)
     }
 }
 
 #[cfg(Py_3_12)]
-impl<T: PyClassImpl> InternalPyClassObjectLayout<T> for PyVariableClassObject<T> {
+impl<T: PyClassImpl> PyClassObjectLayout<T> for PyVariableClassObject<T> {
     unsafe fn contents_uninitialised(
         obj: *mut ffi::PyObject,
-    ) -> *mut MaybeUninit<PyClassObjectContents<T>> {
-        Self::get_contents_of_obj(obj) as *mut MaybeUninit<PyClassObjectContents<T>>
+    ) -> *mut MaybeUninit<WrappedPyClassObjectContents<T>> {
+        Self::get_contents_of_obj(obj) as *mut MaybeUninit<WrappedPyClassObjectContents<T>>
     }
 
     fn get_ptr(&self) -> *mut T {
-        self.contents().value.get()
+        self.contents().0.value.get()
     }
 
     fn ob_base(&self) -> &<T::BaseType as PyClassBaseType>::LayoutAsBase {
         &self.ob_base
     }
 
-    fn contents(&self) -> &PyClassObjectContents<T> {
-        unsafe { (self.get_contents_ptr() as *const PyClassObjectContents<T>).as_ref() }
+    fn contents(&self) -> &WrappedPyClassObjectContents<T> {
+        unsafe { self.get_contents_ptr().cast_const().as_ref() }
             .expect("should be able to cast PyClassObjectContents pointer")
     }
 
-    fn contents_mut(&mut self) -> &mut PyClassObjectContents<T> {
+    fn contents_mut(&mut self) -> &mut WrappedPyClassObjectContents<T> {
         unsafe { self.get_contents_ptr().as_mut() }
             .expect("should be able to cast PyClassObjectContents pointer")
     }
@@ -560,24 +588,24 @@ impl<T: PyClassImpl> InternalPyClassObjectLayout<T> for PyVariableClassObject<T>
 unsafe impl<T: PyClassImpl> PyLayout<T> for PyVariableClassObject<T> {}
 
 #[cfg(Py_3_12)]
-impl<T: PyClassImpl> PyClassObjectLayout<T> for PyVariableClassObject<T>
+impl<T: PyClassImpl> PyClassObjectBaseLayout<T> for PyVariableClassObject<T>
 where
-    <T::BaseType as PyClassBaseType>::LayoutAsBase: PyClassObjectLayout<T::BaseType>,
+    <T::BaseType as PyClassBaseType>::LayoutAsBase: PyClassObjectBaseLayout<T::BaseType>,
 {
     fn ensure_threadsafe(&self) {
-        self.contents().thread_checker.ensure();
+        self.contents().0.thread_checker.ensure();
         self.ob_base.ensure_threadsafe();
     }
     fn check_threadsafe(&self) -> Result<(), PyBorrowError> {
-        if !self.contents().thread_checker.check() {
+        if !self.contents().0.thread_checker.check() {
             return Err(PyBorrowError { _private: () });
         }
         self.ob_base.check_threadsafe()
     }
     unsafe fn tp_dealloc(py: Python<'_>, slf: *mut ffi::PyObject) {
         // Safety: Python only calls tp_dealloc when no references to the object remain.
         let class_object = &mut *(slf.cast::<T::Layout>());
-        class_object.contents_mut().dealloc(py, slf);
+        class_object.contents_mut().0.dealloc(py, slf);
         <T::BaseType as PyClassBaseType>::LayoutAsBase::tp_dealloc(py, slf)
     }
 }

src/pyclass/create_type_object.rs

@@ -11,7 +11,7 @@ use crate::{
         trampoline::trampoline,
     },
     internal_tricks::ptr_from_ref,
-    pycell::impl_::InternalPyClassObjectLayout,
+    pycell::impl_::PyClassObjectLayout,
     types::{typeobject::PyTypeMethods, PyType},
     Py, PyClass, PyResult, PyTypeInfo, Python,
 };