Stop leaking in new_closure

Author: davidhewitt

src/impl_/pymethods.rs

@@ -1,7 +1,9 @@
-use crate::internal_tricks::{extract_cstr_or_leak_cstring, NulByteInString};
+use crate::exceptions::PyValueError;
 use crate::{ffi, IntoPy, Py, PyAny, PyErr, PyObject, PyResult, PyTraverseError, Python};
-use std::ffi::CStr;
+use std::borrow::Cow;
+use std::ffi::{CStr, CString};
 use std::fmt;
+use std::mem::ManuallyDrop;
 use std::os::raw::c_int;
 
 /// Python 3.8 and up - __ipow__ has modulo argument correctly populated.
@@ -95,6 +97,12 @@ pub struct PyClassAttributeDef {
     pub(crate) meth: PyClassAttributeFactory,
 }
 
+impl PyClassAttributeDef {
+    pub(crate) fn attribute_c_string(&self) -> PyResult<Cow<'static, CStr>> {
+        extract_c_string(self.name, "class attribute name cannot contain nul bytes")
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct PyGetterDef {
     pub(crate) name: &'static str,
@@ -161,7 +169,9 @@ impl PyMethodDef {
     }
 
     /// Convert `PyMethodDef` to Python method definition struct `ffi::PyMethodDef`
-    pub(crate) fn as_method_def(&self) -> Result<ffi::PyMethodDef, NulByteInString> {
+    pub(crate) fn as_method_def(
+        &self,
+    ) -> PyResult<(ffi::PyMethodDef, ManuallyDrop<PyMethodDefDestructor>)> {
         let meth = match self.ml_meth {
             PyMethodType::PyCFunction(meth) => ffi::PyMethodDefPointer {
                 PyCFunction: meth.0,
@@ -175,12 +185,19 @@ impl PyMethodDef {
             },
         };
 
-        Ok(ffi::PyMethodDef {
-            ml_name: get_name(self.ml_name)?.as_ptr(),
+        let name = get_name(self.ml_name)?;
+        let doc = get_doc(self.ml_doc)?;
+        let def = ffi::PyMethodDef {
+            ml_name: name.as_ptr(),
             ml_meth: meth,
             ml_flags: self.ml_flags,
-            ml_doc: get_doc(self.ml_doc)?.as_ptr(),
-        })
+            ml_doc: doc.as_ptr(),
+        };
+        let destructor = ManuallyDrop::new(PyMethodDefDestructor {
+            name: Some(name),
+            doc: Some(doc),
+        });
+        Ok((def, destructor))
     }
 }
 
@@ -214,10 +231,16 @@ impl PyGetterDef {
     /// Copy descriptor information to `ffi::PyGetSetDef`
     pub fn copy_to(&self, dst: &mut ffi::PyGetSetDef) {
         if dst.name.is_null() {
-            dst.name = get_name(self.name).unwrap().as_ptr() as _;
+            let name = get_name(self.name).unwrap();
+            dst.name = name.as_ptr() as _;
+            // FIXME: stop leaking name
+            std::mem::forget(name);
         }
         if dst.doc.is_null() {
-            dst.doc = get_doc(self.doc).unwrap().as_ptr() as _;
+            let doc = get_doc(self.doc).unwrap();
+            dst.doc = doc.as_ptr() as _;
+            // FIXME: stop leaking doc
+            std::mem::forget(doc);
         }
         dst.get = Some(self.meth.0);
     }
@@ -236,21 +259,27 @@ impl PySetterDef {
     /// Copy descriptor information to `ffi::PyGetSetDef`
     pub fn copy_to(&self, dst: &mut ffi::PyGetSetDef) {
         if dst.name.is_null() {
-            dst.name = get_name(self.name).unwrap().as_ptr() as _;
+            let name = get_name(self.name).unwrap();
+            dst.name = name.as_ptr() as _;
+            // FIXME: stop leaking name
+            std::mem::forget(name);
         }
         if dst.doc.is_null() {
-            dst.doc = get_doc(self.doc).unwrap().as_ptr() as _;
+            let doc = get_doc(self.doc).unwrap();
+            dst.doc = doc.as_ptr() as _;
+            // FIXME: stop leaking doc
+            std::mem::forget(doc);
         }
         dst.set = Some(self.meth.0);
     }
 }
 
-fn get_name(name: &'static str) -> Result<&'static CStr, NulByteInString> {
-    extract_cstr_or_leak_cstring(name, "Function name cannot contain NUL byte.")
+fn get_name(name: &'static str) -> PyResult<Cow<'static, CStr>> {
+    extract_c_string(name, "Function name cannot contain NUL byte.")
 }
 
-fn get_doc(doc: &'static str) -> Result<&'static CStr, NulByteInString> {
-    extract_cstr_or_leak_cstring(doc, "Document cannot contain NUL byte.")
+fn get_doc(doc: &'static str) -> PyResult<Cow<'static, CStr>> {
+    extract_c_string(doc, "Document cannot contain NUL byte.")
 }
 
 /// Unwraps the result of __traverse__ for tp_traverse
@@ -263,6 +292,14 @@ pub fn unwrap_traverse_result(result: Result<(), PyTraverseError>) -> c_int {
     }
 }
 
+pub(crate) struct PyMethodDefDestructor {
+    // These members are just to avoid leaking CStrings when possible
+    #[allow(dead_code)]
+    name: Option<Cow<'static, CStr>>,
+    #[allow(dead_code)]
+    doc: Option<Cow<'static, CStr>>,
+}
+
 // The macros need to Ok-wrap the output of user defined functions; i.e. if they're not a result, make them into one.
 pub trait OkWrap<T> {
     type Error;
@@ -288,3 +325,16 @@ where
         self.map(|o| o.into_py(py))
     }
 }
+
+fn extract_c_string(src: &'static str, err_msg: &'static str) -> PyResult<Cow<'static, CStr>> {
+    let bytes = src.as_bytes();
+    if bytes.last() == Some(&0) {
+        // Last byte is a nul; try to create as a CStr
+        let c_str = CStr::from_bytes_with_nul(bytes).map_err(|_| PyValueError::new_err(err_msg))?;
+        Ok(Cow::Borrowed(c_str))
+    } else {
+        // Allocate a new CString for this
+        let c_string = CString::new(bytes).map_err(|_| PyValueError::new_err(err_msg))?;
+        Ok(Cow::Owned(c_string))
+    }
+}

src/internal_tricks.rs

@@ -1,6 +1,4 @@
 use crate::ffi::{Py_ssize_t, PY_SSIZE_T_MAX};
-use std::ffi::{CStr, CString};
-
 pub struct PrivateMarker;
 
 macro_rules! private_decl {
@@ -32,20 +30,6 @@ macro_rules! pyo3_exception {
     };
 }
 
-#[derive(Debug)]
-pub(crate) struct NulByteInString(pub(crate) &'static str);
-
-pub(crate) fn extract_cstr_or_leak_cstring(
-    src: &'static str,
-    err_msg: &'static str,
-) -> Result<&'static CStr, NulByteInString> {
-    CStr::from_bytes_with_nul(src.as_bytes())
-        .or_else(|_| {
-            CString::new(src.as_bytes()).map(|c_string| &*Box::leak(c_string.into_boxed_c_str()))
-        })
-        .map_err(|_| NulByteInString(err_msg))
-}
-
 /// Convert an usize index into a Py_ssize_t index, clamping overflow to
 /// PY_SSIZE_T_MAX.
 pub(crate) fn get_ssize_index(index: usize) -> Py_ssize_t {

src/pyclass.rs

@@ -158,7 +158,12 @@ impl PyTypeBuilder {
             }
             PyMethodDefType::Method(def)
             | PyMethodDefType::Class(def)
-            | PyMethodDefType::Static(def) => self.method_defs.push(def.as_method_def().unwrap()),
+            | PyMethodDefType::Static(def) => {
+                let (def, destructor) = def.as_method_def().unwrap();
+                // FIXME: stop leaking destructor
+                std::mem::forget(destructor);
+                self.method_defs.push(def);
+            }
             // These class attributes are added after the type gets created by LazyStaticType
             PyMethodDefType::ClassAttribute(_) => {}
         }

src/type_object.rs

@@ -2,14 +2,15 @@
 //! Python type object information
 
 use crate::impl_::pyclass::PyClassItemsIter;
-use crate::internal_tricks::extract_cstr_or_leak_cstring;
 use crate::once_cell::GILOnceCell;
 use crate::pyclass::create_type_object;
 use crate::pyclass::PyClass;
 use crate::types::{PyAny, PyType};
 use crate::{conversion::IntoPyPointer, PyMethodDefType};
 use crate::{ffi, AsPyPointer, PyNativeType, PyObject, PyResult, Python};
 use parking_lot::{const_mutex, Mutex};
+use std::borrow::Cow;
+use std::ffi::CStr;
 use std::thread::{self, ThreadId};
 
 /// `T: PyLayout<U>` represents that `T` is a concrete representation of `U` in the Python heap.
@@ -180,11 +181,7 @@ impl LazyStaticType {
         for class_items in items_iter {
             for def in class_items.methods {
                 if let PyMethodDefType::ClassAttribute(attr) = def {
-                    let key = extract_cstr_or_leak_cstring(
-                        attr.name,
-                        "class attribute name cannot contain nul bytes",
-                    )
-                    .unwrap();
+                    let key = attr.attribute_c_string().unwrap();
 
                     match (attr.meth.0)(py) {
                         Ok(val) => items.push((key, val)),
@@ -221,7 +218,7 @@ impl LazyStaticType {
 fn initialize_tp_dict(
     py: Python<'_>,
     type_object: *mut ffi::PyObject,
-    items: Vec<(&'static std::ffi::CStr, PyObject)>,
+    items: Vec<(Cow<'static, CStr>, PyObject)>,
 ) -> PyResult<()> {
     // We hold the GIL: the dictionary update can be considered atomic from
     // the POV of other threads.

src/types/function.rs

@@ -1,13 +1,14 @@
 use crate::derive_utils::PyFunctionArguments;
-use crate::exceptions::PyValueError;
 use crate::impl_::panic::PanicTrap;
+use crate::methods::PyMethodDefDestructor;
 use crate::panic::PanicException;
 use crate::{
     ffi,
     impl_::pymethods::{self, PyMethodDef},
     types, AsPyPointer,
 };
 use crate::{prelude::*, GILPool};
+use std::mem::ManuallyDrop;
 use std::os::raw::c_void;
 
 /// Represents a builtin Python function object.
@@ -32,17 +33,25 @@ where
         args,
         kwargs,
         |py, capsule_ptr, args, kwargs| {
-            let boxed_fn: &F = &*(ffi::PyCapsule_GetPointer(
+            let boxed_fn: &ClosureDestructor<F> = &*(ffi::PyCapsule_GetPointer(
                 capsule_ptr,
                 CLOSURE_CAPSULE_NAME.as_ptr() as *const _,
-            ) as *mut F);
+            ) as *mut ClosureDestructor<F>);
             let args = py.from_borrowed_ptr::<types::PyTuple>(args);
             let kwargs = py.from_borrowed_ptr_or_opt::<types::PyDict>(kwargs);
-            crate::callback::convert(py, boxed_fn(args, kwargs))
+            crate::callback::convert(py, (boxed_fn.closure)(args, kwargs))
         },
     )
 }
 
+struct ClosureDestructor<F> {
+    closure: F,
+    def: ffi::PyMethodDef,
+    // Used to destroy the cstrings in `def`, if necessary.
+    #[allow(dead_code)]
+    def_destructor: PyMethodDefDestructor,
+}
+
 unsafe extern "C" fn drop_closure<F, R>(capsule_ptr: *mut ffi::PyObject)
 where
     F: Fn(&types::PyTuple, Option<&types::PyDict>) -> R + Send + 'static,
@@ -51,11 +60,12 @@ where
     let trap = PanicTrap::new("uncaught panic during drop_closure");
     let pool = GILPool::new();
     if let Err(payload) = std::panic::catch_unwind(|| {
-        let boxed_fn: Box<F> = Box::from_raw(ffi::PyCapsule_GetPointer(
+        let destructor: Box<ClosureDestructor<F>> = Box::from_raw(ffi::PyCapsule_GetPointer(
             capsule_ptr,
             CLOSURE_CAPSULE_NAME.as_ptr() as *const _,
-        ) as *mut F);
-        drop(boxed_fn)
+        )
+            as *mut ClosureDestructor<F>);
+        drop(destructor)
     }) {
         let py = pool.python();
         let err = PanicException::from_panic_payload(payload);
@@ -116,46 +126,44 @@ impl PyCFunction {
         py: Python<'a>,
         name: Option<&'static str>,
         doc: Option<&'static str>,
-        f: F,
+        closure: F,
     ) -> PyResult<&'a PyCFunction>
     where
         F: Fn(&types::PyTuple, Option<&types::PyDict>) -> R + Send + 'static,
         R: crate::callback::IntoPyCallbackOutput<*mut ffi::PyObject>,
     {
-        let function_ptr = Box::into_raw(Box::new(f));
-        let capsule = unsafe {
+        let method_def = pymethods::PyMethodDef::cfunction_with_keywords(
+            name.unwrap_or("pyo3-closure\0"),
+            pymethods::PyCFunctionWithKeywords(run_closure::<F, R>),
+            doc.unwrap_or("\0"),
+        );
+        let (def, def_destructor) = method_def.as_method_def()?;
+        let ptr = Box::into_raw(Box::new(ClosureDestructor {
+            closure,
+            def,
+            // Disable the `ManuallyDrop`; we do actually want to drop this later.
+            def_destructor: ManuallyDrop::into_inner(def_destructor),
+        }));
+
+        let destructor = unsafe {
             PyObject::from_owned_ptr_or_err(
                 py,
                 ffi::PyCapsule_New(
-                    function_ptr as *mut c_void,
+                    ptr as *mut c_void,
                     CLOSURE_CAPSULE_NAME.as_ptr() as *const _,
                     Some(drop_closure::<F, R>),
                 ),
             )?
         };
-        let method_def = pymethods::PyMethodDef::cfunction_with_keywords(
-            name.unwrap_or("pyo3-closure\0"),
-            pymethods::PyCFunctionWithKeywords(run_closure::<F, R>),
-            doc.unwrap_or("\0"),
-        );
-        Self::internal_new_from_pointers(&method_def, py, capsule.as_ptr(), std::ptr::null_mut())
-    }
 
-    #[doc(hidden)]
-    fn internal_new_from_pointers<'py>(
-        method_def: &PyMethodDef,
-        py: Python<'py>,
-        mod_ptr: *mut ffi::PyObject,
-        module_name: *mut ffi::PyObject,
-    ) -> PyResult<&'py Self> {
-        let def = method_def
-            .as_method_def()
-            .map_err(|err| PyValueError::new_err(err.0))?;
         unsafe {
             py.from_owned_ptr_or_err::<PyCFunction>(ffi::PyCFunction_NewEx(
-                Box::into_raw(Box::new(def)),
-                mod_ptr,
-                module_name,
+                #[cfg(addr_of)]
+                core::ptr::addr_of_mut!((*ptr).def),
+                #[cfg(not(addr_of))]
+                &mut (*ptr).def,
+                destructor.as_ptr(),
+                std::ptr::null_mut(),
             ))
         }
     }
@@ -173,7 +181,19 @@ impl PyCFunction {
         } else {
             (std::ptr::null_mut(), std::ptr::null_mut())
         };
-        Self::internal_new_from_pointers(method_def, py, mod_ptr, module_name)
+        let (def, destructor) = method_def.as_method_def()?;
+
+        // FIXME: stop leaking the def and destructor
+        let def = Box::into_raw(Box::new(def));
+        std::mem::forget(destructor);
+
+        unsafe {
+            py.from_owned_ptr_or_err::<PyCFunction>(ffi::PyCFunction_NewEx(
+                def,
+                mod_ptr,
+                module_name,
+            ))
+        }
     }
 }