diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index b890e2c06..819337f25 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -410,6 +410,40 @@ impl ArbitraryStorage { /// List of transactions that can be broadcasted. pub type BroadcastableTransactions = VecDeque; +/// Setting for migrating the database to zkEVM storage when starting in ZKsync mode. +/// The migration is performed on the DB via the inspector so must only be performed once. +#[derive(Debug, Default, Clone)] +pub enum ZkStartupMigration { + /// Defer database migration to a later execution point. + /// + /// This is required as we need to wait for some baseline deployments + /// to occur before the test/script execution is performed. + #[default] + Defer, + /// Allow database migration. + Allow, + /// Database migration has already been performed. + Done, +} + +impl ZkStartupMigration { + /// Check if startup migration is allowed. Migration is disallowed if it's to be deferred or has + /// already been performed. + pub fn is_allowed(&self) -> bool { + matches!(self, Self::Allow) + } + + /// Allow migrating the the DB to zkEVM storage. + pub fn allow(&mut self) { + *self = Self::Allow + } + + /// Mark the migration as completed. It must not be performed again. + pub fn done(&mut self) { + *self = Self::Done + } +} + /// An EVM inspector that handles calls to various cheatcodes, each with their own behavior. /// /// Cheatcodes can be called by contracts during execution to modify the VM environment, such as @@ -561,9 +595,8 @@ pub struct Cheatcodes { /// Dual compiled contracts pub dual_compiled_contracts: DualCompiledContracts, - /// Starts the cheatcode inspector in ZK mode. - /// This is set to `false`, once the startup migration is completed. - pub startup_zk: bool, + /// The migration status of the database to zkEVM storage, `None` if we start in EVM context. + pub zk_startup_migration: Option, /// Factory deps stored through `zkUseFactoryDep`. These factory deps are used in the next /// CREATE or CALL, and cleared after. @@ -631,13 +664,14 @@ impl Cheatcodes { let mut persisted_factory_deps = HashMap::new(); persisted_factory_deps.insert(zk_bytecode_hash, zk_deployed_bytecode); - let startup_zk = config.use_zk; + let zk_startup_migration = config.use_zk.then_some(ZkStartupMigration::Defer); + Self { fs_commit: true, labels: config.labels.clone(), config, dual_compiled_contracts, - startup_zk, + zk_startup_migration, block: Default::default(), gas_price: Default::default(), prank: Default::default(), @@ -876,6 +910,7 @@ impl Cheatcodes { let mut known_codes_storage: rHashMap = Default::default(); let mut deployed_codes: HashMap = Default::default(); + let test_contract = data.db.get_test_contract_address(); for address in data.db.persistent_accounts().into_iter().chain([data.env.tx.caller]) { info!(?address, "importing to zk state"); @@ -891,6 +926,12 @@ impl Cheatcodes { let nonce_key = get_nonce_key(address); nonce_storage.insert(nonce_key, EvmStorageSlot::new(full_nonce.to_ru256())); + if test_contract.map(|test_address| address == test_address).unwrap_or_default() { + // avoid migrating test contract code + tracing::trace!(?address, "ignoring code translation for test contract"); + continue; + } + if let Some(contract) = self.dual_compiled_contracts.iter().find(|contract| { info.code_hash != KECCAK_EMPTY && info.code_hash == contract.evm_bytecode_hash }) { @@ -941,17 +982,12 @@ impl Cheatcodes { journaled_account(data, known_codes_addr).expect("failed to load account"); known_codes_account.storage.extend(known_codes_storage.clone()); - let test_contract = data.db.get_test_contract_address(); for (address, info) in deployed_codes { let account = journaled_account(data, address).expect("failed to load account"); let _ = std::mem::replace(&mut account.info.balance, info.balance); let _ = std::mem::replace(&mut account.info.nonce, info.nonce); - if test_contract.map(|addr| addr == address).unwrap_or_default() { - tracing::trace!(?address, "ignoring code translation for test contract"); - } else { - account.info.code_hash = info.code_hash; - account.info.code.clone_from(&info.code); - } + account.info.code_hash = info.code_hash; + account.info.code.clone_from(&info.code); } } @@ -1991,9 +2027,17 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes { ecx.env.tx.gas_price = gas_price; } - if self.startup_zk && !self.use_zk_vm { - self.startup_zk = false; // We only do this once. + let migration_allowed = self + .zk_startup_migration + .as_ref() + .map(|migration| migration.is_allowed()) + .unwrap_or(false); + if migration_allowed && !self.use_zk_vm { self.select_zk_vm(ecx, None); + if let Some(zk_startup_migration) = &mut self.zk_startup_migration { + zk_startup_migration.done(); + } + debug!("startup zkEVM storage migration completed"); } // Record gas for current frame. diff --git a/crates/forge/src/runner.rs b/crates/forge/src/runner.rs index 7d5330fe4..5c1ebc89d 100644 --- a/crates/forge/src/runner.rs +++ b/crates/forge/src/runner.rs @@ -114,6 +114,9 @@ impl ContractRunner<'_> { } let address = self.sender.create(self.executor.get_nonce(self.sender)?); + // NOTE(zk): the test contract is set here instead of where upstream does it as + // the test contract address needs to be retrieved in order to skip + // zkEVM mode for the creation of the test address (and for calls to it later). self.executor.backend_mut().set_test_contract(address); // Set the contracts initial balance before deployment, so it is available during @@ -144,6 +147,15 @@ impl ContractRunner<'_> { self.executor.deploy_create2_deployer()?; + // Test contract has already been deployed so we can migrate the database to zkEVM storage + // in the next runner execution. + if let Some(cheatcodes) = &mut self.executor.inspector.cheatcodes { + if let Some(zk_startup_migration) = &mut cheatcodes.zk_startup_migration { + debug!("test contract deployed, allowing startup storage migration"); + zk_startup_migration.allow(); + } + } + // Optionally call the `setUp` function let result = if call_setup { trace!("calling setUp"); diff --git a/crates/forge/tests/fixtures/zk/test_zksync_keystore b/crates/forge/tests/fixtures/zk/test_zksync_keystore new file mode 100644 index 000000000..1bb13fe59 --- /dev/null +++ b/crates/forge/tests/fixtures/zk/test_zksync_keystore @@ -0,0 +1 @@ +{"crypto":{"cipher":"aes-128-ctr","cipherparams":{"iv":"baf20aa15d09d4ba2fea20a4f45f3b74"},"ciphertext":"19cbaecaf374fcab41fa5899990b904441e0762899ef3727a5a667398db8b7d8","kdf":"scrypt","kdfparams":{"dklen":32,"n":8192,"p":1,"r":8,"salt":"625b00531450f22ae8f74c9df7582150158fe6a88d9a214d61ec80a328e0601a"},"mac":"1c6f106535db170ab8c838abc6c33f019e35663808812f2c793bb00fd23e4baa"},"id":"f651e741-d76b-4583-b5f2-3ca0b4a697a2","version":3} \ No newline at end of file diff --git a/crates/forge/tests/it/zk/mod.rs b/crates/forge/tests/it/zk/mod.rs index 714a33199..11cc1944a 100644 --- a/crates/forge/tests/it/zk/mod.rs +++ b/crates/forge/tests/it/zk/mod.rs @@ -17,4 +17,5 @@ mod ownership; mod paymaster; mod proxy; mod repros; +mod script; mod traces; diff --git a/crates/forge/tests/it/zk/script.rs b/crates/forge/tests/it/zk/script.rs new file mode 100644 index 000000000..0cb9f13d5 --- /dev/null +++ b/crates/forge/tests/it/zk/script.rs @@ -0,0 +1,55 @@ +use foundry_test_utils::{ + forgetest_async, + util::{self, OutputExt}, + ZkSyncNode, +}; +use std::path::Path; + +forgetest_async!(test_zk_can_broadcast_with_keystore_account, |prj, cmd| { + util::initialize(prj.root()); + prj.add_script("Deploy.s.sol", include_str!("../../fixtures/zk/Deploy.s.sol")).unwrap(); + prj.add_source("Greeter.sol", include_str!("../../../../../testdata/zk/Greeter.sol")).unwrap(); + + let node = ZkSyncNode::start(); + let url = node.url(); + + cmd.forge_fuse(); + + let script_path_contract = "./script/Deploy.s.sol:DeployScript"; + let keystore_path = + Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/zk/test_zksync_keystore"); + + let script_args = vec![ + "--zk-startup", + &script_path_contract, + "--broadcast", + "--keystores", + keystore_path.to_str().unwrap(), + "--password", + "password", + "--chain", + "260", + "--gas-estimate-multiplier", + "310", + "--rpc-url", + url.as_str(), + "--slow", + ]; + + cmd.arg("script").args(&script_args); + + cmd.assert_success() + .get_output() + .stdout_lossy() + .contains("ONCHAIN EXECUTION COMPLETE & SUCCESSFUL"); + + let run_latest = foundry_common::fs::json_files(prj.root().join("broadcast").as_path()) + .find(|file| file.ends_with("run-latest.json")) + .expect("No broadcast artifacts"); + + let content = foundry_common::fs::read_to_string(run_latest).unwrap(); + + let json: serde_json::Value = serde_json::from_str(&content).unwrap(); + assert_eq!(json["transactions"].as_array().expect("broadcastable txs").len(), 3); + cmd.forge_fuse(); +}); diff --git a/crates/script/src/runner.rs b/crates/script/src/runner.rs index 52e5efe31..8869b703f 100644 --- a/crates/script/src/runner.rs +++ b/crates/script/src/runner.rs @@ -132,7 +132,6 @@ 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 self.executor.set_balance(address, self.evm_opts.initial_balance)?; @@ -148,6 +147,14 @@ impl ScriptRunner { self.executor.set_nonce(self.evm_opts.sender, u64::MAX / 2)?; } + // NOTE(zk): Address recomputed again after the nonce modification as it will + // differ in case of evm_opts.sender == CALLER + let zk_actual_script_address = CALLER.create(self.executor.get_nonce(CALLER)?); + // NOTE(zk): the test contract is set here instead of where upstream does it as + // the test contract address needs to be retrieved in order to skip + // zkEVM mode for the creation of the test address (and for calls to it later). + self.executor.backend_mut().set_test_contract(zk_actual_script_address); + // Deploy an instance of the contract let DeployResult { address, @@ -163,8 +170,19 @@ impl ScriptRunner { traces.extend(constructor_traces.map(|traces| (TraceKind::Deployment, traces))); + // Script has already been deployed so we can migrate the database to zkEVM storage + // in the next runner execution. + if let Some(cheatcodes) = &mut self.executor.inspector.cheatcodes { + if let Some(zk_startup_migration) = &mut cheatcodes.zk_startup_migration { + debug!("script deployed, allowing startup storage migration"); + zk_startup_migration.allow(); + } + } + // Optionally call the `setUp` function let (success, gas_used, labeled_addresses, transactions) = if !setup { + // NOTE(zk): keeping upstream code for context. Test contract is only set on this branch + // 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/deny.toml b/deny.toml index 303bfec5b..ef0dd7003 100644 --- a/deny.toml +++ b/deny.toml @@ -17,6 +17,11 @@ ignore = [ # Used by alloy # proc-macro-error is unmaintained "RUSTSEC-2024-0370", + # instant is unmaintained + "RUSTSEC-2024-0384", + # Used by boojum + # derivative is unmaintained + "RUSTSEC-2024-0388", ] # This section is considered when running `cargo deny check bans`.