[3.13] gh-148973: fix segfault on mismatch between consts size and oparg in compiler (GH-148974) (#148997)

Author: iritkatriel

Lib/test/test_peepholer.py

@@ -1,3 +1,4 @@
+import ast
 import dis
 from itertools import combinations, product
 import opcode
@@ -6,7 +7,11 @@
 import unittest
 
 from test import support
-from test.support.bytecode_helper import BytecodeTestCase, CfgOptimizationTestCase
+from test.support.bytecode_helper import (
+    BytecodeTestCase,
+    CfgOptimizationTestCase,
+    _testinternalcapi,
+)
 
 
 def compile_pattern_with_fast_locals(pattern):
@@ -962,6 +967,45 @@ def trace(frame, event, arg):
 
 class DirectCfgOptimizerTests(CfgOptimizationTestCase):
 
+    def test_optimize_cfg_const_index_out_of_range(self):
+        insts = [
+            ('LOAD_CONST', 2, 0),
+            ('RETURN_VALUE', None, 0),
+        ]
+        seq = self.seq_from_insts(insts)
+        with self.assertRaisesRegex(ValueError, "out of range"):
+            _testinternalcapi.optimize_cfg(seq, [0, 1], 0)
+
+    def test_optimize_cfg_consts_must_be_list(self):
+        insts = [
+            ('LOAD_CONST', 0, 0),
+            ('RETURN_VALUE', None, 0),
+        ]
+        seq = self.seq_from_insts(insts)
+        with self.assertRaisesRegex(TypeError, "consts must be a list"):
+            _testinternalcapi.optimize_cfg(seq, (0,), 0)
+
+    def test_compiler_codegen_metadata_consts_roundtrips_optimize_cfg(self):
+        tree = ast.parse("x = (1, 2)", mode="exec", optimize=1)
+        insts, meta = _testinternalcapi.compiler_codegen(tree, "<s>", 0)
+        consts = meta["consts"]
+        self.assertIsInstance(consts, list)
+        _testinternalcapi.optimize_cfg(insts, consts, 0)
+
+    def test_compiler_codegen_consts_include_none_required_for_implicit_return(self):
+        tree = ast.parse("pass", mode="exec", optimize=1)
+        insts, meta = _testinternalcapi.compiler_codegen(tree, "<s>", 0)
+        consts = meta["consts"]
+        self.assertEqual(consts, [None])
+
+        load_const = opcode.opmap["LOAD_CONST"]
+        self.assertEqual(
+            [t[1] for t in insts.get_instructions() if t[0] == load_const],
+            [0],
+        )
+
+        _testinternalcapi.optimize_cfg(insts, list(consts), 0)
+
     def cfg_optimization_test(self, insts, expected_insts,
                               consts=None, expected_consts=None,
                               nlocals=0):

Python/compile.c

@@ -7844,6 +7844,7 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
 {
     PyObject *res = NULL;
     PyObject *metadata = NULL;
+    PyObject *consts_list = NULL;
 
     if (!PyAST_Check(ast)) {
         PyErr_SetString(PyExc_TypeError, "expected an AST");
@@ -7889,7 +7890,6 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
 
     SET_MATADATA_ITEM("name", umd->u_name);
     SET_MATADATA_ITEM("qualname", umd->u_qualname);
-    SET_MATADATA_ITEM("consts", umd->u_consts);
     SET_MATADATA_ITEM("names", umd->u_names);
     SET_MATADATA_ITEM("varnames", umd->u_varnames);
     SET_MATADATA_ITEM("cellvars", umd->u_cellvars);
@@ -7915,12 +7915,21 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
     }
 
     if (_PyInstructionSequence_ApplyLabelMap(INSTR_SEQUENCE(c)) < 0) {
-        return NULL;
+        goto finally;
+    }
+    /* After add_return_at_end: const indices match final instruction stream. */
+    consts_list = consts_dict_keys_inorder(umd->u_consts);
+    if (consts_list == NULL) {
+        goto finally;
+    }
+    if (PyDict_SetItemString(metadata, "consts", consts_list) < 0) {
+        goto finally;
     }
     /* Allocate a copy of the instruction sequence on the heap */
     res = PyTuple_Pack(2, INSTR_SEQUENCE(c), metadata);
 
 finally:
+    Py_XDECREF(consts_list);
     Py_XDECREF(metadata);
     compiler_exit_scope(c);
     compiler_free(c);
@@ -7935,6 +7944,10 @@ _PyCompile_OptimizeCfg(PyObject *seq, PyObject *consts, int nlocals)
         PyErr_SetString(PyExc_ValueError, "expected an instruction sequence");
         return NULL;
     }
+    if (!PyList_Check(consts)) {
+        PyErr_SetString(PyExc_TypeError, "consts must be a list");
+        return NULL;
+    }
     PyObject *const_cache = PyDict_New();
     if (const_cache == NULL) {
         return NULL;

Python/flowgraph.c

@@ -1244,6 +1244,14 @@ get_const_value(int opcode, int oparg, PyObject *co_consts)
     PyObject *constant = NULL;
     assert(OPCODE_HAS_CONST(opcode));
     if (opcode == LOAD_CONST) {
+        assert(PyList_Check(co_consts));
+        Py_ssize_t n = PyList_GET_SIZE(co_consts);
+        if (oparg < 0 || oparg >= n) {
+            PyErr_Format(PyExc_ValueError,
+                         "LOAD_CONST index %d is out of range for consts (len=%zd)",
+                         oparg, n);
+            return NULL;
+        }
         constant = PyList_GET_ITEM(co_consts, oparg);
     }
 
@@ -2045,6 +2053,12 @@ remove_unused_consts(basicblock *entryblock, PyObject *consts)
         for (int i = 0; i < b->b_iused; i++) {
             if (OPCODE_HAS_CONST(b->b_instr[i].i_opcode)) {
                 int index = b->b_instr[i].i_oparg;
+                if (index < 0 || index >= nconsts) {
+                    PyErr_Format(PyExc_ValueError,
+                                 "LOAD_CONST index %d is out of range for consts (len=%zd)",
+                                 index, nconsts);
+                    goto end;
+                }
                 index_map[index] = index;
             }
         }