|
| 1 | +use std::collections::{HashMap, VecDeque}; |
| 2 | + |
| 3 | +use starknet_api::hash::HashOutput; |
| 4 | +use thiserror::Error; |
| 5 | + |
| 6 | +use crate::patricia_merkle_tree::node_data::inner_node::{ |
| 7 | + BinaryData, |
| 8 | + EdgeData, |
| 9 | + NodeData, |
| 10 | + Preimage, |
| 11 | + PreimageMap, |
| 12 | +}; |
| 13 | +use crate::patricia_merkle_tree::node_data::leaf::Leaf; |
| 14 | +use crate::patricia_merkle_tree::types::NodeIndex; |
| 15 | +use crate::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunction; |
| 16 | + |
| 17 | +#[derive(Debug, Error, PartialEq, Eq)] |
| 18 | +pub enum ProofVerificationError { |
| 19 | + #[error("missing leaf at index {index:?}")] |
| 20 | + MissingLeaf { index: NodeIndex }, |
| 21 | + #[error("hash mismatch at index {index:?}: proof {proof_value:?}, actual {actual:?}")] |
| 22 | + HashMismatch { index: NodeIndex, proof_value: HashOutput, actual: HashOutput }, |
| 23 | + #[error("duplicate parent for index {index:?}")] |
| 24 | + DuplicateParent { index: NodeIndex }, |
| 25 | +} |
| 26 | + |
| 27 | +/// Verifies that `preimages` form valid Patricia paths from each accessed leaf to `root_hash`. |
| 28 | +pub fn verify_patricia_proof<L: Leaf, TH: TreeHashFunction<L>>( |
| 29 | + root_hash: HashOutput, |
| 30 | + preimages: &PreimageMap, |
| 31 | + requested_leaves: &HashMap<NodeIndex, HashOutput>, |
| 32 | +) -> Result<(), ProofVerificationError> { |
| 33 | + if requested_leaves.is_empty() { |
| 34 | + return Ok(()); |
| 35 | + } |
| 36 | + |
| 37 | + let hash_by_index = build_proof_index_maps::<L, TH>(root_hash, preimages)?; |
| 38 | + |
| 39 | + // Once an index-->hash map is built, it suffices to verify that the expected leaf hashes are |
| 40 | + // present in the map, since we already verified the validity of the path from the root to the |
| 41 | + // leaf during the map construction. |
| 42 | + for (leaf_index, actual_leaf_hash) in requested_leaves { |
| 43 | + match hash_by_index.get(leaf_index) { |
| 44 | + None => return Err(ProofVerificationError::MissingLeaf { index: *leaf_index }), |
| 45 | + Some(proof_value) if *proof_value != *actual_leaf_hash => { |
| 46 | + return Err(ProofVerificationError::HashMismatch { |
| 47 | + index: *leaf_index, |
| 48 | + proof_value: *proof_value, |
| 49 | + actual: *actual_leaf_hash, |
| 50 | + }); |
| 51 | + } |
| 52 | + Some(_) => {} |
| 53 | + } |
| 54 | + } |
| 55 | + |
| 56 | + Ok(()) |
| 57 | +} |
| 58 | + |
| 59 | +/// Builds an `index -> hash` map by expanding a Patricia proof from the root. |
| 60 | +/// |
| 61 | +/// Verifies that the supplied hashes are consistent with the supplied preimages. |
| 62 | +pub fn build_proof_index_maps<L: Leaf, TH: TreeHashFunction<L>>( |
| 63 | + root_hash: HashOutput, |
| 64 | + preimages: &PreimageMap, |
| 65 | +) -> Result<HashMap<NodeIndex, HashOutput>, ProofVerificationError> { |
| 66 | + let mut hash_by_index = HashMap::from([(NodeIndex::ROOT, root_hash)]); |
| 67 | + let mut queue = VecDeque::from([NodeIndex::ROOT]); |
| 68 | + |
| 69 | + while let Some(index) = queue.pop_front() { |
| 70 | + let hash = hash_by_index[&index]; |
| 71 | + |
| 72 | + let Some(preimage) = preimages.get(&hash) else { |
| 73 | + // We reach a child-node of a node in `preimages` that is supposedly not required in |
| 74 | + // the proof (e.g. sibling of a node with no requested leaves in its subtree). |
| 75 | + continue; |
| 76 | + }; |
| 77 | + |
| 78 | + let computed_hash = TH::compute_node_hash(&preimage_to_node_data::<L>(preimage)); |
| 79 | + if hash != computed_hash { |
| 80 | + return Err(ProofVerificationError::HashMismatch { |
| 81 | + index, |
| 82 | + proof_value: hash, |
| 83 | + actual: computed_hash, |
| 84 | + }); |
| 85 | + } |
| 86 | + |
| 87 | + match preimage { |
| 88 | + Preimage::Binary(binary) => { |
| 89 | + let [left_index, right_index] = index.get_children_indices(); |
| 90 | + register_child(&mut hash_by_index, &mut queue, left_index, binary.left_data)?; |
| 91 | + register_child(&mut hash_by_index, &mut queue, right_index, binary.right_data)?; |
| 92 | + } |
| 93 | + Preimage::Edge(edge) => { |
| 94 | + let bottom_index = edge.path_to_bottom.bottom_index(index); |
| 95 | + register_child(&mut hash_by_index, &mut queue, bottom_index, edge.bottom_data)?; |
| 96 | + } |
| 97 | + } |
| 98 | + } |
| 99 | + |
| 100 | + Ok(hash_by_index) |
| 101 | +} |
| 102 | + |
| 103 | +fn register_child( |
| 104 | + hash_by_index: &mut HashMap<NodeIndex, HashOutput>, |
| 105 | + queue: &mut VecDeque<NodeIndex>, |
| 106 | + child_index: NodeIndex, |
| 107 | + child_hash: HashOutput, |
| 108 | +) -> Result<(), ProofVerificationError> { |
| 109 | + if hash_by_index.contains_key(&child_index) { |
| 110 | + return Err(ProofVerificationError::DuplicateParent { index: child_index }); |
| 111 | + } |
| 112 | + hash_by_index.insert(child_index, child_hash); |
| 113 | + queue.push_back(child_index); |
| 114 | + Ok(()) |
| 115 | +} |
| 116 | + |
| 117 | +fn preimage_to_node_data<L: Leaf>(preimage: &Preimage) -> NodeData<L, HashOutput> { |
| 118 | + match preimage { |
| 119 | + Preimage::Binary(binary) => NodeData::Binary(BinaryData { |
| 120 | + left_data: binary.left_data, |
| 121 | + right_data: binary.right_data, |
| 122 | + }), |
| 123 | + Preimage::Edge(edge) => NodeData::Edge(EdgeData { |
| 124 | + bottom_data: edge.bottom_data, |
| 125 | + path_to_bottom: edge.path_to_bottom, |
| 126 | + }), |
| 127 | + } |
| 128 | +} |
0 commit comments