fix: retrieve zks_getBytecodeByHash from SharedBackend when forking (#…

…815)

* retrieve zks_getBytecodeByHash from SharedBackend when forking

crates/evm/core/src/backend/mod.rs

@@ -9,10 +9,13 @@ use crate::{
 };
 use alloy_genesis::GenesisAccount;
 use alloy_network::{AnyRpcBlock, AnyTxEnvelope, TransactionResponse};
-use alloy_primitives::{keccak256, map::HashMap, uint, Address, B256, U256};
+use alloy_primitives::{keccak256, map::HashMap, uint, Address, Bytes, B256, U256};
+use alloy_provider::Provider;
 use alloy_rpc_types::{BlockNumberOrTag, Transaction, TransactionRequest};
 use eyre::Context;
-use foundry_common::{is_known_system_sender, SYSTEM_TRANSACTION_TYPE};
+use foundry_common::{
+    is_known_system_sender, provider::try_get_zksync_http_provider, SYSTEM_TRANSACTION_TYPE,
+};
 pub use foundry_fork_db::{cache::BlockchainDbMeta, BlockchainDb, SharedBackend};
 use itertools::Itertools;
 use revm::{
@@ -28,6 +31,7 @@ use revm::{
 use std::{
     any::Any,
     collections::{BTreeMap, HashSet},
+    sync::Arc,
     time::Instant,
 };
 use strategy::{BackendStrategy, BackendStrategyForkInfo};
@@ -1612,6 +1616,27 @@ impl Database for Backend {
     }
 
     fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
+        // Try obtaining code by hash via zks_getBytecodeByHash for zksync forks.
+        let maybe_zk_fork = self
+            .active_fork_id()
+            .and_then(|id| self.get_fork_info(id).ok())
+            .map(|info| info.fork_type.is_zk())
+            .and_then(|is_zk| if is_zk { self.active_fork_url() } else { None });
+        if let (Some(fork_url), Some(db)) = (maybe_zk_fork, self.active_fork_db_mut()) {
+            let provider = try_get_zksync_http_provider(fork_url)
+                .map_err(|err| DatabaseError::AnyRequest(Arc::new(err)))?;
+            let result = db.db.do_any_request(async move {
+                let bytes = provider
+                    .raw_request::<_, Bytes>("zks_getBytecodeByHash".into(), vec![code_hash])
+                    .await?;
+                Ok(Bytecode::new_raw(bytes))
+            });
+
+            if let Ok(bytes) = result {
+                return Ok(bytes);
+            }
+        }
+
         if let Some(db) = self.active_fork_db_mut() {
             Ok(db.code_by_hash(code_hash)?)
         } else {

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

@@ -20,3 +20,11 @@ async fn test_zk_immutable_vars_persist_after_fork() {
 
     TestConfig::with_filter(runner, filter).spec_id(SpecId::SHANGHAI).run().await;
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_zk_consistent_storage_migration_after_fork() {
+    let runner = TEST_DATA_DEFAULT.runner_zksync();
+    let filter = Filter::new(".*", "ZkForkStorageMigrationTest", ".*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::SHANGHAI).run().await;
+}

testdata/zk/Fork.t.sol

@@ -0,0 +1,124 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import "ds-test/test.sol";
+import "../cheats/Vm.sol";
+import {Globals} from "./Globals.sol";
+
+contract Store {
+    uint256 a;
+
+    constructor() payable {}
+
+    function set(uint256 _a) public {
+        a = _a;
+    }
+
+    function get() public view returns (uint256) {
+        return a;
+    }
+}
+
+interface ERC20 {
+    function decimals() external returns (uint8);
+}
+
+contract ZkForkStorageMigrationTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    uint256 constant ETH_FORK_BLOCK = 19225195;
+    uint256 constant ERA_FORK_BLOCK = 48517149;
+
+    uint256 forkEth;
+    uint256 forkEra;
+
+    // Wrapped native token addresses from https://docs.uniswap.org/contracts/v3/reference/deployments/
+    address uniswapEth = 0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2;
+    address uniswapEra = 0x5AEa5775959fBC2557Cc8789bC1bf90A239D9a91;
+
+    function setUp() public {
+        forkEra = vm.createFork(Globals.ZKSYNC_MAINNET_URL, ERA_FORK_BLOCK);
+        forkEth = vm.createFork(Globals.ETHEREUM_MAINNET_URL, ETH_FORK_BLOCK);
+    }
+
+    function testForkMigrationExecutesChainNativeCalls() public {
+        // assert we have switched to ethereum
+        vm.selectFork(forkEth);
+        assertEq(18, ERC20(uniswapEth).decimals());
+
+        // assert we have switched to era
+        vm.selectFork(forkEra);
+        assertEq(18, ERC20(uniswapEra).decimals());
+    }
+
+    function testForkMigrationConsistentBalanceAfterForkToZkEvm() public {
+        // assert we have switched to ethereum
+        vm.selectFork(forkEth);
+        assertEq(18, ERC20(uniswapEth).decimals());
+
+        // deploy on EVM
+        Store store = new Store{value: 1 ether}();
+        store.set(10);
+        vm.makePersistent(address(store));
+
+        // assert we have switched to era
+        vm.selectFork(forkEra);
+        assertEq(18, ERC20(uniswapEra).decimals());
+
+        // assert balance on zkEVM
+        assertEq(1 ether, address(store).balance);
+    }
+
+    function testForkMigrationConsistentBalanceAfterForkToEvm() public {
+        // assert we have switched to era
+        vm.selectFork(forkEra);
+        assertEq(18, ERC20(uniswapEra).decimals());
+
+        // deploy on zkEVM
+        Store store = new Store{value: 1 ether}();
+        store.set(10);
+        vm.makePersistent(address(store));
+
+        // assert we have switched to ethereum
+        vm.selectFork(forkEth);
+        assertEq(18, ERC20(uniswapEth).decimals());
+
+        // assert balance on EVM
+        assertEq(1 ether, address(store).balance);
+    }
+
+    function testForkMigrationConsistentContractCallsAfterForkToZkEvm() public {
+        // assert we have switched to ethereum
+        vm.selectFork(forkEth);
+        assertEq(18, ERC20(uniswapEth).decimals());
+
+        // deploy on EVM
+        Store store = new Store{value: 1 ether}();
+        store.set(10);
+        vm.makePersistent(address(store));
+
+        // assert we have switched to era
+        vm.selectFork(forkEra);
+        assertEq(18, ERC20(uniswapEra).decimals());
+
+        // assert contract calls on zkEVM
+        assertEq(10, store.get());
+    }
+
+    function testForkMigrationConsistentContractCallsAfterForkToEvm() public {
+        // assert we have switched to era
+        vm.selectFork(forkEra);
+        assertEq(18, ERC20(uniswapEra).decimals());
+
+        // deploy on zkEVM
+        Store store = new Store{value: 1 ether}();
+        store.set(10);
+        vm.makePersistent(address(store));
+
+        // assert we have switched to ethereum
+        vm.selectFork(forkEth);
+        assertEq(18, ERC20(uniswapEth).decimals());
+
+        // assert contract calls on EVM
+        assertEq(10, store.get());
+    }
+}