feat(vm): Support missed storage invocation limit in fast VM

core/lib/multivm/src/versions/testonly/mod.rs

@@ -22,8 +22,8 @@ use zksync_types::{
 };
 
 pub(super) use self::tester::{
-    validation_params, TestedVm, TestedVmForValidation, TestedVmWithCallTracer, VmTester,
-    VmTesterBuilder,
+    validation_params, TestedVm, TestedVmForValidation, TestedVmWithCallTracer,
+    TestedVmWithStorageLimit, VmTester, VmTesterBuilder,
 };
 use crate::{
     interface::{

core/lib/multivm/src/versions/testonly/storage.rs

@@ -1,9 +1,15 @@
+use assert_matches::assert_matches;
 use ethabi::Token;
-use zksync_test_contracts::TestContract;
-use zksync_types::{Address, Execute, U256};
+use zksync_test_contracts::{LoadnextContractExecutionParams, TestContract, TxType};
+use zksync_types::{AccountTreeId, Address, Execute, StorageKey, H256, U256};
 
-use super::{tester::VmTesterBuilder, ContractToDeploy, TestedVm};
-use crate::interface::{InspectExecutionMode, TxExecutionMode, VmInterfaceExt};
+use super::{
+    get_empty_storage, tester::VmTesterBuilder, ContractToDeploy, TestedVm,
+    TestedVmWithStorageLimit, VmTester,
+};
+use crate::interface::{
+    ExecutionResult, Halt, InspectExecutionMode, TxExecutionMode, VmInterfaceExt,
+};
 
 fn test_storage<VM: TestedVm>(first_tx_calldata: Vec<u8>, second_tx_calldata: Vec<u8>) -> u32 {
     let bytecode = TestContract::storage_test().bytecode.to_vec();
@@ -107,3 +113,93 @@ pub(crate) fn test_transient_storage_behavior<VM: TestedVm>() {
 
     test_storage::<VM>(first_tstore_test, second_tstore_test);
 }
+
+pub(crate) fn test_limiting_storage_writes<VM: TestedVmWithStorageLimit>(should_stop: bool) {
+    let bytecode = TestContract::expensive().bytecode.to_vec();
+    let test_address = Address::repeat_byte(1);
+
+    let mut vm: VmTester<VM> = VmTesterBuilder::new()
+        .with_empty_in_memory_storage()
+        .with_execution_mode(TxExecutionMode::VerifyExecute)
+        .with_rich_accounts(1)
+        .with_custom_contracts(vec![ContractToDeploy::new(bytecode, test_address)])
+        .build();
+
+    let account = &mut vm.rich_accounts[0];
+    let test_fn = TestContract::expensive().function("expensive");
+    let (executed_writes, limit) = if should_stop {
+        (1_000, 100)
+    } else {
+        (100, 1_000)
+    };
+    let tx = account.get_l2_tx_for_execute(
+        Execute {
+            contract_address: Some(test_address),
+            calldata: test_fn
+                .encode_input(&[Token::Uint(executed_writes.into())])
+                .unwrap(),
+            value: 0.into(),
+            factory_deps: vec![],
+        },
+        None,
+    );
+    vm.vm.push_transaction(tx);
+
+    let result = vm.vm.execute_with_storage_limit(limit);
+    if should_stop {
+        assert_matches!(
+            &result.result,
+            ExecutionResult::Halt {
+                reason: Halt::TracerCustom(_)
+            }
+        );
+    } else {
+        assert!(!result.result.is_failed(), "{:?}", result.result);
+    }
+}
+
+pub(crate) fn test_limiting_storage_reads<VM: TestedVmWithStorageLimit>(should_stop: bool) {
+    let bytecode = TestContract::load_test().bytecode.to_vec();
+    let test_address = Address::repeat_byte(1);
+    let (executed_reads, limit) = if should_stop {
+        (1_000, 100)
+    } else {
+        (100, 1_000)
+    };
+    let mut storage = get_empty_storage();
+    // Set the read array length in the load test contract.
+    storage.set_value(
+        StorageKey::new(AccountTreeId::new(test_address), H256::zero()),
+        H256::from_low_u64_be(executed_reads),
+    );
+
+    let mut vm: VmTester<VM> = VmTesterBuilder::new()
+        .with_storage(storage)
+        .with_execution_mode(TxExecutionMode::VerifyExecute)
+        .with_rich_accounts(1)
+        .with_custom_contracts(vec![ContractToDeploy::new(bytecode, test_address)])
+        .build();
+
+    let account = &mut vm.rich_accounts[0];
+    let tx = account.get_loadnext_transaction(
+        test_address,
+        LoadnextContractExecutionParams {
+            reads: executed_reads as usize,
+            ..LoadnextContractExecutionParams::empty()
+        },
+        TxType::L2,
+    );
+    vm.vm.push_transaction(tx);
+
+    let result = vm.vm.execute_with_storage_limit(limit);
+    if should_stop {
+        assert_matches!(
+            &result.result,
+            ExecutionResult::Halt {
+                reason: Halt::TracerCustom(_)
+            }
+        );
+    } else {
+        assert!(!result.result.is_failed(), "{:?}", result.result);
+    }
+}

core/lib/multivm/src/versions/testonly/tester/mod.rs

@@ -257,3 +257,7 @@ pub(crate) fn validation_params(tx: &L2Tx, system: &SystemEnv) -> ValidationPara
 pub(crate) trait TestedVmWithCallTracer: TestedVm {
     fn inspect_with_call_tracer(&mut self) -> (VmExecutionResultAndLogs, Vec<Call>);
 }
+
+pub(crate) trait TestedVmWithStorageLimit: TestedVm {
+    fn execute_with_storage_limit(&mut self, limit: usize) -> VmExecutionResultAndLogs;
+}

core/lib/multivm/src/versions/vm_fast/mod.rs

@@ -2,7 +2,7 @@ pub use zksync_vm2::interface;
 
 pub(crate) use self::version::FastVmVersion;
 pub use self::{
-    tracers::{CallTracer, FullValidationTracer, ValidationTracer},
+    tracers::{CallTracer, FullValidationTracer, StorageInvocationsTracer, ValidationTracer},
     vm::Vm,
 };
 

core/lib/multivm/src/versions/vm_fast/tests/mod.rs

@@ -4,21 +4,21 @@ use zksync_types::{
     h256_to_u256, l2::L2Tx, writes::StateDiffRecord, StorageKey, Transaction, H160, H256, U256,
 };
 use zksync_vm2::interface::{Event, HeapId, StateInterface, Tracer};
-use zksync_vm_interface::{
-    pubdata::{PubdataBuilder, PubdataInput},
-    storage::ReadStorage,
-    tracer::ViolatedValidationRule,
-    Call, CurrentExecutionState, InspectExecutionMode, L2BlockEnv, VmExecutionMode,
-    VmExecutionResultAndLogs, VmInterface,
-};
 
 use super::{FullValidationTracer, ValidationTracer, Vm};
 use crate::{
-    interface::storage::{ImmutableStorageView, InMemoryStorage},
+    interface::{
+        pubdata::{PubdataBuilder, PubdataInput},
+        storage::{ImmutableStorageView, InMemoryStorage, ReadStorage, StorageView},
+        tracer::ViolatedValidationRule,
+        Call, CurrentExecutionState, InspectExecutionMode, L2BlockEnv, VmExecutionMode,
+        VmExecutionResultAndLogs, VmInterface,
+    },
     versions::testonly::{
         validation_params, TestedVm, TestedVmForValidation, TestedVmWithCallTracer,
+        TestedVmWithStorageLimit,
     },
-    vm_fast::{tracers::WithBuiltinTracers, CallTracer},
+    vm_fast::{tracers::WithBuiltinTracers, CallTracer, StorageInvocationsTracer},
 };
 
 mod account_validation_rules;
@@ -200,3 +200,13 @@ impl TestedVmWithCallTracer for TestedFastVm<CallTracer, ()> {
         (result, tracer.0.into_result())
     }
 }
+
+type TestStorageLimiter = StorageInvocationsTracer<StorageView<InMemoryStorage>>;
+
+impl TestedVmWithStorageLimit for TestedFastVm<TestStorageLimiter, ()> {
+    fn execute_with_storage_limit(&mut self, limit: usize) -> VmExecutionResultAndLogs {
+        let storage = self.world.storage.to_rc_ptr();
+        let mut tracer = (StorageInvocationsTracer::new(storage, limit), ());
+        self.inspect(&mut tracer, InspectExecutionMode::OneTx)
+    }
+}

core/lib/multivm/src/versions/vm_fast/tests/storage.rs

@@ -1,5 +1,8 @@
 use crate::{
-    versions::testonly::storage::{test_storage_behavior, test_transient_storage_behavior},
+    versions::testonly::storage::{
+        test_limiting_storage_reads, test_limiting_storage_writes, test_storage_behavior,
+        test_transient_storage_behavior,
+    },
     vm_fast::Vm,
 };
 
@@ -12,3 +15,19 @@ fn storage_behavior() {
 fn transient_storage_behavior() {
     test_transient_storage_behavior::<Vm<_>>();
 }
+
+#[test]
+fn limiting_storage_writes() {
+    println!("Sanity check");
+    test_limiting_storage_writes::<Vm<_, _>>(false);
+    println!("Storage limiter check");
+    test_limiting_storage_writes::<Vm<_, _>>(true);
+}
+
+#[test]
+fn limiting_storage_reads() {
+    println!("Sanity check");
+    test_limiting_storage_reads::<Vm<_, _>>(false);
+    println!("Storage limiter check");
+    test_limiting_storage_reads::<Vm<_, _>>(true);
+}

core/lib/multivm/src/versions/vm_fast/tracers/mod.rs

@@ -5,6 +5,7 @@ use zksync_vm2::interface::{CycleStats, GlobalStateInterface, OpcodeType, Should
 pub(super) use self::evm_deploy::DynamicBytecodes;
 pub use self::{
     calls::CallTracer,
+    storage::StorageInvocationsTracer,
     validation::{FullValidationTracer, ValidationTracer},
 };
 use self::{circuits::CircuitsTracer, evm_deploy::EvmDeployTracer};
@@ -13,6 +14,7 @@ use crate::interface::CircuitStatistic;
 mod calls;
 mod circuits;
 mod evm_deploy;
+mod storage;
 mod validation;
 
 #[derive(Debug)]

core/lib/multivm/src/versions/vm_fast/tracers/storage.rs

@@ -0,0 +1,46 @@
+use zksync_vm2::interface::{GlobalStateInterface, Opcode, OpcodeType, ShouldStop, Tracer};
+
+use crate::interface::storage::{StoragePtr, WriteStorage};
+
+/// Tracer that stops VM execution once too many storage slots were read from the underlying storage.
+/// Useful as an anti-DoS measure.
+#[derive(Debug)]
+pub struct StorageInvocationsTracer<ST> {
+    // FIXME: this is very awkward; logically, storage should be obtainable from `GlobalStateInterface` (but how?)
+    storage: Option<StoragePtr<ST>>,
+    limit: usize,
+}
+
+impl<ST: WriteStorage> Default for StorageInvocationsTracer<ST> {
+    fn default() -> Self {
+        Self {
+            storage: None,
+            limit: usize::MAX,
+        }
+    }
+}
+
+impl<ST: WriteStorage> StorageInvocationsTracer<ST> {
+    pub fn new(storage: StoragePtr<ST>, limit: usize) -> Self {
+        Self {
+            storage: Some(storage),
+            limit,
+        }
+    }
+}
+
+impl<ST: WriteStorage> Tracer for StorageInvocationsTracer<ST> {
+    #[inline(always)]
+    fn after_instruction<OP: OpcodeType, S: GlobalStateInterface>(
+        &mut self,
+        _state: &mut S,
+    ) -> ShouldStop {
+        if matches!(OP::VALUE, Opcode::StorageRead | Opcode::StorageWrite) {
+            let storage = self.storage.as_ref().unwrap();
+            if storage.borrow().missed_storage_invocations() >= self.limit {
+                return ShouldStop::Stop;
+            }
+        }
+        ShouldStop::Continue
+    }
+}

core/lib/multivm/src/versions/vm_latest/tests/mod.rs

@@ -24,11 +24,11 @@ use crate::{
         tracer::ViolatedValidationRule,
         CurrentExecutionState, L2BlockEnv, VmExecutionMode, VmExecutionResultAndLogs,
     },
-    tracers::{CallTracer, ValidationTracer},
+    tracers::{CallTracer, StorageInvocations, ValidationTracer},
     utils::bytecode::bytes_to_be_words,
     versions::testonly::{
         filter_out_base_system_contracts, validation_params, TestedVm, TestedVmForValidation,
-        TestedVmWithCallTracer,
+        TestedVmWithCallTracer, TestedVmWithStorageLimit,
     },
     vm_latest::{
         constants::BOOTLOADER_HEAP_PAGE,
@@ -350,3 +350,10 @@ impl TestedVmWithCallTracer for TestedLatestVm {
         (res, traces)
     }
 }
+
+impl TestedVmWithStorageLimit for TestedLatestVm {
+    fn execute_with_storage_limit(&mut self, limit: usize) -> VmExecutionResultAndLogs {
+        let tracer = StorageInvocations::new(limit).into_tracer_pointer();
+        self.inspect(&mut tracer.into(), InspectExecutionMode::OneTx)
+    }
+}

core/lib/multivm/src/versions/vm_latest/tests/storage.rs

@@ -1,5 +1,8 @@
 use crate::{
-    versions::testonly::storage::{test_storage_behavior, test_transient_storage_behavior},
+    versions::testonly::storage::{
+        test_limiting_storage_reads, test_limiting_storage_writes, test_storage_behavior,
+        test_transient_storage_behavior,
+    },
     vm_latest::{HistoryEnabled, Vm},
 };
 
@@ -12,3 +15,19 @@ fn storage_behavior() {
 fn transient_storage_behavior() {
     test_transient_storage_behavior::<Vm<_, HistoryEnabled>>();
 }
+
+#[test]
+fn limiting_storage_writes() {
+    println!("Sanity check");
+    test_limiting_storage_writes::<Vm<_, HistoryEnabled>>(false);
+    println!("Storage limiter check");
+    test_limiting_storage_writes::<Vm<_, HistoryEnabled>>(true);
+}
+
+#[test]
+fn limiting_storage_reads() {
+    println!("Sanity check");
+    test_limiting_storage_reads::<Vm<_, HistoryEnabled>>(false);
+    println!("Storage limiter check");
+    test_limiting_storage_reads::<Vm<_, HistoryEnabled>>(true);
+}