reth_trie_common/
proofs.rs

Help
1//! Merkle trie proofs.
2
3use crate::{Nibbles, TrieAccount};
4use alloc::{borrow::Cow, vec::Vec};
5use alloy_consensus::constants::KECCAK_EMPTY;
6use alloy_primitives::{
7    keccak256,
8    map::{hash_map, B256Map, B256Set, HashMap},
9    Address, Bytes, B256, U256,
10};
11use alloy_rlp::{encode_fixed_size, Decodable, EMPTY_STRING_CODE};
12use alloy_trie::{
13    nodes::TrieNode,
14    proof::{verify_proof, DecodedProofNodes, ProofNodes, ProofVerificationError},
15    TrieMask, EMPTY_ROOT_HASH,
16};
17use derive_more::{Deref, DerefMut, IntoIterator};
18use itertools::Itertools;
19use reth_primitives_traits::Account;
20
21/// Proof targets map.
22#[derive(Deref, DerefMut, IntoIterator, Clone, PartialEq, Eq, Default, Debug)]
23pub struct MultiProofTargets(B256Map<B256Set>);
24
25impl FromIterator<(B256, B256Set)> for MultiProofTargets {
26    fn from_iter<T: IntoIterator<Item = (B256, B256Set)>>(iter: T) -> Self {
27        Self(B256Map::from_iter(iter))
28    }
29}
30
31impl MultiProofTargets {
32    /// Creates an empty `MultiProofTargets` with at least the specified capacity.
33    pub fn with_capacity(capacity: usize) -> Self {
34        Self(B256Map::with_capacity_and_hasher(capacity, Default::default()))
35    }
36
37    /// Create `MultiProofTargets` with a single account as a target.
38    pub fn account(hashed_address: B256) -> Self {
39        Self::accounts([hashed_address])
40    }
41
42    /// Create `MultiProofTargets` with a single account and slots as targets.
43    pub fn account_with_slots<I: IntoIterator<Item = B256>>(
44        hashed_address: B256,
45        slots_iter: I,
46    ) -> Self {
47        Self(B256Map::from_iter([(hashed_address, slots_iter.into_iter().collect())]))
48    }
49
50    /// Create `MultiProofTargets` only from accounts.
51    pub fn accounts<I: IntoIterator<Item = B256>>(iter: I) -> Self {
52        Self(iter.into_iter().map(|hashed_address| (hashed_address, Default::default())).collect())
53    }
54
55    /// Retains the targets representing the difference,
56    /// i.e., the values that are in `self` but not in `other`.
57    pub fn retain_difference(&mut self, other: &Self) {
58        self.0.retain(|hashed_address, hashed_slots| {
59            if let Some(other_hashed_slots) = other.get(hashed_address) {
60                hashed_slots.retain(|hashed_slot| !other_hashed_slots.contains(hashed_slot));
61                !hashed_slots.is_empty()
62            } else {
63                true
64            }
65        });
66    }
67
68    /// Extend multi proof targets with contents of other.
69    pub fn extend(&mut self, other: Self) {
70        self.extend_inner(Cow::Owned(other));
71    }
72
73    /// Extend multi proof targets with contents of other.
74    ///
75    /// Slightly less efficient than [`Self::extend`], but preferred to `extend(other.clone())`.
76    pub fn extend_ref(&mut self, other: &Self) {
77        self.extend_inner(Cow::Borrowed(other));
78    }
79
80    fn extend_inner(&mut self, other: Cow<'_, Self>) {
81        for (hashed_address, hashed_slots) in other.iter() {
82            self.entry(*hashed_address).or_default().extend(hashed_slots);
83        }
84    }
85
86    /// Returns an iterator that yields chunks of the specified size.
87    ///
88    /// See [`ChunkedMultiProofTargets`] for more information.
89    pub fn chunks(self, size: usize) -> ChunkedMultiProofTargets {
90        ChunkedMultiProofTargets::new(self, size)
91    }
92}
93
94/// An iterator that yields chunks of the proof targets of at most `size` account and storage
95/// targets.
96///
97/// For example, for the following proof targets:
98/// ```text
99/// - 0x1: [0x10, 0x20, 0x30]
100/// - 0x2: [0x40]
101/// - 0x3: []
102/// ```
103///
104/// and `size = 2`, the iterator will yield the following chunks:
105/// ```text
106/// - { 0x1: [0x10, 0x20] }
107/// - { 0x1: [0x30], 0x2: [0x40] }
108/// - { 0x3: [] }
109/// ```
110///
111/// It follows two rules:
112/// - If account has associated storage slots, each storage slot is counted towards the chunk size.
113/// - If account has no associated storage slots, the account is counted towards the chunk size.
114#[derive(Debug)]
115pub struct ChunkedMultiProofTargets {
116    flattened_targets: alloc::vec::IntoIter<(B256, Option<B256>)>,
117    size: usize,
118}
119
120impl ChunkedMultiProofTargets {
121    fn new(targets: MultiProofTargets, size: usize) -> Self {
122        let flattened_targets = targets
123            .into_iter()
124            .flat_map(|(address, slots)| {
125                if slots.is_empty() {
126                    // If the account has no storage slots, we still need to yield the account
127                    // address with empty storage slots. `None` here means that
128                    // there's no storage slot to fetch.
129                    itertools::Either::Left(core::iter::once((address, None)))
130                } else {
131                    itertools::Either::Right(
132                        slots.into_iter().map(move |slot| (address, Some(slot))),
133                    )
134                }
135            })
136            .sorted();
137        Self { flattened_targets, size }
138    }
139}
140
141impl Iterator for ChunkedMultiProofTargets {
142    type Item = MultiProofTargets;
143
144    fn next(&mut self) -> Option<Self::Item> {
145        let chunk = self.flattened_targets.by_ref().take(self.size).fold(
146            MultiProofTargets::default(),
147            |mut acc, (address, slot)| {
148                let entry = acc.entry(address).or_default();
149                if let Some(slot) = slot {
150                    entry.insert(slot);
151                }
152                acc
153            },
154        );
155
156        if chunk.is_empty() {
157            None
158        } else {
159            Some(chunk)
160        }
161    }
162}
163
164/// The state multiproof of target accounts and multiproofs of their storage tries.
165/// Multiproof is effectively a state subtrie that only contains the nodes
166/// in the paths of target accounts.
167#[derive(Clone, Default, Debug, PartialEq, Eq)]
168pub struct MultiProof {
169    /// State trie multiproof for requested accounts.
170    pub account_subtree: ProofNodes,
171    /// The hash masks of the branch nodes in the account proof.
172    pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
173    /// The tree masks of the branch nodes in the account proof.
174    pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
175    /// Storage trie multiproofs.
176    pub storages: B256Map<StorageMultiProof>,
177}
178
179impl MultiProof {
180    /// Returns true if the multiproof is empty.
181    pub fn is_empty(&self) -> bool {
182        self.account_subtree.is_empty() &&
183            self.branch_node_hash_masks.is_empty() &&
184            self.branch_node_tree_masks.is_empty() &&
185            self.storages.is_empty()
186    }
187
188    /// Return the account proof nodes for the given account path.
189    pub fn account_proof_nodes(&self, path: &Nibbles) -> Vec<(Nibbles, Bytes)> {
190        self.account_subtree.matching_nodes_sorted(path)
191    }
192
193    /// Return the storage proof nodes for the given storage slots of the account path.
194    pub fn storage_proof_nodes(
195        &self,
196        hashed_address: B256,
197        slots: impl IntoIterator<Item = B256>,
198    ) -> Vec<(B256, Vec<(Nibbles, Bytes)>)> {
199        self.storages
200            .get(&hashed_address)
201            .map(|storage_mp| {
202                slots
203                    .into_iter()
204                    .map(|slot| {
205                        let nibbles = Nibbles::unpack(slot);
206                        (slot, storage_mp.subtree.matching_nodes_sorted(&nibbles))
207                    })
208                    .collect()
209            })
210            .unwrap_or_default()
211    }
212
213    /// Construct the account proof from the multiproof.
214    pub fn account_proof(
215        &self,
216        address: Address,
217        slots: &[B256],
218    ) -> Result<AccountProof, alloy_rlp::Error> {
219        let hashed_address = keccak256(address);
220        let nibbles = Nibbles::unpack(hashed_address);
221
222        // Retrieve the account proof.
223        let proof = self
224            .account_proof_nodes(&nibbles)
225            .into_iter()
226            .map(|(_, node)| node)
227            .collect::<Vec<_>>();
228
229        // Inspect the last node in the proof. If it's a leaf node with matching suffix,
230        // then the node contains the encoded trie account.
231        let info = 'info: {
232            if let Some(last) = proof.last() {
233                if let TrieNode::Leaf(leaf) = TrieNode::decode(&mut &last[..])? {
234                    if nibbles.ends_with(&leaf.key) {
235                        let account = TrieAccount::decode(&mut &leaf.value[..])?;
236                        break 'info Some(Account {
237                            balance: account.balance,
238                            nonce: account.nonce,
239                            bytecode_hash: (account.code_hash != KECCAK_EMPTY)
240                                .then_some(account.code_hash),
241                        })
242                    }
243                }
244            }
245            None
246        };
247
248        // Retrieve proofs for requested storage slots.
249        let storage_multiproof = self.storages.get(&hashed_address);
250        let storage_root = storage_multiproof.map(|m| m.root).unwrap_or(EMPTY_ROOT_HASH);
251        let mut storage_proofs = Vec::with_capacity(slots.len());
252        for slot in slots {
253            let proof = if let Some(multiproof) = &storage_multiproof {
254                multiproof.storage_proof(*slot)?
255            } else {
256                StorageProof::new(*slot)
257            };
258            storage_proofs.push(proof);
259        }
260        Ok(AccountProof { address, info, proof, storage_root, storage_proofs })
261    }
262
263    /// Extends this multiproof with another one, merging both account and storage
264    /// proofs.
265    pub fn extend(&mut self, other: Self) {
266        self.account_subtree.extend_from(other.account_subtree);
267
268        self.branch_node_hash_masks.extend(other.branch_node_hash_masks);
269        self.branch_node_tree_masks.extend(other.branch_node_tree_masks);
270
271        for (hashed_address, storage) in other.storages {
272            match self.storages.entry(hashed_address) {
273                hash_map::Entry::Occupied(mut entry) => {
274                    debug_assert_eq!(entry.get().root, storage.root);
275                    let entry = entry.get_mut();
276                    entry.subtree.extend_from(storage.subtree);
277                    entry.branch_node_hash_masks.extend(storage.branch_node_hash_masks);
278                    entry.branch_node_tree_masks.extend(storage.branch_node_tree_masks);
279                }
280                hash_map::Entry::Vacant(entry) => {
281                    entry.insert(storage);
282                }
283            }
284        }
285    }
286}
287
288/// This is a type of [`MultiProof`] that uses decoded proofs, meaning these proofs are stored as a
289/// collection of [`TrieNode`]s instead of RLP-encoded bytes.
290#[derive(Clone, Default, Debug, PartialEq, Eq)]
291pub struct DecodedMultiProof {
292    /// State trie multiproof for requested accounts.
293    pub account_subtree: DecodedProofNodes,
294    /// The hash masks of the branch nodes in the account proof.
295    pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
296    /// The tree masks of the branch nodes in the account proof.
297    pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
298    /// Storage trie multiproofs.
299    pub storages: B256Map<DecodedStorageMultiProof>,
300}
301
302impl DecodedMultiProof {
303    /// Return the account proof nodes for the given account path.
304    pub fn account_proof_nodes(&self, path: &Nibbles) -> Vec<(Nibbles, TrieNode)> {
305        self.account_subtree.matching_nodes_sorted(path)
306    }
307
308    /// Return the storage proof nodes for the given storage slots of the account path.
309    pub fn storage_proof_nodes(
310        &self,
311        hashed_address: B256,
312        slots: impl IntoIterator<Item = B256>,
313    ) -> Vec<(B256, Vec<(Nibbles, TrieNode)>)> {
314        self.storages
315            .get(&hashed_address)
316            .map(|storage_mp| {
317                slots
318                    .into_iter()
319                    .map(|slot| {
320                        let nibbles = Nibbles::unpack(slot);
321                        (slot, storage_mp.subtree.matching_nodes_sorted(&nibbles))
322                    })
323                    .collect()
324            })
325            .unwrap_or_default()
326    }
327
328    /// Construct the account proof from the multiproof.
329    pub fn account_proof(
330        &self,
331        address: Address,
332        slots: &[B256],
333    ) -> Result<DecodedAccountProof, alloy_rlp::Error> {
334        let hashed_address = keccak256(address);
335        let nibbles = Nibbles::unpack(hashed_address);
336
337        // Retrieve the account proof.
338        let proof = self
339            .account_proof_nodes(&nibbles)
340            .into_iter()
341            .map(|(_, node)| node)
342            .collect::<Vec<_>>();
343
344        // Inspect the last node in the proof. If it's a leaf node with matching suffix,
345        // then the node contains the encoded trie account.
346        let info = 'info: {
347            if let Some(TrieNode::Leaf(leaf)) = proof.last() {
348                if nibbles.ends_with(&leaf.key) {
349                    let account = TrieAccount::decode(&mut &leaf.value[..])?;
350                    break 'info Some(Account {
351                        balance: account.balance,
352                        nonce: account.nonce,
353                        bytecode_hash: (account.code_hash != KECCAK_EMPTY)
354                            .then_some(account.code_hash),
355                    })
356                }
357            }
358            None
359        };
360
361        // Retrieve proofs for requested storage slots.
362        let storage_multiproof = self.storages.get(&hashed_address);
363        let storage_root = storage_multiproof.map(|m| m.root).unwrap_or(EMPTY_ROOT_HASH);
364        let mut storage_proofs = Vec::with_capacity(slots.len());
365        for slot in slots {
366            let proof = if let Some(multiproof) = &storage_multiproof {
367                multiproof.storage_proof(*slot)?
368            } else {
369                DecodedStorageProof::new(*slot)
370            };
371            storage_proofs.push(proof);
372        }
373        Ok(DecodedAccountProof { address, info, proof, storage_root, storage_proofs })
374    }
375
376    /// Extends this multiproof with another one, merging both account and storage
377    /// proofs.
378    pub fn extend(&mut self, other: Self) {
379        self.account_subtree.extend_from(other.account_subtree);
380
381        self.branch_node_hash_masks.extend(other.branch_node_hash_masks);
382        self.branch_node_tree_masks.extend(other.branch_node_tree_masks);
383
384        for (hashed_address, storage) in other.storages {
385            match self.storages.entry(hashed_address) {
386                hash_map::Entry::Occupied(mut entry) => {
387                    debug_assert_eq!(entry.get().root, storage.root);
388                    let entry = entry.get_mut();
389                    entry.subtree.extend_from(storage.subtree);
390                    entry.branch_node_hash_masks.extend(storage.branch_node_hash_masks);
391                    entry.branch_node_tree_masks.extend(storage.branch_node_tree_masks);
392                }
393                hash_map::Entry::Vacant(entry) => {
394                    entry.insert(storage);
395                }
396            }
397        }
398    }
399}
400
401/// The merkle multiproof of storage trie.
402#[derive(Clone, Debug, PartialEq, Eq)]
403pub struct StorageMultiProof {
404    /// Storage trie root.
405    pub root: B256,
406    /// Storage multiproof for requested slots.
407    pub subtree: ProofNodes,
408    /// The hash masks of the branch nodes in the storage proof.
409    pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
410    /// The tree masks of the branch nodes in the storage proof.
411    pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
412}
413
414impl StorageMultiProof {
415    /// Create new storage multiproof for empty trie.
416    pub fn empty() -> Self {
417        Self {
418            root: EMPTY_ROOT_HASH,
419            subtree: ProofNodes::from_iter([(
420                Nibbles::default(),
421                Bytes::from([EMPTY_STRING_CODE]),
422            )]),
423            branch_node_hash_masks: HashMap::default(),
424            branch_node_tree_masks: HashMap::default(),
425        }
426    }
427
428    /// Return storage proofs for the target storage slot (unhashed).
429    pub fn storage_proof(&self, slot: B256) -> Result<StorageProof, alloy_rlp::Error> {
430        let nibbles = Nibbles::unpack(keccak256(slot));
431
432        // Retrieve the storage proof.
433        let proof = self
434            .subtree
435            .matching_nodes_iter(&nibbles)
436            .sorted_by(|a, b| a.0.cmp(b.0))
437            .map(|(_, node)| node.clone())
438            .collect::<Vec<_>>();
439
440        // Inspect the last node in the proof. If it's a leaf node with matching suffix,
441        // then the node contains the encoded slot value.
442        let value = 'value: {
443            if let Some(last) = proof.last() {
444                if let TrieNode::Leaf(leaf) = TrieNode::decode(&mut &last[..])? {
445                    if nibbles.ends_with(&leaf.key) {
446                        break 'value U256::decode(&mut &leaf.value[..])?
447                    }
448                }
449            }
450            U256::ZERO
451        };
452
453        Ok(StorageProof { key: slot, nibbles, value, proof })
454    }
455}
456
457/// The decoded merkle multiproof for a storage trie.
458#[derive(Clone, Debug, PartialEq, Eq)]
459pub struct DecodedStorageMultiProof {
460    /// Storage trie root.
461    pub root: B256,
462    /// Storage multiproof for requested slots.
463    pub subtree: DecodedProofNodes,
464    /// The hash masks of the branch nodes in the storage proof.
465    pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
466    /// The tree masks of the branch nodes in the storage proof.
467    pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
468}
469
470impl DecodedStorageMultiProof {
471    /// Create new storage multiproof for empty trie.
472    pub fn empty() -> Self {
473        Self {
474            root: EMPTY_ROOT_HASH,
475            subtree: DecodedProofNodes::from_iter([(Nibbles::default(), TrieNode::EmptyRoot)]),
476            branch_node_hash_masks: HashMap::default(),
477            branch_node_tree_masks: HashMap::default(),
478        }
479    }
480
481    /// Return storage proofs for the target storage slot (unhashed).
482    pub fn storage_proof(&self, slot: B256) -> Result<DecodedStorageProof, alloy_rlp::Error> {
483        let nibbles = Nibbles::unpack(keccak256(slot));
484
485        // Retrieve the storage proof.
486        let proof = self
487            .subtree
488            .matching_nodes_iter(&nibbles)
489            .sorted_by(|a, b| a.0.cmp(b.0))
490            .map(|(_, node)| node.clone())
491            .collect::<Vec<_>>();
492
493        // Inspect the last node in the proof. If it's a leaf node with matching suffix,
494        // then the node contains the encoded slot value.
495        let value = 'value: {
496            if let Some(TrieNode::Leaf(leaf)) = proof.last() {
497                if nibbles.ends_with(&leaf.key) {
498                    break 'value U256::decode(&mut &leaf.value[..])?
499                }
500            }
501            U256::ZERO
502        };
503
504        Ok(DecodedStorageProof { key: slot, nibbles, value, proof })
505    }
506}
507
508/// The merkle proof with the relevant account info.
509#[derive(Clone, PartialEq, Eq, Debug)]
510#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
511#[cfg_attr(any(test, feature = "serde"), serde(rename_all = "camelCase"))]
512pub struct AccountProof {
513    /// The address associated with the account.
514    pub address: Address,
515    /// Account info.
516    pub info: Option<Account>,
517    /// Array of rlp-serialized merkle trie nodes which starting from the root node and
518    /// following the path of the hashed address as key.
519    pub proof: Vec<Bytes>,
520    /// The storage trie root.
521    pub storage_root: B256,
522    /// Array of storage proofs as requested.
523    pub storage_proofs: Vec<StorageProof>,
524}
525
526#[cfg(feature = "eip1186")]
527impl AccountProof {
528    /// Convert into an EIP-1186 account proof response
529    pub fn into_eip1186_response(
530        self,
531        slots: Vec<alloy_serde::JsonStorageKey>,
532    ) -> alloy_rpc_types_eth::EIP1186AccountProofResponse {
533        let info = self.info.unwrap_or_default();
534        alloy_rpc_types_eth::EIP1186AccountProofResponse {
535            address: self.address,
536            balance: info.balance,
537            code_hash: info.get_bytecode_hash(),
538            nonce: info.nonce,
539            storage_hash: self.storage_root,
540            account_proof: self.proof,
541            storage_proof: self
542                .storage_proofs
543                .into_iter()
544                .filter_map(|proof| {
545                    let input_slot = slots.iter().find(|s| s.as_b256() == proof.key)?;
546                    Some(proof.into_eip1186_proof(*input_slot))
547                })
548                .collect(),
549        }
550    }
551}
552
553impl Default for AccountProof {
554    fn default() -> Self {
555        Self::new(Address::default())
556    }
557}
558
559impl AccountProof {
560    /// Create new account proof entity.
561    pub const fn new(address: Address) -> Self {
562        Self {
563            address,
564            info: None,
565            proof: Vec::new(),
566            storage_root: EMPTY_ROOT_HASH,
567            storage_proofs: Vec::new(),
568        }
569    }
570
571    /// Verify the storage proofs and account proof against the provided state root.
572    pub fn verify(&self, root: B256) -> Result<(), ProofVerificationError> {
573        // Verify storage proofs.
574        for storage_proof in &self.storage_proofs {
575            storage_proof.verify(self.storage_root)?;
576        }
577
578        // Verify the account proof.
579        let expected = if self.info.is_none() && self.storage_root == EMPTY_ROOT_HASH {
580            None
581        } else {
582            Some(alloy_rlp::encode(
583                self.info.unwrap_or_default().into_trie_account(self.storage_root),
584            ))
585        };
586        let nibbles = Nibbles::unpack(keccak256(self.address));
587        verify_proof(root, nibbles, expected, &self.proof)
588    }
589}
590
591/// The merkle proof with the relevant account info.
592#[derive(Clone, PartialEq, Eq, Debug)]
593pub struct DecodedAccountProof {
594    /// The address associated with the account.
595    pub address: Address,
596    /// Account info.
597    pub info: Option<Account>,
598    /// Array of merkle trie nodes which starting from the root node and following the path of the
599    /// hashed address as key.
600    pub proof: Vec<TrieNode>,
601    /// The storage trie root.
602    pub storage_root: B256,
603    /// Array of storage proofs as requested.
604    pub storage_proofs: Vec<DecodedStorageProof>,
605}
606
607impl Default for DecodedAccountProof {
608    fn default() -> Self {
609        Self::new(Address::default())
610    }
611}
612
613impl DecodedAccountProof {
614    /// Create new account proof entity.
615    pub const fn new(address: Address) -> Self {
616        Self {
617            address,
618            info: None,
619            proof: Vec::new(),
620            storage_root: EMPTY_ROOT_HASH,
621            storage_proofs: Vec::new(),
622        }
623    }
624}
625
626/// The merkle proof of the storage entry.
627#[derive(Clone, PartialEq, Eq, Default, Debug)]
628#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
629pub struct StorageProof {
630    /// The raw storage key.
631    pub key: B256,
632    /// The hashed storage key nibbles.
633    pub nibbles: Nibbles,
634    /// The storage value.
635    pub value: U256,
636    /// Array of rlp-serialized merkle trie nodes which starting from the storage root node and
637    /// following the path of the hashed storage slot as key.
638    pub proof: Vec<Bytes>,
639}
640
641impl StorageProof {
642    /// Convert into an EIP-1186 storage proof
643    #[cfg(feature = "eip1186")]
644    pub fn into_eip1186_proof(
645        self,
646        slot: alloy_serde::JsonStorageKey,
647    ) -> alloy_rpc_types_eth::EIP1186StorageProof {
648        alloy_rpc_types_eth::EIP1186StorageProof { key: slot, value: self.value, proof: self.proof }
649    }
650}
651
652impl StorageProof {
653    /// Create new storage proof from the storage slot.
654    pub fn new(key: B256) -> Self {
655        let nibbles = Nibbles::unpack(keccak256(key));
656        Self { key, nibbles, ..Default::default() }
657    }
658
659    /// Create new storage proof from the storage slot and its pre-hashed image.
660    pub fn new_with_hashed(key: B256, hashed_key: B256) -> Self {
661        Self { key, nibbles: Nibbles::unpack(hashed_key), ..Default::default() }
662    }
663
664    /// Create new storage proof from the storage slot and its pre-hashed image.
665    pub fn new_with_nibbles(key: B256, nibbles: Nibbles) -> Self {
666        Self { key, nibbles, ..Default::default() }
667    }
668
669    /// Set proof nodes on storage proof.
670    pub fn with_proof(mut self, proof: Vec<Bytes>) -> Self {
671        self.proof = proof;
672        self
673    }
674
675    /// Verify the proof against the provided storage root.
676    pub fn verify(&self, root: B256) -> Result<(), ProofVerificationError> {
677        let expected =
678            if self.value.is_zero() { None } else { Some(encode_fixed_size(&self.value).to_vec()) };
679        verify_proof(root, self.nibbles.clone(), expected, &self.proof)
680    }
681}
682
683/// The merkle proof of the storage entry, using decoded proofs.
684#[derive(Clone, PartialEq, Eq, Default, Debug)]
685pub struct DecodedStorageProof {
686    /// The raw storage key.
687    pub key: B256,
688    /// The hashed storage key nibbles.
689    pub nibbles: Nibbles,
690    /// The storage value.
691    pub value: U256,
692    /// Array of merkle trie nodes which starting from the storage root node and following the path
693    /// of the hashed storage slot as key.
694    pub proof: Vec<TrieNode>,
695}
696
697impl DecodedStorageProof {
698    /// Create new storage proof from the storage slot.
699    pub fn new(key: B256) -> Self {
700        let nibbles = Nibbles::unpack(keccak256(key));
701        Self { key, nibbles, ..Default::default() }
702    }
703
704    /// Create new storage proof from the storage slot and its pre-hashed image.
705    pub fn new_with_hashed(key: B256, hashed_key: B256) -> Self {
706        Self { key, nibbles: Nibbles::unpack(hashed_key), ..Default::default() }
707    }
708
709    /// Create new storage proof from the storage slot and its pre-hashed image.
710    pub fn new_with_nibbles(key: B256, nibbles: Nibbles) -> Self {
711        Self { key, nibbles, ..Default::default() }
712    }
713
714    /// Set proof nodes on storage proof.
715    pub fn with_proof(mut self, proof: Vec<TrieNode>) -> Self {
716        self.proof = proof;
717        self
718    }
719}
720
721/// Implementation of hasher using our keccak256 hashing function
722/// for compatibility with `triehash` crate.
723#[cfg(any(test, feature = "test-utils"))]
724pub mod triehash {
725    use alloy_primitives::{keccak256, B256};
726    use alloy_rlp::RlpEncodable;
727    use hash_db::Hasher;
728    use plain_hasher::PlainHasher;
729
730    /// A [Hasher] that calculates a keccak256 hash of the given data.
731    #[derive(Default, Debug, Clone, PartialEq, Eq, RlpEncodable)]
732    #[non_exhaustive]
733    pub struct KeccakHasher;
734
735    #[cfg(any(test, feature = "test-utils"))]
736    impl Hasher for KeccakHasher {
737        type Out = B256;
738        type StdHasher = PlainHasher;
739
740        const LENGTH: usize = 32;
741
742        fn hash(x: &[u8]) -> Self::Out {
743            keccak256(x)
744        }
745    }
746}
747
748#[cfg(test)]
749mod tests {
750    use super::*;
751
752    #[test]
753    fn test_multiproof_extend_account_proofs() {
754        let mut proof1 = MultiProof::default();
755        let mut proof2 = MultiProof::default();
756
757        let addr1 = B256::random();
758        let addr2 = B256::random();
759
760        proof1.account_subtree.insert(
761            Nibbles::unpack(addr1),
762            alloy_rlp::encode_fixed_size(&U256::from(42)).to_vec().into(),
763        );
764        proof2.account_subtree.insert(
765            Nibbles::unpack(addr2),
766            alloy_rlp::encode_fixed_size(&U256::from(43)).to_vec().into(),
767        );
768
769        proof1.extend(proof2);
770
771        assert!(proof1.account_subtree.contains_key(&Nibbles::unpack(addr1)));
772        assert!(proof1.account_subtree.contains_key(&Nibbles::unpack(addr2)));
773    }
774
775    #[test]
776    fn test_multiproof_extend_storage_proofs() {
777        let mut proof1 = MultiProof::default();
778        let mut proof2 = MultiProof::default();
779
780        let addr = B256::random();
781        let root = B256::random();
782
783        let mut subtree1 = ProofNodes::default();
784        subtree1.insert(
785            Nibbles::from_nibbles(vec![0]),
786            alloy_rlp::encode_fixed_size(&U256::from(42)).to_vec().into(),
787        );
788        proof1.storages.insert(
789            addr,
790            StorageMultiProof {
791                root,
792                subtree: subtree1,
793                branch_node_hash_masks: HashMap::default(),
794                branch_node_tree_masks: HashMap::default(),
795            },
796        );
797
798        let mut subtree2 = ProofNodes::default();
799        subtree2.insert(
800            Nibbles::from_nibbles(vec![1]),
801            alloy_rlp::encode_fixed_size(&U256::from(43)).to_vec().into(),
802        );
803        proof2.storages.insert(
804            addr,
805            StorageMultiProof {
806                root,
807                subtree: subtree2,
808                branch_node_hash_masks: HashMap::default(),
809                branch_node_tree_masks: HashMap::default(),
810            },
811        );
812
813        proof1.extend(proof2);
814
815        let storage = proof1.storages.get(&addr).unwrap();
816        assert_eq!(storage.root, root);
817        assert!(storage.subtree.contains_key(&Nibbles::from_nibbles(vec![0])));
818        assert!(storage.subtree.contains_key(&Nibbles::from_nibbles(vec![1])));
819    }
820
821    #[test]
822    fn test_multi_proof_retain_difference() {
823        let mut empty = MultiProofTargets::default();
824        empty.retain_difference(&Default::default());
825        assert!(empty.is_empty());
826
827        let targets = MultiProofTargets::accounts((0..10).map(B256::with_last_byte));
828
829        let mut diffed = targets.clone();
830        diffed.retain_difference(&MultiProofTargets::account(B256::with_last_byte(11)));
831        assert_eq!(diffed, targets);
832
833        diffed.retain_difference(&MultiProofTargets::accounts((0..5).map(B256::with_last_byte)));
834        assert_eq!(diffed, MultiProofTargets::accounts((5..10).map(B256::with_last_byte)));
835
836        diffed.retain_difference(&targets);
837        assert!(diffed.is_empty());
838
839        let mut targets = MultiProofTargets::default();
840        let (account1, account2, account3) =
841            (1..=3).map(B256::with_last_byte).collect_tuple().unwrap();
842        let account2_slots = (1..5).map(B256::with_last_byte).collect::<B256Set>();
843        targets.insert(account1, B256Set::from_iter([B256::with_last_byte(1)]));
844        targets.insert(account2, account2_slots.clone());
845        targets.insert(account3, B256Set::from_iter([B256::with_last_byte(1)]));
846
847        let mut diffed = targets.clone();
848        diffed.retain_difference(&MultiProofTargets::accounts((1..=3).map(B256::with_last_byte)));
849        assert_eq!(diffed, targets);
850
851        // remove last 3 slots for account 2
852        let mut account2_slots_expected_len = account2_slots.len();
853        for slot in account2_slots.iter().skip(1) {
854            diffed.retain_difference(&MultiProofTargets::account_with_slots(account2, [*slot]));
855            account2_slots_expected_len -= 1;
856            assert_eq!(
857                diffed.get(&account2).map(|slots| slots.len()),
858                Some(account2_slots_expected_len)
859            );
860        }
861
862        diffed.retain_difference(&targets);
863        assert!(diffed.is_empty());
864    }
865
866    #[test]
867    fn test_multi_proof_retain_difference_no_overlap() {
868        let mut targets = MultiProofTargets::default();
869
870        // populate some targets
871        let (addr1, addr2) = (B256::random(), B256::random());
872        let (slot1, slot2) = (B256::random(), B256::random());
873        targets.insert(addr1, vec![slot1].into_iter().collect());
874        targets.insert(addr2, vec![slot2].into_iter().collect());
875
876        let mut retained = targets.clone();
877        retained.retain_difference(&Default::default());
878        assert_eq!(retained, targets);
879
880        // add a different addr and slot to fetched proof targets
881        let mut other_targets = MultiProofTargets::default();
882        let addr3 = B256::random();
883        let slot3 = B256::random();
884        other_targets.insert(addr3, B256Set::from_iter([slot3]));
885
886        // check that the prefetch proof targets are the same because the fetched proof targets
887        // don't overlap with the prefetch targets
888        let mut retained = targets.clone();
889        retained.retain_difference(&other_targets);
890        assert_eq!(retained, targets);
891    }
892
893    #[test]
894    fn test_get_prefetch_proof_targets_remove_subset() {
895        // populate some targets
896        let mut targets = MultiProofTargets::default();
897        let (addr1, addr2) = (B256::random(), B256::random());
898        let (slot1, slot2) = (B256::random(), B256::random());
899        targets.insert(addr1, B256Set::from_iter([slot1]));
900        targets.insert(addr2, B256Set::from_iter([slot2]));
901
902        // add a subset of the first target to other proof targets
903        let other_targets = MultiProofTargets::account_with_slots(addr1, [slot1]);
904
905        let mut retained = targets.clone();
906        retained.retain_difference(&other_targets);
907
908        // check that the prefetch proof targets do not include the subset
909        assert_eq!(retained.len(), 1);
910        assert!(!retained.contains_key(&addr1));
911        assert!(retained.contains_key(&addr2));
912
913        // now add one more slot to the prefetch targets
914        let slot3 = B256::random();
915        targets.get_mut(&addr1).unwrap().insert(slot3);
916
917        let mut retained = targets.clone();
918        retained.retain_difference(&other_targets);
919
920        // check that the prefetch proof targets do not include the subset
921        // but include the new slot
922        assert_eq!(retained.len(), 2);
923        assert!(retained.contains_key(&addr1));
924        assert_eq!(retained.get(&addr1), Some(&B256Set::from_iter([slot3])));
925        assert!(retained.contains_key(&addr2));
926        assert_eq!(retained.get(&addr2), Some(&B256Set::from_iter([slot2])));
927    }
928}
reth_trie_common/proofs.rs

reth_trie_common/
proofs.rs
