Skip to content

Commit

Permalink
feat: vm.dumpState (#6827)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
tynes authored Jan 17, 2024
1 parent 02f1c1e commit 9e2830d
Show file tree
Hide file tree
Showing 5 changed files with 210 additions and 3 deletions.
20 changes: 20 additions & 0 deletions crates/cheatcodes/assets/cheatcodes.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions crates/cheatcodes/spec/src/vm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
65 changes: 62 additions & 3 deletions crates/cheatcodes/src/evm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -86,6 +95,56 @@ impl Cheatcode for loadAllocsCall {
}
}

impl Cheatcode for dumpStateCall {
fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> 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::<HashMap<_, _>>();

write_json_file(path, &alloc)?;
Ok(Default::default())
}
}

impl Cheatcode for sign_0Call {
fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
let Self { privateKey, digest } = self;
Expand Down
1 change: 1 addition & 0 deletions testdata/cheats/Vm.sol

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

123 changes: 123 additions & 0 deletions testdata/cheats/dumpState.t.sol
Original file line number Diff line number Diff line change
@@ -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);
}
}

0 comments on commit 9e2830d

Please sign in to comment.