chore(blockifier): cap max_validate_gas and max_execute_gas (#2397)

* chore(blockifier): cap max_validate_gas and max_execute_gas

* fix(blockifier): fix tests for sierra_gas usage

crates/blockifier/src/blockifier/stateful_validator.rs

@@ -74,7 +74,7 @@ impl<S: StateReader> StatefulValidator<S> {
 
         // `__validate__` call.
         let (_optional_call_info, actual_cost) =
-            self.validate(&tx, tx_context.initial_sierra_gas())?;
+            self.validate(&tx, tx_context.initial_sierra_gas().0)?;
 
         // Post validations.
         PostValidationReport::verify(&tx_context, &actual_cost)?;

crates/blockifier/src/context.rs

@@ -5,13 +5,15 @@ use papyrus_config::{ParamPath, ParamPrivacyInput, SerializedParam};
 use serde::{Deserialize, Serialize};
 use starknet_api::block::{BlockInfo, FeeType, GasPriceVector};
 use starknet_api::core::{ChainId, ContractAddress};
+use starknet_api::execution_resources::GasAmount;
 use starknet_api::transaction::fields::{
     AllResourceBounds,
     GasVectorComputationMode,
     ValidResourceBounds,
 };
 
 use crate::bouncer::BouncerConfig;
+use crate::execution::call_info::CallInfo;
 use crate::transaction::objects::{
     CurrentTransactionInfo,
     HasRelatedFeeType,
@@ -42,22 +44,49 @@ impl TransactionContext {
 
     /// Returns the initial Sierra gas of the transaction.
     /// This value is used to limit the transaction's run.
-    // TODO(tzahi): replace returned value from u64 to GasAmount.
-    pub fn initial_sierra_gas(&self) -> u64 {
+    pub fn initial_sierra_gas(&self) -> GasAmount {
         match &self.tx_info {
             TransactionInfo::Deprecated(_)
             | TransactionInfo::Current(CurrentTransactionInfo {
                 resource_bounds: ValidResourceBounds::L1Gas(_),
                 ..
-            }) => self.block_context.versioned_constants.default_initial_gas_cost(),
+            }) => self.block_context.versioned_constants.initial_gas_no_user_l2_bound(),
             TransactionInfo::Current(CurrentTransactionInfo {
                 resource_bounds: ValidResourceBounds::AllResources(AllResourceBounds { l2_gas, .. }),
                 ..
-            }) => l2_gas.max_amount.0,
+            }) => l2_gas.max_amount,
         }
     }
 }
 
