From e154b30232a9ca5db31cd1d4331d17a49a840ff7 Mon Sep 17 00:00:00 2001 From: Michael Sproul Date: Tue, 25 Jun 2019 18:46:57 +1000 Subject: [PATCH] merkle_proof: implement tree construction Plus QuickCheck tests! --- eth2/utils/merkle_proof/Cargo.toml | 5 + eth2/utils/merkle_proof/src/lib.rs | 193 ++++++++++++++++++++++++++++- 2 files changed, 193 insertions(+), 5 deletions(-) diff --git a/eth2/utils/merkle_proof/Cargo.toml b/eth2/utils/merkle_proof/Cargo.toml index 6ef6cc0aac..5ffb6af532 100644 --- a/eth2/utils/merkle_proof/Cargo.toml +++ b/eth2/utils/merkle_proof/Cargo.toml @@ -7,3 +7,8 @@ edition = "2018" [dependencies] ethereum-types = "0.6" eth2_hashing = { path = "../eth2_hashing" } +lazy_static = "1.3.0" + +[dev-dependencies] +quickcheck = "0.8" +quickcheck_macros = "0.8" diff --git a/eth2/utils/merkle_proof/src/lib.rs b/eth2/utils/merkle_proof/src/lib.rs index bc8bcea127..73a972c759 100644 --- a/eth2/utils/merkle_proof/src/lib.rs +++ b/eth2/utils/merkle_proof/src/lib.rs @@ -1,6 +1,138 @@ +#[macro_use] +extern crate lazy_static; + use eth2_hashing::hash; use ethereum_types::H256; +const MAX_TREE_DEPTH: usize = 32; +const EMPTY_SLICE: &[H256] = &[]; + +lazy_static! { + /// Cached zero hashes where `ZERO_HASHES[i]` is the hash of a Merkle tree with 2^i zero leaves. + static ref ZERO_HASHES: Vec = { + let mut hashes = vec![H256::from([0; 32]); MAX_TREE_DEPTH + 1]; + + for i in 0..MAX_TREE_DEPTH { + hashes[i + 1] = hash_concat(hashes[i], hashes[i]); + } + + hashes + }; + + /// Zero nodes to act as "synthetic" left and right subtrees of other zero nodes. + static ref ZERO_NODES: Vec = { + (0..MAX_TREE_DEPTH + 1).map(MerkleTree::Zero).collect() + }; +} + +/// Right-sparse Merkle tree. +/// +/// Efficiently represents a Merkle tree of fixed depth where only the first N +/// indices are populated by non-zero leaves (perfect for the deposit contract tree). +#[derive(Debug)] +pub enum MerkleTree { + /// Leaf node with the hash of its content. + Leaf(H256), + /// Internal node with hash, left subtree and right subtree. + Node(H256, Box, Box), + /// Zero subtree of a given depth. + /// + /// It represents a Merkle tree of 2^depth zero leaves. + Zero(usize), +} + +impl MerkleTree { + /// Create a new Merkle tree from a list of leaves and a fixed depth. + pub fn create(leaves: &[H256], depth: usize) -> Self { + use MerkleTree::*; + + if leaves.is_empty() { + return Zero(depth); + } + + match depth { + 0 => { + debug_assert_eq!(leaves.len(), 1); + Leaf(leaves[0]) + } + _ => { + // Split leaves into left and right subtrees + let subtree_capacity = 2usize.pow(depth as u32 - 1); + let (left_leaves, right_leaves) = if leaves.len() <= subtree_capacity { + (leaves, EMPTY_SLICE) + } else { + leaves.split_at(subtree_capacity) + }; + + let left_subtree = MerkleTree::create(left_leaves, depth - 1); + let right_subtree = MerkleTree::create(right_leaves, depth - 1); + let hash = hash_concat(left_subtree.hash(), right_subtree.hash()); + + Node(hash, Box::new(left_subtree), Box::new(right_subtree)) + } + } + } + + /// Retrieve the root hash of this Merkle tree. + pub fn hash(&self) -> H256 { + match *self { + MerkleTree::Leaf(h) => h, + MerkleTree::Node(h, _, _) => h, + MerkleTree::Zero(depth) => ZERO_HASHES[depth], + } + } + + /// Get a reference to the left and right subtrees if they exist. + pub fn left_and_right_branches(&self) -> Option<(&Self, &Self)> { + match *self { + MerkleTree::Leaf(_) | MerkleTree::Zero(0) => None, + MerkleTree::Node(_, ref l, ref r) => Some((l, r)), + MerkleTree::Zero(depth) => Some((&ZERO_NODES[depth - 1], &ZERO_NODES[depth - 1])), + } + } + + /// Is this Merkle tree a leaf? + pub fn is_leaf(&self) -> bool { + match self { + MerkleTree::Leaf(_) => true, + _ => false, + } + } + + /// Return the leaf at `index` and a Merkle proof of its inclusion. + /// + /// The Merkle proof is in "bottom-up" order, starting with a leaf node + /// and moving up the tree. Its length will be exactly equal to `depth`. + pub fn generate_proof(&self, index: usize, depth: usize) -> (H256, Vec) { + let mut proof = vec![]; + let mut current_node = self; + let mut current_depth = depth; + while current_depth > 0 { + let ith_bit = (index >> (current_depth - 1)) & 0x01; + // Note: unwrap is safe because leaves are only ever constructed at depth == 0. + let (left, right) = current_node.left_and_right_branches().unwrap(); + + // Go right, include the left branch in the proof. + if ith_bit == 1 { + proof.push(left.hash()); + current_node = right; + } else { + proof.push(right.hash()); + current_node = left; + } + current_depth -= 1; + } + + debug_assert_eq!(proof.len(), depth); + debug_assert!(current_node.is_leaf()); + + // Put proof in bottom-up order. + proof.reverse(); + + (current_node.hash(), proof) + } +} + /// Verify a proof that `leaf` exists at `index` in a Merkle tree rooted at `root`. /// /// The `branch` argument is the main component of the proof: it should be a list of internal @@ -46,15 +178,66 @@ fn concat(mut vec1: Vec, mut vec2: Vec) -> Vec { vec1 } +/// Compute the hash of two other hashes concatenated. +fn hash_concat(h1: H256, h2: H256) -> H256 { + H256::from_slice(&hash(&concat( + h1.as_bytes().to_vec(), + h2.as_bytes().to_vec(), + ))) +} + #[cfg(test)] mod tests { use super::*; + use quickcheck::TestResult; + use quickcheck_macros::quickcheck; - fn hash_concat(h1: H256, h2: H256) -> H256 { - H256::from_slice(&hash(&concat( - h1.as_bytes().to_vec(), - h2.as_bytes().to_vec(), - ))) + /// Check that we can: + /// 1. Build a MerkleTree from arbitrary leaves and an arbitrary depth. + /// 2. Generate valid proofs for all of the leaves of this MerkleTree. + #[quickcheck] + fn quickcheck_create_and_verify(int_leaves: Vec, depth: usize) -> TestResult { + if depth > MAX_TREE_DEPTH || int_leaves.len() > 2usize.pow(depth as u32) { + return TestResult::discard(); + } + + let leaves: Vec<_> = int_leaves.into_iter().map(H256::from_low_u64_be).collect(); + let merkle_tree = MerkleTree::create(&leaves, depth); + let merkle_root = merkle_tree.hash(); + + let proofs_ok = (0..leaves.len()).into_iter().all(|i| { + let (leaf, branch) = merkle_tree.generate_proof(i, depth); + leaf == leaves[i] && verify_merkle_proof(leaf, &branch, depth, i, merkle_root) + }); + + TestResult::from_bool(proofs_ok) + } + + #[test] + fn sparse_zero_correct() { + let depth = 2; + let zero = H256::from([0x00; 32]); + let dense_tree = MerkleTree::create(&[zero, zero, zero, zero], depth); + let sparse_tree = MerkleTree::create(&[], depth); + assert_eq!(dense_tree.hash(), sparse_tree.hash()); + } + + #[test] + fn create_small_example() { + // Construct a small merkle tree manually and check that it's consistent with + // the MerkleTree type. + let leaf_b00 = H256::from([0xAA; 32]); + let leaf_b01 = H256::from([0xBB; 32]); + let leaf_b10 = H256::from([0xCC; 32]); + let leaf_b11 = H256::from([0xDD; 32]); + + let node_b0x = hash_concat(leaf_b00, leaf_b01); + let node_b1x = hash_concat(leaf_b10, leaf_b11); + + let root = hash_concat(node_b0x, node_b1x); + + let tree = MerkleTree::create(&[leaf_b00, leaf_b01, leaf_b10, leaf_b11], 2); + assert_eq!(tree.hash(), root); } #[test]