fix: guard __class__ reassignment layout

Co-authored-by: youknowone <69878+youknowone@users.noreply.github.com>
RustPython · a3fb6fed740a862e216e674b52932663f8dde1f6
commit a3fb6fed740a862e216e674b52932663f8dde1f6
diff --git a/crates/vm/src/builtins/object.rs b/crates/vm/src/builtins/object.rs
@@ -497,16 +497,29 @@ impl PyBaseObject {
     ) -> PyResult<()> {
         match value.downcast::<PyType>() {
             Ok(cls) => {
-                let both_module = instance.class().fast_issubclass(vm.ctx.types.module_type)
+                let current_cls = instance.class();
+                let both_module = current_cls.fast_issubclass(vm.ctx.types.module_type)
                     && cls.fast_issubclass(vm.ctx.types.module_type);
-                let both_mutable = !instance
-                    .class()
+                let both_mutable = !current_cls
                     .slots
                     .flags
                     .has_feature(PyTypeFlags::IMMUTABLETYPE)
                     && !cls.slots.flags.has_feature(PyTypeFlags::IMMUTABLETYPE);
                 // FIXME(#1979) cls instances might have a payload
                 if both_mutable || both_module {
+                    let has_dict = |typ: &Py<PyType>| typ.slots.flags.has_feature(PyTypeFlags::HAS_DICT);
+                    if current_cls.slots.basicsize != cls.slots.basicsize
+                        || current_cls.slots.member_count != cls.slots.member_count
+                        || has_dict(current_cls) != has_dict(&cls)
+                        || (cls.slots.flags.has_feature(PyTypeFlags::HEAPTYPE)
+                            && instance.typeid() != PyBaseObject::payload_type_id())
+                    {
+                        return Err(vm.new_type_error(format!(
+                            "__class__ assignment: '{}' object layout differs from '{}'",
+                            cls.name(),
+                            current_cls.name()
+                        )));
+                    }
                     instance.set_class(cls, vm);
                     Ok(())
                 } else {

diff --git a/crates/vm/src/vm/interpreter.rs b/crates/vm/src/vm/interpreter.rs
@@ -169,4 +169,31 @@ mod tests {
             assert_eq!(value.as_str(), "Hello Hello Hello Hello ")
         })
     }
+
+    #[test]
+    fn class_assignment_layout_mismatch_error() {
+        Interpreter::without_stdlib(Default::default()).enter(|vm| {
+            let scope = vm.new_scope_with_builtins();
+            let source = r#"
+class TypeA:
+    def __init__(self):
+        self.a = 1
+
+class TypeB:
+    __slots__ = "b"
+    def __init__(self):
+        self.b = 2
+
+obj = TypeA()
+try:
+    obj.__class__ = TypeB
+except TypeError as e:
+    assert str(e) == "__class__ assignment: 'TypeB' object layout differs from 'TypeA'"
+else:
+    raise AssertionError("TypeError not raised")
+"#;
+            vm.run_code_string(scope, source, "<test>".to_owned())
+                .expect("script should complete without uncaught exceptions");
+        })
+    }
 }