From 16abef41bd58051032030176a904586d6ba66135 Mon Sep 17 00:00:00 2001
From: Chayim Refael Friedman <chayimfr@gmail.com>
Date: Mon, 19 Aug 2024 03:02:18 +0300
Subject: [PATCH] Fix a soundness bug with `PyClassInitializer`

From now you cannot initialize a `PyClassInitializer<SubClass>` with `PyClassInitializer::from(Py<BaseClass>).add_subclass(SubClass)`.

This was out of bounds write. Now it panics. See details at https://github.com/PyO3/pyo3/issues/4452.
---
 newsfragments/4454.fixed.md |  1 +
 src/pyclass_init.rs         | 25 +++++++++++++++++++++++++
 tests/add_subclass_to_py.rs | 21 +++++++++++++++++++++
 3 files changed, 47 insertions(+)
 create mode 100644 newsfragments/4454.fixed.md
 create mode 100644 tests/add_subclass_to_py.rs

diff --git a/newsfragments/4454.fixed.md b/newsfragments/4454.fixed.md
new file mode 100644
index 00000000000..e7faf3e690d
--- /dev/null
+++ b/newsfragments/4454.fixed.md
@@ -0,0 +1 @@
+Fix a soundness bug with `PyClassInitializer`: from now you cannot initialize a `PyClassInitializer<SubClass>` with `PyClassInitializer::from(Py<BaseClass>).add_subclass(SubClass)`.
\ No newline at end of file
diff --git a/src/pyclass_init.rs b/src/pyclass_init.rs
index 01983c79b13..228e134d9e9 100644
--- a/src/pyclass_init.rs
+++ b/src/pyclass_init.rs
@@ -27,6 +27,10 @@ pub trait PyObjectInit<T>: Sized {
         py: Python<'_>,
         subtype: *mut PyTypeObject,
     ) -> PyResult<*mut ffi::PyObject>;
+
+    #[doc(hidden)]
+    fn can_be_subclassed(&self) -> bool;
+
     private_decl! {}
 }
 
@@ -81,6 +85,11 @@ impl<T: PyTypeInfo> PyObjectInit<T> for PyNativeTypeInitializer<T> {
         inner(py, type_object, subtype)
     }
 
+    #[inline]
+    fn can_be_subclassed(&self) -> bool {
+        true
+    }
+
     private_impl! {}
 }
 
@@ -147,7 +156,14 @@ impl<T: PyClass> PyClassInitializer<T> {
     /// Constructs a new initializer from value `T` and base class' initializer.
     ///
     /// It is recommended to use `add_subclass` instead of this method for most usage.
+    #[track_caller]
+    #[inline]
     pub fn new(init: T, super_init: <T::BaseType as PyClassBaseType>::Initializer) -> Self {
+        // This is unsound; see https://github.com/PyO3/pyo3/issues/4452.
+        assert!(
+            super_init.can_be_subclassed(),
+            "you cannot add a subclass to an existing value",
+        );
         Self(PyClassInitializerImpl::New { init, super_init })
     }
 
@@ -197,6 +213,8 @@ impl<T: PyClass> PyClassInitializer<T> {
     ///     })
     /// }
     /// ```
+    #[track_caller]
+    #[inline]
     pub fn add_subclass<S>(self, subclass_value: S) -> PyClassInitializer<S>
     where
         S: PyClass<BaseType = T>,
@@ -268,6 +286,11 @@ impl<T: PyClass> PyObjectInit<T> for PyClassInitializer<T> {
             .map(Bound::into_ptr)
     }
 
+    #[inline]
+    fn can_be_subclassed(&self) -> bool {
+        !matches!(self.0, PyClassInitializerImpl::Existing(..))
+    }
+
     private_impl! {}
 }
 
@@ -288,6 +311,8 @@ where
     B: PyClass,
     B::BaseType: PyClassBaseType<Initializer = PyNativeTypeInitializer<B::BaseType>>,
 {
+    #[track_caller]
+    #[inline]
     fn from(sub_and_base: (S, B)) -> PyClassInitializer<S> {
         let (sub, base) = sub_and_base;
         PyClassInitializer::from(base).add_subclass(sub)
diff --git a/tests/add_subclass_to_py.rs b/tests/add_subclass_to_py.rs
new file mode 100644
index 00000000000..5d8ee01cb1f
--- /dev/null
+++ b/tests/add_subclass_to_py.rs
@@ -0,0 +1,21 @@
+//! See https://github.com/PyO3/pyo3/issues/4452.
+#![cfg(feature = "macros")]
+
+use pyo3::prelude::*;
+
+#[pyclass(subclass)]
+struct BaseClass {}
+
+#[pyclass(extends=BaseClass)]
+struct SubClass {
+    _data: i32,
+}
+
+#[test]
+#[should_panic]
+fn add_subclass_to_py_is_unsound() {
+    Python::with_gil(|py| {
+        let base = Py::new(py, BaseClass {}).unwrap();
+        let _subclass = PyClassInitializer::from(base).add_subclass(SubClass { _data: 42 });
+    });
+}