feat(forge): account access cheatcode accounts for extcode* and balance opcodes

crates/cheatcodes/spec/src/vm.rs

@@ -53,6 +53,14 @@ interface Vm {
  SelfDestruct,
  /// Synthetic access indicating the current context has resumed after a previous sub-context (AccountAccess).
  Resume,
+ /// The account's balance was read.
+ Balance,
+ /// The account's codesize was read.
+ Extcodesize,
+ /// The account's code was copied.
+ Extcodecopy,
+ /// The account's codehash was read.
+ Extcodehash,
  }
 
  /// An Ethereum log. Returned by `getRecordedLogs`.

crates/cheatcodes/src/inspector.rs

@@ -505,6 +505,60 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
  data.journaled_state.depth(),
  );
  }
+ // Record account accesses via the EXT family of opcodes
+ opcode::EXTCODECOPY |
+ opcode::EXTCODESIZE |
+ opcode::EXTCODEHASH |
+ opcode::BALANCE => {
+ let kind = match interpreter.current_opcode() {
+ opcode::EXTCODECOPY => crate::Vm::AccountAccessKind::Extcodecopy,
+ opcode::EXTCODESIZE => crate::Vm::AccountAccessKind::Extcodesize,
+ opcode::EXTCODEHASH => crate::Vm::AccountAccessKind::Extcodehash,
+ opcode::BALANCE => crate::Vm::AccountAccessKind::Balance,
+ _ => unreachable!(),
+ };
+ let address = Address::from_word(B256::from(try_or_continue!(interpreter
+ .stack()
+ .peek(0))));
+ let balance;
+ let initialized;
+ if let Ok((acc, _)) = data.journaled_state.load_account(address, data.db) {
+ initialized = acc.info.exists();
+ balance = acc.info.balance;
+ } else {
+ initialized = false;
+ balance = U256::ZERO;
+ }
+ let account_access = crate::Vm::AccountAccess {
+ chainInfo: crate::Vm::ChainInfo {
+ forkId: data.db.active_fork_id().unwrap_or_default(),
+ chainId: U256::from(data.env.cfg.chain_id),
+ },
+ accessor: interpreter.contract().address,
+ account: address,
+ kind,
+ initialized,
+ oldBalance: balance,
+ newBalance: balance,
+ value: U256::ZERO,
+ data: vec![],
+ reverted: false,
+ deployedCode: vec![],
+ storageAccesses: vec![],
+ };
+ let access = AccountAccess {
+ access: account_access,
+ // use current depth; EXT* opcodes are not creating new contexts
+ depth: data.journaled_state.depth(),
+ };
+ // Record the EXT* call as an account access at the current depth
+ // (future storage accesses will be recorded in a new "Resume" context)
+ if let Some(last) = recorded_account_diffs_stack.last_mut() {
+ last.push(access);
+ } else {
+ recorded_account_diffs_stack.push(vec![access]);
+ }
+ }
  _ => (),
  }
  }

testdata/cheats/RecordAccountAccesses.t.sol

@@ -121,6 +121,19 @@ contract NestedRunner {
  }
 }
 
