From 204b3d7415f191652d03a41abd6a4576b5204c34 Mon Sep 17 00:00:00 2001
From: Mathieu <60658558+enitrat@users.noreply.github.com>
Date: Wed, 20 Nov 2024 16:45:45 +0700
Subject: [PATCH] [KGA-51] fix: underconstrained hint in initialize_jumpdests
 (#1626)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fixes an underconstrained hint in initialize_jumpdests. Avoids a
malicious prover to continue iterating over the bytecode limit and
consuming extra resources, but has no impact over code execution

https://github.com/code-423n4/2024-09-kakarot-findings/issues/26

---------

Co-authored-by: Clément Walter <clement0walter@gmail.com>
---
 cairo_zero/tests/src/utils/test_utils.py | 13 +++++++++++++
 cairo_zero/utils/utils.cairo             |  5 +++++
 tests/utils/hints.py                     | 19 ++++++++++++++++++-
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/cairo_zero/tests/src/utils/test_utils.py b/cairo_zero/tests/src/utils/test_utils.py
index 42930b926..171c75564 100644
--- a/cairo_zero/tests/src/utils/test_utils.py
+++ b/cairo_zero/tests/src/utils/test_utils.py
@@ -161,6 +161,19 @@ async def test_should_return_same_as_execution_specs(self, cairo_run):
         output = cairo_run("test__initialize_jumpdests", bytecode=bytecode)
         assert set(output) == get_valid_jump_destinations(bytecode)
 
+    async def test_should_err_on_malicious_prover(self, cairo_program, cairo_run):
+        with (
+            patch_hint(
+                cairo_program,
+                "memory[ap] = 0 if 0 <= (ids.a % PRIME) < range_check_builtin.bound else 1",
+                "memory[ap] = 1",
+                "initialize_jumpdests",
+            ),
+            cairo_error(message="Reading out of bounds bytecode"),
+        ):
+            bytecode = (await get_contract("PlainOpcodes", "Counter")).bytecode_runtime
+            cairo_run("test__initialize_jumpdests", bytecode=bytecode)
+
 
 class TestLoadPackedBytes:
     def test_should_load_packed_bytes(self, cairo_run):
diff --git a/cairo_zero/utils/utils.cairo b/cairo_zero/utils/utils.cairo
index 8095eeb08..ee136ec17 100644
--- a/cairo_zero/utils/utils.cairo
+++ b/cairo_zero/utils/utils.cairo
@@ -910,6 +910,11 @@ namespace Helpers {
         let valid_jumpdests = cast([ap - 2], DictAccess*);
         let i = [ap - 1];
 
+        with_attr error_message("Reading out of bounds bytecode") {
+            assert [range_check_ptr] = bytecode_len - 1 - i;
+        }
+        let range_check_ptr = range_check_ptr + 1;
+
         tempvar opcode = [bytecode + i];
         let is_opcode_ge_0x5f = Helpers.is_le_unchecked(0x5f, opcode);
         let is_opcode_le_0x7f = Helpers.is_le_unchecked(opcode, 0x7f);
diff --git a/tests/utils/hints.py b/tests/utils/hints.py
index 735e5b748..217b75eb2 100644
--- a/tests/utils/hints.py
+++ b/tests/utils/hints.py
@@ -1,5 +1,6 @@
 from collections import defaultdict
 from contextlib import contextmanager
+from typing import Optional
 from unittest.mock import patch
 
 from starkware.cairo.common.dict import DictTracker
@@ -33,12 +34,28 @@ def new_default_dict(
 
 
 @contextmanager
-def patch_hint(program, hint, new_hint):
+def patch_hint(program, hint, new_hint, scope: Optional[str] = None):
+    """
+    Patch a Cairo hint in a program with a new hint.
+
+    Args:
+        program: The Cairo program containing the hints
+        hint: The original hint code to replace
+        new_hint: The new hint code to use instead
+        scope: Optional scope name to restrict which hints are patched. If provided,
+              only hints in scope containing this string will be patched.
+
+    Example:
+        with patch_hint(program, "old_hint", "new_hint", "initialize_jumpdests"):
+            # Code that runs with the patched hint
+
+    """
     patched_hints = {
         k: [
             (
                 hint_
                 if hint_.code != hint
+                or (scope is not None and scope not in str(hint_.accessible_scopes[-1]))
                 else CairoHint(
                     accessible_scopes=hint_.accessible_scopes,
                     flow_tracking_data=hint_.flow_tracking_data,