gh-141976: Check stack bounds in JIT optimizer (GH-142201)

Author: Fidget-Spinner

Lib/test/test_generated_cases.py

@@ -2115,6 +2115,7 @@ def test_validate_uop_unused_input(self):
         """
         output = """
         case OP: {
+            CHECK_STACK_BOUNDS(-1);
             stack_pointer += -1;
             ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
             break;
@@ -2132,6 +2133,7 @@ def test_validate_uop_unused_input(self):
         """
         output = """
         case OP: {
+            CHECK_STACK_BOUNDS(-1);
             stack_pointer += -1;
             ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
             break;
@@ -2153,6 +2155,7 @@ def test_validate_uop_unused_output(self):
         case OP: {
             JitOptRef foo;
             foo = NULL;
+            CHECK_STACK_BOUNDS(1);
             stack_pointer[0] = foo;
             stack_pointer += 1;
             ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
@@ -2172,6 +2175,7 @@ def test_validate_uop_unused_output(self):
         """
         output = """
         case OP: {
+            CHECK_STACK_BOUNDS(1);
             stack_pointer += 1;
             ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
             break;

Misc/NEWS.d/next/Core_and_Builtins/2025-12-02-21-11-46.gh-issue-141976.yu7pDV.rst

@@ -0,0 +1 @@
+Check against abstract stack overflow in the JIT optimizer.

Python/optimizer_analysis.c

@@ -144,10 +144,6 @@ incorrect_keys(PyObject *obj, uint32_t version)
 
 #define CURRENT_FRAME_IS_INIT_SHIM() (ctx->frame->code == ((PyCodeObject *)&_Py_InitCleanup))
 
-#define WITHIN_STACK_BOUNDS() \
-    (CURRENT_FRAME_IS_INIT_SHIM() || (STACK_LEVEL() >= 0 && STACK_LEVEL() <= STACK_SIZE()))
-
-
 #define GETLOCAL(idx)          ((ctx->frame->locals[idx]))
 
 #define REPLACE_OP(INST, OP, ARG, OPERAND)    \
@@ -192,6 +188,27 @@ incorrect_keys(PyObject *obj, uint32_t version)
 
 #define JUMP_TO_LABEL(label) goto label;
 
+static int
+check_stack_bounds(JitOptContext *ctx, JitOptRef *stack_pointer, int offset, int opcode)
+{
+    int stack_level = (int)(stack_pointer + (offset) - ctx->frame->stack);
+    int should_check = !CURRENT_FRAME_IS_INIT_SHIM() ||
+        (opcode == _RETURN_VALUE) ||
+        (opcode == _RETURN_GENERATOR) ||
+        (opcode == _YIELD_VALUE);
+    if (should_check && (stack_level < 0 || stack_level > STACK_SIZE())) {
+        ctx->contradiction = true;
+        ctx->done = true;
+        return 1;
+    }
+    return 0;
+}
+
+#define CHECK_STACK_BOUNDS(offset) \
+    if (check_stack_bounds(ctx, stack_pointer, offset, opcode)) { \
+        break; \
+    } \
+
 static int
 optimize_to_bool(
     _PyUOpInstruction *this_instr,