fix: high offset memory expansion computations (#239)

Close #152 Note to reviewer: this is a fix imported from C4 mitigation, ensure the fix was correctly ported by looking at the corresponding issue and PR.
kkrt-labs · Dec 9, 2024 · bec0337 · bec0337
1 parent 5e773bf
commit bec0337
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 12 deletions.
diff --git a/cairo/src/gas.cairo b/cairo/src/gas.cairo
@@ -1,7 +1,7 @@
 from starkware.cairo.common.math import split_felt
 from starkware.cairo.common.math_cmp import is_not_zero, is_nn, is_le_felt
 from starkware.cairo.common.bool import FALSE
-from starkware.cairo.common.uint256 import Uint256, uint256_lt
+from starkware.cairo.common.uint256 import Uint256, uint256_lt, uint256_add
 
 from src.model import model
 from src.utils.uint256 import uint256_eq
@@ -150,19 +150,41 @@ namespace Gas {
     ) -> model.MemoryExpansion {
         alloc_locals;
 
-        let (is_zero_1) = uint256_eq([size_1], Uint256(0, 0));
-        let (is_zero_2) = uint256_eq([size_2], Uint256(0, 0));
-        tempvar both_zero = is_zero_1 * is_zero_2;
+        let (is_zero_size_1) = uint256_eq([size_1], Uint256(0, 0));
+        let (is_zero_size_2) = uint256_eq([size_2], Uint256(0, 0));
+        tempvar both_zero = is_zero_size_1 * is_zero_size_2;
         jmp no_expansion if both_zero != 0;
 
-        tempvar is_not_saturated = Helpers.is_zero(offset_1.high) * Helpers.is_zero(size_1.high) *
-            Helpers.is_zero(offset_2.high) * Helpers.is_zero(size_2.high);
-        tempvar is_saturated = 1 - is_not_saturated;
-        tempvar range_check_ptr = range_check_ptr;
-        jmp expansion_cost_saturated if is_saturated != 0;
+        if (is_zero_size_1 == FALSE) {
+            let (max_offset_1_res, carry_1) = uint256_add([offset_1], [size_1]);
+            tempvar is_chunk_1_saturated = carry_1 + is_not_zero(max_offset_1_res.high);
+            tempvar max_offset_1 = max_offset_1_res.low;
+            tempvar range_check_ptr = range_check_ptr;
+
+            // Early jump if the chunk is saturated
+            jmp expansion_cost_saturated if is_chunk_1_saturated != 0;
+        } else {
+            tempvar max_offset_1 = 0;
+            tempvar range_check_ptr = range_check_ptr;
+        }
+        let max_offset_1 = [ap - 2];
+        let range_check_ptr = [ap - 1];
+
+        if (is_zero_size_2 == FALSE) {
+            let (max_offset_2_res, carry_2) = uint256_add([offset_2], [size_2]);
+            tempvar is_chunk_2_saturated = carry_2 + is_not_zero(max_offset_2_res.high);
+            tempvar max_offset_2 = max_offset_2_res.low;
+            tempvar range_check_ptr = range_check_ptr;
+
+            // Early jump if the chunk is saturated
+            jmp expansion_cost_saturated if is_chunk_2_saturated != 0;
+        } else {
+            tempvar max_offset_2 = 0;
+            tempvar range_check_ptr = range_check_ptr;
+        }
+        let max_offset_2 = [ap - 2];
+        let range_check_ptr = [ap - 1];
 
-        let max_offset_1 = (1 - is_zero_1) * (offset_1.low + size_1.low);
-        let max_offset_2 = (1 - is_zero_2) * (offset_2.low + size_2.low);
         let max_expansion_is_2 = is_le_felt(max_offset_1, max_offset_2);
         let max_offset = max_offset_1 * (1 - max_expansion_is_2) + max_offset_2 *
             max_expansion_is_2;

diff --git a/cairo/tests/src/test_gas.py b/cairo/tests/src/test_gas.py
@@ -1,5 +1,5 @@
 import pytest
-from hypothesis import given
+from hypothesis import example, given
 from hypothesis.strategies import integers
 
 from ethereum.cancun.vm.gas import (
@@ -44,6 +44,7 @@ def test_should_return_correct_expansion_cost(
                 min_value=0, max_value=0x3C0000
             ),  # upper bound reaching 30M gas limit on expansion
         )
+        @example(offset_1=2**200, size_1=0, offset_2=0, size_2=1, words_len=1)
         def test_should_return_max_expansion_cost(
             self, cairo_run, offset_1, size_1, offset_2, size_2, words_len
         ):