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

10 files changed, 203 insertions, 32 deletions

diff --git a/.gitattributes b/.gitattributes
index c66e765..fd30380 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -57,6 +57,7 @@ Doc/library/token-list.inc  linguist-generated=true
 Include/token.h             linguist-generated=true
 Lib/token.py                linguist-generated=true
 Parser/token.c              linguist-generated=true
+Programs/test_frozenmain.h  linguist-generated=true
 
 # Language aware diff headers
 # https://tekin.co.uk/2020/10/better-git-diff-output-for-ruby-python-elixir-and-more
diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
index 23cf297..c68a662 100644
--- a/Lib/test/test_embed.py
+++ b/Lib/test/test_embed.py
@@ -1480,6 +1480,21 @@ class MiscTests(EmbeddingTestsMixin, unittest.TestCase):
         # when Python is initialized multiples times.
         self.run_embedded_interpreter("test_unicode_id_init")
 
+    # See bpo-44133
+    @unittest.skipIf(os.name == 'nt',
+                     'Py_FrozenMain is not exported on Windows')
+    def test_frozenmain(self):
+        out, err = self.run_embedded_interpreter("test_frozenmain")
+        exe = os.path.realpath('./argv0')
+        expected = textwrap.dedent(f"""
+            Frozen Hello World
+            sys.argv ['./argv0', '-E', 'arg1', 'arg2']
+            config program_name: ./argv0
+            config executable: {exe}
+            config use_environment: 1
+        """).lstrip()
+        self.assertEqual(out, expected)
+
 
 class StdPrinterTests(EmbeddingTestsMixin, unittest.TestCase):
     # Test PyStdPrinter_Type which is used by _PySys_SetPreliminaryStderr():
diff --git a/Makefile.pre.in b/Makefile.pre.in
index 080318b..83788a6 100644
--- a/Makefile.pre.in
+++ b/Makefile.pre.in
@@ -720,6 +720,14 @@ Makefile Modules/config.c: Makefile.pre \
 	@mv config.c Modules
 	@echo "The Makefile was updated, you may need to re-run make."
 
+regen-test-frozenmain: $(BUILDPYTHON)
+	# Regenerate Programs/test_frozenmain.h
+	# from Programs/test_frozenmain.py
+	# using Programs/freeze_test_frozenmain.py
+	$(RUNSHARED) ./$(BUILDPYTHON) Programs/freeze_test_frozenmain.py Programs/test_frozenmain.h
+
+Programs/test_frozenmain.h: Programs/freeze_test_frozenmain.py Programs/test_frozenmain.py
+	$(MAKE) regen-test-frozenmain
 
 Programs/_testembed: Programs/_testembed.o $(LIBRARY_DEPS)
 	$(LINKCC) $(PY_CORE_LDFLAGS) $(LINKFORSHARED) -o $@ Programs/_testembed.o $(BLDLIBRARY) $(LIBS) $(MODLIBS) $(SYSLIBS)
@@ -763,7 +771,7 @@ regen-limited-abi: all
 
 regen-all: regen-opcode regen-opcode-targets regen-typeslots \
 	regen-token regen-ast regen-keyword regen-importlib clinic \
-	regen-pegen-metaparser regen-pegen regen-frozen
+	regen-pegen-metaparser regen-pegen regen-frozen regen-test-frozenmain
 	@echo
 	@echo "Note: make regen-stdlib-module-names and autoconf should be run manually"
 
@@ -794,7 +802,7 @@ Modules/getpath.o: $(srcdir)/Modules/getpath.c Makefile
 Programs/python.o: $(srcdir)/Programs/python.c
 	$(MAINCC) -c $(PY_CORE_CFLAGS) -o $@ $(srcdir)/Programs/python.c
 
-Programs/_testembed.o: $(srcdir)/Programs/_testembed.c
+Programs/_testembed.o: $(srcdir)/Programs/_testembed.c Programs/test_frozenmain.h
 	$(MAINCC) -c $(PY_CORE_CFLAGS) -o $@ $(srcdir)/Programs/_testembed.c
 
 Modules/_sre.o: $(srcdir)/Modules/_sre.c $(srcdir)/Modules/sre.h $(srcdir)/Modules/sre_constants.h $(srcdir)/Modules/sre_lib.h
