bpo-40600: Disallow loading extension modules load more than once in per interpreter.

Doc/whatsnew/3.10.rst

@@ -402,6 +402,11 @@ Optimizations
   bytecode level.  It is now around 100% faster to create a function with parameter
   annotations.  (Contributed by Yurii Karabas and Inada Naoki in :issue:`42202`)
 
+* Add ``Py_MODFLAGS_SINGLE`` of module flags to disallow loading the extension
+  module more than once in per interpreter.
+  (Contributed by Hai Shi in :issue:`40600`.)
+
+
 Deprecated
 ==========
 

Include/internal/pycore_interp.h

@@ -193,6 +193,7 @@ struct _is {
     // sys.modules dictionary
     PyObject *modules;
     PyObject *modules_by_index;
+    PyObject *single_modules;
     // Dictionary of the sys module
     PyObject *sysdict;
     // Dictionary of the builtins module

Include/internal/pycore_pystate.h

@@ -144,6 +144,9 @@ PyAPI_FUNC(int) _PyState_AddModule(
     PyObject* module,
     struct PyModuleDef* def);
 
+/* New in 3.10 */
+PyAPI_FUNC(int) _PyState_AddSingleModule(PyObject *, struct PyModuleDef *);
+PyAPI_FUNC(PyObject *) _PyState_FindSingleModule(struct PyModuleDef *);
 
 PyAPI_FUNC(int) _PyOS_InterruptOccurred(PyThreadState *tstate);
 

Include/moduleobject.h

@@ -82,6 +82,7 @@ typedef struct PyModuleDef{
   traverseproc m_traverse;
   inquiry m_clear;
   freefunc m_free;
+  unsigned int m_flags;
 } PyModuleDef;
 
 #ifdef __cplusplus

Include/object.h

@@ -382,6 +382,8 @@ given type object has a specified feature.
 /* Type structure has tp_finalize member (3.4) */
 #define Py_TPFLAGS_HAVE_FINALIZE (1UL << 0)
 
+/* Set if the module object is only allowed loading once in per interpreter. */
+#define Py_MODFLAGS_SINGLE 1UL
 
 /*
 The macros Py_INCREF(op) and Py_DECREF(op) are used to increment or decrement

Lib/test/test_atexit.py

@@ -222,6 +222,28 @@ def callback():
         self.assertEqual(os.read(r, len(expected)), expected)
         os.close(r)
 
+    def test_multiple_loading_on_subinterpreter(self):
+        expected = b"atexit2 callback\natexit1 callback\n"
+        r, w = os.pipe()
+
+        code = r"""if 1:
+            import sys
+            import os
+            def callback(msg):
+                os.write({:d}, msg)
+            atexit1 = sys.modules.pop('atexit', None)
+            if atexit1 is None:
+                import atexit as atexit1
+                del sys.modules['atexit']
+            import atexit as atexit2
+            atexit1.register(callback, b"atexit1 callback\n")
+            atexit2.register(callback, b"atexit2 callback\n")
+        """.format(w)
+        ret = support.run_in_subinterp(code)
+        os.close(w)
+        self.assertEqual(os.read(r, len(expected)), expected)
+        os.close(r)
+
 
 if __name__ == "__main__":
     unittest.main()

Misc/NEWS.d/next/C API/2020-12-07-18-45-55.bpo-40600.9po0nr.rst

@@ -0,0 +1,3 @@
+Add :c:func:`_PyState_AddSingleModule` and :c:func:`_PyState_FindSingleModule`
+and ``Py_MODFLAGS_SINGLE`` of module flags to disallow loading the extension
+module more than once in per interpreter.

Modules/atexitmodule.c

@@ -338,14 +338,15 @@ static PyModuleDef_Slot atexit_slots[] = {
 
 static struct PyModuleDef atexitmodule = {
     PyModuleDef_HEAD_INIT,
-    "atexit",
-    atexit__doc__,
-    sizeof(atexitmodule_state),
-    atexit_methods,
-    atexit_slots,
-    atexit_m_traverse,
-    atexit_m_clear,
-    (freefunc)atexit_free
+    .m_name = "atexit",
+    .m_doc = atexit__doc__,
+    .m_size = sizeof(atexitmodule_state),
+    .m_methods = atexit_methods,
+    .m_slots = atexit_slots,
+    .m_traverse = atexit_m_traverse,
+    .m_clear = atexit_m_clear,
+    .m_free = (freefunc)atexit_free,
+    .m_flags = Py_MODFLAGS_SINGLE
 };
 
 PyMODINIT_FUNC

Objects/moduleobject.c

@@ -38,8 +38,15 @@ PyTypeObject PyModuleDef_Type = {
 PyObject*
 PyModuleDef_Init(struct PyModuleDef* def)
 {
-    if (PyType_Ready(&PyModuleDef_Type) < 0)
-         return NULL;
+    if (PyType_Ready(&PyModuleDef_Type) < 0) {
+        return NULL;
+    }
+    if (def->m_flags & Py_MODFLAGS_SINGLE) {
+        PyObject *mod = _PyState_FindSingleModule(def);
+        if (mod != NULL) {
+            return mod;
+        }
+    }
     if (def->m_base.m_index == 0) {
         max_module_number++;
         Py_SET_REFCNT(def, 1);

PC/python3dll.c

@@ -30,6 +30,8 @@ EXPORT_FUNC(_PyObject_GC_Resize)
 EXPORT_FUNC(_PyObject_New)
 EXPORT_FUNC(_PyObject_NewVar)
 EXPORT_FUNC(_PyState_AddModule)
+EXPORT_FUNC(_PyState_AddSingleModule)
+EXPORT_FUNC(_PyState_FindSingleModule)
 EXPORT_FUNC(_PyThreadState_Init)
 EXPORT_FUNC(_PyThreadState_Prealloc)
 EXPORT_FUNC(_PyTrash_deposit_object)

Python/import.c

@@ -1002,7 +1002,25 @@ create_builtin(PyThreadState *tstate, PyObject *name, PyObject *spec)
             }
 
             if (PyObject_TypeCheck(mod, &PyModuleDef_Type)) {
-                return PyModule_FromDefAndSpec((PyModuleDef*)mod, spec);
+                PyObject *module = PyModule_FromDefAndSpec((PyModuleDef *)mod,
+                                                           spec);
+                if (module == NULL) {
+                    return NULL;
+                }
+                if (((PyModuleDef *)mod)->m_flags & Py_MODFLAGS_SINGLE) {
+                    if (_PyState_AddSingleModule(module,
+                                                 (PyModuleDef *)mod) < 0) {
+                        Py_DECREF(module);
+                        return NULL;
+                    }
+                }
+                return module;
+            }
+            else if (PyModule_GetDef(mod)->m_slots != NULL) {
+                PyModuleDef *def = PyModule_GetDef(mod);
+                PyObject *module = _PyState_FindSingleModule(def);
+                Py_XINCREF(module);
+                return module;
             }
             else {
                 /* Remember pointer to module init function. */

Python/pystate.c

@@ -291,6 +291,7 @@ interpreter_clear(PyInterpreterState *interp, PyThreadState *tstate)
     Py_CLEAR(interp->codec_error_registry);
     Py_CLEAR(interp->modules);
     Py_CLEAR(interp->modules_by_index);
+    Py_CLEAR(interp->single_modules);
     Py_CLEAR(interp->builtins_copy);
     Py_CLEAR(interp->importlib);
     Py_CLEAR(interp->import_func);
@@ -750,6 +751,56 @@ PyState_AddModule(PyObject* module, struct PyModuleDef* def)
     return _PyState_AddModule(tstate, module, def);
 }
 
