From 40faed6ea05a778642a703627ce912dfb8894cfb Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Wed, 9 Oct 2024 17:59:28 +0200 Subject: [PATCH] fix(witness): destroyed slots as proof targets (#11596) --- crates/trie/db/tests/witness.rs | 38 +++++++++++++++++++++++++++- crates/trie/trie/src/witness.rs | 45 ++++++++++++++++++++++++--------- 2 files changed, 70 insertions(+), 13 deletions(-) diff --git a/crates/trie/db/tests/witness.rs b/crates/trie/db/tests/witness.rs index cc921f6570..59656383de 100644 --- a/crates/trie/db/tests/witness.rs +++ b/crates/trie/db/tests/witness.rs @@ -6,7 +6,7 @@ use alloy_primitives::{ Address, Bytes, B256, U256, }; use alloy_rlp::EMPTY_STRING_CODE; -use reth_primitives::{constants::EMPTY_ROOT_HASH, Account}; +use reth_primitives::{constants::EMPTY_ROOT_HASH, Account, StorageEntry}; use reth_provider::{test_utils::create_test_provider_factory, HashingWriter}; use reth_trie::{proof::Proof, witness::TrieWitness, HashedPostState, HashedStorage, StateRoot}; use reth_trie_db::{DatabaseProof, DatabaseStateRoot, DatabaseTrieWitness}; @@ -55,3 +55,39 @@ fn includes_empty_node_preimage() { // witness includes empty state trie root node assert_eq!(witness.get(&EMPTY_ROOT_HASH), Some(&Bytes::from([EMPTY_STRING_CODE]))); } + +#[test] +fn includes_nodes_for_destroyed_storage_nodes() { + let factory = create_test_provider_factory(); + let provider = factory.provider_rw().unwrap(); + + let address = Address::random(); + let hashed_address = keccak256(address); + let slot = B256::random(); + let hashed_slot = keccak256(slot); + + // Insert account and slot into database + provider.insert_account_for_hashing([(address, Some(Account::default()))]).unwrap(); + provider + .insert_storage_for_hashing([(address, [StorageEntry { key: slot, value: U256::from(1) }])]) + .unwrap(); + + let state_root = StateRoot::from_tx(provider.tx_ref()).root().unwrap(); + let multiproof = Proof::from_tx(provider.tx_ref()) + .multiproof(HashMap::from_iter([(hashed_address, HashSet::from_iter([hashed_slot]))])) + .unwrap(); + + let witness = TrieWitness::from_tx(provider.tx_ref()) + .compute(HashedPostState { + accounts: HashMap::from([(hashed_address, Some(Account::default()))]), + storages: HashMap::from([(hashed_address, HashedStorage::from_iter(true, []))]), // destroyed + }) + .unwrap(); + assert!(witness.contains_key(&state_root)); + for node in multiproof.account_subtree.values() { + assert_eq!(witness.get(&keccak256(node)), Some(node)); + } + for node in multiproof.storages.iter().flat_map(|(_, storage)| storage.subtree.values()) { + assert_eq!(witness.get(&keccak256(node)), Some(node)); + } +} diff --git a/crates/trie/trie/src/witness.rs b/crates/trie/trie/src/witness.rs index d668946e62..c042a0d821 100644 --- a/crates/trie/trie/src/witness.rs +++ b/crates/trie/trie/src/witness.rs @@ -1,7 +1,7 @@ use std::collections::BTreeMap; use crate::{ - hashed_cursor::HashedCursorFactory, + hashed_cursor::{HashedCursor, HashedCursorFactory}, prefix_set::TriePrefixSetsMut, proof::{Proof, StorageProof}, trie_cursor::TrieCursorFactory, @@ -14,7 +14,7 @@ use alloy_primitives::{ }; use alloy_rlp::{BufMut, Decodable, Encodable}; use itertools::{Either, Itertools}; -use reth_execution_errors::TrieWitnessError; +use reth_execution_errors::{StateProofError, TrieWitnessError}; use reth_primitives::constants::EMPTY_ROOT_HASH; use reth_trie_common::{ BranchNode, HashBuilder, Nibbles, StorageMultiProof, TrieAccount, TrieNode, CHILD_INDEX_RANGE, @@ -90,16 +90,7 @@ where return Ok(self.witness) } - let proof_targets = HashMap::from_iter( - state - .accounts - .keys() - .map(|hashed_address| (*hashed_address, HashSet::default())) - .chain(state.storages.iter().map(|(hashed_address, storage)| { - (*hashed_address, storage.storage.keys().copied().collect()) - })), - ); - + let proof_targets = self.get_proof_targets(&state)?; let mut account_multiproof = Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone()) .with_prefix_sets_mut(self.prefix_sets.clone()) @@ -257,6 +248,36 @@ where Ok(trie_nodes) } + /// Retrieve proof targets for incoming hashed state. + /// This method will aggregate all accounts and slots present in the hash state as well as + /// select all existing slots from the database for the accounts that have been destroyed. + fn get_proof_targets( + &self, + state: &HashedPostState, + ) -> Result>, StateProofError> { + let mut proof_targets = HashMap::default(); + for hashed_address in state.accounts.keys() { + proof_targets.insert(*hashed_address, HashSet::default()); + } + for (hashed_address, storage) in &state.storages { + let mut storage_keys = storage.storage.keys().copied().collect::>(); + if storage.wiped { + // storage for this account was destroyed, gather all slots from the current state + let mut storage_cursor = + self.hashed_cursor_factory.hashed_storage_cursor(*hashed_address)?; + // position cursor at the start + if let Some((hashed_slot, _)) = storage_cursor.seek(B256::ZERO)? { + storage_keys.insert(hashed_slot); + } + while let Some((hashed_slot, _)) = storage_cursor.next()? { + storage_keys.insert(hashed_slot); + } + } + proof_targets.insert(*hashed_address, storage_keys); + } + Ok(proof_targets) + } + fn next_root_from_proofs( trie_nodes: BTreeMap>>, mut trie_node_provider: impl FnMut(Nibbles) -> Result,