Stop leaking in new_closure

Cargo.toml

@@ -18,6 +18,7 @@ cfg-if = "1.0"
 libc = "0.2.62"
 parking_lot = ">= 0.11, < 0.13"
 memoffset = "0.6.5"
+memchr = "1"
 
 # ffi bindings to the python interpreter, split into a separate crate so they can be used independently
 pyo3-ffi = { path = "pyo3-ffi", version = "=0.17.2" }

src/impl_/pymethods.rs

@@ -1,7 +1,8 @@
-use crate::internal_tricks::{extract_cstr_or_leak_cstring, NulByteInString};
+use crate::internal_tricks::{extract_cstr_or_leak_cstring, MaybeLeaked, NulByteInString};
 use crate::{ffi, IntoPy, Py, PyAny, PyErr, PyObject, PyResult, PyTraverseError, Python};
-use std::ffi::CStr;
+use std::ffi::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.
@@ -161,7 +162,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,
+    ) -> Result<(ffi::PyMethodDef, ManuallyDrop<PyMethodDefDestructor>), NulByteInString> {
         let meth = match self.ml_meth {
             PyMethodType::PyCFunction(meth) => ffi::PyMethodDefPointer {
                 PyCFunction: meth.0,
@@ -175,12 +178,28 @@ impl PyMethodDef {
             },
         };
 
-        Ok(ffi::PyMethodDef {
-            ml_name: get_name(self.ml_name)?.as_ptr(),
+        let name = get_name(self.ml_name)?;
+        let doc = match get_doc(self.ml_doc) {
+            Ok(d) => d,
+            Err(e) => {
+                // Free `name` if it got leaked and doc failed.
+                if let MaybeLeaked::Leaked(ptr) = name {
+                    let _ = unsafe { CString::from_raw(ptr) };
+                }
+                return Err(e);
+            }
+        };
+        let destructor = ManuallyDrop::new(PyMethodDefDestructor {
+            name: Some(name),
+            doc: Some(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(),
+        };
+        Ok((def, destructor))
     }
 }
 
@@ -245,11 +264,11 @@ impl PySetterDef {
     }
 }
 
-fn get_name(name: &'static str) -> Result<&'static CStr, NulByteInString> {
+fn get_name(name: &'static str) -> Result<MaybeLeaked, NulByteInString> {
     extract_cstr_or_leak_cstring(name, "Function name cannot contain NUL byte.")
 }
 
-fn get_doc(doc: &'static str) -> Result<&'static CStr, NulByteInString> {
+fn get_doc(doc: &'static str) -> Result<MaybeLeaked, NulByteInString> {
     extract_cstr_or_leak_cstring(doc, "Document cannot contain NUL byte.")
 }
 
@@ -263,6 +282,22 @@ pub fn unwrap_traverse_result(result: Result<(), PyTraverseError>) -> c_int {
     }
 }
 
+pub(crate) struct PyMethodDefDestructor {
+    name: Option<MaybeLeaked>,
+    doc: Option<MaybeLeaked>,
+}
+
+impl Drop for PyMethodDefDestructor {
+    fn drop(&mut self) {
+        if let Some(MaybeLeaked::Leaked(ptr)) = self.name.take() {
+            let _ = unsafe { CString::from_raw(ptr) };
+        }
+        if let Some(MaybeLeaked::Leaked(ptr)) = self.doc.take() {
+            let _ = unsafe { CString::from_raw(ptr) };
+        }
+    }
+}
+
 // 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;

src/internal_tricks.rs

@@ -1,6 +1,6 @@
 use crate::ffi::{Py_ssize_t, PY_SSIZE_T_MAX};
 use std::ffi::{CStr, CString};
-
+use std::os::raw::c_char;
 pub struct PrivateMarker;
 
 macro_rules! private_decl {
@@ -35,15 +35,45 @@ macro_rules! pyo3_exception {
 #[derive(Debug)]
 pub(crate) struct NulByteInString(pub(crate) &'static str);
 
+#[derive(Copy, Clone)]
+pub(crate) enum MaybeLeaked {
+    Static(*const c_char),
+    Leaked(*mut c_char),
+}
+
+impl MaybeLeaked {
+    pub(crate) fn as_ptr(&self) -> *const c_char {
+        match *self {
+            Self::Static(ptr) => ptr,
+            Self::Leaked(ptr) => ptr as *const c_char,
+        }
+    }
+}
+
 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))
+) -> Result<MaybeLeaked, NulByteInString> {
+    let bytes = src.as_bytes();
+    let nul_pos = memchr::memchr(0, bytes);
+    match nul_pos {
+        Some(nul_pos) if nul_pos + 1 == bytes.len() => {
+            // SAFETY: We know there is only one nul byte, at the end
+            // of the byte slice.
+            let ptr = unsafe { CStr::from_bytes_with_nul_unchecked(bytes).as_ptr() };
+            Ok(MaybeLeaked::Static(ptr))
+        }
+        Some(_) => Err(NulByteInString(err_msg)),
+        None => {
+            let mut v = Vec::with_capacity(bytes.len() + 1);
+            v.extend_from_slice(bytes);
+
+            // SAFETY: There is no null byte, and `from_vec_unchecked` will append one.
+            let cstring = unsafe { CString::from_vec_unchecked(v) };
+            let ptr = cstring.into_raw();
+            Ok(MaybeLeaked::Leaked(ptr))
+        }
+    }
 }
 
 /// Convert an usize index into a Py_ssize_t index, clamping overflow to

src/pyclass.rs

@@ -158,7 +158,10 @@ 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();
+                self.method_defs.push(def);
+            }
             // These class attributes are added after the type gets created by LazyStaticType
             PyMethodDefType::ClassAttribute(_) => {}
         }

src/type_object.rs

@@ -2,7 +2,7 @@
 //! Python type object information
 
 use crate::impl_::pyclass::PyClassItemsIter;
-use crate::internal_tricks::extract_cstr_or_leak_cstring;
+use crate::internal_tricks::{extract_cstr_or_leak_cstring, MaybeLeaked};
 use crate::once_cell::GILOnceCell;
 use crate::pyclass::create_type_object;
 use crate::pyclass::PyClass;
@@ -221,7 +221,7 @@ impl LazyStaticType {
 fn initialize_tp_dict(
     py: Python<'_>,
     type_object: *mut ffi::PyObject,
-    items: Vec<(&'static std::ffi::CStr, PyObject)>,
+    items: Vec<(MaybeLeaked, 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,15 @@
 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.
@@ -28,15 +30,23 @@ where
     R: crate::callback::IntoPyCallbackOutput<*mut ffi::PyObject>,
 {
     crate::callback_body!(py, {
-        let boxed_fn: &F =
+        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);
-        boxed_fn(args, kwargs)
+        (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,
@@ -45,11 +55,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);
@@ -106,45 +117,45 @@ impl PyCFunction {
     ///     py_run!(py, add_one, "assert add_one(42) == 43");
     /// });
     /// ```
-    pub fn new_closure<F, R>(f: F, py: Python<'_>) -> PyResult<&PyCFunction>
+    pub fn new_closure<F, R>(closure: F, py: Python<'_>) -> PyResult<&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(
+            "pyo3-closure\0",
+            pymethods::PyCFunctionWithKeywords(run_closure::<F, R>),
+            "\0",
+        );
+        let (def, def_destructor) = method_def
+            .as_method_def()
+            .map_err(|err| PyValueError::new_err(err.0))?;
+        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(
-            "pyo3-closure",
-            pymethods::PyCFunctionWithKeywords(run_closure::<F, R>),
-            "",
-        );
-        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(),
             ))
         }
     }
@@ -162,7 +173,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()
+            .map_err(|err| PyValueError::new_err(err.0))?;
+
+        let def = Box::into_raw(Box::new(def));
+
+        unsafe {
+            py.from_owned_ptr_or_err::<PyCFunction>(ffi::PyCFunction_NewEx(
+                def,
+                mod_ptr,
+                module_name,
+            ))
+        }
     }
 }