cfx_storage/impls/merkle_patricia_trie/
trie_proof.rs

1// Copyright 2019 Conflux Foundation. All rights reserved.
2// Conflux is free software and distributed under GNU General Public License.
3// See http://www.gnu.org/licenses/
4
5#[derive(Clone, Debug, Default, PartialEq)]
6pub struct TrieProof {
7    /// The first node must be the root node. Child node must come later than
8    /// one of its parent node.
9    nodes: Vec<TrieProofNode>,
10    merkle_to_node_index: HashMap<MerkleHash, usize>,
11    number_leaf_nodes: u32,
12}
13
14/// The node is the child_index child of the node at parent_node_index.
15#[derive(Clone, Debug, PartialEq)]
16struct ParentInfo {
17    parent_node_index: usize,
18    child_index: u8,
19}
20
21#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
22pub struct TrieProofNode(VanillaTrieNode<MerkleHash>);
23
24impl TrieProofNode {
25    pub fn new(
26        children_table: VanillaChildrenTable<MerkleHash>,
27        maybe_value: Option<Box<[u8]>>, compressed_path: CompressedPathRaw,
28        path_without_first_nibble: bool,
29    ) -> Self {
30        let merkle = compute_merkle(
31            compressed_path.as_ref(),
32            path_without_first_nibble,
33            if children_table.get_children_count() == 0 {
34                None
35            } else {
36                Some(children_table.get_children_table())
37            },
38            maybe_value.as_ref().map(|v| v.as_ref()),
39        );
40        Self(VanillaTrieNode::new(
41            merkle,
42            children_table,
43            maybe_value,
44            compressed_path,
45        ))
46    }
47}
48
49impl TrieProof {
50    pub const MAX_NODES: usize = 1000;
51
52    /// Makes sure that the proof nodes are valid and connected at the time of
53    /// creation.
54    pub fn new(nodes: Vec<TrieProofNode>) -> Result<Self> {
55        let merkle_to_node_index = nodes
56            .iter()
57            .enumerate()
58            .map(|(index, node)| (node.get_merkle().clone(), index))
59            .collect::<HashMap<H256, usize>>();
60        let number_nodes = nodes.len();
61
62        // Connectivity check.
63        let mut connected_child_parent_map =
64            HashMap::<MerkleHash, Vec<ParentInfo>, RandomState>::default();
65        if let Some(node) = nodes.get(0) {
66            connected_child_parent_map
67                .entry(node.get_merkle().clone())
68                .or_insert(vec![]);
69        }
70        for (node_index, node) in nodes.iter().enumerate() {
71            if !connected_child_parent_map.contains_key(node.get_merkle()) {
72                // Not connected.
73                bail!(Error::InvalidTrieProof);
74            }
75            for (child_index, child_merkle) in
76                node.get_children_table_ref().iter()
77            {
78                connected_child_parent_map
79                    .entry(child_merkle.clone())
80                    .or_insert(vec![])
81                    .push(ParentInfo {
82                        parent_node_index: node_index,
83                        child_index,
84                    });
85            }
86        }
87
88        let mut is_non_leaf = vec![false; number_nodes];
89        for node in &nodes {
90            if let Some(parent_infos) =
91                connected_child_parent_map.get(node.get_merkle())
92            {
93                for parent_info in parent_infos {
94                    is_non_leaf[parent_info.parent_node_index] = true;
95                }
96            }
97        }
98
99        let mut number_leaf_nodes = 0;
100        for non_leaf in is_non_leaf {
101            if !non_leaf {
102                number_leaf_nodes += 1;
103            }
104        }
105
106        Ok(TrieProof {
107            nodes,
108            merkle_to_node_index,
109            number_leaf_nodes,
110        })
111    }
112
113    pub fn get_merkle_root(&self) -> &MerkleHash {
114        match self.nodes.get(0) {
115            None => &MERKLE_NULL_NODE,
116            Some(node) => node.get_merkle(),
117        }
118    }
119
120    /// Verify that the trie `root` has `value` under `key`.
121    /// Use `None` for exclusion proofs (i.e. there is no value under `key`).
122    // NOTE: This API cannot be used to prove that there is a value under `key`
123    // but it does not equal a given value. We add this later on if it's needed.
124    pub fn is_valid_kv(
125        &self, key: &[u8], value: Option<&[u8]>, root: &MerkleHash,
126    ) -> bool {
127        self.is_valid(key, root, |node| match node {
128            None => value == None,
129            Some(node) => value == node.value_as_slice().into_option(),
130        })
131    }
132
133    /// Check if the key can be proved. The only reason of inability to prove a
134    /// key is missing nodes.
135    pub fn if_proves_key(&self, key: &[u8]) -> (bool, Option<&TrieProofNode>) {
136        let mut proof_node = None;
137        let proof_node_mut = &mut proof_node;
138        let proves = self.is_valid(key, self.get_merkle_root(), |maybe_node| {
139            *proof_node_mut = maybe_node.clone();
140            true
141        });
142        (proves, proof_node)
143    }
144
145    /// Get the value under `key` starting from `root`.
146    pub fn get_value(
147        &self, key: &[u8], root: &MerkleHash,
148    ) -> (bool, Option<&[u8]>) {
149        let mut proof_node = None;
150        let proof_node_mut = &mut proof_node;
151
152        let proves = self.is_valid(key, root, |maybe_node| {
153            *proof_node_mut = maybe_node.clone();
154            true
155        });
156
157        (
158            proves,
159            proof_node.and_then(|node| node.value_as_slice().into_option()),
160        )
161    }
162
163    #[cfg(test)]
164    /// Verify that the trie `root` has a node with `hash` under `path`.
165    /// Use `MERKLE_NULL_NODE` for exclusion proofs (i.e. `path` does not exist
166    /// or leads to another hash).
167    pub fn is_valid_path_to(
168        &self, path: &[u8], hash: &MerkleHash, root: &MerkleHash,
169    ) -> bool {
170        self.is_valid(path, root, |node| match node {
171            None => hash.eq(&MERKLE_NULL_NODE),
172            Some(node) => hash == node.get_merkle(),
173        })
174    }
175
176    /// Verify that the trie `root` has a node with `node_merkle` under `key`.
177    /// If `node_merkle` is `None`, then it must prove that the trie does not
178    /// contain `key`.
179    /// If `node_merkle` is `Tombstone`, then it must prove that the trie
180    /// contains a tombstone value under `key`.
181    /// If `node_merkle` is `Some(hash)`, then it must prove that the trie
182    /// contains a node under `key` whose node merkle equals `hash`.
183    pub fn is_valid_node_merkle(
184        &self, key: &[u8], node_merkle: &MptValue<MerkleHash>,
185        root: &MerkleHash,
186    ) -> bool {
187        self.is_valid(key, root, |node| match node {
188            None => node_merkle == &MptValue::None,
189            Some(node) if node.value_as_slice().is_tombstone() => {
190                node_merkle == &MptValue::TombStone
191            }
192            Some(node) => {
193                let h = node.get_merkle_hash_wo_compressed_path();
194                node_merkle == &MptValue::Some(h)
195            }
196        })
197    }
198
199    fn is_valid<'this: 'pred_param, 'pred_param>(
200        &'this self, path: &[u8], root: &MerkleHash,
201        pred: impl FnOnce(Option<&'pred_param TrieProofNode>) -> bool,
202    ) -> bool {
203        // empty trie
204        if root == &MERKLE_NULL_NODE {
205            return pred(None);
206        }
207
208        // NOTE: an empty proof is only valid if it is an
209        // exclusion proof for an empty trie, covered above
210
211        // traverse the trie along `path`
212        let mut key = path;
213        let mut hash = root;
214
215        loop {
216            let node = match self.merkle_to_node_index.get(hash) {
217                Some(node_index) => self
218                    .nodes
219                    .get(*node_index)
220                    .expect("Proof node guaranteed to exist"),
221                None => {
222                    // Missing node. The proof can be invalid or incomplete for
223                    // the key.
224                    return false;
225                }
226            };
227
228            match node.walk::<access_mode::Read>(key) {
229                WalkStop::Arrived => {
230                    return pred(Some(node));
231                }
232                WalkStop::PathDiverted { .. } => {
233                    return pred(None);
234                }
235                WalkStop::ChildNotFound { .. } => {
236                    return pred(None);
237                }
238                WalkStop::Descent {
239                    key_remaining,
240                    child_node,
241                    ..
242                } => {
243                    hash = child_node;
244                    key = key_remaining;
245                }
246            }
247        }
248    }
249
250    #[inline]
251    pub fn number_nodes(&self) -> usize { self.nodes.len() }
252
253    #[inline]
254    pub fn number_leaf_nodes(&self) -> u32 { self.number_leaf_nodes }
255
256    #[inline]
257    pub fn get_proof_nodes(&self) -> &Vec<TrieProofNode> { &self.nodes }
258
259    #[inline]
260    pub fn into_proof_nodes(self) -> Vec<TrieProofNode> { self.nodes }
261
262    /// Returns the (snapshot_mpt_key, child_index, trie_node) along the proof
263    /// path of key.
264    pub fn compute_snapshot_mpt_path_for_proof(
265        &self, key: &[u8],
266    ) -> Vec<(CompressedPathRaw, u8, &VanillaTrieNode<MerkleHash>)> {
267        let mut full_paths = Vec::with_capacity(self.nodes.len());
268        let mut keys_child_indices_and_nodes =
269            Vec::with_capacity(self.nodes.len());
270        // At the time of coding, The root node is guaranteed to have empty
271        // compressed_path. But to be on the safe side, we use the root
272        // node's compressed path as its full path.
273        let merkle_root;
274        // root node isn't a child, so we use CHILDREN_COUNT to distinguish
275        let mut child_index = CHILDREN_COUNT as u8;
276        if let Some(node) = self.nodes.get(0) {
277            full_paths.push(node.compressed_path_ref().into());
278            keys_child_indices_and_nodes.push((
279                CompressedPathRaw::default(),
280                child_index,
281                &**node,
282            ));
283            merkle_root = *node.get_merkle();
284        } else {
285            return vec![];
286        }
287
288        let mut key_remaining = key;
289        let mut hash = &merkle_root;
290
291        loop {
292            let node = match self.merkle_to_node_index.get(hash) {
293                Some(node_index) => self
294                    .nodes
295                    .get(*node_index)
296                    .expect("Proof node guaranteed to exist"),
297                None => {
298                    // Missing node. The proof can be invalid or incomplete for
299                    // the key.
300                    return vec![];
301                }
302            };
303
304            if child_index < CHILDREN_COUNT as u8 {
305                let parent_node_full_path = full_paths.last().unwrap();
306                let full_path = CompressedPathRaw::join_connected_paths(
307                    parent_node_full_path,
308                    child_index,
309                    &node.compressed_path_ref(),
310                );
311                keys_child_indices_and_nodes.push((
312                    CompressedPathRaw::extend_path(
313                        parent_node_full_path,
314                        child_index,
315                    ),
316                    child_index,
317                    &**node,
318                ));
319                full_paths.push(full_path);
320            }
321
322            match node.walk::<access_mode::Read>(key_remaining) {
323                WalkStop::Arrived => {
324                    return keys_child_indices_and_nodes;
325                }
326                WalkStop::PathDiverted { .. } => {
327                    return vec![];
328                }
329                WalkStop::ChildNotFound { .. } => {
330                    return vec![];
331                }
332                WalkStop::Descent {
333                    key_remaining: new_key_remaining,
334                    child_node,
335                    child_index: new_child_index,
336                } => {
337                    child_index = new_child_index;
338                    hash = child_node;
339                    key_remaining = new_key_remaining;
340                }
341            }
342        }
343    }
344}
345
346impl Encodable for TrieProof {
347    fn rlp_append(&self, s: &mut RlpStream) { s.append_list(&self.nodes); }
348}
349
350impl Decodable for TrieProof {
351    fn decode(rlp: &Rlp) -> std::result::Result<Self, DecoderError> {
352        if rlp.item_count()? > Self::MAX_NODES {
353            return Err(DecoderError::Custom("TrieProof too long."));
354        }
355        match Self::new(rlp.as_list()?) {
356            Err(_) => Err(DecoderError::Custom("Invalid TrieProof")),
357            Ok(proof) => Ok(proof),
358        }
359    }
360}
361
362impl Serialize for TrieProof {
363    fn serialize<S>(
364        &self, serializer: S,
365    ) -> std::result::Result<S::Ok, S::Error>
366    where S: Serializer {
367        let mut seq = serializer.serialize_seq(Some(self.nodes.len()))?;
368        for element in self.nodes.iter() {
369            seq.serialize_element(element)?;
370        }
371        seq.end()
372    }
373}
374
375impl<'a> Deserialize<'a> for TrieProof {
376    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
377    where D: Deserializer<'a> {
378        let nodes: Vec<TrieProofNode> = Deserialize::deserialize(deserializer)?;
379        if nodes.len() > Self::MAX_NODES {
380            return Err(de::Error::custom("TrieProof too long."));
381        }
382
383        match Self::new(nodes) {
384            Err(_) => Err(de::Error::custom("Invalid TrieProof")),
385            Ok(proof) => Ok(proof),
386        }
387    }
388}
389
390// FIXME: in rlp encode / decode, children_count and merkle_hash should be
391// omitted.
392impl Encodable for TrieProofNode {
393    fn rlp_append(&self, s: &mut RlpStream) { s.append_internal(&self.0); }
394}
395
396impl Decodable for TrieProofNode {
397    fn decode(rlp: &Rlp) -> std::result::Result<Self, DecoderError> {
398        Ok(Self(rlp.as_val()?))
399    }
400}
401
402impl Deref for TrieProofNode {
403    type Target = VanillaTrieNode<MerkleHash>;
404
405    fn deref(&self) -> &Self::Target { &self.0 }
406}
407
408impl DerefMut for TrieProofNode {
409    fn deref_mut(&mut self) -> &mut <Self as Deref>::Target { &mut self.0 }
410}
411
412use crate::{
413    impls::{
414        errors::*,
415        merkle_patricia_trie::{
416            merkle::compute_merkle,
417            walk::{TrieNodeWalkTrait, WalkStop},
418            CompressedPathRaw, CompressedPathTrait, TrieNodeTrait,
419            VanillaChildrenTable, VanillaTrieNode, CHILDREN_COUNT,
420        },
421    },
422    utils::access_mode,
423};
424use cfx_types::H256;
425use primitives::{MerkleHash, MptValue, MERKLE_NULL_NODE};
426use rlp::*;
427use serde::{
428    de, ser::SerializeSeq, Deserialize, Deserializer, Serialize, Serializer,
429};
430use std::{
431    collections::{hash_map::RandomState, HashMap},
432    ops::{Deref, DerefMut},
433};