+/// Helper contract that uses all three EXT* opcodes on a given address
+contract ExtChecker {
+ function checkExts(address a) external {
+ assembly {
+ let x := extcodesize(a)
+ let y := extcodehash(a)
+ extcodecopy(a, x, y, 0)
+ // sstore to check that storage accesses are correctly stored in a new access with a "resume" context
+ sstore(0, balance(a))
+ }
+ }
+}
+
 /// @notice Helper contract that writes to storage in a nested call
 contract NestedStorer {
  mapping(bytes32 key => uint256 value) slots;
@@ -196,13 +209,15 @@ contract RecordAccountAccessesTest is DSTest {
  Create2or create2or;
  StorageAccessor test1;
  StorageAccessor test2;
+ ExtChecker extChecker;
 
  function setUp() public {
  runner = new NestedRunner();
  nestedStorer = new NestedStorer();
  create2or = new Create2or();
  test1 = new StorageAccessor();
  test2 = new StorageAccessor();
+ extChecker = new ExtChecker();
  }
 
  function testStorageAccessDelegateCall() public {
@@ -211,7 +226,7 @@ contract RecordAccountAccessesTest is DSTest {
 
  cheats.startStateDiffRecording();
  address(proxy).call(abi.encodeCall(StorageAccessor.read, bytes32(uint256(1234))));
- Vm.AccountAccess[] memory called = cheats.stopAndReturnStateDiff();
+ Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
 
  assertEq(called.length, 2, "incorrect length");
 
@@ -246,7 +261,7 @@ contract RecordAccountAccessesTest is DSTest {
  two.write(bytes32(uint256(5678)), bytes32(uint256(123469)));
  two.write(bytes32(uint256(5678)), bytes32(uint256(1234)));
 
- Vm.AccountAccess[] memory called = cheats.stopAndReturnStateDiff();
+ Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
  assertEq(called.length, 4, "incorrect length");
 
  assertEq(called[0].storageAccesses.length, 1, "incorrect storage length");
@@ -317,7 +332,7 @@ contract RecordAccountAccessesTest is DSTest {
  // contract calls to self in constructor
  SelfCaller caller = new SelfCaller{value: 2 ether}("hello2 world2");
 
- Vm.AccountAccess[] memory called = cheats.stopAndReturnStateDiff();
+ Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
  assertEq(called.length, 6);
  assertEq(
  called[0],
@@ -430,7 +445,7 @@ contract RecordAccountAccessesTest is DSTest {
  uint256 initBalance = address(this).balance;
  cheats.startStateDiffRecording();
  try this.revertingCall{value: 1 ether}(address(1234), "") {} catch {}
- Vm.AccountAccess[] memory called = cheats.stopAndReturnStateDiff();
+ Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
  assertEq(called.length, 2);
  assertEq(
  called[0],
@@ -485,7 +500,7 @@ contract RecordAccountAccessesTest is DSTest {
  /// @param shouldRevert Whether the first call should revert
  function runNested(bool shouldRevert, bool expectFirstCall) public {
  try runner.run{value: 1 ether}(shouldRevert) {} catch {}
- Vm.AccountAccess[] memory called = cheats.stopAndReturnStateDiff();
+ Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
  assertEq(called.length, 7 + toUint(expectFirstCall), "incorrect length");
 
  uint256 startingIndex = toUint(expectFirstCall);
@@ -737,7 +752,7 @@ contract RecordAccountAccessesTest is DSTest {
  function testNestedStorage() public {
  cheats.startStateDiffRecording();
  nestedStorer.run();
- Vm.AccountAccess[] memory called = cheats.stopAndReturnStateDiff();
+ Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
  assertEq(called.length, 3, "incorrect account access length");
 
  assertEq(called[0].storageAccesses.length, 2, "incorrect run storage length");
@@ -858,7 +873,7 @@ contract RecordAccountAccessesTest is DSTest {
  bytes memory creationCode = abi.encodePacked(type(ConstructorStorer).creationCode, abi.encode(true));
  address hypotheticalStorer = deriveCreate2Address(address(create2or), bytes32(0), keccak256(creationCode));
 
- Vm.AccountAccess[] memory called = cheats.stopAndReturnStateDiff();
+ Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
  assertEq(called.length, 3, "incorrect account access length");
  assertEq(toUint(called[0].kind), toUint(Vm.AccountAccessKind.Create), "incorrect kind");
  assertEq(toUint(called[1].kind), toUint(Vm.AccountAccessKind.Call), "incorrect kind");
@@ -967,7 +982,7 @@ contract RecordAccountAccessesTest is DSTest {
  try create2or.create2(bytes32(0), creationCode) {} catch {}
  address hypotheticalAddress = deriveCreate2Address(address(create2or), bytes32(0), keccak256(creationCode));
 
- Vm.AccountAccess[] memory called = cheats.stopAndReturnStateDiff();
+ Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
  assertEq(called.length, 3, "incorrect length");
  assertEq(
  called[1],
@@ -1013,7 +1028,7 @@ contract RecordAccountAccessesTest is DSTest {
  this.startRecordingFromLowerDepth();
  address a = address(new SelfDestructor{value: 1 ether}(address(this)));
  address b = address(new SelfDestructor{value: 1 ether}(address(bytes20("doesn't exist yet"))));
- Vm.AccountAccess[] memory called = cheats.stopAndReturnStateDiff();
+ Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
  assertEq(called.length, 5, "incorrect length");
  assertEq(
  called[1],
@@ -1093,12 +1108,58 @@ contract RecordAccountAccessesTest is DSTest {
  StorageAccessor a = new StorageAccessor();
 
  cheats.stopBroadcast();
- Vm.AccountAccess[] memory called = cheats.stopAndReturnStateDiff();
+ Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
  assertEq(called.length, 1, "incorrect length");
  assertEq(toUint(called[0].kind), toUint(Vm.AccountAccessKind.Create));
  assertEq(called[0].account, address(a));
  }
 
+ /// @notice Test that EXT* opcodes are recorded as account accesses
+ function testExtOpcodes() public {
+ cheats.startStateDiffRecording();
+ extChecker.checkExts(address(1234));
+ Vm.AccountAccess[] memory called = cheats.stopAndReturnStateDiff();
+ assertEq(called.length, 7, "incorrect length");
+ // initial solidity extcodesize check for calling extChecker
+ assertEq(toUint(called[0].kind), toUint(Vm.AccountAccessKind.Extcodesize));
+ // call to extChecker
+ assertEq(toUint(called[1].kind), toUint(Vm.AccountAccessKind.Call));
+ // extChecker checks
+ assertEq(toUint(called[2].kind), toUint(Vm.AccountAccessKind.Extcodesize));
+ assertEq(toUint(called[3].kind), toUint(Vm.AccountAccessKind.Extcodehash));
+ assertEq(toUint(called[4].kind), toUint(Vm.AccountAccessKind.Extcodecopy));
+ assertEq(toUint(called[5].kind), toUint(Vm.AccountAccessKind.Balance));
+ // resume of extChecker to hold SSTORE access
+ assertEq(toUint(called[6].kind), toUint(Vm.AccountAccessKind.Resume));
+ assertEq(called[6].storageAccesses.length, 1, "incorrect length");
+ }
+
+ /**
+ * @notice Filter out extcodesize account accesses for legacy tests written before
+ * EXT* opcodes were supported.
+ */
+ function filterExtcodesizeForLegacyTests(Vm.AccountAccess[] memory inArr)
+ internal
+ pure
+ returns (Vm.AccountAccess[] memory out)
+ {
+ // allocate max length for out array
+ out = new Vm.AccountAccess[](inArr.length);
+ // track end size
+ uint256 size;
+ for (uint256 i = 0; i < inArr.length; ++i) {
+ // only append if not extcodesize
+ if (inArr[i].kind != Vm.AccountAccessKind.Extcodesize) {
+ out[size] = inArr[i];
+ ++size;
+ }
+ }
+ // manually truncate out array
+ assembly {
+ mstore(out, size)
+ }
+ }
+
  function startRecordingFromLowerDepth() external {
  cheats.startStateDiffRecording();
  assembly {