Replace merkle.rs with custom implementation.

Cargo.toml

@@ -24,7 +24,6 @@ failure = "0.1"
 init_with = "1.1.0"
 itertools = "0.7"
 log = "0.4.1"
-merkle = { git = "https://github.com/afck/merkle.rs", branch = "public-proof", features = [ "serialization-serde" ] }
 pairing = { version = "0.14.2", features = ["u128-support"] }
 rand = "0.4.2"
 rand_derive = "0.3.1"

src/broadcast/broadcast.rs

@@ -1,15 +1,14 @@
 use std::collections::BTreeMap;
 use std::fmt::{self, Debug};
-use std::iter::once;
 use std::sync::Arc;
 
 use byteorder::{BigEndian, ByteOrder};
-use merkle::{MerkleTree, Proof};
+use itertools::Itertools;
 use rand;
 use reed_solomon_erasure as rse;
 use reed_solomon_erasure::ReedSolomon;
-use ring::digest;
 
+use super::merkle::{MerkleTree, Proof};
 use super::{Error, Result};
 use fault_log::{Fault, FaultKind};
 use fmt::{HexBytes, HexList, HexProof};
@@ -36,8 +35,8 @@ impl rand::Rand for Message {
         rng.fill_bytes(&mut buffer);
 
         // Generate a dummy proof to fill broadcast messages with.
-        let tree = MerkleTree::from_vec(&digest::SHA256, vec![buffer.to_vec()]);
-        let proof = tree.gen_proof(buffer.to_vec()).unwrap();
+        let tree = MerkleTree::from_vec(vec![buffer.to_vec()]);
+        let proof = tree.proof(0).unwrap();
 
         match message_type {
             "value" => Message::Value(proof),
@@ -188,28 +187,17 @@ impl<N: NodeUidT> Broadcast<N> {
 
         debug!("Shards: {:?}", HexList(&shards));
 
-        // TODO: `MerkleTree` generates the wrong proof if a leaf occurs more than once, so we
-        // prepend an "index byte" to each shard. Consider using the `merkle_light` crate instead.
-        let shards_t: Vec<Vec<u8>> = shards
-            .into_iter()
-            .enumerate()
-            .map(|(i, s)| once(i as u8).chain(s.iter().cloned()).collect())
-            .collect();
+        // Create a Merkle tree from the shards.
+        let mtree = MerkleTree::from_vec(shards.into_iter().map(|shard| shard.to_vec()).collect());
 
-        // Convert the Merkle tree into a partial binary tree for later
-        // deconstruction into compound branches.
-        let mtree = MerkleTree::from_vec(&digest::SHA256, shards_t);
-
-        // Default result in case of `gen_proof` error.
+        // Default result in case of `proof` error.
         let mut result = Err(Error::ProofConstructionFailed);
-        assert_eq!(self.netinfo.num_nodes(), mtree.iter().count());
+        assert_eq!(self.netinfo.num_nodes(), mtree.values().len());
 
         let mut step = Step::default();
         // Send each proof to a node.
-        for (leaf_value, uid) in mtree.iter().zip(self.netinfo.all_uids()) {
-            let proof = mtree
-                .gen_proof(leaf_value.to_vec())
-                .ok_or(Error::ProofConstructionFailed)?;
+        for (index, uid) in self.netinfo.all_uids().enumerate() {
+            let proof = mtree.proof(index).ok_or(Error::ProofConstructionFailed)?;
             if *uid == *self.netinfo.our_uid() {
                 // The proof is addressed to this node.
                 result = Ok(proof);
@@ -272,7 +260,7 @@ impl<N: NodeUidT> Broadcast<N> {
             return Ok(Fault::new(sender_id.clone(), FaultKind::InvalidProof).into());
         }
 
-        let hash = p.root_hash.clone();
+        let hash = p.root_hash().to_vec();
 
         // Save the proof for reconstructing the tree later.
         self.echos.insert(sender_id.clone(), p);
@@ -352,8 +340,8 @@ impl<N: NodeUidT> Broadcast<N> {
             .all_uids()
             .map(|id| {
                 self.echos.get(id).and_then(|p| {
-                    if p.root_hash.as_slice() == hash {
-                        Some(p.value.clone().into_boxed_slice())
+                    if p.root_hash() == hash {
+                        Some(p.value().clone().into_boxed_slice())
                     } else {
                         None
                     }
@@ -373,16 +361,14 @@ impl<N: NodeUidT> Broadcast<N> {
     /// Returns `true` if the proof is valid and has the same index as the node ID. Otherwise
     /// logs an info message.
     fn validate_proof(&self, p: &Proof<Vec<u8>>, id: &N) -> bool {
-        if !p.validate(&p.root_hash) {
+        if !p.validate(self.netinfo.num_nodes()) {
             info!(
                 "Node {:?} received invalid proof: {:?}",
                 self.netinfo.our_uid(),
                 HexProof(&p)
             );
             false
-        } else if self.netinfo.node_index(id) != Some(p.value[0] as usize)
-            || p.index(self.netinfo.num_nodes()) != p.value[0] as usize
-        {
+        } else if self.netinfo.node_index(id) != Some(p.index()) {
             info!(
                 "Node {:?} received proof for wrong position: {:?}.",
                 self.netinfo.our_uid(),
@@ -398,7 +384,7 @@ impl<N: NodeUidT> Broadcast<N> {
     fn count_echos(&self, hash: &[u8]) -> usize {
         self.echos
             .values()
-            .filter(|p| p.root_hash.as_slice() == hash)
+            .filter(|p| p.root_hash() == hash)
             .count()
     }
 
@@ -500,10 +486,10 @@ fn decode_from_shards(
     debug!("Reconstructed shards: {:?}", HexList(&shards));
 
     // Construct the Merkle tree.
-    let mtree = MerkleTree::from_vec(&digest::SHA256, shards);
+    let mtree = MerkleTree::from_vec(shards);
     // If the root hash of the reconstructed tree does not match the one
     // received with proofs then abort.
-    if &mtree.root_hash()[..] != root_hash {
+    if mtree.root_hash() != root_hash {
         None // The proposer is faulty.
     } else {
         // Reconstruct the value from the data shards.
@@ -516,7 +502,7 @@ fn decode_from_shards(
 /// and forgetting the leaves that contain parity information.
 fn glue_shards(m: MerkleTree<Vec<u8>>, n: usize) -> Option<Vec<u8>> {
     // Create an iterator over the shard payload, drop the index bytes.
-    let mut bytes = m.into_iter().take(n).flat_map(|s| s.into_iter().skip(1));
+    let mut bytes = Itertools::flatten(m.into_values().into_iter().take(n));
     let payload_len = match (bytes.next(), bytes.next(), bytes.next(), bytes.next()) {
         (Some(b0), Some(b1), Some(b2), Some(b3)) => BigEndian::read_u32(&[b0, b1, b2, b3]) as usize,
         _ => return None, // The proposing node is faulty: no payload size.

src/broadcast/merkle.rs

@@ -0,0 +1,161 @@
+use std::mem;
+
+use ring::digest::{self, Digest, SHA256};
+
+type DigestBytes = Vec<u8>;
+
+/// A Merkle tree: The leaves are values and their hashes. Each level consists of the hashes of
+/// pairs of values on the previous level. The root is the value in the first level with only one
+/// entry.
+#[derive(Debug)]
+pub struct MerkleTree<T> {
+    levels: Vec<Vec<Digest>>,
+    values: Vec<T>,
+    root_hash: Digest,
+}
+
+impl<T: AsRef<[u8]> + Clone> MerkleTree<T> {
+    /// Creates a new Merkle tree with the given values.
+    pub fn from_vec(values: Vec<T>) -> Self {
+        let mut levels = Vec::new();
+        let mut cur_lvl: Vec<Digest> = values.iter().map(hash).collect();
+        while cur_lvl.len() > 1 {
+            let next_lvl = cur_lvl.chunks(2).map(hash_chunk).collect();
+            levels.push(mem::replace(&mut cur_lvl, next_lvl));
+        }
+        let root_hash = cur_lvl[0];
+        MerkleTree {
+            levels,
+            values,
+            root_hash,
+        }
+    }
+
+    /// Returns the proof for entry `index`, if that is a valid index.
+    pub fn proof(&self, index: usize) -> Option<Proof<T>> {
+        let value = self.values.get(index)?.clone();
+        let mut lvl_i = index;
+        let mut digests = Vec::new();
+        for level in &self.levels {
+            // Insert the sibling hash if there is one.
+            if let Some(digest) = level.get(lvl_i ^ 1) {
+                digests.push(digest.as_ref().to_vec());
+            }
+            lvl_i /= 2;
+        }
+        Some(Proof {
+            index,
+            digests,
+            value,
+            root_hash: self.root_hash.as_ref().to_vec(),
+        })
+    }
+
+    /// Returns the root hash of the tree.
+    pub fn root_hash(&self) -> &[u8] {
+        self.root_hash.as_ref()
+    }
+
+    /// Returns a the slice containing all leaf values.
+    pub fn values(&self) -> &[T] {
+        &self.values
+    }
+
+    /// Consumes the tree, and returns the vector of leaf values.
+    pub fn into_values(self) -> Vec<T> {
+        self.values
+    }
+}
+
+/// A proof that a value is at a particular index in the Merkle tree specified by its root hash.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct Proof<T> {
+    value: T,
+    index: usize,
+    digests: Vec<DigestBytes>,
+    root_hash: DigestBytes,
+}
+
+impl<T: AsRef<[u8]>> Proof<T> {
+    /// Returns `true` if the digests in this proof constitute a valid branch in a Merkle tree with
+    /// the root hash.
+    pub fn validate(&self, n: usize) -> bool {
+        let mut digest = hash(&self.value);
+        let mut lvl_i = self.index;
+        let mut lvl_n = n;
+        let mut digest_itr = self.digests.iter();
+        while lvl_n > 1 {
+            if lvl_i ^ 1 < lvl_n {
+                digest = match digest_itr.next() {
+                    None => return false, // Not enough levels in the proof.
+                    Some(sibling) if lvl_i & 1 == 1 => hash_pair(&sibling, &digest),
+                    Some(sibling) => hash_pair(&digest, &sibling),
+                };
+            }
+            lvl_i /= 2; // Our index on the next level.
+            lvl_n = (lvl_n + 1) / 2; // The next level's size.
+        }
+        if digest_itr.next().is_some() {
+            return false; // Too many levels in the proof.
+        }
+        digest.as_ref() == &self.root_hash[..]
+    }
+
+    /// Returns the index of this proof's value in the tree.
+    pub fn index(&self) -> usize {
+        self.index
+    }
+
+    /// Returns the tree's root hash.
+    pub fn root_hash(&self) -> &[u8] {
+        self.root_hash.as_ref()
+    }
+
+    /// Returns the leaf value.
+    pub fn value(&self) -> &T {
+        &self.value
+    }
+
+    /// Consumes the proof and returns the leaf value.
+    pub fn into_value(self) -> T {
+        self.value
+    }
+}
+
+/// Takes a chunk of one or two digests. In the former case, returns the digest itself, in the
+/// latter, it returns the hash of the two digests.
+fn hash_chunk(chunk: &[Digest]) -> Digest {
+    if chunk.len() == 1 {
+        chunk[0]
+    } else {
+        hash_pair(&chunk[0], &chunk[1])
+    }
+}
+
+/// Returns the hash of the concatenated bytes of `d0` and `d1`.
+fn hash_pair<T0: AsRef<[u8]>, T1: AsRef<[u8]>>(v0: &T0, v1: &T1) -> Digest {
+    let bytes: Vec<u8> = v0.as_ref().iter().chain(v1.as_ref()).cloned().collect();
+    hash(&bytes)
+}
+
+/// Returns the SHA-256 hash of the value's `[u8]` representation.
+fn hash<T: AsRef<[u8]>>(value: T) -> Digest {
+    digest::digest(&SHA256, value.as_ref())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::MerkleTree;
+
+    #[test]
+    fn test_merkle() {
+        for &n in &[4, 7, 8, 9, 17] {
+            let tree = MerkleTree::from_vec((0..n).map(|i| vec![i as u8]).collect());
+            for i in 0..n {
+                let proof = tree.proof(i).expect("couldn't get proof");
+                assert!(proof.validate(n));
+            }
+            assert!(tree.proof(n).is_none());
+        }
+    }
+}

src/broadcast/mod.rs

@@ -151,6 +151,7 @@
 
 mod broadcast;
 mod error;
+pub mod merkle;
 
 pub use self::broadcast::{Broadcast, Message, Step};
 pub use self::error::{Error, Result};

src/fmt.rs

@@ -1,4 +1,4 @@
-use merkle::Proof;
+use broadcast::merkle::Proof;
 use std::fmt;
 
 /// Wrapper for a byte array, whose `Debug` implementation outputs shortened hexadecimal strings.
@@ -40,10 +40,10 @@ impl<'a, T: AsRef<[u8]>> fmt::Debug for HexProof<'a, T> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(
             f,
-            "Proof {{ algorithm: {:?}, root_hash: {:?}, value: {:?}, .. }}",
-            self.0.algorithm,
-            HexBytes(&self.0.root_hash),
-            HexBytes(&self.0.value.as_ref())
+            "Proof {{ #{}, root_hash: {:?}, value: {:?}, .. }}",
+            &self.0.index(),
+            HexBytes(&self.0.root_hash()),
+            HexBytes(&self.0.value().as_ref())
         )
     }
 }

src/lib.rs

@@ -111,7 +111,6 @@ extern crate init_with;
 #[macro_use]
 extern crate log;
 extern crate itertools;
-extern crate merkle;
 extern crate pairing;
 extern crate rand;
 #[macro_use]