From 1f39c555a99ded293d12af93fe20f1894a1c2033 Mon Sep 17 00:00:00 2001
From: Linda Guiga <lindaguiga3@gmail.com>
Date: Thu, 13 Apr 2023 18:49:18 +0800
Subject: [PATCH] Address overflow-related TODOs in ASM code, using a macro
 add_or_fault. This is related to
 https://github.com/mir-protocol/plonky2/pull/930/files/a4ea0965d79561c345e2f77836c07949c7e0bc69

---
 evm/src/cpu/kernel/asm/account_code.asm    |  2 +-
 evm/src/cpu/kernel/asm/core/call.asm       |  2 +-
 evm/src/cpu/kernel/asm/core/terminate.asm  |  4 ++--
 evm/src/cpu/kernel/asm/memory/metadata.asm | 12 ++++++++++++
 evm/src/cpu/kernel/asm/memory/syscalls.asm |  2 +-
 evm/src/cpu/kernel/asm/util/keccak.asm     |  2 +-
 6 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/evm/src/cpu/kernel/asm/account_code.asm b/evm/src/cpu/kernel/asm/account_code.asm
index ba19b606..9f950d77 100644
--- a/evm/src/cpu/kernel/asm/account_code.asm
+++ b/evm/src/cpu/kernel/asm/account_code.asm
@@ -106,7 +106,7 @@ global sys_extcodecopy:
     %charge_gas
 
     %stack (kexit_info, address, dest_offset, offset, size) -> (dest_offset, size, kexit_info, address, dest_offset, offset, size)
-    ADD // TODO: check for overflow, see discussion here https://github.com/mir-protocol/plonky2/pull/930/files/a4ea0965d79561c345e2f77836c07949c7e0bc69#r1143630253
+    %add_or_fault
     // stack: expanded_num_bytes, kexit_info, address, dest_offset, offset, size
     DUP1 %ensure_reasonable_offset
     %update_mem_bytes
diff --git a/evm/src/cpu/kernel/asm/core/call.asm b/evm/src/cpu/kernel/asm/core/call.asm
index e9e9d13b..6a80015f 100644
--- a/evm/src/cpu/kernel/asm/core/call.asm
+++ b/evm/src/cpu/kernel/asm/core/call.asm
@@ -401,7 +401,7 @@ global after_call_instruction:
 %macro checked_mem_expansion
     // stack: size, offset, kexit_info
     DUP1 ISZERO %jumpi(%%zero)
-    ADD // TODO: check for overflow
+    %add_or_fault
     // stack: expanded_num_bytes, kexit_info
     DUP1 %ensure_reasonable_offset
     %update_mem_bytes
diff --git a/evm/src/cpu/kernel/asm/core/terminate.asm b/evm/src/cpu/kernel/asm/core/terminate.asm
index a46bbf0f..e078d9b8 100644
--- a/evm/src/cpu/kernel/asm/core/terminate.asm
+++ b/evm/src/cpu/kernel/asm/core/terminate.asm
@@ -14,7 +14,7 @@ global sys_stop:
 global sys_return:
     // stack: kexit_info, offset, size
     %stack (kexit_info, offset, size) -> (offset, size, kexit_info, offset, size)
-    ADD // TODO: Check for overflow?
+    %add_or_fault 
     DUP1 %ensure_reasonable_offset
     %update_mem_bytes
 
@@ -109,7 +109,7 @@ sys_selfdestruct_same_addr:
 global sys_revert:
     // stack: kexit_info, offset, size
     %stack (kexit_info, offset, size) -> (offset, size, kexit_info, offset, size)
-    ADD // TODO: Check for overflow?
+    %add_or_fault
     DUP1 %ensure_reasonable_offset
     %update_mem_bytes
 
diff --git a/evm/src/cpu/kernel/asm/memory/metadata.asm b/evm/src/cpu/kernel/asm/memory/metadata.asm
index 42a88957..70ae7f12 100644
--- a/evm/src/cpu/kernel/asm/memory/metadata.asm
+++ b/evm/src/cpu/kernel/asm/memory/metadata.asm
@@ -312,3 +312,15 @@ global sys_basefee:
     %mload_context_metadata(@CTX_METADATA_STATIC)
     %jumpi(fault_exception)
 %endmacro
+
+%macro add_or_fault
+    // stack: offset, size, kexit_info, offset, size
+    DUP1
+    %ensure_reasonable_offset
+    // stack: offset, size, kexit_info, offset, size
+    DUP2
+    // stack: size, offset, size, kexit_info, offset, size
+    %ensure_reasonable_offset
+    // stack: offset, size, kexit_info, offset, size
+    ADD 
+%endmacro
diff --git a/evm/src/cpu/kernel/asm/memory/syscalls.asm b/evm/src/cpu/kernel/asm/memory/syscalls.asm
index 206f5735..1f70ba91 100644
--- a/evm/src/cpu/kernel/asm/memory/syscalls.asm
+++ b/evm/src/cpu/kernel/asm/memory/syscalls.asm
@@ -131,7 +131,7 @@ sys_calldataload_after_mload_packing:
     DUP5 %num_bytes_to_num_words %mul_const(@GAS_COPY) ADD %charge_gas
 
     %stack (kexit_info, dest_offset, offset, size) -> (dest_offset, size, kexit_info, dest_offset, offset, size)
-    ADD // TODO: check for overflow, see discussion here https://github.com/mir-protocol/plonky2/pull/930/files/a4ea0965d79561c345e2f77836c07949c7e0bc69#r1143630253
+    %add_or_fault
     // stack: expanded_num_bytes, kexit_info, dest_offset, offset, size, kexit_info
     DUP1 %ensure_reasonable_offset
     %update_mem_bytes
diff --git a/evm/src/cpu/kernel/asm/util/keccak.asm b/evm/src/cpu/kernel/asm/util/keccak.asm
index 1d3f985b..1a1f4372 100644
--- a/evm/src/cpu/kernel/asm/util/keccak.asm
+++ b/evm/src/cpu/kernel/asm/util/keccak.asm
@@ -11,7 +11,7 @@ global sys_keccak256:
     // stack: kexit_info, offset, len
 
     %stack (kexit_info, offset, len) -> (offset, len, kexit_info, offset, len)
-    ADD // TODO: need to check for overflow?
+    %add_or_fault
     DUP1 %ensure_reasonable_offset
     %update_mem_bytes