added tests for merkleblock message

Cargo.lock

@@ -662,6 +662,7 @@ dependencies = [
  "p2p 0.1.0",
  "parking_lot 0.3.6 (registry+https://github.com/rust-lang/crates.io-index)",
  "primitives 0.1.0",
+ "rand 0.3.14 (registry+https://github.com/rust-lang/crates.io-index)",
  "script 0.1.0",
  "serialization 0.1.0",
  "test-data 0.1.0",

chain/src/lib.rs

@@ -19,6 +19,7 @@ pub use primitives::{hash, bytes};
 pub use self::block::Block;
 pub use self::block_header::BlockHeader;
 pub use self::merkle_root::merkle_root;
+pub use self::merkle_root::merkle_node_hash;
 pub use self::transaction::{
 	Transaction, TransactionInput, TransactionOutput, OutPoint,
 	SEQUENCE_LOCKTIME_DISABLE_FLAG, SEQUENCE_FINAL,

chain/src/merkle_root.rs

@@ -19,19 +19,24 @@ pub fn merkle_root(hashes: &[H256]) -> H256 {
 	let mut row = vec![];
 	let mut i = 0;
 	while i + 1 < hashes.len() {
-		row.push(dhash256(&*concat(&hashes[i], &hashes[i + 1])));
+		row.push(merkle_node_hash(&hashes[i], &hashes[i + 1]));
 		i += 2
 	}
 
 	// duplicate the last element if len is not even
 	if hashes.len() % 2 == 1 {
 		let last = &hashes[hashes.len() - 1];
-		row.push(dhash256(&*concat(last, last)));
+		row.push(merkle_node_hash(last, last));
 	}
 
 	merkle_root(&row)
 }
 
+/// Calculate merkle tree node hash
+pub fn merkle_node_hash(left: &H256, right: &H256) -> H256 {
+	dhash256(&*concat(left, right))
+}
+
 #[cfg(test)]
 mod tests {
 	use hash::H256;

sync/Cargo.toml

@@ -14,6 +14,7 @@ linked-hash-map = "0.3"
 ethcore-devtools = { path = "../devtools" }
 bit-vec = "0.4.3"
 murmur3 = "0.3"
+rand = "0.3"
 
 chain = { path = "../chain" }
 db = { path = "../db" }

sync/src/connection_filter.rs

@@ -2,7 +2,7 @@ use std::cmp::min;
 use linked_hash_map::LinkedHashMap;
 use bit_vec::BitVec;
 use murmur3::murmur3_32;
-use chain::{Block, Transaction, OutPoint, merkle_root};
+use chain::{Block, Transaction, OutPoint, merkle_root, merkle_node_hash};
 use ser::serialize;
 use message::types;
 use primitives::bytes::Bytes;
@@ -41,6 +41,7 @@ struct ConnectionBloom {
 }
 
 /// `merkleblock` build artefacts
+#[derive(Debug, PartialEq)]
 pub struct MerkleBlockArtefacts {
 	/// `merkleblock` message
 	pub merkleblock: types::MerkleBlock,
@@ -299,18 +300,55 @@ impl PartialMerkleTree {
 		(partial_merkle_tree.hashes, partial_merkle_tree.matches)
 	}
 
-	fn build_tree(&mut self) {
-		let tree_height = {
-			let mut height = 0usize;
-			while self.level_width(height) > 1 {
-				height += 1;
-			}
-			height
+	#[cfg(test)]
+	/// Parse partial merkle tree as described here:
+	/// https://bitcoin.org/en/developer-reference#parsing-a-merkleblock-message
+	pub fn parse(all_len: usize, hashes: Vec<H256>, matches: BitVec) -> Result<(H256, Vec<H256>, BitVec), String> {
+		let mut partial_merkle_tree = PartialMerkleTree {
+			all_len: all_len,
+			all_hashes: Vec::new(),
+			all_matches: BitVec::from_elem(all_len, false),
+			hashes: hashes,
+			matches: matches,
 		};
+		let merkle_root = try!(partial_merkle_tree.parse_tree());
+		Ok((merkle_root, partial_merkle_tree.all_hashes, partial_merkle_tree.all_matches))
+	}
+
+	fn build_tree(&mut self) {
+		let tree_height = self.tree_height();
 		self.build_branch(tree_height, 0)
 	}
 
-	fn build_branch(&mut self, pos: usize, height: usize) {
+	#[cfg(test)]
+	fn parse_tree(&mut self) -> Result<H256, String> {
+		if self.all_len == 0 {
+			return Err("no transactions".into());
+		}
+		if self.hashes.len() > self.all_len {
+			return Err("too many hashes".into());
+		}
+		if self.matches.len() < self.hashes.len() {
+			return Err("too few matches".into());
+		}
+
+		// parse tree
+		let mut matches_used = 0usize;
+		let mut hashes_used = 0usize;
+		let tree_height = self.tree_height();
+		let merkle_root = try!(self.parse_branch(tree_height, 0, &mut matches_used, &mut hashes_used));
+
+		if matches_used != self.matches.len() {
+			return Err("not all matches used".into());
+		}
+		if hashes_used != self.hashes.len() {
+			return Err("not all hashes used".into());
+		}
+
+		Ok(merkle_root)
+	}
+
+	fn build_branch(&mut self, height: usize, pos: usize) {
 		// determine whether this node is the parent of at least one matched txid
 		let transactions_begin = pos << height;
 		let transactions_end = min(self.all_len, (pos + 1) << height);
@@ -320,26 +358,88 @@ impl PartialMerkleTree {
 		// proceeed with descendants
 		if height == 0 || !flag {
 			// we're at the leaf level || there is no match
-			let hash = self.merkle_hash(height, pos);
+			let hash = self.branch_hash(height, pos);
 			self.hashes.push(hash);
 		} else {
 			// proceed with left child
-			self.build_branch(pos << 1, height - 1);
+			self.build_branch(height - 1, pos << 1);
 			// proceed with right child if any
 			if (pos << 1) + 1 < self.level_width(height - 1) {
-				self.build_branch((pos << 1) + 1, height - 1);
+				self.build_branch(height - 1, (pos << 1) + 1);
 			}
 		}
 	}
 
+	#[cfg(test)]
+	fn parse_branch(&mut self, height: usize, pos: usize, matches_used: &mut usize, hashes_used: &mut usize) -> Result<H256, String> {
+		if *matches_used >= self.matches.len() {
+			return Err("all matches used".into());
+		}
+
+		let flag = self.matches[*matches_used];
+		*matches_used += 1;
+
+		if height == 0 || !flag {
+			// we're at the leaf level || there is no match
+			if *hashes_used > self.hashes.len() {
+				return Err("all hashes used".into());
+			}
+
+			// get node hash
+			let ref hash = self.hashes[*hashes_used];
+			*hashes_used += 1;
+
+			// on leaf level && matched flag set => mark transaction as matched
+			if height == 0 && flag {
+				self.all_hashes.push(hash.clone());
+				self.all_matches.set(pos, true);
+			}
+
+			Ok(hash.clone())
+		} else {
+			// proceed with left child
+			let left = try!(self.parse_branch(height - 1, pos << 1, matches_used, hashes_used));
+			// proceed with right child if any
+			let has_right_child = (pos << 1) + 1 < self.level_width(height - 1);
+			let right = if has_right_child {
+				try!(self.parse_branch(height - 1, (pos << 1) + 1, matches_used, hashes_used))
+			} else {
+				left.clone()
+			};
+
+			if has_right_child && left == right {
+				Err("met same hash twice".into())
+			} else {
+				Ok(merkle_node_hash(&left, &right))
+			}
+		}
+	}
+
+	fn tree_height(&self) -> usize {
+		let mut height = 0usize;
+		while self.level_width(height) > 1 {
+			height += 1;
+		}
+		height
+	}
+
 	fn level_width(&self, height: usize) -> usize {
 		(self.all_len + (1 << height) - 1) >> height
 	}
 
-	fn merkle_hash(&self, height: usize, pos: usize) -> H256 {
-		let transactions_begin = pos << height;
-		let transactions_end = min(self.all_len, (pos + 1) << height);
-		merkle_root(&self.all_hashes[transactions_begin..transactions_end])
+	fn branch_hash(&self, height: usize, pos: usize) -> H256 {
+		if height == 0 {
+			self.all_hashes[pos].clone()
+		} else {
+			let left = self.branch_hash(height - 1, pos << 1);
+			let right = if (pos << 1) + 1 < self.level_width(height - 1) {
+				self.branch_hash(height - 1, (pos << 1) + 1)
+			} else {
+				left.clone()
+			};
+
+			merkle_node_hash(&left, &right)
+		}
 	}
 }
 
@@ -348,11 +448,11 @@ pub mod tests {
 	use std::iter::{Iterator, repeat};
 	use test_data;
 	use message::types;
-	use chain::Transaction;
+	use chain::{merkle_root, Transaction};
 	use primitives::hash::H256;
 	use primitives::bytes::Bytes;
 	use ser::serialize;
-	use super::{ConnectionFilter, ConnectionBloom};
+	use super::{ConnectionFilter, ConnectionBloom, PartialMerkleTree};
 
 	pub fn default_filterload() -> types::FilterLoad {
 		types::FilterLoad {
@@ -458,4 +558,46 @@ pub mod tests {
 		assert!(filter.filter_transaction(&tx1.hash(), &tx1));
 		assert!(!filter.filter_transaction(&tx2.hash(), &tx2));
 	}
+
+	#[test]
+	// test from core implementation
+	fn test_build_merkle_block() {
+		use bit_vec::BitVec;
+		use rand::{Rng, SeedableRng, StdRng};
+
+		let rng_seed: &[_] = &[0, 0, 0, 0];
+		let mut rng: StdRng = SeedableRng::from_seed(rng_seed);
+
+		// for some transactions counts
+		let tx_counts: Vec<usize> = vec![1, 4, 7, 17, 56, 100, 127, 256, 312, 513, 1000, 4095];
+		for tx_count in tx_counts {
+			// build block with given transactions number
+			let transactions: Vec<Transaction> = (0..tx_count).map(|n| test_data::TransactionBuilder::with_version(n as i32).into()).collect();
+			let hashes: Vec<_> = transactions.iter().map(|t| t.hash()).collect();
+			let merkle_root = merkle_root(&hashes);
+
+			// mark different transactions as matched
+			for seed_tweak in 1..15 {
+				let mut matches: BitVec = BitVec::with_capacity(tx_count);
+				let mut matched_hashes: Vec<H256> = Vec::with_capacity(tx_count);
+				for i in 0usize..tx_count {
+					let is_match = (rng.gen::<u32>() & ((1 << (seed_tweak / 2)) - 1)) == 0;
+					matches.push(is_match);
+					if is_match {
+						matched_hashes.push(hashes[i].clone());
+					}
+				}
+
+				// build partial merkle tree
+				let (built_hashes, built_flags) = PartialMerkleTree::build(hashes.clone(), matches.clone());
+				// parse tree back
+				let (parsed_root, parsed_hashes, parsed_positions) = PartialMerkleTree::parse(tx_count, built_hashes, built_flags)
+					.expect("no error");
+
+				assert_eq!(matched_hashes, parsed_hashes);
+				assert_eq!(matches, parsed_positions);
+				assert_eq!(merkle_root, parsed_root);
+			}
+		}
+	}
 }

sync/src/lib.rs

@@ -20,6 +20,7 @@ extern crate script;
 extern crate serialization as ser;
 #[cfg(test)]
 extern crate ethcore_devtools as devtools;
+extern crate rand;
 
 mod best_headers_chain;
 mod blocks_writer;