bpo-44131: Test Py_FrozenMain() (GH-26126)

* Add test_frozenmain to test_embed
* Add Programs/test_frozenmain.py
* Add Programs/freeze_test_frozenmain.py
* Add Programs/test_frozenmain.h
* Add make regen-test-frozenmain
* Add test_frozenmain command to Programs/_testembed
* _testembed.c: add error(msg) function

4 files changed, 150 insertions, 9 deletions

diff --git a/Programs/_testembed.c b/Programs/_testembed.c
index 21b24f7..a5ae7c1 100644
--- a/Programs/_testembed.c
+++ b/Programs/_testembed.c
@@ -27,6 +27,14 @@
 _Py_COMP_DIAG_PUSH
 _Py_COMP_DIAG_IGNORE_DEPR_DECLS
 
+
+static void error(const char *msg)
+{
+    fprintf(stderr, "ERROR: %s\n", msg);
+    fflush(stderr);
+}
+
+
 static void _testembed_Py_Initialize(void)
 {
     Py_SetProgramName(PROGRAM_NAME);
@@ -239,7 +247,7 @@ static void bpo20891_thread(void *lockp)
 
     PyGILState_STATE state = PyGILState_Ensure();
     if (!PyGILState_Check()) {
-        fprintf(stderr, "PyGILState_Check failed!");
+        error("PyGILState_Check failed!");
         abort();
     }
 
@@ -259,7 +267,7 @@ static int test_bpo20891(void)
        crash. */
     PyThread_type_lock lock = PyThread_allocate_lock();
     if (!lock) {
-        fprintf(stderr, "PyThread_allocate_lock failed!");
+        error("PyThread_allocate_lock failed!");
         return 1;
     }
 
@@ -267,7 +275,7 @@ static int test_bpo20891(void)
 
     unsigned long thrd = PyThread_start_new_thread(bpo20891_thread, &lock);
     if (thrd == PYTHREAD_INVALID_THREAD_ID) {
-        fprintf(stderr, "PyThread_start_new_thread failed!");
+        error("PyThread_start_new_thread failed!");
         return 1;
     }
     PyThread_acquire_lock(lock, WAIT_LOCK);
@@ -1397,12 +1405,12 @@ static int test_init_setpath(void)
 {
     char *env = getenv("TESTPATH");
     if (!env) {
-        fprintf(stderr, "missing TESTPATH env var\n");
+        error("missing TESTPATH env var");
         return 1;
     }
     wchar_t *path = Py_DecodeLocale(env, NULL);
     if (path == NULL) {
-        fprintf(stderr, "failed to decode TESTPATH\n");
+        error("failed to decode TESTPATH");
         return 1;
     }
     Py_SetPath(path);
@@ -1430,12 +1438,12 @@ static int test_init_setpath_config(void)
 
     char *env = getenv("TESTPATH");
     if (!env) {
-        fprintf(stderr, "missing TESTPATH env var\n");
+        error("missing TESTPATH env var");
         return 1;
     }
     wchar_t *path = Py_DecodeLocale(env, NULL);
     if (path == NULL) {
-        fprintf(stderr, "failed to decode TESTPATH\n");
+        error("failed to decode TESTPATH");
         return 1;
     }
     Py_SetPath(path);
@@ -1459,12 +1467,12 @@ static int test_init_setpythonhome(void)
 {
     char *env = getenv("TESTHOME");
     if (!env) {
-        fprintf(stderr, "missing TESTHOME env var\n");
+        error("missing TESTHOME env var");
         return 1;
     }
     wchar_t *home = Py_DecodeLocale(env, NULL);
     if (home == NULL) {
-        fprintf(stderr, "failed to decode TESTHOME\n");
+        error("failed to decode TESTHOME");
         return 1;
     }
     Py_SetPythonHome(home);
@@ -1726,6 +1734,48 @@ static int test_unicode_id_init(void)
 }
 
 
+#ifndef MS_WINDOWS
+#include "test_frozenmain.h"      // M_test_frozenmain
+
+static int test_frozenmain(void)
+{
+    // Get "_frozen_importlib" and "_frozen_importlib_external"
+    // from PyImport_FrozenModules
+    const struct _frozen *importlib = NULL, *importlib_external = NULL;
+    for (const struct _frozen *mod = PyImport_FrozenModules; mod->name != NULL; mod++) {
+        if (strcmp(mod->name, "_frozen_importlib") == 0) {
+            importlib = mod;
+        }
+        else if (strcmp(mod->name, "_frozen_importlib_external") == 0) {
+            importlib_external = mod;
+        }
+    }
+    if (importlib == NULL || importlib_external == NULL) {
+        error("cannot find frozen importlib and importlib_external");
+        return 1;
+    }
+
+    static struct _frozen frozen_modules[4] = {
+        {0, 0, 0},  // importlib
+        {0, 0, 0},  // importlib_external
+        {"__main__", M_test_frozenmain, sizeof(M_test_frozenmain)},
+        {0, 0, 0}   // sentinel
+    };
+    frozen_modules[0] = *importlib;
+    frozen_modules[1] = *importlib_external;
+
+    char* argv[] = {
+        "./argv0",
+        "-E",
+        "arg1",
+        "arg2",
+    };
+    PyImport_FrozenModules = frozen_modules;
+    return Py_FrozenMain(Py_ARRAY_LENGTH(argv), argv);
+}
+#endif  // !MS_WINDOWS
+
+
 // List frozen modules.
 // Command used by Tools/scripts/generate_stdlib_module_names.py script.
 static int list_frozen(void)
@@ -1811,11 +1861,15 @@ static struct TestCase TestCases[] = {
     {"test_audit_run_stdin", test_audit_run_stdin},
 
     {"test_unicode_id_init", test_unicode_id_init},
+#ifndef MS_WINDOWS
+    {"test_frozenmain", test_frozenmain},
+#endif
 
     {"list_frozen", list_frozen},
     {NULL, NULL}
 };
 
+
 int main(int argc, char *argv[])
 {
     if (argc > 1) {
diff --git a/Programs/freeze_test_frozenmain.py b/Programs/freeze_test_frozenmain.py
new file mode 100644
index 0000000..848fc31
--- /dev/null
+++ b/Programs/freeze_test_frozenmain.py
@@ -0,0 +1,48 @@
+import marshal
+import tokenize
+import os.path
+import sys
+
+PROGRAM_DIR = os.path.dirname(__file__)
+SRC_DIR = os.path.dirname(PROGRAM_DIR)
+
+
+def writecode(fp, mod, data):
+    print('unsigned char M_%s[] = {' % mod, file=fp)
+    indent = ' ' * 4
+    for i in range(0, len(data), 16):
+        print(indent, file=fp, end='')
+        for c in bytes(data[i:i+16]):
+            print('%d,' % c, file=fp, end='')
+        print('', file=fp)
+    print('};', file=fp)
+
+
+def dump(fp, filename, name):
+    # Strip the directory to get reproducible marshal dump
+    code_filename = os.path.basename(filename)
+
+    with tokenize.open(filename) as source_fp:
+        source = source_fp.read()
+        code = compile(source, code_filename, 'exec')
+
+    data = marshal.dumps(code)
+    writecode(fp, name, data)
+
+
+def main():
+    if len(sys.argv) < 2:
+        print(f"usage: {sys.argv[0]} filename")
+        sys.exit(1)
+    filename = sys.argv[1]
+
+    with open(filename, "w") as fp:
+        print("// Auto-generated by Programs/freeze_test_frozenmain.py", file=fp)
+        frozenmain = os.path.join(PROGRAM_DIR, 'test_frozenmain.py')
+        dump(fp, frozenmain, 'test_frozenmain')
+
+    print(f"{filename} written")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/Programs/test_frozenmain.h b/Programs/test_frozenmain.h
new file mode 100644
index 0000000..ac3dfd3
--- /dev/null
+++ b/Programs/test_frozenmain.h
@@ -0,0 +1,30 @@
+// Auto-generated by Programs/freeze_test_frozenmain.py
+unsigned char M_test_frozenmain[] = {
+    227,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
+    0,4,0,0,0,64,0,0,0,115,106,0,0,0,100,0,
+    100,1,108,0,90,0,100,0,100,1,108,1,90,1,101,2,
+    100,2,131,1,1,0,101,2,100,3,101,0,106,3,131,2,
+    1,0,101,1,160,4,161,0,100,4,25,0,90,5,101,2,
+    100,5,101,5,100,6,25,0,155,0,157,2,131,1,1,0,
+    101,2,100,7,101,5,100,8,25,0,155,0,157,2,131,1,
+    1,0,101,2,100,9,101,5,100,10,25,0,155,0,157,2,
+    131,1,1,0,100,1,83,0,41,11,233,0,0,0,0,78,
+    122,18,70,114,111,122,101,110,32,72,101,108,108,111,32,87,
+    111,114,108,100,122,8,115,121,115,46,97,114,103,118,218,6,
+    99,111,110,102,105,103,122,21,99,111,110,102,105,103,32,112,
+    114,111,103,114,97,109,95,110,97,109,101,58,32,90,12,112,
+    114,111,103,114,97,109,95,110,97,109,101,122,19,99,111,110,
+    102,105,103,32,101,120,101,99,117,116,97,98,108,101,58,32,
+    218,10,101,120,101,99,117,116,97,98,108,101,122,24,99,111,
+    110,102,105,103,32,117,115,101,95,101,110,118,105,114,111,110,
+    109,101,110,116,58,32,90,15,117,115,101,95,101,110,118,105,
+    114,111,110,109,101,110,116,41,6,218,3,115,121,115,90,17,
+    95,116,101,115,116,105,110,116,101,114,110,97,108,99,97,112,
+    105,218,5,112,114,105,110,116,218,4,97,114,103,118,90,11,
+    103,101,116,95,99,111,110,102,105,103,115,114,2,0,0,0,
+    169,0,114,7,0,0,0,114,7,0,0,0,250,18,116,101,
+    115,116,95,102,114,111,122,101,110,109,97,105,110,46,112,121,
+    218,8,60,109,111,100,117,108,101,62,1,0,0,0,115,16,
+    0,0,0,8,0,8,1,8,2,12,1,12,1,18,1,18,
+    1,22,1,243,0,0,0,0,
+};
diff --git a/Programs/test_frozenmain.py b/Programs/test_frozenmain.py
new file mode 100644
index 0000000..aa79106
--- /dev/null
+++ b/Programs/test_frozenmain.py
@@ -0,0 +1,9 @@
+import sys
+import _testinternalcapi
+
+print("Frozen Hello World")
+print("sys.argv", sys.argv)
+config = _testinternalcapi.get_configs()['config']
+print(f"config program_name: {config['program_name']}")
+print(f"config executable: {config['executable']}")
+print(f"config use_environment: {config['use_environment']}")