diff --git a/Misc/NEWS.d/next/Tests/2021-05-14-14-13-44.bpo-44131.YPizoC.rst b/Misc/NEWS.d/next/Tests/2021-05-14-14-13-44.bpo-44131.YPizoC.rst
new file mode 100644
index 0000000..a646acf8
--- /dev/null
+++ b/Misc/NEWS.d/next/Tests/2021-05-14-14-13-44.bpo-44131.YPizoC.rst
@@ -0,0 +1,2 @@
+Add test_frozenmain to test_embed to test the :c:func:`Py_FrozenMain` C
+function. Patch by Victor Stinner.
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']}")
diff --git a/Python/frozenmain.c b/Python/frozenmain.c
index 5eb9e31..c3104da 100644
--- a/Python/frozenmain.c
+++ b/Python/frozenmain.c
@@ -1,4 +1,3 @@
-
 /* Python interpreter main program for frozen scripts */
 
 #include "Python.h"
@@ -43,10 +42,12 @@ Py_FrozenMain(int argc, char **argv)
     PyConfig_InitPythonConfig(&config);
     config.pathconfig_warnings = 0;   /* Suppress errors from getpath.c */
 
-    if ((p = Py_GETENV("PYTHONINSPECT")) && *p != '\0')
+    if ((p = Py_GETENV("PYTHONINSPECT")) && *p != '\0') {
         inspect = 1;
-    if ((p = Py_GETENV("PYTHONUNBUFFERED")) && *p != '\0')
+    }
+    if ((p = Py_GETENV("PYTHONUNBUFFERED")) && *p != '\0') {
         unbuffered = 1;
+    }
 
     if (unbuffered) {
         setbuf(stdin, (char *)NULL);
@@ -65,8 +66,9 @@ Py_FrozenMain(int argc, char **argv)
         argv_copy[i] = Py_DecodeLocale(argv[i], NULL);
         argv_copy2[i] = argv_copy[i];
         if (!argv_copy[i]) {
-            fprintf(stderr, "Unable to decode the command line argument #%i\n",
-                            i + 1);
+            fprintf(stderr,
+                    "Unable to decode the command line argument #%i\n",
+                    i + 1);
             argc = i;
             goto error;
         }
@@ -97,24 +99,28 @@ Py_FrozenMain(int argc, char **argv)
     PyWinFreeze_ExeInit();
 #endif
 
-    if (Py_VerboseFlag)
+    if (Py_VerboseFlag) {
         fprintf(stderr, "Python %s\n%s\n",
-            Py_GetVersion(), Py_GetCopyright());
+                Py_GetVersion(), Py_GetCopyright());
+    }
 
     PySys_SetArgv(argc, argv_copy);
 
     n = PyImport_ImportFrozenModule("__main__");
-    if (n == 0)
+    if (n == 0) {
         Py_FatalError("the __main__ module is not frozen");
+    }
     if (n < 0) {
         PyErr_Print();
         sts = 1;
     }
-    else
+    else {
         sts = 0;
+    }
 
-    if (inspect && isatty((int)fileno(stdin)))
+    if (inspect && isatty((int)fileno(stdin))) {
         sts = PyRun_AnyFile(stdin, "<stdin>") != 0;
+    }
 
 #ifdef MS_WINDOWS
     PyWinFreeze_ExeTerm();
diff --git a/Tools/freeze/makefreeze.py b/Tools/freeze/makefreeze.py
index 64e3e6b..d7d05db 100644
--- a/Tools/freeze/makefreeze.py
+++ b/Tools/freeze/makefreeze.py
@@ -74,14 +74,12 @@ def makefreeze(base, dict, debug=0, entry_point=None, fail_import=()):
 # Write a C initializer for a module containing the frozen python code.
 # The array is called M_<mod>.
 
-def writecode(outfp, mod, str):
-    outfp.write('unsigned char M_%s[] = {' % mod)
-    for i in range(0, len(str), 16):
-        outfp.write('\n\t')
-        for c in bytes(str[i:i+16]):
-            outfp.write('%d,' % c)
-    outfp.write('\n};\n')
-
-## def writecode(outfp, mod, str):
-##     outfp.write('unsigned char M_%s[%d] = "%s";\n' % (mod, len(str),
-##     '\\"'.join(map(lambda s: repr(s)[1:-1], str.split('"')))))
+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)