diff --git a/.github/workflows/pr-main_l2.yaml b/.github/workflows/pr-main_l2.yaml index 9f81dec502f..feb97d0c8d1 100644 --- a/.github/workflows/pr-main_l2.yaml +++ b/.github/workflows/pr-main_l2.yaml @@ -45,6 +45,8 @@ jobs: - "fixtures/contracts/**" - "crates/blockchain/dev/**" - "crates/vm/levm/**" + - "crates/common/types/block_execution_witness.rs" + - "crates/networking/rpc/debug/execution_witness.rs" - ".github/workflows/pr-main_l2.yaml" - "cmd/ethrex/l2/**" non_docs: diff --git a/crates/blockchain/blockchain.rs b/crates/blockchain/blockchain.rs index d68e6117050..8e44587303c 100644 --- a/crates/blockchain/blockchain.rs +++ b/crates/blockchain/blockchain.rs @@ -1552,15 +1552,6 @@ impl Blockchain { block_headers_bytes.push(current_header.encode_to_vec()); } - // Create a list of all read/write addresses and storage slots - let mut keys = Vec::new(); - for (address, touched_storage_slots) in touched_account_storage_slots { - keys.push(address.as_bytes().to_vec()); - for slot in touched_storage_slots.iter() { - keys.push(slot.as_bytes().to_vec()); - } - } - // Get initial state trie root and embed the rest of the trie into it let nodes: BTreeMap = used_trie_nodes .into_iter() @@ -1581,12 +1572,9 @@ impl Blockchain { Trie::new_temp() }; let mut storage_trie_roots = BTreeMap::new(); - for key in &keys { - if key.len() != 20 { - continue; // not an address - } - let address = Address::from_slice(key); - let hashed_address = hash_address(&address); + for address in touched_account_storage_slots.keys() { + let hashed_address = hash_address(address); + let hashed_address_h256 = H256::from_slice(&hashed_address); let Some(encoded_account) = state_trie.get(&hashed_address)? else { continue; // empty account, doesn't have a storage trie }; @@ -1603,7 +1591,7 @@ impl Blockchain { "execution witness does not contain non-empty storage trie".to_string(), )); }; - storage_trie_roots.insert(address, (*node).clone()); + storage_trie_roots.insert(hashed_address_h256, (*node).clone()); } Ok(ExecutionWitness { @@ -1613,7 +1601,6 @@ impl Blockchain { chain_config: self.storage.get_chain_config(), state_trie_root, storage_trie_roots, - keys, }) } @@ -1794,15 +1781,6 @@ impl Blockchain { block_headers_bytes.push(current_header.encode_to_vec()); } - // Create a list of all read/write addresses and storage slots - let mut keys = Vec::new(); - for (address, touched_storage_slots) in touched_account_storage_slots { - keys.push(address.as_bytes().to_vec()); - for slot in touched_storage_slots.iter() { - keys.push(slot.as_bytes().to_vec()); - } - } - // Get initial state trie root and embed the rest of the trie into it let nodes: BTreeMap = used_trie_nodes .into_iter() @@ -1823,12 +1801,9 @@ impl Blockchain { Trie::new_temp() }; let mut storage_trie_roots = BTreeMap::new(); - for key in &keys { - if key.len() != 20 { - continue; // not an address - } - let address = Address::from_slice(key); - let hashed_address = hash_address(&address); + for address in touched_account_storage_slots.keys() { + let hashed_address = hash_address(address); + let hashed_address_h256 = H256::from_slice(&hashed_address); let Some(encoded_account) = state_trie.get(&hashed_address)? else { continue; // empty account, doesn't have a storage trie }; @@ -1845,7 +1820,7 @@ impl Blockchain { "execution witness does not contain non-empty storage trie".to_string(), )); }; - storage_trie_roots.insert(address, (*node).clone()); + storage_trie_roots.insert(hashed_address_h256, (*node).clone()); } Ok(ExecutionWitness { @@ -1855,7 +1830,6 @@ impl Blockchain { chain_config: self.storage.get_chain_config(), state_trie_root, storage_trie_roots, - keys, }) } diff --git a/crates/common/types/block_execution_witness.rs b/crates/common/types/block_execution_witness.rs index 28b4f1a913a..09be4a3d8d0 100644 --- a/crates/common/types/block_execution_witness.rs +++ b/crates/common/types/block_execution_witness.rs @@ -2,7 +2,7 @@ use std::collections::{BTreeMap, BTreeSet}; use bytes::Bytes; -use crate::rkyv_utils::H160Wrapper; +use crate::rkyv_utils::H256Wrapper; use crate::serde_utils; use crate::types::{Block, Code, CodeMetadata}; use crate::{ @@ -42,17 +42,17 @@ pub struct GuestProgramState { pub first_block_number: u64, /// The chain configuration. pub chain_config: ChainConfig, - /// Map of storage root hashes to their corresponding storage tries. - pub storage_tries: BTreeMap, + /// Map of hashed addresses to their corresponding storage tries. + pub storage_tries: BTreeMap, /// Map of account addresses to their corresponding hashed addresses. /// This is a convenience map to avoid recomputing the hashed address /// multiple times during guest program execution. /// It is built on-demand during guest program execution, inside the zkVM. - pub account_hashes_by_address: BTreeMap>, - /// Map of account addresses to booleans, indicating whose account's storage tries were + pub account_hashes_by_address: BTreeMap, + /// Map of hashed addresses to booleans, indicating whose account's storage tries were /// verified. /// Verification is done by hashing the trie and comparing the root hash with the account's storage root. - pub verified_storage_roots: BTreeMap, + pub verified_storage_roots: BTreeMap, } /// Witness data produced by the client and consumed by the guest program @@ -76,13 +76,10 @@ pub struct ExecutionWitness { pub chain_config: ChainConfig, /// Root node embedded with the rest of the trie's nodes pub state_trie_root: Option, - /// Root nodes per account storage embedded with the rest of the trie's nodes - #[rkyv(with = MapKV)] - pub storage_trie_roots: BTreeMap, - /// Flattened map of account addresses and storage keys whose values - /// are needed for stateless execution. - #[rkyv(with = crate::rkyv_utils::VecVecWrapper)] - pub keys: Vec>, + /// Root nodes per account storage embedded with the rest of the trie's nodes, + /// keyed by the keccak256 hash of the account address. + #[rkyv(with = MapKV)] + pub storage_trie_roots: BTreeMap, } /// RPC-friendly representation of an execution witness. @@ -98,6 +95,7 @@ pub struct RpcExecutionWitness { )] pub state: Vec, #[serde( + default, serialize_with = "serde_utils::bytes::vec::serialize", deserialize_with = "serde_utils::bytes::vec::deserialize" )] @@ -127,7 +125,7 @@ impl TryFrom for RpcExecutionWitness { } Ok(Self { state: nodes.into_iter().map(Bytes::from).collect(), - keys: value.keys.into_iter().map(Bytes::from).collect(), + keys: Vec::new(), codes: value.codes.into_iter().map(Bytes::from).collect(), headers: value .block_headers_bytes @@ -199,11 +197,11 @@ impl TryFrom for GuestProgramState { state_trie.hash_no_commit(); let mut storage_tries = BTreeMap::new(); - for (address, storage_trie_root) in value.storage_trie_roots { + for (hashed_address, storage_trie_root) in value.storage_trie_roots { // hash storage trie nodes let storage_trie = Trie::new_temp_with_root(storage_trie_root.into()); storage_trie.hash_no_commit(); - storage_tries.insert(address, storage_trie); + storage_tries.insert(hashed_address, storage_trie); } // hash codes @@ -242,18 +240,18 @@ impl GuestProgramState { account_updates: &[AccountUpdate], ) -> Result<(), GuestProgramStateError> { for update in account_updates.iter() { - let hashed_address = self + let hashed_address = *self .account_hashes_by_address .entry(update.address) .or_insert_with(|| hash_address(&update.address)); if update.removed { // Remove account from trie - self.state_trie.remove(hashed_address)?; + self.state_trie.remove(hashed_address.as_bytes())?; } else { // Add or update AccountState in the trie // Fetch current state or create a new state to be inserted - let mut account_state = match self.state_trie.get(hashed_address)? { + let mut account_state = match self.state_trie.get(hashed_address.as_bytes())? { Some(encoded_state) => AccountState::decode(&encoded_state)?, None => AccountState::default(), }; @@ -271,7 +269,7 @@ impl GuestProgramState { } // Store the added storage in the account's storage trie and compute its new root if !update.added_storage.is_empty() { - let storage_trie = self.storage_tries.entry(update.address).or_default(); + let storage_trie = self.storage_tries.entry(hashed_address).or_default(); // Inserts must come before deletes, otherwise deletes might require extra nodes // Example: @@ -295,8 +293,10 @@ impl GuestProgramState { account_state.storage_root = storage_root; } - self.state_trie - .insert(hashed_address.clone(), account_state.encode_to_vec())?; + self.state_trie.insert( + hashed_address.as_bytes().to_vec(), + account_state.encode_to_vec(), + )?; } } Ok(()) @@ -362,12 +362,12 @@ impl GuestProgramState { &mut self, address: Address, ) -> Result, GuestProgramStateError> { - let hashed_address = self + let hashed_address = *self .account_hashes_by_address .entry(address) .or_insert_with(|| hash_address(&address)); - let Ok(Some(encoded_state)) = self.state_trie.get(hashed_address) else { + let Ok(Some(encoded_state)) = self.state_trie.get(hashed_address.as_bytes()) else { return Ok(None); }; let state = AccountState::decode(&encoded_state).map_err(|_| { @@ -500,16 +500,24 @@ impl GuestProgramState { &mut self, address: Address, ) -> Result, GuestProgramStateError> { - let is_storage_verified = *self.verified_storage_roots.get(&address).unwrap_or(&false); + let hashed_address = *self + .account_hashes_by_address + .entry(address) + .or_insert_with(|| hash_address(&address)); + + let is_storage_verified = *self + .verified_storage_roots + .get(&hashed_address) + .unwrap_or(&false); if is_storage_verified { - Ok(self.storage_tries.get(&address)) + Ok(self.storage_tries.get(&hashed_address)) } else { let Some(storage_root) = self.get_account_state(address)?.map(|a| a.storage_root) else { // empty account return Ok(None); }; - let storage_trie = match self.storage_tries.get(&address) { + let storage_trie = match self.storage_tries.get(&hashed_address) { None if storage_root == *EMPTY_TRIE_HASH => return Ok(None), Some(trie) if trie.hash_no_commit() == storage_root => trie, _ => { @@ -518,14 +526,14 @@ impl GuestProgramState { ))); } }; - self.verified_storage_roots.insert(address, true); + self.verified_storage_roots.insert(hashed_address, true); Ok(Some(storage_trie)) } } } -fn hash_address(address: &Address) -> Vec { - keccak_hash(address.to_fixed_bytes()).to_vec() +fn hash_address(address: &Address) -> H256 { + H256(keccak_hash(address.to_fixed_bytes())) } pub fn hash_key(key: &H256) -> Vec { diff --git a/crates/networking/rpc/debug/execution_witness.rs b/crates/networking/rpc/debug/execution_witness.rs index c7459416bf8..39efc39e9db 100644 --- a/crates/networking/rpc/debug/execution_witness.rs +++ b/crates/networking/rpc/debug/execution_witness.rs @@ -2,7 +2,7 @@ use std::collections::BTreeMap; use bytes::Bytes; use ethrex_common::{ - Address, H256, + H256, types::{ AccountState, BlockHeader, ChainConfig, block_execution_witness::{ExecutionWitness, GuestProgramStateError, RpcExecutionWitness}, @@ -10,8 +10,7 @@ use ethrex_common::{ utils::keccak, }; use ethrex_rlp::{decode::RLPDecode, error::RLPDecodeError}; -use ethrex_storage::hash_address; -use ethrex_trie::{EMPTY_TRIE_HASH, Node, NodeRef, Trie}; +use ethrex_trie::{EMPTY_TRIE_HASH, Nibbles, Node, NodeRef, Trie}; use serde_json::Value; use tracing::debug; @@ -64,36 +63,33 @@ pub fn execution_witness_from_rpc_chain_config( None }; - // get all storage trie roots and embed the rest of the trie into it - let state_trie = if let Some(state_trie_root) = &state_trie_root { - Trie::new_temp_with_root(state_trie_root.clone().into()) - } else { - Trie::new_temp() - }; + // Walk the state trie to discover accounts and their storage roots, + // instead of relying on the keys field which is being removed from the RPC spec. let mut storage_trie_roots = BTreeMap::new(); - for key in &rpc_witness.keys { - if key.len() != 20 { - continue; // not an address - } - let address = Address::from_slice(key); - let hashed_address = hash_address(&address); - let Some(encoded_account) = state_trie.get(&hashed_address)? else { - continue; // empty account, doesn't have a storage trie - }; - let storage_root_hash = AccountState::decode(&encoded_account)?.storage_root; - if storage_root_hash == *EMPTY_TRIE_HASH { - continue; // empty storage trie - } - if !nodes.contains_key(&storage_root_hash) { - continue; // storage trie isn't relevant to this execution + if let Some(state_trie_root) = &state_trie_root { + let mut accounts = Vec::new(); + collect_accounts_from_node( + state_trie_root, + Nibbles::from_raw(&[], false), + &mut accounts, + &nodes, + ); + + for (hashed_address, storage_root_hash) in accounts { + if storage_root_hash == *EMPTY_TRIE_HASH { + continue; // empty storage trie + } + if !nodes.contains_key(&storage_root_hash) { + continue; // storage trie isn't relevant to this execution + } + let node = Trie::get_embedded_root(&nodes, storage_root_hash)?; + let NodeRef::Node(node, _) = node else { + return Err(GuestProgramStateError::Custom( + "execution witness does not contain non-empty storage trie".to_string(), + )); + }; + storage_trie_roots.insert(hashed_address, (*node).clone()); } - let node = Trie::get_embedded_root(&nodes, storage_root_hash)?; - let NodeRef::Node(node, _) = node else { - return Err(GuestProgramStateError::Custom( - "execution witness does not contain non-empty storage trie".to_string(), - )); - }; - storage_trie_roots.insert(address, (*node).clone()); } let witness = ExecutionWitness { @@ -107,12 +103,74 @@ pub fn execution_witness_from_rpc_chain_config( .collect(), state_trie_root, storage_trie_roots, - keys: rpc_witness.keys.into_iter().map(|b| b.to_vec()).collect(), }; Ok(witness) } +/// Recursively walks an embedded state trie node and collects +/// `(hashed_address, storage_root)` pairs from leaf nodes. +/// Also resolves `NodeRef::Hash` references using the flat `nodes` map, +/// in case some children weren't fully embedded by `get_embedded_root`. +fn collect_accounts_from_node( + node: &Node, + path: Nibbles, + accounts: &mut Vec<(H256, H256)>, + nodes: &BTreeMap, +) { + match node { + Node::Branch(branch) => { + for (i, child) in branch.choices.iter().enumerate() { + let child_node: Option<&Node> = match child { + NodeRef::Node(n, _) => Some(n), + NodeRef::Hash(hash) if hash.is_valid() => nodes.get(&hash.finalize()), + _ => None, + }; + if let Some(child_node) = child_node { + collect_accounts_from_node( + child_node, + path.append_new(i as u8), + accounts, + nodes, + ); + } + } + } + Node::Extension(ext) => { + let child_node: Option<&Node> = match &ext.child { + NodeRef::Node(n, _) => Some(n), + NodeRef::Hash(hash) if hash.is_valid() => nodes.get(&hash.finalize()), + _ => None, + }; + if let Some(child_node) = child_node { + collect_accounts_from_node(child_node, path.concat(&ext.prefix), accounts, nodes); + } + } + Node::Leaf(leaf) => { + let full_path = path.concat(&leaf.partial); + let path_bytes = full_path.to_bytes(); + if path_bytes.len() == 32 { + let hashed_address = H256::from_slice(&path_bytes); + match AccountState::decode(&leaf.value) { + Ok(account_state) => { + accounts.push((hashed_address, account_state.storage_root)); + } + Err(e) => { + debug!( + "Failed to decode AccountState from state trie leaf at {hashed_address:?}: {e}" + ); + } + } + } else { + debug!( + "Unexpected state trie leaf path length: {} (expected 32)", + path_bytes.len() + ); + } + } + } +} + pub struct ExecutionWitnessRequest { pub from: BlockIdentifier, pub to: Option,