+pub(crate) struct GasCounter {
+    pub(crate) spent_gas: GasAmount,
+    pub(crate) remaining_gas: GasAmount,
+}
+
+impl GasCounter {
+    pub(crate) fn new(initial_gas: GasAmount) -> Self {
+        GasCounter { spent_gas: GasAmount(0), remaining_gas: initial_gas }
+    }
+
+    fn spend(&mut self, amount: GasAmount) {
+        self.spent_gas = self.spent_gas.checked_add(amount).expect("Gas overflow");
+        self.remaining_gas = self
+            .remaining_gas
+            .checked_sub(amount)
+            .expect("Overuse of gas; should have been caught earlier");
+    }
+
+    /// Limits the amount of gas that can be used (in validate\execute) by the given global limit.
+    pub(crate) fn limit_usage(&self, amount: GasAmount) -> u64 {
+        self.remaining_gas.min(amount).0
+    }
+
+    pub(crate) fn subtract_used_gas(&mut self, call_info: &CallInfo) {
+        self.spend(GasAmount(call_info.execution.gas_consumed));
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct BlockContext {
     // TODO(Yoni, 1/10/2024): consider making these fields public.

crates/blockifier/src/execution/contract_address_test.rs

@@ -34,7 +34,7 @@ fn test_calculate_contract_address() {
             calldata,
             entry_point_selector: selector_from_name("test_contract_address"),
             storage_address: deployer_address,
-            initial_gas: versioned_constants.default_initial_gas_cost(),
+            initial_gas: versioned_constants.inifite_gas_for_vm_mode(),
             ..Default::default()
         };
         let contract_address =

crates/blockifier/src/execution/entry_point.rs

@@ -10,6 +10,7 @@ use starknet_api::abi::abi_utils::selector_from_name;
 use starknet_api::abi::constants::CONSTRUCTOR_ENTRY_POINT_NAME;
 use starknet_api::contract_class::EntryPointType;
 use starknet_api::core::{ClassHash, ContractAddress, EntryPointSelector};
+use starknet_api::execution_resources::GasAmount;
 use starknet_api::state::StorageKey;
 use starknet_api::transaction::fields::{
     AllResourceBounds,
@@ -375,6 +376,10 @@ impl EntryPointExecutionContext {
         &self.versioned_constants().os_constants.gas_costs
     }
 
+    pub fn mode_sierra_gas_limit(&self) -> GasAmount {
+        self.tx_context.block_context.versioned_constants.sierra_gas_limit(&self.execution_mode)
+    }
+
     /// Reverts the state back to the way it was when self.revert_infos.0['revert_idx'] was created.
     pub fn revert(&mut self, revert_idx: usize, state: &mut dyn State) -> StateResult<()> {
         for contract_revert_info in self.revert_infos.0.drain(revert_idx..).rev() {

crates/blockifier/src/execution/entry_point_test.rs

@@ -200,7 +200,7 @@ fn run_security_test(
         entry_point_selector: selector_from_name(entry_point_name),
         calldata,
         storage_address: security_contract.get_instance_address(0),
-        initial_gas: versioned_constants.default_initial_gas_cost(),
+        initial_gas: versioned_constants.inifite_gas_for_vm_mode(),
         ..Default::default()
     };
     let error = match entry_point_call.execute_directly(state) {

crates/blockifier/src/execution/execution_utils.rs

@@ -63,7 +63,7 @@ pub fn execute_entry_point_call_wrapper(
     );
     if current_tracked_resource == TrackedResource::CairoSteps {
         // Override the initial gas with a high value so it won't limit the run.
-        call.initial_gas = context.versioned_constants().default_initial_gas_cost();
+        call.initial_gas = context.versioned_constants().inifite_gas_for_vm_mode();
     }
     let orig_call = call.clone();
 

crates/blockifier/src/fee/fee_test.rs

@@ -361,7 +361,7 @@ fn test_get_fee_by_gas_vector_overflow(
 
 #[rstest]
 #[case::default(
-    VersionedConstants::create_for_account_testing().default_initial_gas_cost(),
+    VersionedConstants::create_for_account_testing().initial_gas_no_user_l2_bound().0,
     GasVectorComputationMode::NoL2Gas
 )]
 #[case::from_l2_gas(4321, GasVectorComputationMode::All)]
@@ -384,6 +384,6 @@ fn test_initial_sierra_gas(
         }),
     };
     let account_tx = invoke_tx_with_default_flags(invoke_tx_args!(resource_bounds));
-    let actual = block_context.to_tx_context(&account_tx).initial_sierra_gas();
+    let actual = block_context.to_tx_context(&account_tx).initial_sierra_gas().0;
     assert_eq!(actual, expected)
 }

crates/blockifier/src/transaction/account_transaction.rs

@@ -21,7 +21,7 @@ use starknet_api::transaction::fields::{
 use starknet_api::transaction::{constants, TransactionHash, TransactionVersion};
 use starknet_types_core::felt::Felt;
 
-use crate::context::{BlockContext, TransactionContext};
+use crate::context::{BlockContext, GasCounter, TransactionContext};
 use crate::execution::call_info::CallInfo;
 use crate::execution::contract_class::RunnableCompiledClass;
 use crate::execution::entry_point::{CallEntryPoint, CallType, EntryPointExecutionContext};
@@ -372,11 +372,19 @@ impl AccountTransaction {
         &self,
         state: &mut dyn State,
         tx_context: Arc<TransactionContext>,
-        remaining_gas: &mut u64,
+        remaining_gas: &mut GasCounter,
     ) -> TransactionExecutionResult<Option<CallInfo>> {
-        let limit_steps_by_resources = self.execution_flags.charge_fee;
         if self.execution_flags.validate {
-            self.validate_tx(state, tx_context, remaining_gas, limit_steps_by_resources)
+            let limit_steps_by_resources = self.execution_flags.charge_fee;
+            // TODO(Aner): cap the gas for validation.
+            let remaining_validation_gas = &mut remaining_gas
+                .limit_usage(tx_context.block_context.versioned_constants.validate_max_sierra_gas);
+            Ok(self
+                .validate_tx(state, tx_context, remaining_validation_gas, limit_steps_by_resources)?
+                .inspect(|call_info| {
+                    // TODO(Aner): Update the gas counter.
+                    remaining_gas.subtract_used_gas(call_info);
+                }))
         } else {
             Ok(None)
         }
@@ -503,20 +511,29 @@ impl AccountTransaction {
         &self,
         state: &mut S,
         context: &mut EntryPointExecutionContext,
-        remaining_gas: &mut u64,
+        remaining_gas: &mut GasCounter,
     ) -> TransactionExecutionResult<Option<CallInfo>> {
-        match &self.tx {
-            Transaction::Declare(tx) => tx.run_execute(state, context, remaining_gas),
-            Transaction::DeployAccount(tx) => tx.run_execute(state, context, remaining_gas),
-            Transaction::Invoke(tx) => tx.run_execute(state, context, remaining_gas),
-        }
+        // TODO(Aner): cap the gas usage for execution.
+        let remaining_execution_gas =
+            &mut remaining_gas.limit_usage(context.mode_sierra_gas_limit());
+        Ok(match &self.tx {
+            Transaction::Declare(tx) => tx.run_execute(state, context, remaining_execution_gas),
+            Transaction::DeployAccount(tx) => {
+                tx.run_execute(state, context, remaining_execution_gas)
+            }
+            Transaction::Invoke(tx) => tx.run_execute(state, context, remaining_execution_gas),
+        }?
+        .inspect(|call_info| {
+            // TODO(Aner): Update the gas counter.
+            remaining_gas.subtract_used_gas(call_info);
+        }))
     }
 
     fn run_non_revertible<S: StateReader>(
         &self,
         state: &mut TransactionalState<'_, S>,
         tx_context: Arc<TransactionContext>,
-        remaining_gas: &mut u64,
+        remaining_gas: &mut GasCounter,
     ) -> TransactionExecutionResult<ValidateExecuteCallInfo> {
         let validate_call_info: Option<CallInfo>;
         let execute_call_info: Option<CallInfo>;
@@ -572,7 +589,7 @@ impl AccountTransaction {
         &self,
         state: &mut TransactionalState<'_, S>,
         tx_context: Arc<TransactionContext>,
-        remaining_gas: &mut u64,
+        remaining_gas: &mut GasCounter,
     ) -> TransactionExecutionResult<ValidateExecuteCallInfo> {
         let mut execution_context = EntryPointExecutionContext::new_invoke(
             tx_context.clone(),
@@ -717,7 +734,7 @@ impl AccountTransaction {
     fn run_or_revert<S: StateReader>(
         &self,
         state: &mut TransactionalState<'_, S>,
-        remaining_gas: &mut u64,
+        remaining_gas: &mut GasCounter,
         tx_context: Arc<TransactionContext>,
     ) -> TransactionExecutionResult<ValidateExecuteCallInfo> {
         if self.is_non_revertible(&tx_context.tx_info) {
@@ -743,7 +760,7 @@ impl<U: UpdatableState> ExecutableTransaction<U> for AccountTransaction {
         self.perform_pre_validation_stage(state, &tx_context, strict_nonce_check)?;
 
         // Run validation and execution.
-        let mut remaining_gas = tx_context.initial_sierra_gas();
+        let initial_gas = tx_context.initial_sierra_gas();
         let ValidateExecuteCallInfo {
             validate_call_info,
             execute_call_info,
@@ -755,7 +772,7 @@ impl<U: UpdatableState> ExecutableTransaction<U> for AccountTransaction {
                     resources: final_resources,
                     gas: total_gas,
                 },
-        } = self.run_or_revert(state, &mut remaining_gas, tx_context.clone())?;
+        } = self.run_or_revert(state, &mut GasCounter::new(initial_gas), tx_context.clone())?;
         let fee_transfer_call_info = Self::handle_fee(
             state,
             tx_context,

crates/blockifier/src/transaction/account_transactions_test.rs

@@ -1691,18 +1691,25 @@ fn test_initial_gas(
         resource_bounds:  default_all_resource_bounds,
         version: TransactionVersion::THREE
     });
+    let user_gas_bound = block_context.to_tx_context(&account_tx).initial_sierra_gas();
 
     let transaction_ex_info = account_tx.execute(state, &block_context).unwrap();
 
     let validate_call_info = &transaction_ex_info.validate_call_info.unwrap();
     let validate_initial_gas = validate_call_info.call.initial_gas;
-    assert_eq!(validate_initial_gas, DEFAULT_L2_GAS_MAX_AMOUNT.0);
+    assert_eq!(validate_initial_gas, block_context.versioned_constants.validate_max_sierra_gas.0);
     let validate_gas_consumed = validate_call_info.execution.gas_consumed;
     assert!(validate_gas_consumed > 0, "New Cairo1 contract should consume gas.");
 
     let default_call_info = CallInfo::default();
-    let mut prev_initial_gas = validate_initial_gas;
     let mut execute_call_info = &transaction_ex_info.execute_call_info.unwrap();
+    // Initial gas for execution is the minimum between the max execution gas and the initial gas
+    // minus the gas consumed by validate. Need to add 1 as the check is strictly less than.
+    let mut prev_initial_gas = block_context
+        .versioned_constants
+        .execute_max_sierra_gas
+        .min(user_gas_bound - GasAmount(validate_gas_consumed) + GasAmount(1))
+        .0;
     let mut curr_initial_gas;
     let mut started_vm_mode = false;
     // The __validate__ call of a the account contract.
@@ -1729,7 +1736,7 @@ fn test_initial_gas(
                 );
                 assert_eq!(
                     curr_initial_gas,
-                    block_context.versioned_constants.default_initial_gas_cost()
+                    block_context.versioned_constants.inifite_gas_for_vm_mode()
                 );
                 started_vm_mode = true;
             }

crates/blockifier/src/transaction/transaction_execution.rs

@@ -144,7 +144,7 @@ impl<U: UpdatableState> ExecutableTransaction<U> for L1HandlerTransaction {
         let limit_steps_by_resources = false;
         let mut context =
             EntryPointExecutionContext::new_invoke(tx_context.clone(), limit_steps_by_resources);
-        let mut remaining_gas = tx_context.initial_sierra_gas();
+        let mut remaining_gas = tx_context.initial_sierra_gas().0;
         let execute_call_info = self.run_execute(state, &mut context, &mut remaining_gas)?;
         let l1_handler_payload_size = self.payload_size();
         let TransactionReceipt {