vm.store: implement store cheatcode (#353)

Signed-off-by: Alexandru Gheorghe <[email protected]>

Author: alexggh

crates/cheatcodes/src/error.rs

@@ -74,7 +74,7 @@ macro_rules! ensure_not_precompile {
 }
 
 #[cold]
-pub(crate) fn precompile_error(address: &Address) -> Error {
+pub fn precompile_error(address: &Address) -> Error {
     fmt_err!("cannot use precompile {address} as an argument")
 }
 

crates/cheatcodes/src/lib.rs

@@ -22,7 +22,7 @@ use spec::Status;
 
 pub use Vm::ForgeContext;
 pub use config::CheatsConfig;
-pub use error::{Error, ErrorKind, Result};
+pub use error::{Error, ErrorKind, Result, precompile_error};
 pub use inspector::{
     BroadcastableTransaction, BroadcastableTransactions, Cheatcodes, CheatcodesExecutor,
 };
@@ -181,7 +181,7 @@ impl std::ops::DerefMut for CheatsCtxt<'_, '_, '_, '_> {
 
 impl CheatsCtxt<'_, '_, '_, '_> {
     #[inline]
-    pub(crate) fn is_precompile(&self, address: &Address) -> bool {
+    pub fn is_precompile(&self, address: &Address) -> bool {
         self.ecx.journaled_state.inner.precompiles.contains(address)
     }
 }

crates/forge/tests/it/revive/cheat_store.rs

@@ -0,0 +1,27 @@
+use crate::{config::*, test_helpers::TEST_DATA_REVIVE};
+use foundry_test_utils::Filter;
+use revm::primitives::hardfork::SpecId;
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_store_works() {
+    let runner: forge::MultiContractRunner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testStoreWorks", "Store", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_store_fuzzed() {
+    let runner: forge::MultiContractRunner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testStoreFuzzed", "Store", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_store_not_available_on_precompiles() {
+    let runner: forge::MultiContractRunner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testStoreNotAvailableOnPrecompiles", "Store", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}

crates/forge/tests/it/revive/mod.rs

@@ -1,3 +1,4 @@
 //! Revive strategy tests
 
+pub mod cheat_store;
 pub mod migration;

crates/revive-strategy/src/cheatcodes/mod.rs

@@ -4,12 +4,12 @@ use alloy_sol_types::SolValue;
 use foundry_cheatcodes::{
     Broadcast, BroadcastableTransactions, CheatcodeInspectorStrategy,
     CheatcodeInspectorStrategyContext, CheatcodeInspectorStrategyRunner, CheatsConfig, CheatsCtxt,
-    CommonCreateInput, DealRecord, Ecx, EvmCheatcodeInspectorStrategyRunner, Result,
+    CommonCreateInput, DealRecord, Ecx, Error, EvmCheatcodeInspectorStrategyRunner, Result,
     Vm::{
         dealCall, getNonce_0Call, loadCall, pvmCall, rollCall, setNonceCall, setNonceUnsafeCall,
-        warpCall,
+        storeCall, warpCall,
     },
-    journaled_account,
+    journaled_account, precompile_error,
 };
 use foundry_common::sh_err;
 use foundry_compilers::resolc::dual_compiled_contracts::DualCompiledContracts;
@@ -299,6 +299,26 @@ impl CheatcodeInspectorStrategyRunner for PvmCheatcodeInspectorStrategyRunner {
                     .unwrap_or(B256::ZERO);
                 Ok(result.abi_encode())
             }
+            t if using_pvm && is::<storeCall>(t) => {
+                tracing::info!(cheatcode = ?cheatcode.as_debug() , using_pvm = ?using_pvm);
+                let &storeCall { target, slot, value } = cheatcode.as_any().downcast_ref().unwrap();
+                if ccx.is_precompile(&target) {
+                    return Err(precompile_error(&target));
+                }
+
+                let target_address_h160 = H160::from_slice(target.as_slice());
+                let _ = execute_with_externalities(|externalities| {
+                    externalities.execute_with(|| {
+                        Pallet::<Runtime>::set_storage(
+                            target_address_h160,
+                            slot.into(),
+                            Some(value.to_vec()),
+                        )
+                    })
+                })
+                .map_err(|_| <&str as Into<Error>>::into("Could not set storage"))?;
+                Ok(Default::default())
+            }
             // Not custom, just invoke the default behavior
             _ => cheatcode.dyn_apply(ccx, executor),
         }

testdata/default/revive/Store.t.sol

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity ^0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+import "../../default/logs/console.sol";
+
+contract Storage {
+    uint256 public slot0 = 10;
+    uint256 public slot1 = 20;
+}
+
+contract StoreTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    Storage store;
+
+    function setUp() public {
+        vm.pvm(true);
+        store = new Storage();
+    }
+
+    function testStoreWorks() public {
+        assertEq(store.slot0(), 10, "initial value for slot 0 is incorrect");
+        assertEq(store.slot1(), 20, "initial value for slot 1 is incorrect");
+
+        vm.store(address(store), bytes32(0), bytes32(uint256(1)));
+        assertEq(store.slot0(), 1, "store failed");
+        assertEq(store.slot1(), 20, "store failed");
+    }
+
+    function testStoreNotAvailableOnPrecompiles() public {
+        assertEq(store.slot0(), 10, "initial value for slot 0 is incorrect");
+        assertEq(store.slot1(), 20, "initial value for slot 1 is incorrect");
+
+        vm._expectCheatcodeRevert("cannot use precompile 0x0000000000000000000000000000000000000001 as an argument");
+        vm.store(address(1), bytes32(0), bytes32(uint256(1)));
+    }
+
+    function testStoreFuzzed(uint256 slot0, uint256 slot1) public {
+        assertEq(store.slot0(), 10, "initial value for slot 0 is incorrect");
+        assertEq(store.slot1(), 20, "initial value for slot 1 is incorrect");
+
+        vm.store(address(store), bytes32(0), bytes32(slot0));
+        vm.store(address(store), bytes32(uint256(1)), bytes32(slot1));
+        assertEq(store.slot0(), slot0, "store failed");
+        assertEq(store.slot1(), slot1, "store failed");
+    }
+}