Merge branch 'jrigada-fix-nonce-mismatch-scripts' into nish-jrigada-f…

…ix-nonce-mismatch-scripts

crates/cheatcodes/src/inspector.rs

@@ -443,6 +443,40 @@ impl ZkPersistNonceUpdate {
     }
 }
 
+/// 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
@@ -594,9 +628,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<ZkStartupMigration>,
 
     /// Factory deps stored through `zkUseFactoryDep`. These factory deps are used in the next
     /// CREATE or CALL, and cleared after.
@@ -663,13 +696,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(),
@@ -908,6 +942,7 @@ impl Cheatcodes {
         let mut known_codes_storage: rHashMap<U256, EvmStorageSlot> = Default::default();
         let mut deployed_codes: HashMap<Address, AccountInfo> = 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");
 
@@ -923,6 +958,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
             }) {
@@ -973,17 +1014,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);
         }
     }
 
@@ -2018,9 +2054,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.

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");

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}

crates/forge/tests/it/zk/mod.rs

@@ -17,4 +17,5 @@ mod ownership;
 mod paymaster;
 mod proxy;
 mod repros;
+mod script;
 mod traces;

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();
+});

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) {

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`.