From 9e2830d7f0566e0d00b1104eeaedd5032a4e556e Mon Sep 17 00:00:00 2001 From: Mark Tyneway Date: Wed, 17 Jan 2024 18:40:39 +0300 Subject: [PATCH] feat: vm.dumpState (#6827) * feat: vm.dumpState Implements a cheatcode `vm.dumpState(string)` that dumps the current revm state to disk in the same format as the "allocs" field in a geth style `genesis.json`. This can dump state that can be read in by `vm.loadAllocs(string)`. The implementation of the cheatcode skips dumping system contracts. It includes various test coverage. * solidity: forge fmt * spec: update Run `cargo test` from within the specs crate * dumpstate: cleanup Better implementation based on review --- crates/cheatcodes/assets/cheatcodes.json | 20 ++++ crates/cheatcodes/spec/src/vm.rs | 4 + crates/cheatcodes/src/evm.rs | 65 +++++++++++- testdata/cheats/Vm.sol | 1 + testdata/cheats/dumpState.t.sol | 123 +++++++++++++++++++++++ 5 files changed, 210 insertions(+), 3 deletions(-) create mode 100644 testdata/cheats/dumpState.t.sol diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index 7819914cba8c..a88ebf3b5232 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -1153,6 +1153,26 @@ "status": "stable", "safety": "unsafe" }, + { + "func": { + "id": "dumpState", + "description": "Dump a genesis JSON file's `allocs` to disk.", + "declaration": "function dumpState(string calldata pathToStateJson) external;", + "visibility": "external", + "mutability": "", + "signature": "dumpState(string)", + "selector": "0x709ecd3f", + "selectorBytes": [ + 112, + 158, + 205, + 63 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, { "func": { "id": "envAddress_0", diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index d1611222d54b..5178bd83d824 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -227,6 +227,10 @@ interface Vm { #[cheatcode(group = Evm, safety = Safe)] function addr(uint256 privateKey) external pure returns (address keyAddr); + /// Dump a genesis JSON file's `allocs` to disk. + #[cheatcode(group = Evm, safety = Unsafe)] + function dumpState(string calldata pathToStateJson) external; + /// Gets the nonce of an account. #[cheatcode(group = Evm, safety = Safe)] function getNonce(address account) external view returns (uint64 nonce); diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index 24e26def0efd..94a15e6cbb9d 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -3,10 +3,19 @@ use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Result, Vm::*}; use alloy_primitives::{Address, Bytes, U256}; use alloy_sol_types::SolValue; -use ethers_core::utils::{Genesis, GenesisAccount}; +use ethers_core::{ + types::H256, + utils::{Genesis, GenesisAccount}, +}; use ethers_signers::Signer; -use foundry_common::{fs::read_json_file, types::ToAlloy}; -use foundry_evm_core::backend::{DatabaseExt, RevertSnapshotAction}; +use foundry_common::{ + fs::{read_json_file, write_json_file}, + types::{ToAlloy, ToEthers}, +}; +use foundry_evm_core::{ + backend::{DatabaseExt, RevertSnapshotAction}, + constants::{CALLER, CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, TEST_CONTRACT_ADDRESS}, +}; use revm::{ primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY}, EVMData, @@ -86,6 +95,56 @@ impl Cheatcode for loadAllocsCall { } } +impl Cheatcode for dumpStateCall { + fn apply_full(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { pathToStateJson } = self; + let path = Path::new(pathToStateJson); + + // Do not include system accounts in the dump. + let skip = |key: &Address| { + key == &CHEATCODE_ADDRESS || + key == &CALLER || + key == &HARDHAT_CONSOLE_ADDRESS || + key == &TEST_CONTRACT_ADDRESS || + key == &ccx.caller || + key == &ccx.state.config.evm_opts.sender + }; + + let alloc = ccx + .data + .journaled_state + .state() + .into_iter() + .filter(|(key, _)| !skip(key)) + .map(|(key, val)| { + ( + key, + GenesisAccount { + nonce: Some(val.info.nonce), + balance: val.info.balance.to_ethers(), + code: Some( + val.info.code.clone().unwrap_or_default().original_bytes().to_ethers(), + ), + storage: Some( + val.storage + .iter() + .map(|(k, v)| { + let key = k.to_be_bytes::<32>(); + let val = v.present_value().to_be_bytes::<32>(); + (H256::from_slice(&key), H256::from_slice(&val)) + }) + .collect(), + ), + }, + ) + }) + .collect::>(); + + write_json_file(path, &alloc)?; + Ok(Default::default()) + } +} + impl Cheatcode for sign_0Call { fn apply_full(&self, ccx: &mut CheatsCtxt) -> Result { let Self { privateKey, digest } = self; diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index 78ca56eda835..237b68db8835 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -55,6 +55,7 @@ interface Vm { function deriveKey(string calldata mnemonic, uint32 index, string calldata language) external pure returns (uint256 privateKey); function deriveKey(string calldata mnemonic, string calldata derivationPath, uint32 index, string calldata language) external pure returns (uint256 privateKey); function difficulty(uint256 newDifficulty) external; + function dumpState(string calldata pathToStateJson) external; function envAddress(string calldata name) external view returns (address value); function envAddress(string calldata name, string calldata delim) external view returns (address[] memory value); function envBool(string calldata name) external view returns (bool value); diff --git a/testdata/cheats/dumpState.t.sol b/testdata/cheats/dumpState.t.sol new file mode 100644 index 000000000000..51d6eb38ab8f --- /dev/null +++ b/testdata/cheats/dumpState.t.sol @@ -0,0 +1,123 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +pragma solidity 0.8.18; + +import "ds-test/test.sol"; +import "./Vm.sol"; + +contract SimpleContract { + constructor() { + assembly { + sstore(1, 2) + } + } +} + +contract DumpStateTest is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + + function testDumpStateCheatAccount() public { + // Path to temporary file that is deleted after the test + string memory path = string.concat(vm.projectRoot(), "/fixtures/Json/test_dump_state_cheat.json"); + + // Define some values to set in the state using cheatcodes + address target = address(1001); + bytes memory bytecode = hex"11223344"; + uint256 balance = 1.2 ether; + uint64 nonce = 45; + + vm.etch(target, bytecode); + vm.deal(target, balance); + vm.setNonce(target, nonce); + vm.store(target, bytes32(uint256(0x20)), bytes32(uint256(0x40))); + vm.store(target, bytes32(uint256(0x40)), bytes32(uint256(0x60))); + + // Write the state to disk + vm.dumpState(path); + + string memory json = vm.readFile(path); + string[] memory keys = vm.parseJsonKeys(json, ""); + assertEq(keys.length, 1); + + string memory key = keys[0]; + assertEq(nonce, vm.parseJsonUint(json, string.concat(".", key, ".nonce"))); + assertEq(balance, vm.parseJsonUint(json, string.concat(".", key, ".balance"))); + assertEq(bytecode, vm.parseJsonBytes(json, string.concat(".", key, ".code"))); + + string[] memory slots = vm.parseJsonKeys(json, string.concat(".", key, ".storage")); + assertEq(slots.length, 2); + + assertEq( + bytes32(uint256(0x40)), + vm.parseJsonBytes32(json, string.concat(".", key, ".storage.", vm.toString(bytes32(uint256(0x20))))) + ); + assertEq( + bytes32(uint256(0x60)), + vm.parseJsonBytes32(json, string.concat(".", key, ".storage.", vm.toString(bytes32(uint256(0x40))))) + ); + + vm.removeFile(path); + } + + function testDumpStateMultipleAccounts() public { + string memory path = string.concat(vm.projectRoot(), "/fixtures/Json/test_dump_state_multiple_accounts.json"); + + vm.setNonce(address(0x100), 1); + vm.deal(address(0x200), 1 ether); + vm.store(address(0x300), bytes32(uint256(1)), bytes32(uint256(2))); + vm.etch(address(0x400), hex"af"); + + vm.dumpState(path); + + string memory json = vm.readFile(path); + string[] memory keys = vm.parseJsonKeys(json, ""); + assertEq(keys.length, 4); + + assertEq(4, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x100)))).length); + assertEq(1, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x100)), ".nonce"))); + assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x100)), ".balance"))); + assertEq(hex"", vm.parseJsonBytes(json, string.concat(".", vm.toString(address(0x100)), ".code"))); + assertEq(0, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x100)), ".storage")).length); + + assertEq(4, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x200)))).length); + assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x200)), ".nonce"))); + assertEq(1 ether, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x200)), ".balance"))); + assertEq(hex"", vm.parseJsonBytes(json, string.concat(".", vm.toString(address(0x200)), ".code"))); + assertEq(0, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x200)), ".storage")).length); + + assertEq(4, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x300)))).length); + assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x300)), ".nonce"))); + assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x300)), ".balance"))); + assertEq(hex"", vm.parseJsonBytes(json, string.concat(".", vm.toString(address(0x300)), ".code"))); + assertEq(1, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x300)), ".storage")).length); + assertEq( + 2, + vm.parseJsonUint( + json, string.concat(".", vm.toString(address(0x300)), ".storage.", vm.toString(bytes32(uint256(1)))) + ) + ); + + assertEq(4, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x400)))).length); + assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x400)), ".nonce"))); + assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x400)), ".balance"))); + assertEq(hex"af", vm.parseJsonBytes(json, string.concat(".", vm.toString(address(0x400)), ".code"))); + assertEq(0, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x400)), ".storage")).length); + + vm.removeFile(path); + } + + function testDumpStateDeployment() public { + string memory path = string.concat(vm.projectRoot(), "/fixtures/Json/test_dump_state_deployment.json"); + + SimpleContract s = new SimpleContract(); + vm.dumpState(path); + + string memory json = vm.readFile(path); + string[] memory keys = vm.parseJsonKeys(json, ""); + assertEq(keys.length, 1); + assertEq(address(s), vm.parseAddress(keys[0])); + assertEq(1, vm.parseJsonKeys(json, string.concat(".", keys[0], ".storage")).length); + assertEq(2, vm.parseJsonUint(json, string.concat(".", keys[0], ".storage.", vm.toString(bytes32(uint256(1)))))); + + vm.removeFile(path); + } +}