feat: Allow bytecode replacement (#336)

* feat: add option to replace bytecodes

* apply earlier

* updated comments

* cleaned up error handling

* clippy

README.md

@@ -170,6 +170,19 @@ If you wish to replay a remote transaction locally for deep debugging, use the f
 era_test_node replay_tx <network> <transaction_hash>
 ```
 
+## Replacing bytecodes
+
+You can also replace / override the contract bytecode with the local version. This is especially useful if you are replaying some mainnet transactions and would like to see how they would behave on the different bytecode. Or when you want to fork mainnet to see how your code would
+behave on mainnet state.
+
+You have to prepare a directory, with files in format `0xabc..93f.json` that contain the json outputs that you can get from zkout directories from your compiler.
+
+Then you have to add `--override-bytecodes-dir=XX` flag to point at that directory. See the `example_override` dir for more details.
+
+```bash
+cargo run -- --override-bytecodes-dir=example_override --show-storage-logs all fork mainnet
+```
+
 ## 📞 Sending Network Calls
 
 You can send network calls against a running `era-test-node`. For example, to check the testnet LINK balance or mainnet USDT, use `curl` or `foundry-zksync`.

example_override/README.md

@@ -0,0 +1,6 @@
+# Example override directory
+
+This directory is an example on how to use the `override-bytecodes-dir` flag.
+
+You put your bytecodes in this directory, under a proper file name - that should match the address that you're overriding, and run
+the era-test-node with `--override-bytecodes-dir=example_override` flag.

src/bytecode_override.rs

@@ -0,0 +1,56 @@
+use std::fs;
+
+use crate::fork::ForkSource;
+use crate::node::InMemoryNode;
+use eyre::Context;
+use hex::FromHex;
+use serde::Deserialize;
+use std::str::FromStr;
+use zksync_types::Address;
+
+#[derive(Debug, Deserialize)]
+struct ContractJson {
+    bytecode: Bytecode,
+}
+
+#[derive(Debug, Deserialize)]
+struct Bytecode {
+    object: String,
+}
+
+// Loads a list of bytecodes and addresses from the directory and then inserts them directly
+// into the Node's storage.
+pub fn override_bytecodes<T: Clone + ForkSource + std::fmt::Debug>(
+    node: &InMemoryNode<T>,
+    bytecodes_dir: String,
+) -> eyre::Result<()> {
+    for entry in fs::read_dir(bytecodes_dir)? {
+        let entry = entry?;
+        let path = entry.path();
+
+        if path.is_file() {
+            let filename = match path.file_name().and_then(|name| name.to_str()) {
+                Some(name) => name,
+                None => eyre::bail!("Invalid filename {}", path.display().to_string()),
+            };
+
+            // Look only at .json files.
+            if let Some(filename) = filename.strip_suffix(".json") {
+                let address = Address::from_str(filename)
+                    .wrap_err(format!("Cannot parse {} as address", filename))?;
+
+                let file_content = fs::read_to_string(&path)?;
+                let contract: ContractJson = serde_json::from_str(&file_content)
+                    .wrap_err(format!("Failed to  parse json file {:?}", path))?;
+
+                let bytecode = Vec::from_hex(contract.bytecode.object)
+                    .wrap_err(format!("Failed to parse hex from {:?}", path))?;
+
+                node.override_bytecode(&address, &bytecode)
+                    .expect("Failed to override bytecode");
+                tracing::info!("+++++ Replacing bytecode at address {:?} +++++", address);
+            }
+        }
+    }
+    Ok(())
+}

src/config/cli.rs

@@ -93,6 +93,9 @@ pub struct Cli {
     /// Cache directory location for `disk` cache - default: ".cache"
     #[arg(long)]
     pub cache_dir: Option<String>,
+
+    #[arg(long)]
+    pub override_bytecodes_dir: Option<String>,
 }
 
 #[derive(Debug, Subcommand)]

src/main.rs

@@ -1,5 +1,6 @@
 use crate::observability::Observability;
 use anyhow::anyhow;
+use bytecode_override::override_bytecodes;
 use clap::Parser;
 use colored::Colorize;
 use config::cli::{Cli, Command};
@@ -10,6 +11,7 @@ use logging_middleware::LoggingMiddleware;
 use tracing_subscriber::filter::LevelFilter;
 
 mod bootloader_debug;
+mod bytecode_override;
 mod cache;
 mod config;
 mod console_log;
@@ -171,6 +173,10 @@ async fn main() -> anyhow::Result<()> {
     let node: InMemoryNode<HttpForkSource> =
         InMemoryNode::new(fork_details, Some(observability), config.node, config.gas);
 
+    if let Some(bytecodes_dir) = opt.override_bytecodes_dir {
+        override_bytecodes(&node, bytecodes_dir).unwrap();
+    }
+
     if !transactions_to_replay.is_empty() {
         let _ = node.apply_txs(transactions_to_replay);
     }

src/node/in_memory.rs

@@ -60,16 +60,15 @@ use zksync_types::{
     block::{unpack_block_info, L2BlockHasher},
     fee::Fee,
     fee_model::{BatchFeeInput, PubdataIndependentBatchFeeModelInput},
-    get_nonce_key,
-    l2::L2Tx,
-    l2::TransactionType,
+    get_code_key, get_nonce_key,
+    l2::{L2Tx, TransactionType},
     utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
     vm_trace::Call,
     PackedEthSignature, StorageKey, StorageLogQueryType, StorageValue, Transaction,
     ACCOUNT_CODE_STORAGE_ADDRESS, MAX_L2_TX_GAS_LIMIT, SYSTEM_CONTEXT_ADDRESS,
     SYSTEM_CONTEXT_BLOCK_INFO_POSITION,
 };
-use zksync_utils::{h256_to_account_address, h256_to_u256, u256_to_h256};
+use zksync_utils::{bytecode::hash_bytecode, h256_to_account_address, h256_to_u256, u256_to_h256};
 use zksync_web3_decl::error::Web3Error;
 
 /// Max possible size of an ABI encoded tx (in bytes).
@@ -1699,6 +1698,26 @@ impl<S: ForkSource + std::fmt::Debug + Clone> InMemoryNode<S> {
 
         Ok(())
     }
+
+    // Forcefully stores the given bytecode at a given account.
+    pub fn override_bytecode(&self, address: &Address, bytecode: &[u8]) -> Result<(), String> {
+        let mut inner = self
+            .inner
+            .write()
+            .map_err(|e| format!("Failed to acquire write lock: {}", e))?;
+
+        let code_key = get_code_key(address);
+
+        let bytecode_hash = hash_bytecode(bytecode);
+
+        inner
+            .fork_storage
+            .store_factory_dep(bytecode_hash, bytecode.to_owned());
+
+        inner.fork_storage.set_value(code_key, bytecode_hash);
+
+        Ok(())
+    }
 }
 
 /// Keeps track of a block's batch number, miniblock number and timestamp.