gh-81057: Move PyImport_Inittab to _PyRuntimeState (gh-99402)

We actually don't move PyImport_Inittab.  Instead, we make a copy that we keep on _PyRuntimeState and use only that after Py_Initialize().  We also prevent folks from modifying PyImport_Inittab (the best we can) after that point.

https://github.com/python/cpython/issues/81057

3 files changed, 68 insertions, 6 deletions

diff --git a/Python/import.c b/Python/import.c
index d161613..daee16e 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -30,7 +30,11 @@ static PyObject *import_add_module(PyThreadState *tstate, PyObject *name);
 /* This table is defined in config.c: */
 extern struct _inittab _PyImport_Inittab[];
 
+// This is not used after Py_Initialize() is called.
+// (See _PyRuntimeState.imports.inittab.)
 struct _inittab *PyImport_Inittab = _PyImport_Inittab;
+// When we dynamically allocate a larger table for PyImport_ExtendInittab(),
+// we track the pointer here so we can deallocate it during finalization.
 static struct _inittab *inittab_copy = NULL;
 
 /*[clinic input]
@@ -218,6 +222,38 @@ _imp_release_lock_impl(PyObject *module)
     Py_RETURN_NONE;
 }
 
+PyStatus
+_PyImport_Init(void)
+{
+    if (_PyRuntime.imports.inittab != NULL) {
+        return _PyStatus_ERR("global import state already initialized");
+    }
+    PyStatus status = _PyStatus_OK();
+
+    size_t size;
+    for (size = 0; PyImport_Inittab[size].name != NULL; size++)
+        ;
+    size++;
+
+    /* Force default raw memory allocator to get a known allocator to be able
+       to release the memory in _PyImport_Fini() */
+    PyMemAllocatorEx old_alloc;
+    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
+
+    /* Make the copy. */
+    struct _inittab *copied = PyMem_RawMalloc(size * sizeof(struct _inittab));
+    if (copied == NULL) {
+        status = PyStatus_NoMemory();
+        goto done;
+    }
+    memcpy(copied, PyImport_Inittab, size * sizeof(struct _inittab));
+    _PyRuntime.imports.inittab = copied;
+
+done:
+    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
+    return status;
+}
+
 static inline void _extensions_cache_clear(void);
 
 void
@@ -228,6 +264,17 @@ _PyImport_Fini(void)
         PyThread_free_lock(import_lock);
         import_lock = NULL;
     }
+
+    /* Use the same memory allocator as _PyImport_Init(). */
+    PyMemAllocatorEx old_alloc;
+    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
+
+    /* Free memory allocated by _PyImport_Init() */
+    struct _inittab *inittab = _PyRuntime.imports.inittab;
+    _PyRuntime.imports.inittab = NULL;
+    PyMem_RawFree(inittab);
+
+    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 }
 
 void
@@ -889,9 +936,10 @@ static int
 is_builtin(PyObject *name)
 {
     int i;
-    for (i = 0; PyImport_Inittab[i].name != NULL; i++) {
-        if (_PyUnicode_EqualToASCIIString(name, PyImport_Inittab[i].name)) {
-            if (PyImport_Inittab[i].initfunc == NULL)
+    struct _inittab *inittab = _PyRuntime.imports.inittab;
+    for (i = 0; inittab[i].name != NULL; i++) {
+        if (_PyUnicode_EqualToASCIIString(name, inittab[i].name)) {
+            if (inittab[i].initfunc == NULL)
                 return -1;
             else
                 return 1;
@@ -984,7 +1032,7 @@ create_builtin(PyThreadState *tstate, PyObject *name, PyObject *spec)
     }
 
     PyObject *modules = tstate->interp->modules;
-    for (struct _inittab *p = PyImport_Inittab; p->name != NULL; p++) {
+    for (struct _inittab *p = _PyRuntime.imports.inittab; p->name != NULL; p++) {
         if (_PyUnicode_EqualToASCIIString(name, p->name)) {
             if (p->initfunc == NULL) {
                 /* Cannot re-init internal module ("sys" or "builtins") */
@@ -2592,6 +2640,10 @@ PyImport_ExtendInittab(struct _inittab *newtab)
     size_t i, n;
     int res = 0;
 
+    if (_PyRuntime.imports.inittab != NULL) {
+        Py_FatalError("PyImport_ExtendInittab() may be be called after Py_Initialize()");
+    }
+
     /* Count the number of entries in both tables */
     for (n = 0; newtab[n].name != NULL; n++)
         ;
@@ -2636,6 +2688,10 @@ PyImport_AppendInittab(const char *name, PyObject* (*initfunc)(void))
 {
     struct _inittab newtab[2];
 
+    if (_PyRuntime.imports.inittab != NULL) {
+        Py_FatalError("PyImport_AppendInittab() may be be called after Py_Initialize()");
+    }
+
     memset(newtab, '\0', sizeof newtab);
 
     newtab[0].name = name;
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index 39e54c1..44f8442 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -605,6 +605,11 @@ pycore_init_runtime(_PyRuntimeState *runtime,
         return status;
     }
 
+    status = _PyImport_Init();
+    if (_PyStatus_EXCEPTION(status)) {
+        return status;
+    }
+
     status = _PyInterpreterState_Enable(runtime);
     if (_PyStatus_EXCEPTION(status)) {
         return status;
diff --git a/Python/sysmodule.c b/Python/sysmodule.c
index 5de684e..1090b12 100644
--- a/Python/sysmodule.c
+++ b/Python/sysmodule.c
@@ -2252,8 +2252,9 @@ list_builtin_module_names(void)
     if (list == NULL) {
         return NULL;
     }
-    for (Py_ssize_t i = 0; PyImport_Inittab[i].name != NULL; i++) {
-        PyObject *name = PyUnicode_FromString(PyImport_Inittab[i].name);
+    struct _inittab *inittab = _PyRuntime.imports.inittab;
+    for (Py_ssize_t i = 0; inittab[i].name != NULL; i++) {
+        PyObject *name = PyUnicode_FromString(inittab[i].name);
         if (name == NULL) {
             goto error;
         }