From df011e0d7caa2c81eb54fb24d10c7ac5fbccb713 Mon Sep 17 00:00:00 2001 From: Nisheeth Barthwal Date: Mon, 16 Sep 2024 16:00:00 +0200 Subject: [PATCH] feat: add cheatcode to skip zkEVM execution (#569) * panic with bytecode hash * add zkVmSkip cheatcode * sticky addresses deployed in skip --- crates/cheatcodes/assets/cheatcodes.json | 82 +++++---- crates/cheatcodes/spec/src/vm.rs | 6 + crates/cheatcodes/src/inspector.rs | 212 +++++++++++++++++------ crates/cheatcodes/src/test.rs | 8 + crates/forge/tests/it/zk/cheats.rs | 8 + crates/zksync/compiler/src/zksolc/mod.rs | 6 +- testdata/cheats/Vm.sol | 1 + testdata/zk/Cheatcodes.t.sol | 57 ++++++ 8 files changed, 298 insertions(+), 82 deletions(-) diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index f9456dc8a..9b4dad514 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -8631,26 +8631,6 @@ "status": "stable", "safety": "safe" }, - { - "func": { - "id": "zkRegisterContract", - "description": "Registers bytecodes for ZK-VM for transact/call and create instructions.", - "declaration": "function zkRegisterContract(string calldata name, bytes32 evmBytecodeHash, bytes calldata evmDeployedBytecode, bytes calldata evmBytecode, bytes32 zkBytecodeHash, bytes calldata zkDeployedBytecode) external pure;", - "visibility": "external", - "mutability": "pure", - "signature": "zkRegisterContract(string,bytes32,bytes,bytes,bytes32,bytes)", - "selector": "0x428cb039", - "selectorBytes": [ - 66, - 140, - 176, - 57 - ] - }, - "group": "testing", - "status": "stable", - "safety": "safe" - }, { "func": { "id": "writeToml_0", @@ -8671,6 +8651,46 @@ "status": "stable", "safety": "safe" }, + { + "func": { + "id": "writeToml_1", + "description": "Takes serialized JSON, converts to TOML and write a serialized TOML table to an **existing** TOML file, replacing a value with key = \nThis is useful to replace a specific value of a TOML file, without having to parse the entire thing.", + "declaration": "function writeToml(string calldata json, string calldata path, string calldata valueKey) external;", + "visibility": "external", + "mutability": "", + "signature": "writeToml(string,string,string)", + "selector": "0x51ac6a33", + "selectorBytes": [ + 81, + 172, + 106, + 51 + ] + }, + "group": "toml", + "status": "stable", + "safety": "safe" + }, + { + "func": { + "id": "zkRegisterContract", + "description": "Registers bytecodes for ZK-VM for transact/call and create instructions.", + "declaration": "function zkRegisterContract(string calldata name, bytes32 evmBytecodeHash, bytes calldata evmDeployedBytecode, bytes calldata evmBytecode, bytes32 zkBytecodeHash, bytes calldata zkDeployedBytecode) external pure;", + "visibility": "external", + "mutability": "pure", + "signature": "zkRegisterContract(string,bytes32,bytes,bytes,bytes32,bytes)", + "selector": "0x428cb039", + "selectorBytes": [ + 66, + 140, + 176, + 57 + ] + }, + "group": "testing", + "status": "stable", + "safety": "safe" + }, { "func": { "id": "zkVm", @@ -8693,21 +8713,21 @@ }, { "func": { - "id": "writeToml_1", - "description": "Takes serialized JSON, converts to TOML and write a serialized TOML table to an **existing** TOML file, replacing a value with key = \nThis is useful to replace a specific value of a TOML file, without having to parse the entire thing.", - "declaration": "function writeToml(string calldata json, string calldata path, string calldata valueKey) external;", + "id": "zkVmSkip", + "description": "When running in zkEVM context, skips the next CREATE or CALL, executing it on the EVM instead.\nAll `CREATE`s executed within this skip, will automatically have `CALL`s to their target addresses\nexecuted in the EVM, and need not be marked with this cheatcode at every usage location.", + "declaration": "function zkVmSkip() external pure;", "visibility": "external", - "mutability": "", - "signature": "writeToml(string,string,string)", - "selector": "0x51ac6a33", + "mutability": "pure", + "signature": "zkVmSkip()", + "selector": "0x99c48bb9", "selectorBytes": [ - 81, - 172, - 106, - 51 + 153, + 196, + 139, + 185 ] }, - "group": "toml", + "group": "testing", "status": "stable", "safety": "safe" } diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index 7fe6812a7..5e4d567f2 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -694,6 +694,12 @@ interface Vm { #[cheatcode(group = Testing, safety = Safe)] function zkVm(bool enable) external pure; + /// When running in zkEVM context, skips the next CREATE or CALL, executing it on the EVM instead. + /// All `CREATE`s executed within this skip, will automatically have `CALL`s to their target addresses + /// executed in the EVM, and need not be marked with this cheatcode at every usage location. + #[cheatcode(group = Testing, safety = Safe)] + function zkVmSkip() external pure; + /// Registers bytecodes for ZK-VM for transact/call and create instructions. #[cheatcode(group = Testing, safety = Safe)] function zkRegisterContract(string calldata name, bytes32 evmBytecodeHash, bytes calldata evmDeployedBytecode, bytes calldata evmBytecode, bytes32 zkBytecodeHash, bytes calldata zkDeployedBytecode) external pure; diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index 616787055..bbfe7cfd6 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -51,7 +51,7 @@ use revm::{ use rustc_hash::FxHashMap; use serde_json::Value; use std::{ - collections::{BTreeMap, HashMap, VecDeque}, + collections::{BTreeMap, HashMap, HashSet, VecDeque}, fs::File, io::BufReader, ops::Range, @@ -348,6 +348,17 @@ pub struct Cheatcodes { /// Use ZK-VM to execute CALLs and CREATEs. pub use_zk_vm: bool, + /// When in zkEVM context, execute the next CALL or CREATE in the EVM instead. + pub skip_zk_vm: bool, + + /// Any contracts that were deployed in `skip_zk_vm` step. + /// This makes it easier to dispatch calls to any of these addresses in zkEVM context, directly + /// to EVM. Alternatively, we'd need to add `vm.zkVmSkip()` to these calls manually. + pub skip_zk_vm_addresses: HashSet
, + + /// Records the next create address for `skip_zk_vm_addresses`. + pub record_next_create_address: bool, + /// Dual compiled contracts pub dual_compiled_contracts: DualCompiledContracts, @@ -444,6 +455,9 @@ impl Cheatcodes { pc: Default::default(), breakpoints: Default::default(), use_zk_vm: Default::default(), + skip_zk_vm: Default::default(), + skip_zk_vm_addresses: Default::default(), + record_next_create_address: Default::default(), persisted_factory_deps: Default::default(), } } @@ -893,40 +907,74 @@ impl Cheatcodes { } if self.use_zk_vm { - info!("running create in zk vm"); - if input.init_code().0 == DEFAULT_CREATE2_DEPLOYER_CODE { - info!("ignoring DEFAULT_CREATE2_DEPLOYER_CODE for zk"); - return None + if let Some(result) = self.try_create_in_zk(ecx, input, executor) { + return Some(result); } + } - let zk_contract = self - .dual_compiled_contracts - .find_by_evm_bytecode(&input.init_code().0) - .unwrap_or_else(|| panic!("failed finding contract for {:?}", input.init_code())); + None + } - let factory_deps = self.dual_compiled_contracts.fetch_all_factory_deps(zk_contract); - tracing::debug!(contract = zk_contract.name, "using dual compiled contract"); + /// Try handling the `CREATE` within zkEVM. + /// If `Some` is returned then the result must be returned immediately, else the call must be + /// handled in EVM. + fn try_create_in_zk( + &mut self, + ecx: &mut EvmContext, + input: Input, + executor: &mut impl CheatcodesExecutor, + ) -> Option + where + DB: DatabaseExt, + Input: CommonCreateInput, + { + if self.skip_zk_vm { + self.skip_zk_vm = false; // handled the skip, reset flag + self.record_next_create_address = true; + info!("running create in EVM, instead of zkEVM (skipped)"); + return None + } - let ccx = foundry_zksync_core::vm::CheatcodeTracerContext { - mocked_calls: self.mocked_calls.clone(), - expected_calls: Some(&mut self.expected_calls), - accesses: self.accesses.as_mut(), - persisted_factory_deps: Some(&mut self.persisted_factory_deps), - }; - let create_inputs = CreateInputs { - scheme: input.scheme().unwrap_or(CreateScheme::Create), - init_code: input.init_code(), - value: input.value(), - caller: input.caller(), - gas_limit: input.gas_limit(), - }; - if let Ok(result) = foundry_zksync_core::vm::create::<_, DatabaseError>( - &create_inputs, - zk_contract, - factory_deps, - ecx, - ccx, - ) { + if input.init_code().0 == DEFAULT_CREATE2_DEPLOYER_CODE { + info!("running create in EVM, instead of zkEVM (DEFAULT_CREATE2_DEPLOYER_CODE)"); + return None + } + + info!("running create in zkEVM"); + + let zk_contract = self + .dual_compiled_contracts + .find_by_evm_bytecode(&input.init_code().0) + .unwrap_or_else(|| panic!("failed finding contract for {:?}", input.init_code())); + + let factory_deps = self.dual_compiled_contracts.fetch_all_factory_deps(zk_contract); + tracing::debug!(contract = zk_contract.name, "using dual compiled contract"); + + let ccx = foundry_zksync_core::vm::CheatcodeTracerContext { + mocked_calls: self.mocked_calls.clone(), + expected_calls: Some(&mut self.expected_calls), + accesses: self.accesses.as_mut(), + persisted_factory_deps: Some(&mut self.persisted_factory_deps), + }; + let create_inputs = CreateInputs { + scheme: input.scheme().unwrap_or(CreateScheme::Create), + init_code: input.init_code(), + value: input.value(), + caller: input.caller(), + gas_limit: input.gas_limit(), + }; + + // We currently exhaust the entire gas for the call as zkEVM returns a very high + // amount of gas that OOGs revm. + let gas = Gas::new(input.gas_limit()); + match foundry_zksync_core::vm::create::<_, DatabaseError>( + &create_inputs, + zk_contract, + factory_deps, + ecx, + ccx, + ) { + Ok(result) => { if let Some(recorded_logs) = &mut self.recorded_logs { recorded_logs.extend(result.logs.clone().into_iter().map(|log| Vm::Log { topics: log.data.topics().to_vec(), @@ -959,7 +1007,7 @@ impl Cheatcodes { } } - return match result.execution_result { + match result.execution_result { ExecutionResult::Success { output, .. } => match output { Output::Create(bytes, address) => Some(CreateOutcome { result: InterpreterResult { @@ -996,9 +1044,20 @@ impl Cheatcodes { }), } } + Err(err) => { + error!("error inspecting zkEVM: {err:?}"); + Some(CreateOutcome { + result: InterpreterResult { + result: InstructionResult::Revert, + output: Bytes::from_iter( + format!("error inspecting zkEVM: {err:?}").as_bytes(), + ), + gas, + }, + address: None, + }) + } } - - None } // common create_end functionality for both legacy and EOF. @@ -1111,6 +1170,14 @@ impl Cheatcodes { } } } + + if self.record_next_create_address { + self.record_next_create_address = false; + if let Some(address) = outcome.address { + self.skip_zk_vm_addresses.insert(address); + } + } + outcome } @@ -1400,22 +1467,56 @@ impl Cheatcodes { } if self.use_zk_vm { - if let TransactTo::Call(test_contract) = ecx.env.tx.transact_to { - if call.bytecode_address == test_contract { - info!("using evm for calls to test contract {:?}", ecx.env); - return None - } + if let Some(result) = self.try_call_in_zk(ecx, call, executor) { + return Some(result); } + } - info!("running call in zk vm {:#?}", call); + None + } - let ccx = foundry_zksync_core::vm::CheatcodeTracerContext { - mocked_calls: self.mocked_calls.clone(), - expected_calls: Some(&mut self.expected_calls), - accesses: self.accesses.as_mut(), - persisted_factory_deps: Some(&mut self.persisted_factory_deps), - }; - if let Ok(result) = foundry_zksync_core::vm::call::<_, DatabaseError>(call, ecx, ccx) { + /// Try handling the `CALL` within zkEVM. + /// If `Some` is returned then the result must be returned immediately, else the call must be + /// handled in EVM. + fn try_call_in_zk( + &mut self, + ecx: &mut EvmContext, + call: &mut CallInputs, + executor: &mut impl CheatcodesExecutor, + ) -> Option + where + DB: DatabaseExt, + { + // also skip if the target was created during a zkEVM skip + self.skip_zk_vm = + self.skip_zk_vm || self.skip_zk_vm_addresses.contains(&call.target_address); + if self.skip_zk_vm { + self.skip_zk_vm = false; // handled the skip, reset flag + info!("running create in EVM, instead of zkEVM (skipped) {:#?}", call); + 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 + } + } + + info!("running call in zkEVM {:#?}", call); + + let ccx = foundry_zksync_core::vm::CheatcodeTracerContext { + mocked_calls: self.mocked_calls.clone(), + expected_calls: Some(&mut self.expected_calls), + accesses: self.accesses.as_mut(), + persisted_factory_deps: Some(&mut self.persisted_factory_deps), + }; + + // We currently exhaust the entire gas for the call as zkEVM returns a very high amount + // of gas that OOGs revm. + let gas = Gas::new(call.gas_limit); + match foundry_zksync_core::vm::call::<_, DatabaseError>(call, ecx, ccx) { + Ok(result) => { // append console logs from zkEVM to the current executor's LogTracer result.logs.iter().filter_map(decode_console_log).for_each(|decoded_log| { executor.console_log( @@ -1451,7 +1552,7 @@ impl Cheatcodes { } } - return match result.execution_result { + match result.execution_result { ExecutionResult::Success { output, .. } => match output { Output::Call(bytes) => Some(CallOutcome { result: InterpreterResult { @@ -1488,9 +1589,20 @@ impl Cheatcodes { }), } } + Err(err) => { + error!("error inspecting zkEVM: {err:?}"); + Some(CallOutcome { + result: InterpreterResult { + result: InstructionResult::Revert, + output: Bytes::from_iter( + format!("error inspecting zkEVM: {err:?}").as_bytes(), + ), + gas, + }, + memory_offset: call.return_memory_offset.clone(), + }) + } } - - None } } diff --git a/crates/cheatcodes/src/test.rs b/crates/cheatcodes/src/test.rs index 015e1e917..deaaa71ba 100644 --- a/crates/cheatcodes/src/test.rs +++ b/crates/cheatcodes/src/test.rs @@ -23,6 +23,14 @@ impl Cheatcode for zkVmCall { } } +impl Cheatcode for zkVmSkipCall { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + ccx.state.skip_zk_vm = ccx.state.use_zk_vm; + + Ok(Default::default()) + } +} + impl Cheatcode for zkRegisterContractCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { diff --git a/crates/forge/tests/it/zk/cheats.rs b/crates/forge/tests/it/zk/cheats.rs index 135e3c333..ca33723a0 100644 --- a/crates/forge/tests/it/zk/cheats.rs +++ b/crates/forge/tests/it/zk/cheats.rs @@ -121,3 +121,11 @@ async fn test_zk_cheatcodes_in_zkvm() { TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await; } + +#[tokio::test(flavor = "multi_thread")] +async fn test_zk_zk_vm_skip_works() { + let runner = TEST_DATA_DEFAULT.runner_zksync(); + let filter = Filter::new(".*", "ZkCheatcodeZkVmSkipTest", ".*"); + + TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await; +} diff --git a/crates/zksync/compiler/src/zksolc/mod.rs b/crates/zksync/compiler/src/zksolc/mod.rs index a5b0dc7e2..db4c9b6e1 100644 --- a/crates/zksync/compiler/src/zksolc/mod.rs +++ b/crates/zksync/compiler/src/zksolc/mod.rs @@ -133,7 +133,11 @@ impl DualCompiledContracts { let bytecode_vec = bytecode.object.clone().into_bytes().unwrap().to_vec(); let mut factory_deps_vec: Vec> = factory_deps_map .keys() - .map(|factory_hash| zksolc_all_bytecodes.get(factory_hash).unwrap()) + .map(|factory_hash| { + zksolc_all_bytecodes.get(factory_hash).unwrap_or_else(|| { + panic!("failed to find zksolc artifact with hash {factory_hash:?}") + }) + }) .cloned() .collect(); diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index c76457b23..6ce385a91 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -697,4 +697,5 @@ interface Vm { bytes calldata zkDeployedBytecode ) external pure; function zkVm(bool enable) external pure; + function zkVmSkip() external pure; } diff --git a/testdata/zk/Cheatcodes.t.sol b/testdata/zk/Cheatcodes.t.sol index f8a3e5fa9..d19736726 100644 --- a/testdata/zk/Cheatcodes.t.sol +++ b/testdata/zk/Cheatcodes.t.sol @@ -314,3 +314,60 @@ contract ZkCheatcodesInZkVmTest is DSTest { assertEq(expected, got); } } + +contract Calculator { + event Added(uint8 indexed sum); + + function add(uint8 a, uint8 b) public returns (uint8) { + uint8 sum = a + b; + emit Added(sum); + return sum; + } +} + +contract EvmTargetContract is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + + event Added(uint8 indexed sum); + + function exec() public { + // We emit the event we expect to see. + vm.expectEmit(); + emit Added(3); + + Calculator calc = new Calculator(); // deployed on zkEVM + uint8 sum = calc.add(1, 2); // deployed on zkEVM + assertEq(3, sum); + } +} + +contract ZkCheatcodeZkVmSkipTest is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + EvmTargetContract helper; + + function setUp() external { + vm.zkVm(true); + helper = new EvmTargetContract(); + // ensure we can call cheatcodes from the helper + vm.allowCheatcodes(address(helper)); + // and that the contract is kept between vm switches + vm.makePersistent(address(helper)); + } + + function testFail_UseCheatcodesInZkVmWithoutSkip() external { + helper.exec(); + } + + function testUseCheatcodesInEvmWithSkip() external { + vm.zkVmSkip(); + helper.exec(); + } + + function testAutoSkipAfterDeployInEvmWithSkip() external { + vm.zkVmSkip(); + EvmTargetContract helper2 = new EvmTargetContract(); + + // this should auto execute in EVM + helper2.exec(); + } +}