From bf19a840952cb364b5448f033213a8f6ee9eb2c5 Mon Sep 17 00:00:00 2001 From: Daniyar Itegulov Date: Wed, 15 Jan 2025 17:05:47 +1100 Subject: [PATCH] reduce usage of `ForkStorage` outside `inner` module --- crates/api_server/src/impls/anvil.rs | 18 ++- crates/cli/src/bytecode_override.rs | 2 +- crates/cli/src/main.rs | 3 +- crates/core/src/node/debug.rs | 2 +- crates/core/src/node/eth.rs | 22 +-- crates/core/src/node/in_memory.rs | 42 +++-- crates/core/src/node/in_memory_ext.rs | 78 ++++------ crates/core/src/node/inner/fork.rs | 21 +++ crates/core/src/node/inner/in_memory_inner.rs | 38 ++--- crates/core/src/node/inner/mod.rs | 6 +- crates/core/src/node/inner/node_executor.rs | 143 +++++++++++++++++- crates/core/src/node/inner/storage.rs | 23 +++ crates/core/src/node/zks.rs | 25 +-- crates/core/src/utils.rs | 24 --- 14 files changed, 286 insertions(+), 161 deletions(-) create mode 100644 crates/core/src/node/inner/storage.rs diff --git a/crates/api_server/src/impls/anvil.rs b/crates/api_server/src/impls/anvil.rs index 58829fea..6e20f3ad 100644 --- a/crates/api_server/src/impls/anvil.rs +++ b/crates/api_server/src/impls/anvil.rs @@ -155,11 +155,19 @@ impl AnvilNamespaceServer for AnvilNamespace { } async fn set_balance(&self, address: Address, balance: U256) -> RpcResult { - Ok(self.node.set_balance(address, balance).await) + Ok(self + .node + .set_balance(address, balance) + .await + .map_err(RpcError::from)?) } async fn set_nonce(&self, address: Address, nonce: U256) -> RpcResult { - Ok(self.node.set_nonce(address, nonce).await) + Ok(self + .node + .set_nonce(address, nonce) + .await + .map_err(RpcError::from)?) } async fn anvil_mine(&self, num_blocks: Option, interval: Option) -> RpcResult<()> { @@ -203,7 +211,11 @@ impl AnvilNamespaceServer for AnvilNamespace { } async fn set_storage_at(&self, address: Address, slot: U256, value: U256) -> RpcResult { - Ok(self.node.set_storage_at(address, slot, value).await) + Ok(self + .node + .set_storage_at(address, slot, value) + .await + .map_err(RpcError::from)?) } async fn set_chain_id(&self, id: u32) -> RpcResult<()> { diff --git a/crates/cli/src/bytecode_override.rs b/crates/cli/src/bytecode_override.rs index f6dbe91e..cc71cd6f 100644 --- a/crates/cli/src/bytecode_override.rs +++ b/crates/cli/src/bytecode_override.rs @@ -41,7 +41,7 @@ pub async fn override_bytecodes(node: &InMemoryNode, bytecodes_dir: String) -> a let bytecode = Vec::from_hex(contract.bytecode.object) .with_context(|| format!("Failed to parse hex from {:?}", path))?; - node.override_bytecode(&address, &bytecode) + node.override_bytecode(address, bytecode) .await .expect("Failed to override bytecode"); tracing::info!("+++++ Replacing bytecode at address {:?} +++++", address); diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 1f804c6a..d31e4112 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -221,7 +221,7 @@ async fn main() -> anyhow::Result<()> { let system_contracts = SystemContracts::from_options(&config.system_contracts_options, config.use_evm_emulator); - let (node_inner, _fork_storage, blockchain, time) = InMemoryNodeInner::init( + let (node_inner, storage, blockchain, time) = InMemoryNodeInner::init( fork_details, fee_input_provider.clone(), filters, @@ -245,6 +245,7 @@ async fn main() -> anyhow::Result<()> { let node: InMemoryNode = InMemoryNode::new( node_inner, blockchain, + storage, node_handle, Some(observability), time, diff --git a/crates/core/src/node/debug.rs b/crates/core/src/node/debug.rs index 7c8c68a7..abcb5191 100644 --- a/crates/core/src/node/debug.rs +++ b/crates/core/src/node/debug.rs @@ -68,7 +68,7 @@ impl InMemoryNode { // update the enforced_base_fee within l1_batch_env to match the logic in zksync_core l1_batch_env.enforced_base_fee = Some(l2_tx.common_data.fee.max_fee_per_gas.as_u64()); let system_env = inner.create_system_env(system_contracts.clone(), execution_mode); - let storage = StorageView::new(&inner.fork_storage).into_rc_ptr(); + let storage = StorageView::new(inner.read_storage()).into_rc_ptr(); let mut vm: Vm<_, HistoryDisabled> = Vm::new(l1_batch_env, system_env, storage); // We must inject *some* signature (otherwise bootloader code fails to generate hash). diff --git a/crates/core/src/node/eth.rs b/crates/core/src/node/eth.rs index 89fabfdf..c0640e8d 100644 --- a/crates/core/src/node/eth.rs +++ b/crates/core/src/node/eth.rs @@ -75,7 +75,7 @@ impl InMemoryNode { } pub async fn send_raw_transaction_impl(&self, tx_bytes: Bytes) -> Result { - let chain_id = self.inner.read().await.fork_storage.chain_id; + let chain_id = self.chain_id().await; let (tx_req, hash) = TransactionRequest::from_bytes(&tx_bytes.0, chain_id)?; // Impersonation does not matter in this context so we assume the tx is not impersonated: @@ -104,10 +104,7 @@ impl InMemoryNode { ) -> Result { let (chain_id, l2_gas_price) = { let reader = self.inner.read().await; - ( - reader.fork_storage.chain_id, - reader.fee_input_provider.gas_price(), - ) + (self.chain_id().await, reader.fee_input_provider.gas_price()) }; let mut tx_req = TransactionRequest::from(tx.clone()); @@ -187,9 +184,7 @@ impl InMemoryNode { AccountTreeId::new(L2_BASE_TOKEN_ADDRESS), &address, ); - - let inner_guard = self.inner.read().await; - match inner_guard.fork_storage.read_value_internal(&balance_key) { + match self.storage.read_value_alt(&balance_key).await { Ok(balance) => Ok(h256_to_u256(balance)), Err(error) => Err(anyhow::anyhow!("failed to read account balance: {error}")), } @@ -258,12 +253,9 @@ impl InMemoryNode { // TODO: Support _block: Option, ) -> anyhow::Result { - let inner = self.inner.write().await; - let code_key = get_code_key(&address); - - match inner.fork_storage.read_value_internal(&code_key) { - Ok(code_hash) => match inner.fork_storage.load_factory_dep_internal(code_hash) { + match self.storage.read_value_alt(&code_key).await { + Ok(code_hash) => match self.storage.load_factory_dep_alt(code_hash).await { Ok(raw_code) => { let code = raw_code.unwrap_or_default(); Ok(Bytes::from(code)) @@ -280,10 +272,8 @@ impl InMemoryNode { // TODO: Support _block: Option, ) -> anyhow::Result { - let inner = self.inner.read().await; let nonce_key = get_nonce_key(&address); - - match inner.fork_storage.read_value_internal(&nonce_key) { + match self.storage.read_value_alt(&nonce_key).await { Ok(result) => Ok(h256_to_u64(result).into()), Err(error) => Err(anyhow::anyhow!("failed to read nonce storage: {error}")), } diff --git a/crates/core/src/node/in_memory.rs b/crates/core/src/node/in_memory.rs index 80d809ce..a0a4b03a 100644 --- a/crates/core/src/node/in_memory.rs +++ b/crates/core/src/node/in_memory.rs @@ -2,7 +2,8 @@ use super::inner::fork::ForkDetails; use super::inner::node_executor::NodeExecutorHandle; use super::inner::InMemoryNodeInner; -use crate::deps::{storage_view::StorageView, InMemoryStorage}; +use crate::deps::storage_view::StorageView; +use crate::deps::InMemoryStorage; use crate::filters::EthFilters; use crate::formatter; use crate::node::call_error_tracer::CallErrorTracer; @@ -10,6 +11,7 @@ use crate::node::error::LoadStateError; use crate::node::fee_model::TestNodeFeeInputProvider; use crate::node::impersonate::{ImpersonationManager, ImpersonationState}; use crate::node::inner::blockchain::ReadBlockchain; +use crate::node::inner::storage::ReadStorageDyn; use crate::node::inner::time::ReadTime; use crate::node::sealer::BlockSealerState; use crate::node::state::VersionedState; @@ -46,17 +48,15 @@ use zksync_multivm::vm_latest::{HistoryDisabled, ToTracerPointer, Vm}; use zksync_multivm::VmVersion; use zksync_types::api::{Block, DebugCall, TransactionReceipt, TransactionVariant}; use zksync_types::block::unpack_block_info; -use zksync_types::bytecode::BytecodeHash; use zksync_types::fee_model::BatchFeeInput; use zksync_types::l2::L2Tx; use zksync_types::storage::{ EMPTY_UNCLES_HASH, SYSTEM_CONTEXT_ADDRESS, SYSTEM_CONTEXT_BLOCK_INFO_POSITION, }; use zksync_types::web3::{keccak256, Bytes}; -use zksync_types::{get_code_key, h256_to_u256}; use zksync_types::{ - AccountTreeId, Address, Bloom, L1BatchNumber, L2BlockNumber, PackedEthSignature, StorageKey, - StorageValue, Transaction, H160, H256, H64, U256, U64, + h256_to_u256, AccountTreeId, Address, Bloom, L1BatchNumber, L2BlockNumber, L2ChainId, + PackedEthSignature, StorageKey, StorageValue, Transaction, H160, H256, H64, U256, U64, }; /// Max possible size of an ABI encoded tx (in bytes). @@ -244,6 +244,7 @@ pub struct InMemoryNode { /// A thread safe reference to the [InMemoryNodeInner]. pub(crate) inner: Arc>, pub(crate) blockchain: Box, + pub(crate) storage: Box, pub(crate) node_handle: NodeExecutorHandle, /// List of snapshots of the [InMemoryNodeInner]. This is bounded at runtime by [MAX_SNAPSHOTS]. pub(crate) snapshots: Arc>>, @@ -261,6 +262,7 @@ impl InMemoryNode { pub fn new( inner: Arc>, blockchain: Box, + storage: Box, node_handle: NodeExecutorHandle, observability: Option, time: Box, @@ -272,6 +274,7 @@ impl InMemoryNode { InMemoryNode { inner, blockchain, + storage, node_handle, snapshots: Default::default(), time, @@ -380,7 +383,7 @@ impl InMemoryNode { let (batch_env, _) = inner.create_l1_batch_env().await; let system_env = inner.create_system_env(base_contracts, execution_mode); - let storage = StorageView::new(&inner.fork_storage).into_rc_ptr(); + let storage = StorageView::new(inner.read_storage()).into_rc_ptr(); let mut vm: Vm<_, HistoryDisabled> = Vm::new(batch_env, system_env, storage); // We must inject *some* signature (otherwise bootloader code fails to generate hash). @@ -455,22 +458,10 @@ impl InMemoryNode { // Forcefully stores the given bytecode at a given account. pub async fn override_bytecode( &self, - address: &Address, - bytecode: &[u8], - ) -> Result<(), String> { - let inner = self.inner.write().await; - - let code_key = get_code_key(address); - - let bytecode_hash = BytecodeHash::for_bytecode(bytecode).value(); - - inner - .fork_storage - .store_factory_dep(bytecode_hash, bytecode.to_owned()); - - inner.fork_storage.set_value(code_key, bytecode_hash); - - Ok(()) + address: Address, + bytecode: Vec, + ) -> anyhow::Result<()> { + self.node_handle.set_code_sync(address, bytecode).await } pub async fn dump_state(&self, preserve_historical_states: bool) -> anyhow::Result { @@ -605,6 +596,10 @@ impl InMemoryNode { observability.set_logging(directive)?; Ok(true) } + + pub async fn chain_id(&self) -> L2ChainId { + self.inner.read().await.chain_id() + } } pub fn load_last_l1_batch(storage: StoragePtr) -> Option<(u64, u64)> { @@ -636,7 +631,7 @@ impl InMemoryNode { &config.system_contracts_options, config.use_evm_emulator, ); - let (inner, _, blockchain, time) = InMemoryNodeInner::init( + let (inner, storage, blockchain, time) = InMemoryNodeInner::init( fork, fee_provider, Arc::new(RwLock::new(Default::default())), @@ -661,6 +656,7 @@ impl InMemoryNode { Self::new( inner, blockchain, + storage, node_handle, None, time, diff --git a/crates/core/src/node/in_memory_ext.rs b/crates/core/src/node/in_memory_ext.rs index a07bb238..70daf646 100644 --- a/crates/core/src/node/in_memory_ext.rs +++ b/crates/core/src/node/in_memory_ext.rs @@ -2,18 +2,14 @@ use super::inner::fork::ForkDetails; use super::pool::TxBatch; use super::sealer::BlockSealerMode; use super::InMemoryNode; -use crate::utils::bytecode_to_factory_dep; use anvil_zksync_types::api::{DetailedTransaction, ResetRequest}; -use anyhow::anyhow; +use anyhow::{anyhow, Context}; use std::time::Duration; use zksync_types::api::{Block, TransactionVariant}; +use zksync_types::bytecode::BytecodeHash; use zksync_types::u256_to_h256; -use zksync_types::{ - get_code_key, get_nonce_key, - utils::{nonces_to_full_nonce, storage_key_for_eth_balance}, - L2BlockNumber, StorageKey, -}; use zksync_types::{AccountTreeId, Address, H256, U256, U64}; +use zksync_types::{L2BlockNumber, StorageKey}; type Result = anyhow::Result; @@ -168,33 +164,24 @@ impl InMemoryNode { .map_err(|err| anyhow!("{}", err)) } - pub async fn set_balance(&self, address: Address, balance: U256) -> bool { - let writer = self.inner.write().await; - let balance_key = storage_key_for_eth_balance(&address); - writer - .fork_storage - .set_value(balance_key, u256_to_h256(balance)); + pub async fn set_balance(&self, address: Address, balance: U256) -> anyhow::Result { + self.node_handle.set_balance_sync(address, balance).await?; tracing::info!( "👷 Balance for address {:?} has been manually set to {} Wei", address, balance ); - true + Ok(true) } - pub async fn set_nonce(&self, address: Address, nonce: U256) -> bool { - let writer = self.inner.write().await; - let nonce_key = get_nonce_key(&address); - let enforced_full_nonce = nonces_to_full_nonce(nonce, nonce); + pub async fn set_nonce(&self, address: Address, nonce: U256) -> anyhow::Result { + self.node_handle.set_nonce_sync(address, nonce).await?; tracing::info!( "👷 Nonces for address {:?} have been set to {}", address, nonce ); - writer - .fork_storage - .set_value(nonce_key, u256_to_h256(enforced_full_nonce)); - true + Ok(true) } pub async fn mine_blocks(&self, num_blocks: Option, interval: Option) -> Result<()> { @@ -316,34 +303,24 @@ impl InMemoryNode { } pub async fn set_code(&self, address: Address, code: String) -> Result<()> { - let writer = self.inner.write().await; - let code_key = get_code_key(&address); - tracing::info!("set code for address {address:#x}"); let code_slice = code .strip_prefix("0x") .ok_or_else(|| anyhow!("code must be 0x-prefixed"))?; - let code_bytes = hex::decode(code_slice)?; - let hashcode = bytecode_to_factory_dep(code_bytes)?; - let hash = u256_to_h256(hashcode.0); - let code = hashcode - .1 - .iter() - .flat_map(|entry| { - let mut bytes = vec![0u8; 32]; - entry.to_big_endian(&mut bytes); - bytes.to_vec() - }) - .collect(); - writer.fork_storage.store_factory_dep(hash, code); - writer.fork_storage.set_value(code_key, hash); + let bytecode = hex::decode(code_slice)?; + zksync_types::bytecode::validate_bytecode(&bytecode).context("Invalid bytecode")?; + tracing::info!( + ?address, + bytecode_hash = ?BytecodeHash::for_bytecode(&bytecode).value(), + "set code" + ); + self.node_handle.set_code_sync(address, bytecode).await?; Ok(()) } - pub async fn set_storage_at(&self, address: Address, slot: U256, value: U256) -> bool { - let writer = self.inner.write().await; + pub async fn set_storage_at(&self, address: Address, slot: U256, value: U256) -> Result { let key = StorageKey::new(AccountTreeId::new(address), u256_to_h256(slot)); - writer.fork_storage.set_value(key, u256_to_h256(value)); - true + self.node_handle.set_storage_sync(key, value).await?; + Ok(true) } pub fn set_logging_enabled(&self, enable: bool) -> Result<()> { @@ -470,7 +447,7 @@ mod tests { let balance_before = node.get_balance_impl(address, None).await.unwrap(); - let result = node.set_balance(address, U256::from(1337)).await; + let result = node.set_balance(address, U256::from(1337)).await.unwrap(); assert!(result); let balance_after = node.get_balance_impl(address, None).await.unwrap(); @@ -488,7 +465,7 @@ mod tests { .await .unwrap(); - let result = node.set_nonce(address, U256::from(1337)).await; + let result = node.set_nonce(address, U256::from(1337)).await.unwrap(); assert!(result); let nonce_after = node @@ -498,7 +475,7 @@ mod tests { assert_eq!(nonce_after, U256::from(1337)); assert_ne!(nonce_before, nonce_after); - let result = node.set_nonce(address, U256::from(1336)).await; + let result = node.set_nonce(address, U256::from(1336)).await.unwrap(); assert!(result); let nonce_after = node @@ -592,7 +569,7 @@ mod tests { .get_transaction_count_impl(address, None) .await .unwrap(); - assert!(node.set_nonce(address, U256::from(1337)).await); + assert!(node.set_nonce(address, U256::from(1337)).await.unwrap()); assert!(node.reset_network(None).await.unwrap()); @@ -619,7 +596,10 @@ mod tests { Address::from_str("0xd8da6bf26964af9d7eed9e03e53415d37aa96045").unwrap(); // give impersonated account some balance - let result = node.set_balance(to_impersonate, U256::exp10(18)).await; + let result = node + .set_balance(to_impersonate, U256::exp10(18)) + .await + .unwrap(); assert!(result); // construct a tx @@ -717,7 +697,7 @@ mod tests { let value_before = node.inner.write().await.fork_storage.read_value(&key); assert_eq!(H256::default(), value_before); - let result = node.set_storage_at(address, slot, value).await; + let result = node.set_storage_at(address, slot, value).await.unwrap(); assert!(result); let value_after = node.inner.write().await.fork_storage.read_value(&key); diff --git a/crates/core/src/node/inner/fork.rs b/crates/core/src/node/inner/fork.rs index 6332253d..78c58549 100644 --- a/crates/core/src/node/inner/fork.rs +++ b/crates/core/src/node/inner/fork.rs @@ -3,6 +3,7 @@ //! There is ForkStorage (that is a wrapper over InMemoryStorage) //! And ForkDetails - that parses network address and fork height from arguments. +use crate::node::inner::storage::ReadStorageDyn; use crate::utils::block_on; use crate::{deps::InMemoryStorage, http_fork_source::HttpForkSource}; use anvil_zksync_config::constants::{ @@ -10,6 +11,7 @@ use anvil_zksync_config::constants::{ DEFAULT_FAIR_PUBDATA_PRICE, TEST_NODE_NETWORK_ID, }; use anvil_zksync_config::types::{CacheConfig, SystemContractsOptions}; +use async_trait::async_trait; use eyre::eyre; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; @@ -323,6 +325,25 @@ impl ReadStorage for &ForkStorage { } } +#[async_trait] +impl ReadStorageDyn for ForkStorage { + fn dyn_cloned(&self) -> Box { + Box::new(self.clone()) + } + + async fn read_value_alt(&self, key: &StorageKey) -> anyhow::Result { + // TODO: Get rid of `block_on` inside to propagate asynchronous execution up to this level + self.read_value_internal(key) + .map_err(|e| anyhow::anyhow!("failed reading value: {:?}", e)) + } + + async fn load_factory_dep_alt(&self, hash: H256) -> anyhow::Result>> { + // TODO: Get rid of `block_on` inside to propagate asynchronous execution up to this level + self.load_factory_dep_internal(hash) + .map_err(|e| anyhow::anyhow!("failed to load factory dep: {:?}", e)) + } +} + impl ForkStorage { pub fn set_value(&self, key: StorageKey, value: zksync_types::StorageValue) { let mut mutator = self.inner.write().unwrap(); diff --git a/crates/core/src/node/inner/in_memory_inner.rs b/crates/core/src/node/inner/in_memory_inner.rs index 55daf3ea..ab09c3d9 100644 --- a/crates/core/src/node/inner/in_memory_inner.rs +++ b/crates/core/src/node/inner/in_memory_inner.rs @@ -15,11 +15,12 @@ use crate::node::{ MAX_PREVIOUS_STATES, MAX_TX_SIZE, }; use crate::system_contracts::SystemContracts; -use crate::utils::{bytecode_to_factory_dep, create_debug_output}; +use crate::utils::create_debug_output; use crate::{formatter, utils}; use anvil_zksync_config::constants::NON_FORK_FIRST_BLOCK_TIMESTAMP; use anvil_zksync_config::TestNodeConfig; use anvil_zksync_types::{ShowCalls, ShowGasDetails, ShowStorageLogs, ShowVMDetails}; +use anyhow::Context; use colored::Colorize; use indexmap::IndexMap; use once_cell::sync::OnceCell; @@ -45,6 +46,7 @@ use zksync_multivm::vm_latest::{HistoryDisabled, HistoryEnabled, ToTracerPointer use zksync_multivm::{HistoryMode, VmVersion}; use zksync_types::api::{BlockIdVariant, TransactionVariant}; use zksync_types::block::build_bloom; +use zksync_types::bytecode::BytecodeHash; use zksync_types::fee::Fee; use zksync_types::fee_model::{BatchFeeInput, PubdataIndependentBatchFeeModelInput}; use zksync_types::l2::{L2Tx, TransactionType}; @@ -55,7 +57,7 @@ use zksync_types::utils::{ use zksync_types::web3::{Bytes, Index}; use zksync_types::{ api, get_nonce_key, h256_to_address, h256_to_u256, u256_to_h256, AccountTreeId, Address, Bloom, - BloomInput, L1BatchNumber, L2BlockNumber, StorageKey, StorageValue, Transaction, + BloomInput, L1BatchNumber, L2BlockNumber, L2ChainId, StorageKey, StorageValue, Transaction, ACCOUNT_CODE_STORAGE_ADDRESS, H160, H256, MAX_L2_TX_GAS_LIMIT, U256, U64, }; use zksync_web3_decl::error::Web3Error; @@ -415,11 +417,9 @@ impl InMemoryNodeInner { let mut bytecodes = HashMap::new(); for b in &*compressed_bytecodes { - let (hash, bytecode) = bytecode_to_factory_dep(b.original.clone()).map_err(|err| { - tracing::error!("{}", format!("cannot convert bytecode: {err}").on_red()); - err - })?; - bytecodes.insert(hash, bytecode); + zksync_types::bytecode::validate_bytecode(&b.original).context("Invalid bytecode")?; + let hash = BytecodeHash::for_bytecode(&b.original).value(); + bytecodes.insert(hash, b.original.clone()); } Ok(TxExecutionOutput { @@ -466,17 +466,8 @@ impl InMemoryNodeInner { } // Write all the factory deps. - for (hash, code) in bytecodes.iter() { - self.fork_storage.store_factory_dep( - u256_to_h256(*hash), - code.iter() - .flat_map(|entry| { - let mut bytes = vec![0u8; 32]; - entry.to_big_endian(&mut bytes); - bytes.to_vec() - }) - .collect(), - ) + for (hash, code) in bytecodes { + self.fork_storage.store_factory_dep(hash, code) } let logs = result @@ -1282,13 +1273,22 @@ impl InMemoryNodeInner { } self.rich_accounts.insert(address); } + + pub fn read_storage(&self) -> Box { + Box::new(&self.fork_storage) + } + + // TODO: Remove, this should also be made available from somewhere else + pub fn chain_id(&self) -> L2ChainId { + self.fork_storage.chain_id + } } #[derive(Debug)] pub struct TxExecutionOutput { result: VmExecutionResultAndLogs, call_traces: Vec, - bytecodes: HashMap>, + bytecodes: HashMap>, } /// Keeps track of a block's batch number, miniblock number and timestamp. diff --git a/crates/core/src/node/inner/mod.rs b/crates/core/src/node/inner/mod.rs index 760af367..1aeb0238 100644 --- a/crates/core/src/node/inner/mod.rs +++ b/crates/core/src/node/inner/mod.rs @@ -13,12 +13,14 @@ pub mod blockchain; pub mod fork; mod in_memory_inner; pub mod node_executor; +pub mod storage; pub mod time; pub use in_memory_inner::{InMemoryNodeInner, TxExecutionOutput}; use crate::filters::EthFilters; use crate::node::blockchain::Blockchain; +use crate::node::inner::storage::ReadStorageDyn; use crate::node::{ImpersonationManager, TestNodeFeeInputProvider}; use crate::system_contracts::SystemContracts; use anvil_zksync_config::constants::NON_FORK_FIRST_BLOCK_TIMESTAMP; @@ -41,7 +43,7 @@ impl InMemoryNodeInner { system_contracts: SystemContracts, ) -> ( Arc>, - ForkStorage, + Box, Box, Box, ) { @@ -76,7 +78,7 @@ impl InMemoryNodeInner { ( Arc::new(RwLock::new(node_inner)), - fork_storage, + Box::new(fork_storage), Box::new(blockchain), Box::new(time), ) diff --git a/crates/core/src/node/inner/node_executor.rs b/crates/core/src/node/inner/node_executor.rs index b06bd9d4..2402a832 100644 --- a/crates/core/src/node/inner/node_executor.rs +++ b/crates/core/src/node/inner/node_executor.rs @@ -4,7 +4,11 @@ use crate::system_contracts::SystemContracts; use std::sync::Arc; use tokio::sync::{mpsc, oneshot, RwLock}; use zksync_multivm::interface::TxExecutionMode; -use zksync_types::L2BlockNumber; +use zksync_types::bytecode::BytecodeHash; +use zksync_types::utils::{nonces_to_full_nonce, storage_key_for_eth_balance}; +use zksync_types::{ + get_code_key, get_nonce_key, u256_to_h256, Address, L2BlockNumber, StorageKey, U256, +}; pub struct NodeExecutor { node_inner: Arc>, @@ -36,6 +40,18 @@ impl NodeExecutor { Command::SealBlocks(tx_batches, interval, reply) => { self.seal_blocks(tx_batches, interval, reply).await; } + Command::SetCode(address, code, reply) => { + self.set_code(address, code, reply).await; + } + Command::SetStorage(key, value, reply) => { + self.set_storage(key, value, reply).await; + } + Command::SetBalance(address, balance, reply) => { + self.set_balance(address, balance, reply).await; + } + Command::SetNonce(address, nonce, reply) => { + self.set_nonce(address, nonce, reply).await; + } Command::IncreaseTime(delta, reply) => { self.increase_time(delta, reply).await; } @@ -138,6 +154,64 @@ impl NodeExecutor { } } + async fn set_code(&self, address: Address, bytecode: Vec, reply: oneshot::Sender<()>) { + let code_key = get_code_key(&address); + let bytecode_hash = BytecodeHash::for_bytecode(&bytecode).value(); + // TODO: Likely fork_storage can be moved to `NodeExecutor` instead + let node_inner = self.node_inner.read().await; + node_inner + .fork_storage + .store_factory_dep(bytecode_hash, bytecode); + node_inner.fork_storage.set_value(code_key, bytecode_hash); + drop(node_inner); + // Reply to sender if we can + if reply.send(()).is_err() { + tracing::info!("failed to reply as receiver has been dropped"); + } + } + + async fn set_storage(&self, key: StorageKey, value: U256, reply: oneshot::Sender<()>) { + // TODO: Likely fork_storage can be moved to `NodeExecutor` instead + self.node_inner + .read() + .await + .fork_storage + .set_value(key, u256_to_h256(value)); + // Reply to sender if we can + if reply.send(()).is_err() { + tracing::info!("failed to reply as receiver has been dropped"); + } + } + + async fn set_balance(&self, address: Address, balance: U256, reply: oneshot::Sender<()>) { + let balance_key = storage_key_for_eth_balance(&address); + // TODO: Likely fork_storage can be moved to `NodeExecutor` instead + self.node_inner + .read() + .await + .fork_storage + .set_value(balance_key, u256_to_h256(balance)); + // Reply to sender if we can + if reply.send(()).is_err() { + tracing::info!("failed to reply as receiver has been dropped"); + } + } + + async fn set_nonce(&self, address: Address, nonce: U256, reply: oneshot::Sender<()>) { + let nonce_key = get_nonce_key(&address); + let enforced_full_nonce = nonces_to_full_nonce(nonce, nonce); + // TODO: Likely fork_storage can be moved to `NodeExecutor` instead + self.node_inner + .read() + .await + .fork_storage + .set_value(nonce_key, u256_to_h256(enforced_full_nonce)); + // Reply to sender if we can + if reply.send(()).is_err() { + tracing::info!("failed to reply as receiver has been dropped"); + } + } + async fn increase_time(&self, delta: u64, reply: oneshot::Sender<()>) { self.node_inner.write().await.time.increase_time(delta); // Reply to sender if we can @@ -267,6 +341,68 @@ impl NodeExecutorHandle { } } + /// Request [`NodeExecutor`] to set bytecode for given address. Waits for the change to take place. + pub async fn set_code_sync(&self, address: Address, bytecode: Vec) -> anyhow::Result<()> { + let (response_sender, response_receiver) = oneshot::channel(); + self.command_sender + .send(Command::SetCode(address, bytecode, response_sender)) + .await + .map_err(|_| anyhow::anyhow!("failed to set code as node executor is dropped"))?; + match response_receiver.await { + Ok(()) => Ok(()), + Err(_) => { + anyhow::bail!("failed to set code as node executor is dropped") + } + } + } + + /// Request [`NodeExecutor`] to set storage key-value pair. Waits for the change to take place. + pub async fn set_storage_sync(&self, key: StorageKey, value: U256) -> anyhow::Result<()> { + let (response_sender, response_receiver) = oneshot::channel(); + self.command_sender + .send(Command::SetStorage(key, value, response_sender)) + .await + .map_err(|_| anyhow::anyhow!("failed to set storage as node executor is dropped"))?; + match response_receiver.await { + Ok(()) => Ok(()), + Err(_) => { + anyhow::bail!("failed to set storage as node executor is dropped") + } + } + } + + /// Request [`NodeExecutor`] to set account's balance to the given value. Waits for the change + /// to take place. + pub async fn set_balance_sync(&self, address: Address, balance: U256) -> anyhow::Result<()> { + let (response_sender, response_receiver) = oneshot::channel(); + self.command_sender + .send(Command::SetBalance(address, balance, response_sender)) + .await + .map_err(|_| anyhow::anyhow!("failed to set balance as node executor is dropped"))?; + match response_receiver.await { + Ok(()) => Ok(()), + Err(_) => { + anyhow::bail!("failed to set balance as node executor is dropped") + } + } + } + + /// Request [`NodeExecutor`] to set account's nonce to the given value. Waits for the change + /// to take place. + pub async fn set_nonce_sync(&self, address: Address, nonce: U256) -> anyhow::Result<()> { + let (response_sender, response_receiver) = oneshot::channel(); + self.command_sender + .send(Command::SetNonce(address, nonce, response_sender)) + .await + .map_err(|_| anyhow::anyhow!("failed to set nonce as node executor is dropped"))?; + match response_receiver.await { + Ok(()) => Ok(()), + Err(_) => { + anyhow::bail!("failed to set nonce as node executor is dropped") + } + } + } + /// Request [`NodeExecutor`] to increase time by the given delta (in seconds). Waits for the /// change to take place. pub async fn increase_time_sync(&self, delta: u64) -> anyhow::Result<()> { @@ -357,6 +493,11 @@ enum Command { u64, oneshot::Sender>>, ), + // Storage manipulation commands + SetCode(Address, Vec, oneshot::Sender<()>), + SetStorage(StorageKey, U256, oneshot::Sender<()>), + SetBalance(Address, U256, oneshot::Sender<()>), + SetNonce(Address, U256, oneshot::Sender<()>), // Time manipulation commands. Caveat: reply-able commands can hold user connections alive for // a long time (until the command is processed). IncreaseTime(u64, oneshot::Sender<()>), diff --git a/crates/core/src/node/inner/storage.rs b/crates/core/src/node/inner/storage.rs new file mode 100644 index 00000000..eaa3c5b2 --- /dev/null +++ b/crates/core/src/node/inner/storage.rs @@ -0,0 +1,23 @@ +use async_trait::async_trait; +use zksync_multivm::interface::storage::ReadStorage; +use zksync_types::{StorageKey, StorageValue, H256}; + +#[async_trait] +pub trait ReadStorageDyn: ReadStorage + Send + Sync { + /// Alternative for [`Clone::clone`] that is object safe. + fn dyn_cloned(&self) -> Box; + + /// Alternative version of [`ReadStorage::read_value`] that is fallible, async and does not + /// require `&mut self`. + async fn read_value_alt(&self, key: &StorageKey) -> anyhow::Result; + + /// Alternative version of [`ReadStorage::load_factory_dep`] that is fallible, async and does not + /// require `&mut self`. + async fn load_factory_dep_alt(&self, hash: H256) -> anyhow::Result>>; +} + +impl Clone for Box { + fn clone(&self) -> Self { + self.dyn_cloned() + } +} diff --git a/crates/core/src/node/zks.rs b/crates/core/src/node/zks.rs index f8973875..3074527a 100644 --- a/crates/core/src/node/zks.rs +++ b/crates/core/src/node/zks.rs @@ -125,22 +125,13 @@ impl InMemoryNode { let tokens = self.get_confirmed_tokens_impl(0, 100).await?; let balances = { - let writer = self.inner.write().await; let mut balances = HashMap::new(); for token in tokens { let balance_key = storage_key_for_standard_token_balance( AccountTreeId::new(token.l2_address), &address, ); - let balance = match writer.fork_storage.read_value_internal(&balance_key) { - Ok(balance) => balance, - Err(error) => { - return Err(Web3Error::InternalError(anyhow::anyhow!( - "failed reading value: {:?}", - error - ))); - } - }; + let balance = self.storage.read_value_alt(&balance_key).await?; if !balance.is_zero() { balances.insert(token.l2_address, h256_to_u256(balance)); } @@ -223,19 +214,11 @@ impl InMemoryNode { } pub async fn get_bytecode_by_hash_impl(&self, hash: H256) -> anyhow::Result>> { - let writer = self.inner.write().await; - - let maybe_bytecode = match writer.fork_storage.load_factory_dep_internal(hash) { - Ok(maybe_bytecode) => maybe_bytecode, - Err(error) => { - return Err(anyhow::anyhow!("failed to load factory dep: {:?}", error)); - } - }; - - if maybe_bytecode.is_some() { - return Ok(maybe_bytecode); + if let Some(bytecode) = self.storage.load_factory_dep_alt(hash).await? { + return Ok(Some(bytecode)); } + let writer = self.inner.write().await; let maybe_fork_details = &writer .fork_storage .inner diff --git a/crates/core/src/utils.rs b/crates/core/src/utils.rs index 9b3756f9..f84d2b54 100644 --- a/crates/core/src/utils.rs +++ b/crates/core/src/utils.rs @@ -20,15 +20,6 @@ use zksync_types::{ }; use zksync_web3_decl::error::Web3Error; -pub(crate) fn bytes_to_be_words(bytes: &[u8]) -> Vec { - assert_eq!( - bytes.len() % 32, - 0, - "Bytes must be divisible by 32 to split into chunks" - ); - bytes.chunks(32).map(U256::from_big_endian).collect() -} - /// Takes long integers and returns them in human friendly format with "_". /// For example: 12_334_093 pub fn to_human_size(input: U256) -> String { @@ -48,21 +39,6 @@ pub fn to_human_size(input: U256) -> String { tmp.iter().rev().collect() } -// TODO: Approach to encoding bytecode has changed in the core, so this function is likely no -// longer needed. See `zksync_contracts::SystemContractCode` for general approach. -// -// Use of `bytecode_to_factory_dep` needs to be refactored and vendored `bytes_to_be_words` -// should be removed. -pub fn bytecode_to_factory_dep(bytecode: Vec) -> Result<(U256, Vec), anyhow::Error> { - zksync_types::bytecode::validate_bytecode(&bytecode).context("Invalid bytecode")?; - let bytecode_hash = zksync_types::bytecode::BytecodeHash::for_bytecode(&bytecode).value(); - let bytecode_hash = U256::from_big_endian(bytecode_hash.as_bytes()); - - let bytecode_words = bytes_to_be_words(&bytecode); - - Ok((bytecode_hash, bytecode_words)) -} - /// Returns the actual [U64] block number from [BlockNumber]. /// /// # Arguments