+int
+_PyState_AddSingleModule(PyObject *module, PyModuleDef *def)
+{
+    if (def == NULL) {
+        Py_FatalError("module definition is NULL");
+        return -1;
+    }
+    if (module == NULL) {
+        Py_FatalError("module is NULL");
+        return -1;
+    }
+
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    if (interp->single_modules == NULL) {
+        interp->single_modules = PyList_New(0);
+        if (interp->single_modules == NULL) {
+            return -1;
+        }
+    }
+
+    while (PyList_GET_SIZE(interp->single_modules) <= def->m_base.m_index) {
+        if (PyList_Append(interp->single_modules, Py_None) < 0) {
+            return -1;
+        }
+    }
+
+    Py_INCREF(module);
+    return PyList_SetItem(interp->single_modules,
+                          def->m_base.m_index, module);
+}
+
+PyObject *
+_PyState_FindSingleModule(PyModuleDef *def)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    Py_ssize_t index = def->m_base.m_index;
+
+    if (index == 0) {
+        return NULL;
+    }
+    if (interp->single_modules == NULL) {
+        return NULL;
+    }
+    if (index >= PyList_GET_SIZE(interp->single_modules)) {
+        return NULL;
+    }
+    PyObject *res = PyList_GET_ITEM(interp->single_modules, index);
+    return res == Py_None ? NULL : res;
+}
+
 int
 PyState_RemoveModule(struct PyModuleDef* def)
 {