|
| 1 | +//! C level type slots. |
| 2 | +//! |
| 3 | +//! The slot table keeps its Rust signatures, so a `newfunc` from an extension |
| 4 | +//! goes into the type's [`CSlots`] table and `new` gets the trampoline that |
| 5 | +//! reads it. The table pointer is inherited alongside the trampoline, so a |
| 6 | +//! subclass reaches the C function without a lookup, the way an inherited |
| 7 | +//! `tp_new` pointer does. |
| 8 | +//! |
| 9 | +//! `__new__` is the same wrapper a native type gets, so a call through it is |
| 10 | +//! checked by `PyType::__new__` before it reaches the slot. |
| 11 | +
|
| 12 | +use crate::object::PyTypeObject; |
| 13 | +use core::ffi::{c_int, c_void}; |
| 14 | +use core::ptr; |
| 15 | +use rustpython_vm::builtins::PyType; |
| 16 | +use rustpython_vm::function::PyMethodFlags; |
| 17 | +use rustpython_vm::types::{CNewFunc, CSlotId, CSlots}; |
| 18 | +use rustpython_vm::{Py, PyResult, VirtualMachine, identifier}; |
| 19 | + |
| 20 | +#[allow(non_camel_case_types)] |
| 21 | +pub type newfunc = CNewFunc; |
| 22 | + |
| 23 | +#[allow(non_upper_case_globals)] |
| 24 | +pub const Py_tp_new: c_int = CSlotId::TpNew as c_int; |
| 25 | + |
| 26 | +/// Install a C `newfunc` as the type's tp_new. |
| 27 | +/// |
| 28 | +/// `ty` must be a heap type that does not already define `__new__` itself. |
| 29 | +pub fn set_tp_new(vm: &VirtualMachine, ty: &Py<PyType>, tp_new: newfunc) -> PyResult<()> { |
| 30 | + let c_slots = CSlots::new(); |
| 31 | + c_slots.new.store(Some(tp_new)); |
| 32 | + ty.set_c_slots(c_slots, vm)?; |
| 33 | + |
| 34 | + // The wrapper that reaches the slot the checked way, as a native type gets |
| 35 | + // from `extend_class`. Stored without the attribute protocol so that |
| 36 | + // `update_slot` does not read it back as a definition of `__new__`. |
| 37 | + let def = vm |
| 38 | + .ctx |
| 39 | + .new_method_def("__new__", PyType::__new__, PyMethodFlags::METHOD, None); |
| 40 | + let wrapper = def.build_function(vm, Some(ty.to_owned().into())); |
| 41 | + ty.set_attr(identifier!(vm, __new__), wrapper.into()); |
| 42 | + Ok(()) |
| 43 | +} |
| 44 | + |
| 45 | +/// Only slots installed from C are reported. A slot backed by a Rust function |
| 46 | +/// has no C ABI entry point yet and reads as empty, as does a slot id this |
| 47 | +/// layer does not handle. |
| 48 | +#[unsafe(no_mangle)] |
| 49 | +#[allow(non_upper_case_globals)] |
| 50 | +pub unsafe extern "C" fn PyType_GetSlot(ty: *const PyTypeObject, slot: c_int) -> *mut c_void { |
| 51 | + let ty = unsafe { &*ty }; |
| 52 | + let Some(c_slots) = ty.slots.c_slots() else { |
| 53 | + return ptr::null_mut(); |
| 54 | + }; |
| 55 | + match CSlotId::from_raw(slot) { |
| 56 | + Some(CSlotId::TpNew) => c_slots |
| 57 | + .new |
| 58 | + .load() |
| 59 | + .map_or(ptr::null_mut(), |f| f as *mut c_void), |
| 60 | + None => ptr::null_mut(), |
| 61 | + } |
| 62 | +} |
| 63 | + |
| 64 | +#[cfg(test)] |
| 65 | +mod tests { |
| 66 | + use super::*; |
| 67 | + use crate::PyObject; |
| 68 | + use pyo3::Python; |
| 69 | + use rustpython_vm::builtins::{PyStrRef, PyTuple, PyTypeRef}; |
| 70 | + use rustpython_vm::function::{FuncArgs, KwArgs}; |
| 71 | + use rustpython_vm::vm::thread::{current_vm_is_set, with_current_vm}; |
| 72 | + use rustpython_vm::{AsObject, PyObjectRef, PyRef}; |
| 73 | + |
| 74 | + /// Returns `(subtype.__name__, args, kwds is NULL)`. |
| 75 | + unsafe extern "C" fn echo_new( |
| 76 | + subtype: *mut PyTypeObject, |
| 77 | + args: *mut PyObject, |
| 78 | + kwds: *mut PyObject, |
| 79 | + ) -> *mut PyObject { |
| 80 | + assert!(current_vm_is_set()); |
| 81 | + with_current_vm(|vm| { |
| 82 | + let subtype = unsafe { &*subtype }; |
| 83 | + let args = unsafe { &*args }.to_owned(); |
| 84 | + let echo = vm.ctx.new_tuple(vec![ |
| 85 | + vm.ctx.new_str(subtype.name().to_string()).into(), |
| 86 | + args, |
| 87 | + vm.ctx.new_bool(kwds.is_null()).into(), |
| 88 | + ]); |
| 89 | + PyObjectRef::from(echo).into_raw().as_ptr() |
| 90 | + }) |
| 91 | + } |
| 92 | + |
| 93 | + /// Returns an empty tuple, so it is told apart from `echo_new` by result. |
| 94 | + unsafe extern "C" fn other_new( |
| 95 | + _subtype: *mut PyTypeObject, |
| 96 | + _args: *mut PyObject, |
| 97 | + _kwds: *mut PyObject, |
| 98 | + ) -> *mut PyObject { |
| 99 | + with_current_vm(|vm| { |
| 100 | + PyObjectRef::from(vm.ctx.new_tuple(vec![])) |
| 101 | + .into_raw() |
| 102 | + .as_ptr() |
| 103 | + }) |
| 104 | + } |
| 105 | + |
| 106 | + fn heap_type(name: &str, base: &Py<PyType>, vm: &VirtualMachine) -> PyTypeRef { |
| 107 | + PyType::new_simple_heap(name, base, &vm.ctx).unwrap() |
| 108 | + } |
| 109 | + |
| 110 | + fn call(ty: &Py<PyType>, args: FuncArgs, vm: &VirtualMachine) -> PyRef<PyTuple> { |
| 111 | + ty.as_object() |
| 112 | + .call(args, vm) |
| 113 | + .unwrap() |
| 114 | + .downcast::<PyTuple>() |
| 115 | + .unwrap() |
| 116 | + } |
| 117 | + |
| 118 | + fn echoed_name(echo: &PyTuple) -> String { |
| 119 | + let name: PyStrRef = echo[0].clone().downcast().unwrap(); |
| 120 | + name.to_string() |
| 121 | + } |
| 122 | + |
| 123 | + #[test] |
| 124 | + fn c_tp_new_is_called() { |
| 125 | + Python::attach(|_py| { |
| 126 | + with_current_vm(|vm| { |
| 127 | + let ty = heap_type("CType", vm.ctx.types.object_type, vm); |
| 128 | + set_tp_new(vm, &ty, echo_new).unwrap(); |
| 129 | + |
| 130 | + let echo = call(&ty, FuncArgs::default(), vm); |
| 131 | + assert_eq!(echoed_name(&echo), "CType"); |
| 132 | + assert_eq!(echo[1].clone().downcast::<PyTuple>().unwrap().len(), 0); |
| 133 | + // No keywords were passed, so kwds must be NULL. |
| 134 | + assert!(echo[2].clone().try_to_bool(vm).unwrap()); |
| 135 | + |
| 136 | + let echo = call(&ty, FuncArgs::from(vec![vm.ctx.new_int(1).into()]), vm); |
| 137 | + assert_eq!(echo[1].clone().downcast::<PyTuple>().unwrap().len(), 1); |
| 138 | + |
| 139 | + let kwargs: KwArgs = |
| 140 | + core::iter::once(("k".to_owned(), vm.ctx.new_int(2).into())).collect(); |
| 141 | + let echo = call(&ty, FuncArgs::new(vec![], kwargs), vm); |
| 142 | + assert!(!echo[2].clone().try_to_bool(vm).unwrap()); |
| 143 | + }) |
| 144 | + }) |
| 145 | + } |
| 146 | + |
| 147 | + /// A subclass reaches the C slot through the inherited slot pair, and the |
| 148 | + /// type it is instantiated with is the one handed to the slot. |
| 149 | + #[test] |
| 150 | + fn c_tp_new_is_inherited() { |
| 151 | + Python::attach(|_py| { |
| 152 | + with_current_vm(|vm| { |
| 153 | + let base = heap_type("CBase", vm.ctx.types.object_type, vm); |
| 154 | + set_tp_new(vm, &base, echo_new).unwrap(); |
| 155 | + let sub = heap_type("CSub", &base, vm); |
| 156 | + |
| 157 | + assert!(sub.slots.c_slots().and_then(|c| c.new.load()).is_some()); |
| 158 | + assert_eq!(echoed_name(&call(&sub, FuncArgs::default(), vm)), "CSub"); |
| 159 | + }) |
| 160 | + }) |
| 161 | + } |
| 162 | + |
| 163 | + /// Every C type reaches its slot through one shared trampoline, so the |
| 164 | + /// "is not safe" check has to compare the C functions behind it, not the |
| 165 | + /// trampoline. Without that, `CBase.__new__(CSub)` would silently run |
| 166 | + /// CSub's tp_new where the caller asked for CBase's. |
| 167 | + #[test] |
| 168 | + fn cross_type_dunder_new_call_is_rejected() { |
| 169 | + Python::attach(|_py| { |
| 170 | + with_current_vm(|vm| { |
| 171 | + let base = heap_type("COuter", vm.ctx.types.object_type, vm); |
| 172 | + set_tp_new(vm, &base, echo_new).unwrap(); |
| 173 | + let sub = heap_type("CInner", &base, vm); |
| 174 | + set_tp_new(vm, &sub, other_new).unwrap(); |
| 175 | + |
| 176 | + // Each type reaches its own C function. |
| 177 | + assert_eq!(echoed_name(&call(&base, FuncArgs::default(), vm)), "COuter"); |
| 178 | + assert!(call(&sub, FuncArgs::default(), vm).is_empty()); |
| 179 | + |
| 180 | + // But COuter.__new__(CInner) must not reach CInner's. |
| 181 | + let dunder_new = base |
| 182 | + .as_object() |
| 183 | + .get_attr(identifier!(vm, __new__), vm) |
| 184 | + .unwrap(); |
| 185 | + let err = dunder_new.call((sub,), vm).unwrap_err(); |
| 186 | + let msg = err.as_object().str(vm).unwrap().to_string(); |
| 187 | + assert!( |
| 188 | + msg.contains("is not safe"), |
| 189 | + "expected an is-not-safe error, got {msg}" |
| 190 | + ); |
| 191 | + }) |
| 192 | + }) |
| 193 | + } |
| 194 | + |
| 195 | + /// `__new__` reaches the slot through `PyType::__new__`, so a direct call |
| 196 | + /// is argument-checked instead of handing the raw pointer to the callee. |
| 197 | + #[test] |
| 198 | + fn direct_dunder_new_call_is_checked() { |
| 199 | + Python::attach(|_py| { |
| 200 | + with_current_vm(|vm| { |
| 201 | + let ty = heap_type("CChecked", vm.ctx.types.object_type, vm); |
| 202 | + set_tp_new(vm, &ty, echo_new).unwrap(); |
| 203 | + let dunder_new = ty |
| 204 | + .as_object() |
| 205 | + .get_attr(identifier!(vm, __new__), vm) |
| 206 | + .unwrap(); |
| 207 | + |
| 208 | + // No type argument at all. |
| 209 | + assert!(dunder_new.call((), vm).is_err()); |
| 210 | + // First argument is not a type. |
| 211 | + assert!(dunder_new.call((vm.ctx.new_int(42),), vm).is_err()); |
| 212 | + // First argument is a type, but not a subtype of this one. |
| 213 | + assert!( |
| 214 | + dunder_new |
| 215 | + .call((vm.ctx.types.dict_type.to_owned(),), vm) |
| 216 | + .is_err() |
| 217 | + ); |
| 218 | + |
| 219 | + // The type itself and a subclass of it are accepted. |
| 220 | + let echo = dunder_new |
| 221 | + .call((ty.clone(),), vm) |
| 222 | + .unwrap() |
| 223 | + .downcast::<PyTuple>() |
| 224 | + .unwrap(); |
| 225 | + assert_eq!(echoed_name(&echo), "CChecked"); |
| 226 | + |
| 227 | + let sub = heap_type("CCheckedSub", &ty, vm); |
| 228 | + let echo = dunder_new |
| 229 | + .call((sub,), vm) |
| 230 | + .unwrap() |
| 231 | + .downcast::<PyTuple>() |
| 232 | + .unwrap(); |
| 233 | + assert_eq!(echoed_name(&echo), "CCheckedSub"); |
| 234 | + }) |
| 235 | + }) |
| 236 | + } |
| 237 | + |
| 238 | + /// A Python-level `__new__` on a subclass replaces the inherited C slot. |
| 239 | + #[test] |
| 240 | + fn python_subclass_new_overrides_the_slot() { |
| 241 | + Python::attach(|_py| { |
| 242 | + with_current_vm(|vm| { |
| 243 | + let base = heap_type("COverBase", vm.ctx.types.object_type, vm); |
| 244 | + set_tp_new(vm, &base, echo_new).unwrap(); |
| 245 | + let sub = heap_type("COverSub", &base, vm); |
| 246 | + |
| 247 | + let py_new = vm.ctx.new_method_def( |
| 248 | + "__new__", |
| 249 | + |_args: FuncArgs, vm: &VirtualMachine| -> PyResult { |
| 250 | + Ok(vm.ctx.new_str("from python").into()) |
| 251 | + }, |
| 252 | + PyMethodFlags::STATIC, |
| 253 | + None, |
| 254 | + ); |
| 255 | + let py_new = py_new.build_function(vm, None); |
| 256 | + sub.as_object() |
| 257 | + .set_attr(identifier!(vm, __new__), py_new, vm) |
| 258 | + .unwrap(); |
| 259 | + |
| 260 | + assert!(sub.slots.c_slots().is_none()); |
| 261 | + let obj = sub.as_object().call((), vm).unwrap(); |
| 262 | + let obj: PyStrRef = obj.downcast().unwrap(); |
| 263 | + assert_eq!(obj.to_string(), "from python"); |
| 264 | + }) |
| 265 | + }) |
| 266 | + } |
| 267 | + |
| 268 | + #[test] |
| 269 | + fn get_slot_round_trips() { |
| 270 | + Python::attach(|_py| { |
| 271 | + with_current_vm(|vm| { |
| 272 | + let ty = heap_type("CGet", vm.ctx.types.object_type, vm); |
| 273 | + assert!(unsafe { PyType_GetSlot(&*ty, Py_tp_new) }.is_null()); |
| 274 | + |
| 275 | + set_tp_new(vm, &ty, echo_new).unwrap(); |
| 276 | + assert_eq!( |
| 277 | + unsafe { PyType_GetSlot(&*ty, Py_tp_new) }, |
| 278 | + echo_new as *mut c_void |
| 279 | + ); |
| 280 | + |
| 281 | + // Inherited by pointer, as tp_new is. |
| 282 | + let sub = heap_type("CGetSub", &ty, vm); |
| 283 | + assert_eq!( |
| 284 | + unsafe { PyType_GetSlot(&*sub, Py_tp_new) }, |
| 285 | + echo_new as *mut c_void |
| 286 | + ); |
| 287 | + }) |
| 288 | + }) |
| 289 | + } |
| 290 | +} |
0 commit comments