Implement Merkle proof

src/lib.rs

@@ -19,6 +19,7 @@ extern crate rlp_derive;
 
 mod nibbleslice;
 pub mod node;
+pub mod proof;
 mod skewed;
 #[allow(dead_code)]
 pub mod snapshot;

src/proof.rs

@@ -0,0 +1,277 @@
+// Copyright 2020 Kodebox, Inc.
+// This file is part of CodeChain.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+use crate::nibbleslice::NibbleSlice;
+use crate::node::Node;
+use ccrypto::{blake256, BLAKE_NULL_RLP};
+use primitives::Bytes;
+use primitives::H256;
+
+// Unit of a proof.
+//#[derive(Clone, Eq, PartialEq, Debug, RlpEncodable, RlpDecodable)]
+#[derive(Clone, Eq, PartialEq, Debug)]
+pub struct CryptoProofUnit {
+    pub root: H256,
+    pub key: H256,
+    pub value: Option<Bytes>, // None in case of absence
+}
+
+#[derive(Clone, Eq, PartialEq, Debug)]
+pub struct CryptoProof(pub Vec<Bytes>);
+
+pub trait CryptoStructure {
+    fn make_proof(&self, key: &H256) -> crate::Result<(CryptoProofUnit, CryptoProof)>;
+}
+
+/// A verification logic of TrieDB's Merkle proof.
+/// For the format of proof, check the make_proof() function.
+/// It verifies the proof with a given unit of test.
+/// It should never abort or fail, but only return 'false' as a result of getting an invalid or ill-formed proof.
+pub fn verify(proof: &CryptoProof, test: &CryptoProofUnit) -> bool {
+    // step1: verify the value
+    fn step1(proof: &CryptoProof, test: &CryptoProofUnit) -> bool {
+        match Node::decoded(&proof.0.last().unwrap()) {
+            Some(x) => match x {
+                Node::Leaf(_, value) => test.value.as_ref().unwrap() == &value,
+                _ => false,
+            },
+            _ => false,
+        }
+    };
+
+    // step2: verify the root
+    fn step2(proof: &CryptoProof, test: &CryptoProofUnit) -> bool {
+        blake256(&proof.0[0]) == test.root
+    };
+
+    // step3 (presence): verify the key
+    fn step3_p(proof: &CryptoProof, test: &CryptoProofUnit) -> bool {
+        fn verify_branch(path: &NibbleSlice<'_>, hash: &H256, proof: &[Bytes]) -> bool {
+            if *hash != blake256(&proof[0]) {
+                return false
+            }
+            match Node::decoded(&proof[0]) {
+                Some(Node::Leaf(partial, _)) => path == &partial,
+                Some(Node::Branch(partial, table)) => {
+                    if proof.len() < 2 {
+                        // detect ill-formed proof
+                        return false
+                    }
+                    if !path.starts_with(&partial) {
+                        return false
+                    }
+                    match table[path.at(partial.len()) as usize] {
+                        Some(x) => verify_branch(&path.mid(partial.len() + 1), &x, &proof[1..]),
+                        None => false,
+                    }
+                }
+                None => false,
+            }
+        };
+        verify_branch(&NibbleSlice::new(&test.key), &test.root, &proof.0)
+    };
+
+    // step3 (absence): verify the key.
+    fn step3_a(proof: &CryptoProof, test: &CryptoProofUnit) -> bool {
+        fn verify_branch(path: &NibbleSlice<'_>, hash: &H256, proof: &[Bytes]) -> bool {
+            if *hash != blake256(&proof[0]) {
+                return false
+            }
+            match Node::decoded(&proof[0]) {
+                Some(Node::Leaf(partial, _)) => path != &partial, // special case : there is only one leaf node in the trie,
+                Some(Node::Branch(partial, children)) => {
+                    if !path.starts_with(&partial) {
+                        return false
+                    }
+                    match children[path.at(partial.len()) as usize] {
+                        Some(x) => proof.len() >= 2 && verify_branch(&path.mid(partial.len() + 1), &x, &proof[1..]),
+                        None => proof.len() == 1,
+                    }
+                }
+                None => false,
+            }
+        };
+        verify_branch(&NibbleSlice::new(&test.key), &test.root, &proof.0)
+    };
+
+    if proof.0.is_empty() {
+        return test.root == BLAKE_NULL_RLP && test.value.is_none() // special case of an empty trie.
+    }
+    if test.value.is_some() {
+        step1(proof, test) && step2(proof, test) && step3_p(proof, test)
+    } else {
+        step2(proof, test) && step3_a(proof, test)
+    }
+}
+
+
+#[cfg(test)]
+mod tests {
+    extern crate rand;
+
+    use super::*;
+    use crate::*;
+    use cdb::MemoryDB;
+    use rand::{rngs::StdRng, Rng};
+
+    fn simple_test<'db>(t: &TrieDB<'db>, key: &H256, value: Option<&[u8]>, key_proof: &H256, result: bool) {
+        let unit = CryptoProofUnit {
+            root: *t.root(),
+            key: *key,
+            value: value.map(|x| x.to_vec()),
+        };
+        let proof = t.make_proof(key_proof).unwrap();
+        assert_eq!(verify(&proof.1, &unit), result);
+    }
+
+    #[test]
+    fn empty_trie() {
+        let iteration = 100;
+        let seed = [0 as u8; 32];
+        let mut rng: StdRng = rand::SeedableRng::from_seed(seed);
+
+        for _ in 0..iteration {
+            let mut memdb = MemoryDB::new();
+            let mut root = H256::zero();
+            TrieDBMut::new(&mut memdb, &mut root);
+
+            // unused pair
+            let k1 = format!("{}", rng.gen::<u64>());
+            let v1 = format!("{}", rng.gen::<u64>());
+            let (keyu, valu) = { (blake256(&k1), v1.as_bytes()) };
+
+            let t = TrieDB::try_new(&memdb, &root).unwrap();
+
+            simple_test(&t, &keyu, Some(valu), &keyu, false);
+            simple_test(&t, &keyu, None, &keyu, true);
+        }
+    }
+
+    #[test]
+    fn single_trie() {
+        let iteration = 100;
+        let seed = [0 as u8; 32];
+        let mut rng: StdRng = rand::SeedableRng::from_seed(seed);
+
+        for _ in 0..iteration {
+            let mut memdb = MemoryDB::new();
+            let mut root = H256::zero();
+            let mut mt = TrieDBMut::new(&mut memdb, &mut root);
+
+            // unused pair
+            let ku = format!("{}", rng.gen::<u64>());
+            let vu = format!("{}", rng.gen::<u64>());
+            let (keyu, valu) = { (blake256(&ku), vu.as_bytes()) };
+
+            let k1 = format!("{}", rng.gen::<u64>());
+            let v1 = format!("{}", rng.gen::<u64>());
+            let (key1, val1) = { (blake256(&k1), v1.as_bytes()) };
+            mt.insert(&k1.as_bytes(), val1).unwrap();
+
+            if key1 == keyu || val1 == valu {
+                continue
+            }
+
+            let t = TrieDB::try_new(&memdb, &root).unwrap();
+
+            // Be careful: there are some case where the proof is not unique.
+            simple_test(&t, &key1, Some(val1), &key1, true);
+            simple_test(&t, &key1, Some(val1), &keyu, true); //be careful!
+            simple_test(&t, &key1, Some(valu), &key1, false);
+            simple_test(&t, &key1, Some(valu), &keyu, false);
+            simple_test(&t, &key1, None, &key1, false);
+            simple_test(&t, &key1, None, &keyu, false);
+            simple_test(&t, &keyu, Some(val1), &key1, false);
+            simple_test(&t, &keyu, Some(val1), &keyu, false);
+            simple_test(&t, &keyu, Some(valu), &key1, false);
+            simple_test(&t, &keyu, Some(valu), &keyu, false);
+            simple_test(&t, &keyu, None, &key1, true); //be careful!
+            simple_test(&t, &keyu, None, &keyu, true);
+        }
+    }
+
+    #[test]
+    fn some_trie() {
+        let iteration = 100;
+        let size = 234;
+        let seed = [0 as u8; 32];
+        let mut rng: StdRng = rand::SeedableRng::from_seed(seed);
+
+        for _ in 0..iteration {
+            let mut memdb = MemoryDB::new();
+            let mut root = H256::zero();
+            let mut mt = TrieDBMut::new(&mut memdb, &mut root);
+
+            // unused pair
+            let ku = format!("{}", rng.gen::<u64>());
+            let vu = format!("{}", rng.gen::<u64>());
+            let (keyu, valu) = { (blake256(&ku), vu.as_bytes()) };
+
+            let k1 = format!("{}", rng.gen::<u64>());
+            let v1 = format!("{}", rng.gen::<u64>());
+            let (key1, val1) = { (blake256(&k1), v1.as_bytes()) };
+            mt.insert(&k1.as_bytes(), val1).unwrap();
+
+            let k2 = format!("{}", rng.gen::<u64>());
+            let v2 = format!("{}", rng.gen::<u64>());
+            let (key2, val2) = { (blake256(&k2), v2.as_bytes()) };
+            mt.insert(&k2.as_bytes(), val2).unwrap();
+
+            if key1 == keyu || val1 == valu || key2 == keyu || val2 == valu {
+                continue
+            }
+
+            let mut flag = true;
+            for _ in 0..size {
+                let k = format!("{}", rng.gen::<u64>());
+                let v = format!("{}", rng.gen::<u64>());
+                mt.insert(k.as_bytes(), v.as_bytes()).unwrap();
+                if blake256(k) == keyu || v.as_bytes() == valu {
+                    flag = false;
+                    break
+                }
+            }
+            if !flag {
+                continue // skip this iteration
+            }
+
+            let t = TrieDB::try_new(&memdb, &root).unwrap();
+
+            simple_test(&t, &key1, Some(val1), &key1, true);
+            simple_test(&t, &key1, Some(val1), &key2, false);
+            simple_test(&t, &key1, Some(val1), &keyu, false);
+            simple_test(&t, &key1, Some(val2), &key1, false);
+            simple_test(&t, &key1, Some(val2), &key2, false);
+            simple_test(&t, &key1, Some(val2), &keyu, false);
+            simple_test(&t, &key1, None, &key1, false);
+            simple_test(&t, &key1, None, &key2, false);
+            simple_test(&t, &key1, None, &keyu, false);
+
+            simple_test(&t, &keyu, Some(val1), &key1, false);
+            simple_test(&t, &keyu, Some(val1), &key2, false);
+            simple_test(&t, &keyu, Some(val1), &keyu, false);
+            simple_test(&t, &keyu, None, &key1, false);
+            simple_test(&t, &keyu, None, &key2, false);
+            simple_test(&t, &keyu, None, &keyu, true);
+        }
+    }
+
+    // proof is created manually here
+    #[test]
+    fn some_malicious() {
+        // TODO
+    }
+}

