bpo-45395: Make custom frozen modules additions instead of replacements. (gh-28778)

Currently custom modules (the array set on PyImport_FrozenModules) replace all the frozen stdlib modules. That can be problematic and is unlikely to be what the user wants. This change treats the custom frozen modules as additions instead. They take precedence over all other frozen modules except for those needed to bootstrap the import system. If the "code" field of an entry in the custom array is NULL then that frozen module is treated as disabled, which allows a custom entry to disable a frozen stdlib module.

This change allows us to get rid of is_essential_frozen_module() and simplifies the logic for which frozen modules should be ignored.

https://bugs.python.org/issue45395

2 files changed, 125 insertions, 53 deletions

diff --git a/Python/frozen.c b/Python/frozen.c
index 499b3b9..15baa97 100644
--- a/Python/frozen.c
+++ b/Python/frozen.c
@@ -63,14 +63,15 @@
 
 /* Note that a negative size indicates a package. */
 
-static const struct _frozen _PyImport_FrozenModules[] = {
-    /* import system */
+static const struct _frozen bootstrap_modules[] = {
     {"_frozen_importlib", _Py_M__importlib__bootstrap,
         (int)sizeof(_Py_M__importlib__bootstrap)},
     {"_frozen_importlib_external", _Py_M__importlib__bootstrap_external,
         (int)sizeof(_Py_M__importlib__bootstrap_external)},
     {"zipimport", _Py_M__zipimport, (int)sizeof(_Py_M__zipimport)},
-
+    {0, 0, 0} /* bootstrap sentinel */
+};
+static const struct _frozen stdlib_modules[] = {
     /* stdlib - startup, without site (python -S) */
     {"abc", _Py_M__abc, (int)sizeof(_Py_M__abc)},
     {"codecs", _Py_M__codecs, (int)sizeof(_Py_M__codecs)},
@@ -87,8 +88,9 @@ static const struct _frozen _PyImport_FrozenModules[] = {
     {"os", _Py_M__os, (int)sizeof(_Py_M__os)},
     {"site", _Py_M__site, (int)sizeof(_Py_M__site)},
     {"stat", _Py_M__stat, (int)sizeof(_Py_M__stat)},
-
-    /* Test module */
+    {0, 0, 0} /* stdlib sentinel */
+};
+static const struct _frozen test_modules[] = {
     {"__hello__", _Py_M____hello__, (int)sizeof(_Py_M____hello__)},
     {"__hello_alias__", _Py_M____hello__, (int)sizeof(_Py_M____hello__)},
     {"__phello_alias__", _Py_M____hello__, -(int)sizeof(_Py_M____hello__)},
@@ -103,8 +105,11 @@ static const struct _frozen _PyImport_FrozenModules[] = {
     {"__phello__.spam", _Py_M____phello___spam,
         (int)sizeof(_Py_M____phello___spam)},
     {"__hello_only__", _Py_M__frozen_only, (int)sizeof(_Py_M__frozen_only)},
-    {0, 0, 0} /* modules sentinel */
+    {0, 0, 0} /* test sentinel */
 };
+const struct _frozen *_PyImport_FrozenBootstrap = bootstrap_modules;
+const struct _frozen *_PyImport_FrozenStdlib = stdlib_modules;
+const struct _frozen *_PyImport_FrozenTest = test_modules;
 
 static const struct _module_alias aliases[] = {
     {"_frozen_importlib", "importlib._bootstrap"},
@@ -124,4 +129,4 @@ const struct _module_alias *_PyImport_FrozenAliases = aliases;
 /* Embedding apps may change this pointer to point to their favorite
    collection of frozen modules: */
 
-const struct _frozen *PyImport_FrozenModules = _PyImport_FrozenModules;
+const struct _frozen *PyImport_FrozenModules = NULL;
diff --git a/Python/import.c b/Python/import.c
index 15b1956..48ea912 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -1071,27 +1071,6 @@ resolve_module_alias(const char *name, const struct _module_alias *aliases,
 /* Frozen modules */
 
 static bool
-is_essential_frozen_module(const char *name)
-{
-    /* These modules are necessary to bootstrap the import system. */
-    if (strcmp(name, "_frozen_importlib") == 0) {
-        return true;
-    }
-    if (strcmp(name, "_frozen_importlib_external") == 0) {
-        return true;
-    }
-    if (strcmp(name, "zipimport") == 0) {
-        return true;
-    }
-    /* This doesn't otherwise have anywhere to find the module.
-       See frozenmain.c. */
-    if (strcmp(name, "__main__") == 0) {
-        return true;
-    }
-    return false;
-}
-
-static bool
 use_frozen(void)
 {
     PyInterpreterState *interp = _PyInterpreterState_GET();
@@ -1115,26 +1094,76 @@ list_frozen_module_names()
         return NULL;
     }
     bool enabled = use_frozen();
-    for (const struct _frozen *p = PyImport_FrozenModules; ; p++) {
+    const struct _frozen *p;
+#define ADD_MODULE(name) \
+    do { \
+        PyObject *nameobj = PyUnicode_FromString(name); \
+        if (nameobj == NULL) { \
+            goto error; \
+        } \
+        int res = PyList_Append(names, nameobj); \
+        Py_DECREF(nameobj); \
+        if (res != 0) { \
+            goto error; \
+        } \
+    } while(0)
+    // We always use the bootstrap modules.
+    for (p = _PyImport_FrozenBootstrap; ; p++) {
         if (p->name == NULL) {
             break;
         }
-        if (!enabled && !is_essential_frozen_module(p->name)) {
-            continue;
+        ADD_MODULE(p->name);
+    }
+    // Frozen stdlib modules may be disabled.
+    for (p = _PyImport_FrozenStdlib; ; p++) {
+        if (p->name == NULL) {
+            break;
         }
-        PyObject *name = PyUnicode_FromString(p->name);
-        if (name == NULL) {
-            Py_DECREF(names);
-            return NULL;
+        if (enabled) {
+            ADD_MODULE(p->name);
         }
-        int res = PyList_Append(names, name);
-        Py_DECREF(name);
-        if (res != 0) {
-            Py_DECREF(names);
-            return NULL;
+    }
+    for (p = _PyImport_FrozenTest; ; p++) {
+        if (p->name == NULL) {
+            break;
+        }
+        if (enabled) {
+            ADD_MODULE(p->name);
+        }
+    }
+#undef ADD_MODULE
+    // Add any custom modules.
+    if (PyImport_FrozenModules != NULL) {
+        for (p = PyImport_FrozenModules; ; p++) {
+            if (p->name == NULL) {
+                break;
+            }
+            PyObject *nameobj = PyUnicode_FromString(p->name);
+            if (nameobj == NULL) {
+                goto error;
+            }
+            int found = PySequence_Contains(names, nameobj);
+            if (found < 0) {
+                Py_DECREF(nameobj);
+                goto error;
+            }
+            else if (found) {
+                Py_DECREF(nameobj);
+            }
+            else {
+                int res = PyList_Append(names, nameobj);
+                Py_DECREF(nameobj);
+                if (res != 0) {
+                    goto error;
+                }
+            }
         }
     }
     return names;
+
+error:
+    Py_DECREF(names);
+    return NULL;
 }
 
 typedef enum {
@@ -1180,6 +1209,54 @@ set_frozen_error(frozen_status status, PyObject *modname)
     }
 }
 
+static const struct _frozen *
+look_up_frozen(const char *name)
+{
+    const struct _frozen *p;
+    // We always use the bootstrap modules.
+    for (p = _PyImport_FrozenBootstrap; ; p++) {
+        if (p->name == NULL) {
+            // We hit the end-of-list sentinel value.
+            break;
+        }
+        if (strcmp(name, p->name) == 0) {
+            return p;
+        }
+    }
+    // Prefer custom modules, if any.  Frozen stdlib modules can be
+    // disabled here by setting "code" to NULL in the array entry.
+    if (PyImport_FrozenModules != NULL) {
+        for (p = PyImport_FrozenModules; ; p++) {
+            if (p->name == NULL) {
+                break;
+            }
+            if (strcmp(name, p->name) == 0) {
+                return p;
+            }
+        }
+    }
+    // Frozen stdlib modules may be disabled.
+    if (use_frozen()) {
+        for (p = _PyImport_FrozenStdlib; ; p++) {
+            if (p->name == NULL) {
+                break;
+            }
+            if (strcmp(name, p->name) == 0) {
+                return p;
+            }
+        }
+        for (p = _PyImport_FrozenTest; ; p++) {
+            if (p->name == NULL) {
+                break;
+            }
+            if (strcmp(name, p->name) == 0) {
+                return p;
+            }
+        }
+    }
+    return NULL;
+}
+
 struct frozen_info {
     PyObject *nameobj;
     const char *data;
@@ -1209,19 +1286,9 @@ find_frozen(PyObject *nameobj, struct frozen_info *info)
         return FROZEN_BAD_NAME;
     }
 
-    if (!use_frozen() && !is_essential_frozen_module(name)) {
-        return FROZEN_DISABLED;
-    }
-
-    const struct _frozen *p;
-    for (p = PyImport_FrozenModules; ; p++) {
-        if (p->name == NULL) {
-            // We hit the end-of-list sentinel value.
-            return FROZEN_NOT_FOUND;
-        }
-        if (strcmp(name, p->name) == 0) {
-            break;
-        }
+    const struct _frozen *p = look_up_frozen(name);
+    if (p == NULL) {
+        return FROZEN_NOT_FOUND;
     }
     if (info != NULL) {
         info->nameobj = nameobj;  // borrowed