fix: overflow in message_call_gas

kkrt-labs · Sep 30, 2024 · 69abf6e · 69abf6e
1 parent 8b29b2f
commit 69abf6e
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 12 deletions.
diff --git a/crates/evm/src/gas.cairo b/crates/evm/src/gas.cairo
@@ -103,27 +103,32 @@ pub fn max_message_call_gas(gas: u64) -> u64 {
 /// * `memory_cost`: The amount needed to extend the memory in the current frame.
 /// * `extra_gas`: The amount of gas needed for transferring value + creating a new account inside a
 /// message call.
-/// * `call_stipend`: The amount of stipend provided to a message call to execute code while
-/// transferring value(native token).
 ///
 /// # Returns
 ///
-/// * `message_call_gas`: `MessageCallGas`
+/// * `Result<MessageCallGas, EVMError>`: The calculated MessageCallGas or an error if overflow occurs.
 pub fn calculate_message_call_gas(
     value: u256, gas: u64, gas_left: u64, memory_cost: u64, extra_gas: u64
-) -> MessageCallGas {
+) -> Result<MessageCallGas, EVMError> {
     let call_stipend = if value == 0 {
         0
     } else {
         CALL_STIPEND
     };
-    let gas = if gas_left < extra_gas + memory_cost {
+
+    // Check for overflow when adding extra_gas and memory_cost
+    let total_extra_cost = extra_gas.checked_add(memory_cost).ok_or(EVMError::OutOfGas)?;
+    let gas = if gas_left < total_extra_cost {
         gas
     } else {
-        min(gas, max_message_call_gas(gas_left - memory_cost - extra_gas))
+        let remaining_gas = gas_left - total_extra_cost; // Safe because of the check above
+        min(gas, max_message_call_gas(remaining_gas))
     };
 
-    return MessageCallGas { cost: gas + extra_gas, stipend: gas + call_stipend };
+    let cost = gas.checked_add(extra_gas).ok_or(EVMError::OutOfGas)?;
+    let stipend = gas.checked_add(call_stipend).ok_or(EVMError::OutOfGas)?;
+
+    Result::Ok(MessageCallGas { cost, stipend })
 }
 
 

diff --git a/crates/evm/src/instructions/system_operations.cairo b/crates/evm/src/instructions/system_operations.cairo
@@ -9,6 +9,7 @@ use crate::model::vm::{VM, VMTrait};
 use crate::stack::StackTrait;
 use crate::state::StateTrait;
 use utils::set::SetTrait;
+use core::integer::CheckedAdd;
 
 #[generate_trait]
 pub impl SystemOperations of SystemOperationsTrait {
@@ -64,7 +65,7 @@ pub impl SystemOperations of SystemOperationsTrait {
             self.gas_left(),
             memory_expansion.expansion_cost,
             access_gas_cost + transfer_gas_cost + create_gas_cost
-        );
+        )?;
         self.charge_gas(message_call_gas.cost + memory_expansion.expansion_cost)?;
         // Only the transfer gas is left to charge.
 
@@ -139,7 +140,7 @@ pub impl SystemOperations of SystemOperationsTrait {
             self.gas_left(),
             memory_expansion.expansion_cost,
             access_gas_cost + transfer_gas_cost
-        );
+        )?;
         self.charge_gas(message_call_gas.cost + memory_expansion.expansion_cost)?;
 
         // If sender_balance < value, return early, pushing
@@ -211,7 +212,7 @@ pub impl SystemOperations of SystemOperationsTrait {
 
         let message_call_gas = gas::calculate_message_call_gas(
             0, gas, self.gas_left(), memory_expansion.expansion_cost, access_gas_cost
-        );
+        )?;
         self.charge_gas(message_call_gas.cost + memory_expansion.expansion_cost)?;
 
         self
@@ -264,8 +265,12 @@ pub impl SystemOperations of SystemOperationsTrait {
 
         let message_call_gas = gas::calculate_message_call_gas(
             0, gas, self.gas_left(), memory_expansion.expansion_cost, access_gas_cost
-        );
-        self.charge_gas(message_call_gas.cost + memory_expansion.expansion_cost)?;
+        )?;
+        let gas_to_charge = message_call_gas
+            .cost
+            .checked_add(memory_expansion.expansion_cost)
+            .ok_or(EVMError::OutOfGas)?;
+        self.charge_gas(gas_to_charge)?;
 
         self
             .generic_call(