src/triedb.rs

@@ -16,10 +16,12 @@
 
 use crate::nibbleslice::NibbleSlice;
 use crate::node::Node as RlpNode;
-use crate::{Trie, TrieError};
+use crate::proof::{CryptoProof, CryptoProofUnit, CryptoStructure};
+use crate::{Node, Trie, TrieError};
 use ccrypto::{blake256, BLAKE_NULL_RLP};
 use cdb::HashDB;
 use lru_cache::LruCache;
+use primitives::Bytes;
 use primitives::H256;
 use std::cell::RefCell;
 
@@ -98,11 +100,7 @@ impl<'db> TrieDB<'db> {
                     }
                     Some(RlpNode::Branch(partial, children)) => {
                         if path.starts_with(&partial) {
-                            self.get_aux(
-                                &path.mid(partial.len() + 1),
-                                children[path.mid(partial.len()).at(0) as usize],
-                                query,
-                            )
+                            self.get_aux(&path.mid(partial.len() + 1), children[path.at(partial.len()) as usize], query)
                         } else {
                             Ok(None)
                         }
@@ -147,6 +145,67 @@ impl<'db> Trie for TrieDB<'db> {
     }
 }
 
+impl<'db> CryptoStructure for TrieDB<'db> {
+    /// A proof creation logic for TrieDB.
+    /// A proof is basically a list of serialized trie nodes, Vec<Bytes>.
+    /// It starts from the one closest to the root and to the leaf. (It may not reach the leaf in absence case.)
+    /// Each node can be decoded with RLP. (Note that RLP doesn't guarantee format detail, so you must check our serialization code.)
+    /// In case of precense, the list will contain a path from the root to the leaf with the key.
+    /// In case of absence, the list will contain a path to the last node that matches the key.
+    //
+    //          (A: [nil])
+    //         /         \
+    //      (B, g)        \
+    //      /    \         \
+    // (C, iant) (D, mail)  (E, clang)
+    //
+    // Here, the proof of key 'gmail' will be [(RLP encoding of A), (RLP encoding of B), (RLP encoding of D)]
+    // Here, the proof of key 'galbi' (absence) will be [(RLP encoding of A), (RLP encoding of B)]
+    fn make_proof(&self, key: &H256) -> crate::Result<(CryptoProofUnit, CryptoProof)> {
+        // it creates a reversed proof for the sake of a more efficient push() operation. (than concat)
+        fn make_proof_upto(
+            db: &dyn HashDB,
+            path: &NibbleSlice<'_>,
+            hash: &H256,
+        ) -> crate::Result<(Option<Bytes>, Vec<Bytes>)> {
+            let node_rlp = db.get(&hash).ok_or_else(|| TrieError::IncompleteDatabase(*hash))?;
+
+            match Node::decoded(&node_rlp) {
+                Some(Node::Leaf(partial, value)) => {
+                    if &partial == path {
+                        Ok((Some(value.to_vec()), vec![node_rlp]))
+                    } else {
+                        Ok((None, vec![node_rlp]))
+                    }
+                }
+                Some(Node::Branch(partial, children)) => {
+                    if path.starts_with(&partial) {
+                        match children[path.at(partial.len()) as usize] {
+                            Some(x) => {
+                                let (value, mut reversed_proof) =
+                                    make_proof_upto(db, &path.mid(partial.len() + 1), &x)?;
+                                reversed_proof.push(node_rlp);
+                                Ok((value, reversed_proof))
+                            }
+                            None => Ok((None, vec![node_rlp])),
+                        }
+                    } else {
+                        Ok((None, Vec::new()))
+                    }
+                }
+                None => Ok((None, Vec::new())), // empty trie
+            }
+        }
+        let (value, reversed_proof) = make_proof_upto(self.db, &NibbleSlice::new(&key), self.root())?;
+        let unit = CryptoProofUnit {
+            root: *self.root(),
+            key: *key,
+            value,
+        };
+        Ok((unit, CryptoProof(reversed_proof.iter().rev().cloned().collect())))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use cdb::MemoryDB;