From e2e2d57197cbdd5fac7fc3c197becb8584780d59 Mon Sep 17 00:00:00 2001 From: Karrq Date: Fri, 20 Sep 2024 17:59:43 +0200 Subject: [PATCH] fix(zk): invariant testing (#581) * test(zk): add #565 repro * test(zk:565): fix invariant and use zkVmSkip * fix(zk): use constant test addr for zkvm skip test * test(zk): proper basic invariant test * test(zk:invariant): add direct invariant test * feat: `InZkVm` utility to detect if contract is running in zkVm * chore: lints & formatting * refactor(test:zk): edit runner test opts * fix(test:zk): pre-select target senders for issue565 w/o handler * fix(zk): use `get_test_contract_address` * fix(backend): don't set test contract on each init * chore: lints * fix(executor): set_test_contract even w/o setup * test(zk): filter test to specific contract * refactor(zk): simplify test contract check Co-authored-by: Nisheeth Barthwal * chore: fmt --------- Co-authored-by: Nisheeth Barthwal --- crates/cheatcodes/src/inspector.rs | 18 ++-- crates/evm/core/src/backend/mod.rs | 14 +-- crates/evm/evm/src/executors/mod.rs | 2 +- crates/forge/src/runner.rs | 1 + crates/forge/tests/cli/test_cmd.rs | 2 +- crates/forge/tests/it/zk/invariant.rs | 10 +- crates/forge/tests/it/zk/repros.rs | 7 ++ crates/script/src/runner.rs | 2 +- testdata/zk/Globals.sol | 2 + testdata/zk/InZkVm.sol | 31 ++++++ testdata/zk/InvariantDeposit.t.sol | 53 +++++++-- testdata/zk/repros/Issue565.t.sol | 149 ++++++++++++++++++++++++++ 12 files changed, 256 insertions(+), 35 deletions(-) create mode 100644 testdata/zk/InZkVm.sol create mode 100644 testdata/zk/repros/Issue565.t.sol diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index d3acedfa5..974c3c8c0 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -48,7 +48,7 @@ use revm::{ }, primitives::{ AccountInfo, BlockEnv, Bytecode, CreateScheme, EVMError, Env, EvmStorageSlot, - ExecutionResult, HashMap as rHashMap, Output, TransactTo, KECCAK_EMPTY, + ExecutionResult, HashMap as rHashMap, Output, KECCAK_EMPTY, }, EvmContext, InnerEvmContext, Inspector, }; @@ -1530,11 +1530,17 @@ impl Cheatcodes { return None; } - if let TransactTo::Call(test_contract) = ecx.env.tx.transact_to { - if call.bytecode_address == test_contract { - info!("running call in EVM, instead of zkEVM (Test Contract) {:#?}", ecx.env.tx); - return None - } + if ecx + .db + .get_test_contract_address() + .map(|addr| call.bytecode_address == addr) + .unwrap_or_default() + { + info!( + "running call in EVM, instead of zkEVM (Test Contract) {:#?}", + call.bytecode_address + ); + return None } info!("running call in zkEVM {:#?}", call); diff --git a/crates/evm/core/src/backend/mod.rs b/crates/evm/core/src/backend/mod.rs index 36f184dc2..70d55f048 100644 --- a/crates/evm/core/src/backend/mod.rs +++ b/crates/evm/core/src/backend/mod.rs @@ -24,7 +24,7 @@ use revm::{ precompile::{PrecompileSpecId, Precompiles}, primitives::{ Account, AccountInfo, Bytecode, Env, EnvWithHandlerCfg, EvmState, EvmStorageSlot, - HashMap as Map, Log, ResultAndState, SpecId, TxKind, KECCAK_EMPTY, + HashMap as Map, Log, ResultAndState, SpecId, KECCAK_EMPTY, }, Database, DatabaseCommit, JournaledState, }; @@ -771,18 +771,6 @@ impl Backend { pub(crate) fn initialize(&mut self, env: &EnvWithHandlerCfg) { self.set_caller(env.tx.caller); self.set_spec_id(env.handler_cfg.spec_id); - - let test_contract = match env.tx.transact_to { - TxKind::Call(to) => to, - TxKind::Create => { - let nonce = self - .basic_ref(env.tx.caller) - .map(|b| b.unwrap_or_default().nonce) - .unwrap_or_default(); - env.tx.caller.create(nonce) - } - }; - self.set_test_contract(test_contract); } /// Returns the `EnvWithHandlerCfg` with the current `spec_id` set. diff --git a/crates/evm/evm/src/executors/mod.rs b/crates/evm/evm/src/executors/mod.rs index a75b81639..eb658aed1 100644 --- a/crates/evm/evm/src/executors/mod.rs +++ b/crates/evm/evm/src/executors/mod.rs @@ -328,7 +328,7 @@ impl Executor { trace!(?from, ?to, "setting up contract"); let from = from.unwrap_or(CALLER); - self.backend_mut().set_test_contract(to).set_caller(from); + self.backend_mut().set_caller(from); let calldata = Bytes::from_static(&ITest::setUpCall::SELECTOR); let mut res = self.transact_raw(from, to, calldata, U256::ZERO)?; res = res.into_result(rd)?; diff --git a/crates/forge/src/runner.rs b/crates/forge/src/runner.rs index 55fb9b653..2d27fa29f 100644 --- a/crates/forge/src/runner.rs +++ b/crates/forge/src/runner.rs @@ -118,6 +118,7 @@ impl<'a> ContractRunner<'a> { } let address = self.sender.create(self.executor.get_nonce(self.sender)?); + self.executor.backend_mut().set_test_contract(address); // Set the contracts initial balance before deployment, so it is available during // construction diff --git a/crates/forge/tests/cli/test_cmd.rs b/crates/forge/tests/cli/test_cmd.rs index 205e89737..c06c431c4 100644 --- a/crates/forge/tests/cli/test_cmd.rs +++ b/crates/forge/tests/cli/test_cmd.rs @@ -1048,7 +1048,7 @@ contract CallEmptyCode is Test { "#, ) .unwrap(); - cmd.args(["test", "--zksync", "--evm-version", "shanghai"]); + cmd.args(["test", "--zksync", "--evm-version", "shanghai", "--mc", "CallEmptyCode"]); let output = cmd.stdout_lossy(); assert!(output.contains("call may fail or behave unexpectedly due to empty code")); diff --git a/crates/forge/tests/it/zk/invariant.rs b/crates/forge/tests/it/zk/invariant.rs index b1519b16c..eb7bb23db 100644 --- a/crates/forge/tests/it/zk/invariant.rs +++ b/crates/forge/tests/it/zk/invariant.rs @@ -6,8 +6,14 @@ use foundry_test_utils::Filter; #[tokio::test(flavor = "multi_thread")] async fn test_zk_invariant_deposit() { - let runner = TEST_DATA_DEFAULT.runner_zksync(); - let filter = Filter::new("testZkInvariantDeposit", "ZkInvariantTest", ".*"); + let mut runner = TEST_DATA_DEFAULT.runner_zksync(); + + // FIXME: just use the inline config + runner.test_options.invariant.no_zksync_reserved_addresses = true; + runner.test_options.invariant.fail_on_revert = true; + runner.test_options.invariant.runs = 10; + + let filter = Filter::new(".*", "ZkInvariantTest", ".*"); TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await; } diff --git a/crates/forge/tests/it/zk/repros.rs b/crates/forge/tests/it/zk/repros.rs index 23c9b6ee6..febc28cbd 100644 --- a/crates/forge/tests/it/zk/repros.rs +++ b/crates/forge/tests/it/zk/repros.rs @@ -36,3 +36,10 @@ async fn repro_config( // https://github.com/matter-labs/foundry-zksync/issues/497 test_repro!(497); + +test_repro!(565; |cfg| { + // FIXME: just use the inline config + cfg.runner.test_options.invariant.no_zksync_reserved_addresses = true; + cfg.runner.test_options.invariant.fail_on_revert = true; + cfg.runner.test_options.invariant.runs = 2; +}); diff --git a/crates/script/src/runner.rs b/crates/script/src/runner.rs index 300f57809..a54c5e549 100644 --- a/crates/script/src/runner.rs +++ b/crates/script/src/runner.rs @@ -129,6 +129,7 @@ impl ScriptRunner { }; let address = CALLER.create(self.executor.get_nonce(CALLER)?); + self.executor.backend_mut().set_test_contract(address); // Set the contracts initial balance before deployment, so it is available during the // construction @@ -147,7 +148,6 @@ impl ScriptRunner { // Optionally call the `setUp` function let (success, gas_used, labeled_addresses, transactions) = if !setup { - self.executor.backend_mut().set_test_contract(address); (true, 0, Default::default(), Some(library_transactions)) } else { match self.executor.setup(Some(self.evm_opts.sender), address, None) { diff --git a/testdata/zk/Globals.sol b/testdata/zk/Globals.sol index 5f45e65ca..f8eecc1ac 100644 --- a/testdata/zk/Globals.sol +++ b/testdata/zk/Globals.sol @@ -5,4 +5,6 @@ library Globals { string public constant ETHEREUM_MAINNET_URL = "https://eth-mainnet.alchemyapi.io/v2/Lc7oIGYeL_QvInzI0Wiu_pOZZDEKBrdf"; // trufflehog:ignore string public constant ZKSYNC_MAINNET_URL = "mainnet"; + + address public constant SYSTEM_CONTEXT_ADDR = address(0x000000000000000000000000000000000000800B); } diff --git a/testdata/zk/InZkVm.sol b/testdata/zk/InZkVm.sol new file mode 100644 index 000000000..3c825da90 --- /dev/null +++ b/testdata/zk/InZkVm.sol @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity >=0.8.7 <0.9.0; + +import {Globals} from "./Globals.sol"; + +// excerpt from system-contracts +interface ISystemContext { + function chainId() external view returns (uint256); +} + +library InZkVmLib { + function _inZkVm() internal returns (bool) { + (bool success, bytes memory retdata) = + Globals.SYSTEM_CONTEXT_ADDR.call(abi.encodeWithSelector(ISystemContext.chainId.selector)); + + return success; + } +} + +abstract contract InZkVm { + modifier inZkVm() { + require(InZkVmLib._inZkVm(), "must be executed in zkVM"); + _; + } +} + +abstract contract DeployOnlyInZkVm is InZkVm { + constructor() { + require(InZkVmLib._inZkVm(), "must be deployed in zkVM"); + } +} diff --git a/testdata/zk/InvariantDeposit.t.sol b/testdata/zk/InvariantDeposit.t.sol index b5b1da6fc..e25189935 100644 --- a/testdata/zk/InvariantDeposit.t.sol +++ b/testdata/zk/InvariantDeposit.t.sol @@ -5,25 +5,56 @@ import "ds-test/test.sol"; import "../cheats/Vm.sol"; import "./Deposit.sol"; -contract ZkInvariantTest is DSTest { +// partial from forge-std/StdInvariant.sol +abstract contract StdInvariant { + struct FuzzSelector { + address addr; + bytes4[] selectors; + } + + address[] internal _targetedContracts; + + function targetContracts() public view returns (address[] memory) { + return _targetedContracts; + } + + FuzzSelector[] internal _targetedSelectors; + + function targetSelectors() public view returns (FuzzSelector[] memory) { + return _targetedSelectors; + } + + address[] internal _targetedSenders; + + function targetSenders() public view returns (address[] memory) { + return _targetedSenders; + } +} + +contract ZkInvariantTest is DSTest, StdInvariant { Vm constant vm = Vm(HEVM_ADDRESS); - // forge-config: default.invariant.runs = 2 Deposit deposit; + uint256 constant dealAmount = 1 ether; + function setUp() external { + // to fund for fees + _targetedSenders.push(address(65536 + 1)); + _targetedSenders.push(address(65536 + 12)); + _targetedSenders.push(address(65536 + 123)); + _targetedSenders.push(address(65536 + 1234)); + + for (uint256 i = 0; i < _targetedSenders.length; i++) { + vm.deal(_targetedSenders[i], dealAmount); // to pay fees + } + deposit = new Deposit(); - vm.deal(address(deposit), 100 ether); + _targetedContracts.push(address(deposit)); } + //FIXME: seems to not be detected, forcing values in test config // forge-config: default.invariant.runs = 2 - function testZkInvariantDeposit() external payable { - deposit.deposit{value: 1 ether}(); - uint256 balanceBefore = deposit.balance(address(this)); - assertEq(balanceBefore, 1 ether); - deposit.withdraw(); - uint256 balanceAfter = deposit.balance(address(this)); - assertGt(balanceBefore, balanceAfter); - } + function invariant_itWorks() external payable {} receive() external payable {} } diff --git a/testdata/zk/repros/Issue565.t.sol b/testdata/zk/repros/Issue565.t.sol new file mode 100644 index 000000000..bda54fbd2 --- /dev/null +++ b/testdata/zk/repros/Issue565.t.sol @@ -0,0 +1,149 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +pragma solidity ^0.8.18; + +import "ds-test/test.sol"; +import "cheats/Vm.sol"; +import {Globals} from "../Globals.sol"; +import {DeployOnlyInZkVm} from "../InZkVm.sol"; + +import "../../default/logs/console.sol"; + +contract Counter is DeployOnlyInZkVm { + uint256 public number; + + function inc() public inZkVm { + number += 1; + } + + function reset() public inZkVm { + number = 0; + } +} + +contract CounterHandler is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + + uint256 public incCounter; + uint256 public resetCounter; + bool public isResetLast; + Counter public counter; + + constructor(Counter _counter) { + counter = _counter; + } + + function inc() public { + console.log("inc"); + incCounter += 1; + isResetLast = false; + + vm.deal(tx.origin, 1 ether); // ensure caller has funds + counter.inc(); + } + + function reset() public { + console.log("reset"); + resetCounter += 1; + isResetLast = true; + + vm.deal(tx.origin, 1 ether); // ensure caller has funds + counter.reset(); + } +} + +// partial from forge-std/StdInvariant.sol +abstract contract StdInvariant { + struct FuzzSelector { + address addr; + bytes4[] selectors; + } + + address[] internal _targetedContracts; + + function targetContracts() public view returns (address[] memory) { + return _targetedContracts; + } + + FuzzSelector[] internal _targetedSelectors; + + function targetSelectors() public view returns (FuzzSelector[] memory) { + return _targetedSelectors; + } + + address[] internal _targetedSenders; + + function targetSenders() public view returns (address[] memory) { + return _targetedSenders; + } +} + +// https://github.com/matter-labs/foundry-zksync/issues/565 +contract Issue565 is DSTest, StdInvariant { + Vm constant vm = Vm(HEVM_ADDRESS); + Counter cnt; + CounterHandler handler; + + function setUp() public { + cnt = new Counter(); + + vm.zkVmSkip(); + handler = new CounterHandler(cnt); + + // add the handler selectors to the fuzzing targets + bytes4[] memory selectors = new bytes4[](2); + selectors[0] = CounterHandler.inc.selector; + selectors[1] = CounterHandler.reset.selector; + + _targetedContracts.push(address(handler)); + _targetedSelectors.push(FuzzSelector({addr: address(handler), selectors: selectors})); + } + + //FIXME: seems to not be detected, forcing values in test config + /// forge-config: default.invariant.fail-on-revert = true + /// forge-config: default.invariant.no-zksync-reserved-addresses = true + function invariant_ghostVariables() external { + uint256 num = cnt.number(); + + if (handler.resetCounter() == 0) { + assert(handler.incCounter() == num); + } else if (handler.isResetLast()) { + assert(num == 0); + } else { + assert(num != 0); + } + } +} + +contract Issue565WithoutHandler is DSTest, StdInvariant { + Vm constant vm = Vm(HEVM_ADDRESS); + Counter cnt; + + uint256 constant dealAmount = 1 ether; + + function setUp() public { + cnt = new Counter(); + + // so we can fund them ahead of time for fees + _targetedSenders.push(address(65536 + 1)); + _targetedSenders.push(address(65536 + 12)); + _targetedSenders.push(address(65536 + 123)); + _targetedSenders.push(address(65536 + 1234)); + + for (uint256 i = 0; i < _targetedSenders.length; i++) { + vm.deal(_targetedSenders[i], dealAmount); + } + + // add the handler selectors to the fuzzing targets + bytes4[] memory selectors = new bytes4[](2); + selectors[0] = Counter.inc.selector; + selectors[1] = Counter.reset.selector; + + _targetedContracts.push(address(cnt)); + _targetedSelectors.push(FuzzSelector({addr: address(cnt), selectors: selectors})); + } + + //FIXME: seems to not be detected, forcing values in test config + /// forge-config: default.invariant.fail-on-revert = true + /// forge-config: default.invariant.no-zksync-reserved-addresses = true + function invariant_itWorks() external {} +}