PyModuleDef_Init: Fail if non-free-threading ABI is detected.

Author: encukou

Include/object.h

@@ -71,6 +71,8 @@ whose size is determined when the object is allocated.
  *
  * Statically allocated objects might be shared between
  * interpreters, so must be marked as immortal.
+ *
+ * Before changing this, see the check in PyModuleDef_Init().
  */
 #if defined(Py_GIL_DISABLED)
 #define PyObject_HEAD_INIT(type)    \
@@ -645,6 +647,7 @@ given type object has a specified feature.
 
 // Flag values for ob_flags (16 bits available, if SIZEOF_VOID_P > 4).
 #define _Py_IMMORTAL_FLAGS (1 << 0)
+#define _Py_LEGACY_ABI_CHECK_FLAG (1 << 1) /* see PyModuleDef_Init() */
 #define _Py_STATICALLY_ALLOCATED_FLAG (1 << 2)
 #if defined(Py_GIL_DISABLED) && defined(Py_DEBUG)
 #define _Py_TYPE_REVEALED_FLAG (1 << 3)

Objects/moduleobject.c

@@ -52,6 +52,59 @@ _PyModule_IsExtension(PyObject *obj)
 PyObject*
 PyModuleDef_Init(PyModuleDef* def)
 {
+#if Py_GIL_DISABLED
+    // Check that this def does not come from a non-free-threading ABI.
+    //
+    // This is meant as a "sanity check"; users should never rely on it.
+    // In particular, if we run out of ob_flags bits, or otherwise need to
+    // change some of the internals, this check can go away. Still, it
+    // would be nice to keep it for the free-threading transition.
+    //
+    // A PyModuleDef must be initialized with PyModuleDef_HEAD_INIT,
+    // which (via PyObject_HEAD_INIT) sets _Py_STATICALLY_ALLOCATED_FLAG
+    // and not _Py_LEGACY_ABI_CHECK_FLAG. For PyModuleDef, these flags never
+    // change.
+    // This means that the lower nibble of a valid PyModuleDef's ob_flags is
+    // always `_10_` (in binary; `_` is don't care).
+    //
+    // So, a check for these bits won't reject valid PyModuleDef.
+    // Rejecting incompatible extensions is slightly less important; here's
+    // how that works:
+    //
+    // In the pre-free-threading stable ABI, PyModuleDef_HEAD_INIT is big
+    // enough to overlap with free-threading ABI's ob_flags, is all zeros
+    // except for the refcount field.
+    // The refcount field can be:
+    // - 1 (3.11 and below)
+    // - UINT_MAX >> 2 (32-bit 3.12 & 3.13)
+    // - UINT_MAX (64-bit 3.12 & 3.13)
+    // - 7L << 28 (3.14)
+    //
+    // This means that the lower nibble of *any byte* in PyModuleDef_HEAD_INIT
+    // is not `_10_` -- it can be:
+    // - 0b0000
+    // - 0b0001
+    // - 0b0011 (from UINT_MAX >> 2)
+    // - 0b0111 (from 7L << 28)
+    // - 0b1111 (e.g. from UINT_MAX)
+    // (The values may change at runtime as the PyModuleDef is used, but
+    // PyModuleDef_Init is required before using the def as a Python object,
+    // so this check at least once with the initial values.
+    uint16_t flags = ((PyObject*)def)->ob_flags;
+    uint16_t bits = _Py_STATICALLY_ALLOCATED_FLAG | _Py_LEGACY_ABI_CHECK_FLAG;
+    if ((flags & bits) != _Py_STATICALLY_ALLOCATED_FLAG) {
+        const char *message = "invalid PyModuleDef, extension possibly "
+            "compiled for non-free-threaded Python";
+        // Write the error as unraisable: if the extension tries calling
+        // any API, it's likely to segfault and lose the exception.
+        PyErr_SetString(PyExc_SystemError, message);
+        PyErr_WriteUnraisable(NULL);
+        // But also raise the exception normally -- this is technically
+        // a recoverable state.
+        PyErr_SetString(PyExc_SystemError, message);
+        return NULL;
+    }
+#endif
     assert(PyModuleDef_Type.tp_flags & Py_TPFLAGS_READY);
     if (def->m_base.m_index == 0) {
         Py_SET_REFCNT(def, 1);