From 3a9e67a9fdb4fad13bf42df6eb91126ab2ef45a1 Mon Sep 17 00:00:00 2001
From: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>
Date: Thu, 15 Feb 2024 14:32:21 +0000
Subject: [PATCH] gh-115376: fix segfault in _testinternalcapi.compiler_codegen
 on bad input (#115379)

---
 Lib/test/test_compiler_codegen.py             |  7 +++-
 ...-02-12-22-35-01.gh-issue-115376.n9vubZ.rst |  1 +
 Python/compile.c                              | 42 ++++++++++++-------
 3 files changed, 35 insertions(+), 15 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Tests/2024-02-12-22-35-01.gh-issue-115376.n9vubZ.rst

diff --git a/Lib/test/test_compiler_codegen.py b/Lib/test/test_compiler_codegen.py
index dbeadd9ca47..166294a40c1 100644
--- a/Lib/test/test_compiler_codegen.py
+++ b/Lib/test/test_compiler_codegen.py
@@ -8,7 +8,7 @@ class IsolatedCodeGenTests(CodegenTestCase):
 
     def codegen_test(self, snippet, expected_insts):
         import ast
-        a = ast.parse(snippet, "my_file.py", "exec");
+        a = ast.parse(snippet, "my_file.py", "exec")
         insts = self.generate_code(a)
         self.assertInstructionsMatch(insts, expected_insts)
 
@@ -54,3 +54,8 @@ class IsolatedCodeGenTests(CodegenTestCase):
             ('RETURN_VALUE', None),
         ]
         self.codegen_test(snippet, expected)
+
+    def test_syntax_error__return_not_in_function(self):
+        snippet = "return 42"
+        with self.assertRaisesRegex(SyntaxError, "'return' outside function"):
+            self.codegen_test(snippet, None)
diff --git a/Misc/NEWS.d/next/Tests/2024-02-12-22-35-01.gh-issue-115376.n9vubZ.rst b/Misc/NEWS.d/next/Tests/2024-02-12-22-35-01.gh-issue-115376.n9vubZ.rst
new file mode 100644
index 00000000000..e09d78a9c4b
--- /dev/null
+++ b/Misc/NEWS.d/next/Tests/2024-02-12-22-35-01.gh-issue-115376.n9vubZ.rst
@@ -0,0 +1 @@
+Fix segfault in ``_testinternalcapi.compiler_codegen`` on bad input.
diff --git a/Python/compile.c b/Python/compile.c
index f95e3cd8a79..d857239690e 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1735,16 +1735,10 @@ compiler_body(struct compiler *c, location loc, asdl_stmt_seq *stmts)
 static int
 compiler_codegen(struct compiler *c, mod_ty mod)
 {
-    _Py_DECLARE_STR(anon_module, "<module>");
-    RETURN_IF_ERROR(
-        compiler_enter_scope(c, &_Py_STR(anon_module), COMPILER_SCOPE_MODULE,
-                             mod, 1));
-
     location loc = LOCATION(1, 1, 0, 0);
     switch (mod->kind) {
     case Module_kind:
         if (compiler_body(c, loc, mod->v.Module.body) < 0) {
-            compiler_exit_scope(c);
             return ERROR;
         }
         break;
@@ -1753,10 +1747,10 @@ compiler_codegen(struct compiler *c, mod_ty mod)
             ADDOP(c, loc, SETUP_ANNOTATIONS);
         }
         c->c_interactive = 1;
-        VISIT_SEQ_IN_SCOPE(c, stmt, mod->v.Interactive.body);
+        VISIT_SEQ(c, stmt, mod->v.Interactive.body);
         break;
     case Expression_kind:
-        VISIT_IN_SCOPE(c, expr, mod->v.Expression.body);
+        VISIT(c, expr, mod->v.Expression.body);
         break;
     default:
         PyErr_Format(PyExc_SystemError,
@@ -1767,14 +1761,29 @@ compiler_codegen(struct compiler *c, mod_ty mod)
     return SUCCESS;
 }
 
+static int
+compiler_enter_anonymous_scope(struct compiler* c, mod_ty mod)
+{
+    _Py_DECLARE_STR(anon_module, "<module>");
+    RETURN_IF_ERROR(
+        compiler_enter_scope(c, &_Py_STR(anon_module), COMPILER_SCOPE_MODULE,
+                             mod, 1));
+    return SUCCESS;
+}
+
 static PyCodeObject *
 compiler_mod(struct compiler *c, mod_ty mod)
 {
+    PyCodeObject *co = NULL;
     int addNone = mod->kind != Expression_kind;
-    if (compiler_codegen(c, mod) < 0) {
+    if (compiler_enter_anonymous_scope(c, mod) < 0) {
         return NULL;
     }
-    PyCodeObject *co = optimize_and_assemble(c, addNone);
+    if (compiler_codegen(c, mod) < 0) {
+        goto finally;
+    }
+    co = optimize_and_assemble(c, addNone);
+finally:
     compiler_exit_scope(c);
     return co;
 }
@@ -7920,15 +7929,20 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
         return NULL;
     }
 
+    metadata = PyDict_New();
+    if (metadata == NULL) {
+        return NULL;
+    }
+
+    if (compiler_enter_anonymous_scope(c, mod) < 0) {
+        return NULL;
+    }
     if (compiler_codegen(c, mod) < 0) {
         goto finally;
     }
 
     _PyCompile_CodeUnitMetadata *umd = &c->u->u_metadata;
-    metadata = PyDict_New();
-    if (metadata == NULL) {
-        goto finally;
-    }
+
 #define SET_MATADATA_ITEM(key, value) \
     if (value != NULL) { \
         if (PyDict_SetItemString(metadata, key, value) < 0) goto finally; \