Implement trusted self conversion for all extension-type method and slot wrappers

Author: Copilot

pyo3-macros-backend/src/method.rs

@@ -263,7 +263,7 @@ impl FnType {
         &self,
         cls: Option<&syn::Type>,
         error_mode: ExtractErrorMode,
-        descriptor_slot_receiver: bool,
+        self_conversion: SelfConversionPolicy,
         holders: &mut Holders,
         ctx: &Ctx,
     ) -> Option<TokenStream> {
@@ -273,7 +273,7 @@ impl FnType {
                 Some(st.receiver(
                     cls.expect("no class given for Fn with a \"self\" receiver"),
                     error_mode,
-                    descriptor_slot_receiver,
+                    self_conversion,
                     holders,
                     ctx,
                 ))
@@ -322,6 +322,41 @@ pub enum SelfType {
     },
 }
 
+/// Receiver conversion policy for extension-type method wrappers.
+///
+/// Controls whether the `self` receiver is validated with a runtime type check
+/// (`Checked`) or treated as trusted and cast directly without checking
+/// (`Trusted`).
+///
+/// # Invariant
+///
+/// The `Trusted` path is valid due to CPython's slot/method receiver contract:
+/// when CPython dispatches a method call on an extension type — whether through
+/// a type slot or through `tp_methods` — the receiver is guaranteed to be an
+/// instance of the owning type (or a compatible subtype).  For `tp_methods`
+/// entries, CPython's method-wrapper descriptor enforces this before the C
+/// function is reached.
+///
+/// `Checked` should be used in cases where that guarantee does not hold:
+/// - Standalone `#[pyfunction]`s (no class receiver).
+/// - Number-protocol binary operator fragments (`__add__`, `__radd__`, …,
+///   `__pow__`, `__rpow__`): CPython combines the forward and reflected
+///   fragments into a single `nb_add`/`nb_power` slot, and the runtime helper
+///   may call the reflected fragment with the operands swapped, meaning `_slf`
+///   can arrive with a non-class type.  The existing
+///   `ExtractErrorMode::NotImplemented` behaviour on type mismatch is preserved
+///   by using `Checked` for those fragments.
+#[derive(Clone, Copy, Debug)]
+pub enum SelfConversionPolicy {
+    /// The receiver's type is guaranteed by CPython's slot/method dispatch contract.
+    /// Used for all extension-type method and slot entrypoints.
+    Trusted,
+    /// The receiver's type is verified at runtime. Used for standalone functions
+    /// and number-protocol binary operator fragments where the CPython dispatch
+    /// contract does not guarantee the receiver type.
+    Checked,
+}
+
 #[derive(Clone, Copy)]
 pub enum ExtractErrorMode {
     NotImplemented,
@@ -348,7 +383,7 @@ impl SelfType {
         &self,
         cls: &syn::Type,
         error_mode: ExtractErrorMode,
-        descriptor_slot_receiver: bool,
+        self_conversion: SelfConversionPolicy,
         holders: &mut Holders,
         ctx: &Ctx,
     ) -> TokenStream {
@@ -370,22 +405,47 @@ impl SelfType {
                 };
                 let arg =
                     quote! { unsafe { #pyo3_path::impl_::extract_argument::#cast_fn(#py, #slf) } };
-                let method = if *mutable {
-                    syn::Ident::new("extract_pyclass_ref_mut", *span)
-                } else {
-                    syn::Ident::new("extract_pyclass_ref", *span)
-                };
                 let holder = holders.push_holder(*span);
                 let pyo3_path = pyo3_path.to_tokens_spanned(*span);
-                error_mode.handle_error(
-                    quote_spanned! { *span =>
-                        #pyo3_path::impl_::extract_argument::#method::<#cls>(
-                            #arg,
-                            &mut #holder,
+                match self_conversion {
+                    SelfConversionPolicy::Trusted => {
+                        let method = if *mutable {
+                            syn::Ident::new("extract_pyclass_ref_mut_trusted", *span)
+                        } else {
+                            syn::Ident::new("extract_pyclass_ref_trusted", *span)
+                        };
+                        // Use `quote!` (not `quote_spanned!`) for the `unsafe` block so that
+                        // the `unsafe` keyword has `Span::call_site()` and does not inherit the
+                        // user's code span.  This prevents triggering `#![forbid(unsafe_code)]`
+                        // in user crates (see the analogous comment in `impl_py_getter_def`).
+                        // Safety: slot wrappers are only installed on the extension type itself.
+                        // CPython's slot dispatch contract ensures the receiver is an instance
+                        // of the correct type before invoking the slot.
+                        let trusted_call = quote! {
+                            unsafe { #pyo3_path::impl_::extract_argument::#method::<#cls>(
+                                #arg,
+                                &mut #holder,
+                            ) }
+                        };
+                        error_mode.handle_error(trusted_call, ctx)
+                    }
+                    SelfConversionPolicy::Checked => {
+                        let method = if *mutable {
+                            syn::Ident::new("extract_pyclass_ref_mut", *span)
+                        } else {
+                            syn::Ident::new("extract_pyclass_ref", *span)
+                        };
+                        error_mode.handle_error(
+                            quote_spanned! { *span =>
+                                #pyo3_path::impl_::extract_argument::#method::<#cls>(
+                                    #arg,
+                                    &mut #holder,
+                                )
+                            },
+                            ctx,
                         )
-                    },
-                    ctx,
-                )
+                    }
+                }
             }
             SelfType::TryFromBoundRef { span, non_null } => {
                 let bound_ref = if *non_null {
@@ -394,22 +454,27 @@ impl SelfType {
                     quote! { unsafe { #pyo3_path::Bound::ref_from_ptr(#py, &#slf) } }
                 };
                 let pyo3_path = pyo3_path.to_tokens_spanned(*span);
-                let receiver = if descriptor_slot_receiver {
-                    quote_spanned! { *span =>
-                        // Safety: descriptor slot wrappers are only installed on the descriptor
-                        // type itself. CPython calls those slots with `self` set to the
-                        // descriptor object found during lookup, and explicit Python calls to
-                        // `__get__`, `__set__`, and `__delete__` first pass through CPython's
-                        // slot wrapper, which rejects receivers of the wrong type before
-                        // reaching this generated wrapper.
-                        ::std::result::Result::<_, #pyo3_path::PyErr>::Ok(unsafe {
-                            #bound_ref.cast_unchecked::<#cls>()
-                        })
+                let receiver = match self_conversion {
+                    SelfConversionPolicy::Trusted => {
+                        // Use `quote!` (not `quote_spanned!`) for the inner `unsafe` block so
+                        // that it has `Span::call_site()` and does not trigger
+                        // `#![forbid(unsafe_code)]` in user crates.
+                        // Safety: slot wrappers are only installed on the extension type
+                        // itself. CPython's slot dispatch contract ensures the receiver is
+                        // an instance of the correct type (or a compatible subtype) before
+                        // invoking the slot.
+                        let cast = quote! {
+                            unsafe { #bound_ref.cast_unchecked::<#cls>() }
+                        };
+                        quote_spanned! { *span =>
+                            ::std::result::Result::<_, #pyo3_path::PyErr>::Ok(#cast)
+                        }
                     }
-                } else {
-                    quote_spanned! { *span =>
-                        #bound_ref.cast::<#cls>()
-                            .map_err(::std::convert::Into::<#pyo3_path::PyErr>::into)
+                    SelfConversionPolicy::Checked => {
+                        quote_spanned! { *span =>
+                            #bound_ref.cast::<#cls>()
+                                .map_err(::std::convert::Into::<#pyo3_path::PyErr>::into)
+                        }
                     }
                 };
                 error_mode.handle_error(
@@ -697,6 +762,7 @@ impl<'a> FnSpec<'a> {
         ident: &proc_macro2::Ident,
         cls: Option<&syn::Type>,
         convention: CallingConvention,
+        self_conversion: SelfConversionPolicy,
         ctx: &Ctx,
     ) -> Result<TokenStream> {
         let Ctx {
@@ -719,9 +785,13 @@ impl<'a> FnSpec<'a> {
         }
 
         let rust_call = |args: Vec<TokenStream>, mut holders: Holders| {
-            let self_arg = self
-                .tp
-                .self_arg(cls, ExtractErrorMode::Raise, false, &mut holders, ctx);
+            let self_arg = self.tp.self_arg(
+                cls,
+                ExtractErrorMode::Raise,
+                self_conversion,
+                &mut holders,
+                ctx,
+            );
             let init_holders = holders.init_holders(ctx);
 
             // We must assign the output_span to the return value of the call,

pyo3-macros-backend/src/pyfunction.rs

@@ -10,7 +10,7 @@ use crate::{
         self, get_pyo3_options, take_attributes, take_pyo3_options, CrateAttribute,
         FromPyWithAttribute, NameAttribute, TextSignatureAttribute,
     },
-    method::{self, CallingConvention, FnArg},
+    method::{self, CallingConvention, FnArg, SelfConversionPolicy},
     pymethod::check_generic,
 };
 use proc_macro2::{Span, TokenStream};
@@ -430,7 +430,13 @@ pub fn impl_wrap_pyfunction(
         );
     }
     let calling_convention = CallingConvention::from_signature(&spec.signature);
-    let wrapper = spec.get_wrapper_function(&wrapper_ident, None, calling_convention, ctx)?;
+    let wrapper = spec.get_wrapper_function(
+        &wrapper_ident,
+        None,
+        calling_convention,
+        SelfConversionPolicy::Checked,
+        ctx,
+    )?;
     let methoddef = spec.get_methoddef(
         wrapper_ident,
         spec.get_doc(&func.attrs).as_ref(),