[mypyc] Fix crash and reference leaks on base resolution failures

mypyc/lib-rt/misc_ops.c

@@ -230,6 +230,8 @@ PyObject *CPyType_FromTemplate(PyObject *template,
 
     if (orig_bases) {
         bases = update_bases(orig_bases);
+        if (bases == NULL)
+            goto error;
         // update_bases doesn't increment the refcount if nothing changes,
         // so we do it to make sure we have distinct "references" to both
         if (bases == orig_bases)
@@ -252,17 +254,6 @@ PyObject *CPyType_FromTemplate(PyObject *template,
     if (!name)
         goto error;
 
-    if (template_->tp_doc) {
-        // cpython expects tp_doc to be heap-allocated so convert it here to
-        // avoid segfaults on deallocation.
-        Py_ssize_t size = strlen(template_->tp_doc) + 1;
-        char *doc = (char *)PyMem_Malloc(size);
-        if (!doc)
-            goto error;
-        memcpy(doc, template_->tp_doc, size);
-        template_->tp_doc = doc;
-    }
-
     // Allocate the type and then copy the main stuff in.
     t = (PyHeapTypeObject*)PyType_GenericAlloc(&PyType_Type, 0);
     if (!t)
@@ -271,6 +262,21 @@ PyObject *CPyType_FromTemplate(PyObject *template,
            (char *)template_ + sizeof(PyVarObject),
            sizeof(PyTypeObject) - sizeof(PyVarObject));
 
+    t->ht_type.tp_doc = NULL;
+    // Unlike the tp_doc slots of most other object, a heap type's tp_doc
+    // must be heap allocated.
+    if (template_->tp_doc) {
+        // Silently truncate the docstring if it contains a null byte
+        size_t size = strlen(template_->tp_doc) + 1;
+        char *doc = (char *)PyMem_Malloc(size);
+        if (doc == NULL) {
+            PyErr_NoMemory();
+            goto error;
+        }
+        memcpy(doc, template_->tp_doc, size);
+        t->ht_type.tp_doc = doc;
+    }
+
     if (bases != orig_bases) {
         if (PyObject_SetAttr((PyObject *)t, mypyc_interned_str.__orig_bases__, orig_bases) < 0)
             goto error;
@@ -333,21 +339,6 @@ PyObject *CPyType_FromTemplate(PyObject *template,
 
     Py_XDECREF(dummy_class);
 
-    // Unlike the tp_doc slots of most other object, a heap type's tp_doc
-    // must be heap allocated.
-    if (template_->tp_doc) {
-        // Silently truncate the docstring if it contains a null byte
-        Py_ssize_t size = strlen(template_->tp_doc) + 1;
-        char *tp_doc = (char *)PyMem_Malloc(size);
-        if (tp_doc == NULL) {
-            PyErr_NoMemory();
-            goto error;
-        }
-
-        memcpy(tp_doc, template_->tp_doc, size);
-        t->ht_type.tp_doc = tp_doc;
-    }
-
 #if PY_MINOR_VERSION == 11
     // This is a hack. Python 3.11 doesn't include good public APIs to work with managed
     // dicts, which are the default for heap types. So we try to opt-out until Python 3.12.

mypyc/lib-rt/pythonsupport.c

@@ -130,6 +130,7 @@ update_bases(PyObject *bases)
             /* If this is a first successful replacement, create new_bases list and
                copy previously encountered bases. */
             if (!(new_bases = PyList_New(i))) {
+                Py_DECREF(new_base);
                 goto error;
             }
             for (j = 0; j < i; j++) {
@@ -140,6 +141,7 @@ update_bases(PyObject *bases)
         }
         j = PyList_GET_SIZE(new_bases);
         if (PyList_SetSlice(new_bases, j, j, new_base) < 0) {
+            Py_DECREF(new_base);
             goto error;
         }
         Py_DECREF(new_base);

mypyc/test-data/run-classes.test

@@ -2272,29 +2272,101 @@ native.test()
 ok
 
 [case testNativeClassHiddenMetaclass]
-# The stub hides the metaclass of the base class, so it's only detected at runtime
-from hidden import Base
+# The stub hides the base classes' metaclass and base resolution behavior.
+from hidden import Prefix, Base, Second
 
-class C(Base):
+class C(Prefix, Base, Second):
     pass
 
 [file hidden.pyi]
+class Prefix: ...
 class Base: ...
+class Second: ...
 
 [file hidden.py]
+class Prefix:
+    pass
+
 class Meta(type):
     pass
 
 class Base(metaclass=Meta):
     pass
 
+class Second:
+    pass
+
+class SlotsBase:
+    __slots__ = ("attribute",)
+
+# Resolution fails before type creation, so a sentinel can stand in for a base.
+# Unlike types, its refcount is reliable on free-threaded Python too.
+Resolved = object()
+
+class Proxy:
+    def __init__(self, result):
+        self.result = result
+        self.failure = None
+
+    @property
+    def __mro_entries__(self):
+        if self.failure == "lookup":
+            raise RuntimeError("base lookup failed")
+        return self.resolve
+
+    def resolve(self, bases):
+        if self.failure == "call":
+            raise RuntimeError("base resolution failed")
+        return self.result
+
 [file driver.py]
-try:
-    import native
-except TypeError as e:
-    assert str(e) == "mypyc classes can't have a metaclass"
-else:
-    assert False
+import gc
+import sys
+import hidden
+from testutil import assertRaises
+
+for _ in range(3):
+    with assertRaises(TypeError, "mypyc classes can't have a metaclass"):
+        import native
+    gc.collect()
+
+hidden.Base = hidden.SlotsBase
+for _ in range(3):
+    with assertRaises(TypeError, "mypyc classes can't have __slots__"):
+        import native
+    gc.collect()
+
+second_base = hidden.Second
+hidden.Base = hidden.Proxy((hidden.Resolved,))
+hidden.Second = hidden.Proxy([])
+# Base creates a replacement list containing Prefix and Resolved before
+# Second fails. Neither that list nor a non-tuple result may leak references.
+tracked = (hidden.Resolved, hidden.Base, hidden.Second,
+           hidden.Base.result, hidden.Second.result)
+for failure, error, message in [
+    ("lookup", RuntimeError, "base lookup failed"),
+    ("call", RuntimeError, "base resolution failed"),
+    (None, TypeError, "__mro_entries__ must return a tuple"),
+]:
+    hidden.Second.failure = failure
+    gc.collect()
+    before = [sys.getrefcount(obj) for obj in tracked]
+    for _ in range(3):
+        with assertRaises(error, message):
+            import native
+        gc.collect()
+        after = [sys.getrefcount(obj) for obj in tracked]
+        assert after == before, (failure, before, after)
+
+# Keep the base layout but permit the next initialization to complete.
+hidden.Base = hidden.SlotsBase
+hidden.Base.__slots__ = ()
+hidden.Second = second_base
+import native
+# The template's tp_doc contains only the generated constructor signature.
+# CPython extracts it into __text_signature__, leaving __doc__ empty.
+assert native.C.__doc__ == ""
+assert native.C.__text_signature__ == "()"
 
 [case testMetaclass]
 from meta import Meta