Improve error handling when computing PMMR roots (#2988)

* Improve error handling for invalid PMMR roots

* Update tests that rely on pmmr root

* Fix pmmr store tests

chain/src/chain.rs

@@ -545,12 +545,12 @@ impl Chain {
 				pipe::rewind_and_apply_fork(&previous_header, &header_head, extension)?;
 
 				// Retrieve the header root before we apply the new block
-				let prev_root = extension.header_root();
+				let prev_root = extension.header_root()?;
 
 				// Apply the latest block to the chain state via the extension.
 				extension.apply_block(b)?;
 
-				Ok((prev_root, extension.roots(), extension.sizes()))
+				Ok((prev_root, extension.roots()?, extension.sizes()))
 			})?;
 
 		// Set the prev_root on the header.

chain/src/txhashset/txhashset.rs

@@ -670,7 +670,7 @@ impl<'a> HeaderExtension<'a> {
 	pub fn apply_header(&mut self, header: &BlockHeader) -> Result<Hash, Error> {
 		self.pmmr.push(header).map_err(&ErrorKind::TxHashSetErr)?;
 		self.head = Tip::from_header(header);
-		Ok(self.root())
+		Ok(self.root()?)
 	}
 
 	/// Rewind the header extension to the specified header.
@@ -749,8 +749,8 @@ impl<'a> HeaderExtension<'a> {
 	}
 
 	/// The root of the header MMR for convenience.
-	pub fn root(&self) -> Hash {
-		self.pmmr.root()
+	pub fn root(&self) -> Result<Hash, Error> {
+		Ok(self.pmmr.root().map_err(|_| ErrorKind::InvalidRoot)?)
 	}
 
 	/// Validate the prev_root of the header against the root of the current header MMR.
@@ -760,7 +760,7 @@ impl<'a> HeaderExtension<'a> {
 		if header.height == 0 {
 			return Ok(());
 		}
-		if self.root() != header.prev_root {
+		if self.root()? != header.prev_root {
 			Err(ErrorKind::InvalidRoot.into())
 		} else {
 			Ok(())
@@ -1095,18 +1095,33 @@ impl<'a> Extension<'a> {
 
 	/// Current root hashes and sums (if applicable) for the Output, range proof
 	/// and kernel sum trees.
-	pub fn roots(&self) -> TxHashSetRoots {
-		TxHashSetRoots {
-			header_root: self.header_pmmr.root(),
-			output_root: self.output_pmmr.root(),
-			rproof_root: self.rproof_pmmr.root(),
-			kernel_root: self.kernel_pmmr.root(),
-		}
+	pub fn roots(&self) -> Result<TxHashSetRoots, Error> {
+		Ok(TxHashSetRoots {
+			header_root: self
+				.header_pmmr
+				.root()
+				.map_err(|_| ErrorKind::InvalidRoot)?,
+			output_root: self
+				.output_pmmr
+				.root()
+				.map_err(|_| ErrorKind::InvalidRoot)?,
+			rproof_root: self
+				.rproof_pmmr
+				.root()
+				.map_err(|_| ErrorKind::InvalidRoot)?,
+			kernel_root: self
+				.kernel_pmmr
+				.root()
+				.map_err(|_| ErrorKind::InvalidRoot)?,
+		})
 	}
 
 	/// Get the root of the current header MMR.
-	pub fn header_root(&self) -> Hash {
-		self.header_pmmr.root()
+	pub fn header_root(&self) -> Result<Hash, Error> {
+		Ok(self
+			.header_pmmr
+			.root()
+			.map_err(|_| ErrorKind::InvalidRoot)?)
 	}
 
 	/// Validate the following MMR roots against the latest header applied -
@@ -1126,7 +1141,7 @@ impl<'a> Extension<'a> {
 			return Ok(());
 		}
 		let head_header = self.batch.get_block_header(&self.head.last_block_h)?;
-		let roots = self.roots();
+		let roots = self.roots()?;
 		if roots.output_root != head_header.output_root
 			|| roots.rproof_root != head_header.range_proof_root
 			|| roots.kernel_root != head_header.kernel_root
@@ -1143,7 +1158,7 @@ impl<'a> Extension<'a> {
 		if header.height == 0 {
 			return Ok(());
 		}
-		let roots = self.roots();
+		let roots = self.roots()?;
 		if roots.header_root != header.prev_root {
 			Err(ErrorKind::InvalidRoot.into())
 		} else {

core/src/core/pmmr/pmmr.rs

@@ -129,9 +129,9 @@ where
 
 	/// Computes the root of the MMR. Find all the peaks in the current
 	/// tree and "bags" them to get a single peak.
-	pub fn root(&self) -> Hash {
+	pub fn root(&self) -> Result<Hash, String> {
 		if self.is_empty() {
-			return ZERO_HASH;
+			return Ok(ZERO_HASH);
 		}
 		let mut res = None;
 		for peak in self.peaks().iter().rev() {
@@ -140,7 +140,7 @@ where
 				Some(rhash) => Some((*peak, rhash).hash_with_index(self.unpruned_size())),
 			}
 		}
-		res.expect("no root, invalid tree")
+		res.ok_or_else(|| "no root, invalid tree".to_owned())
 	}
 
 	/// Build a Merkle proof for the element at the given position.

core/tests/merkle_proof.rs

@@ -81,45 +81,45 @@ fn pmmr_merkle_proof() {
 
 	let proof = pmmr.merkle_proof(1).unwrap();
 	assert_eq!(proof.path, vec![]);
-	assert!(proof.verify(pmmr.root(), &elems[0], 1).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[0], 1).is_ok());
 
 	pmmr.push(&elems[1]).unwrap();
 	let pos_1 = elems[1].hash_with_index(1);
 	assert_eq!(pmmr.get_hash(2).unwrap(), pos_1);
 	let pos_2 = (pos_0, pos_1).hash_with_index(2);
 	assert_eq!(pmmr.get_hash(3).unwrap(), pos_2);
 
-	assert_eq!(pmmr.root(), pos_2);
+	assert_eq!(pmmr.root().unwrap(), pos_2);
 	assert_eq!(pmmr.peaks(), [pos_2]);
 
 	// single peak, path with single sibling
 	let proof = pmmr.merkle_proof(1).unwrap();
 	assert_eq!(proof.path, vec![pos_1]);
-	assert!(proof.verify(pmmr.root(), &elems[0], 1).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[0], 1).is_ok());
 
 	let proof = pmmr.merkle_proof(2).unwrap();
 	assert_eq!(proof.path, vec![pos_0]);
-	assert!(proof.verify(pmmr.root(), &elems[1], 2).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[1], 2).is_ok());
 
 	// three leaves, two peaks (one also the right-most leaf)
 	pmmr.push(&elems[2]).unwrap();
 	let pos_3 = elems[2].hash_with_index(3);
 	assert_eq!(pmmr.get_hash(4).unwrap(), pos_3);
 
-	assert_eq!(pmmr.root(), (pos_2, pos_3).hash_with_index(4));
+	assert_eq!(pmmr.root().unwrap(), (pos_2, pos_3).hash_with_index(4));
 	assert_eq!(pmmr.peaks(), [pos_2, pos_3]);
 
 	let proof = pmmr.merkle_proof(1).unwrap();
 	assert_eq!(proof.path, vec![pos_1, pos_3]);
-	assert!(proof.verify(pmmr.root(), &elems[0], 1).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[0], 1).is_ok());
 
 	let proof = pmmr.merkle_proof(2).unwrap();
 	assert_eq!(proof.path, vec![pos_0, pos_3]);
-	assert!(proof.verify(pmmr.root(), &elems[1], 2).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[1], 2).is_ok());
 
 	let proof = pmmr.merkle_proof(4).unwrap();
 	assert_eq!(proof.path, vec![pos_2]);
-	assert!(proof.verify(pmmr.root(), &elems[2], 4).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[2], 4).is_ok());
 
 	// 7 leaves, 3 peaks, 11 pos in total
 	pmmr.push(&elems[3]).unwrap();
@@ -152,38 +152,38 @@ fn pmmr_merkle_proof() {
 		proof.path,
 		vec![pos_1, pos_5, (pos_9, pos_10).hash_with_index(11)]
 	);
-	assert!(proof.verify(pmmr.root(), &elems[0], 1).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[0], 1).is_ok());
 
 	let proof = pmmr.merkle_proof(2).unwrap();
 	assert_eq!(
 		proof.path,
 		vec![pos_0, pos_5, (pos_9, pos_10).hash_with_index(11)]
 	);
-	assert!(proof.verify(pmmr.root(), &elems[1], 2).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[1], 2).is_ok());
 
 	let proof = pmmr.merkle_proof(4).unwrap();
 	assert_eq!(
 		proof.path,
 		vec![pos_4, pos_2, (pos_9, pos_10).hash_with_index(11)]
 	);
-	assert!(proof.verify(pmmr.root(), &elems[2], 4).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[2], 4).is_ok());
 
 	let proof = pmmr.merkle_proof(5).unwrap();
 	assert_eq!(
 		proof.path,
 		vec![pos_3, pos_2, (pos_9, pos_10).hash_with_index(11)]
 	);
-	assert!(proof.verify(pmmr.root(), &elems[3], 5).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[3], 5).is_ok());
 
 	let proof = pmmr.merkle_proof(8).unwrap();
 	assert_eq!(proof.path, vec![pos_8, pos_10, pos_6]);
-	assert!(proof.verify(pmmr.root(), &elems[4], 8).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[4], 8).is_ok());
 
 	let proof = pmmr.merkle_proof(9).unwrap();
 	assert_eq!(proof.path, vec![pos_7, pos_10, pos_6]);
-	assert!(proof.verify(pmmr.root(), &elems[5], 9).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[5], 9).is_ok());
 
 	let proof = pmmr.merkle_proof(11).unwrap();
 	assert_eq!(proof.path, vec![pos_9, pos_6]);
-	assert!(proof.verify(pmmr.root(), &elems[6], 11).is_ok());
+	assert!(proof.verify(pmmr.root().unwrap(), &elems[6], 11).is_ok());
 }

core/tests/pmmr.rs

@@ -280,7 +280,7 @@ fn pmmr_push_root() {
 	pmmr.dump(false);
 	let pos_0 = elems[0].hash_with_index(0);
 	assert_eq!(pmmr.peaks(), vec![pos_0]);
-	assert_eq!(pmmr.root(), pos_0);
+	assert_eq!(pmmr.root().unwrap(), pos_0);
 	assert_eq!(pmmr.unpruned_size(), 1);
 
 	// two elements
@@ -289,15 +289,15 @@ fn pmmr_push_root() {
 	let pos_1 = elems[1].hash_with_index(1);
 	let pos_2 = (pos_0, pos_1).hash_with_index(2);
 	assert_eq!(pmmr.peaks(), vec![pos_2]);
-	assert_eq!(pmmr.root(), pos_2);
+	assert_eq!(pmmr.root().unwrap(), pos_2);
 	assert_eq!(pmmr.unpruned_size(), 3);
 
 	// three elements
 	pmmr.push(&elems[2]).unwrap();
 	pmmr.dump(false);
 	let pos_3 = elems[2].hash_with_index(3);
 	assert_eq!(pmmr.peaks(), vec![pos_2, pos_3]);
-	assert_eq!(pmmr.root(), (pos_2, pos_3).hash_with_index(4));
+	assert_eq!(pmmr.root().unwrap(), (pos_2, pos_3).hash_with_index(4));
 	assert_eq!(pmmr.unpruned_size(), 4);
 
 	// four elements
@@ -307,31 +307,31 @@ fn pmmr_push_root() {
 	let pos_5 = (pos_3, pos_4).hash_with_index(5);
 	let pos_6 = (pos_2, pos_5).hash_with_index(6);
 	assert_eq!(pmmr.peaks(), vec![pos_6]);
-	assert_eq!(pmmr.root(), pos_6);
+	assert_eq!(pmmr.root().unwrap(), pos_6);
 	assert_eq!(pmmr.unpruned_size(), 7);
 
 	// five elements
 	pmmr.push(&elems[4]).unwrap();
 	pmmr.dump(false);
 	let pos_7 = elems[4].hash_with_index(7);
 	assert_eq!(pmmr.peaks(), vec![pos_6, pos_7]);
-	assert_eq!(pmmr.root(), (pos_6, pos_7).hash_with_index(8));
+	assert_eq!(pmmr.root().unwrap(), (pos_6, pos_7).hash_with_index(8));
 	assert_eq!(pmmr.unpruned_size(), 8);
 
 	// six elements
 	pmmr.push(&elems[5]).unwrap();
 	let pos_8 = elems[5].hash_with_index(8);
 	let pos_9 = (pos_7, pos_8).hash_with_index(9);
 	assert_eq!(pmmr.peaks(), vec![pos_6, pos_9]);
-	assert_eq!(pmmr.root(), (pos_6, pos_9).hash_with_index(10));
+	assert_eq!(pmmr.root().unwrap(), (pos_6, pos_9).hash_with_index(10));
 	assert_eq!(pmmr.unpruned_size(), 10);
 
 	// seven elements
 	pmmr.push(&elems[6]).unwrap();
 	let pos_10 = elems[6].hash_with_index(10);
 	assert_eq!(pmmr.peaks(), vec![pos_6, pos_9, pos_10]);
 	assert_eq!(
-		pmmr.root(),
+		pmmr.root().unwrap(),
 		(pos_6, (pos_9, pos_10).hash_with_index(11)).hash_with_index(11)
 	);
 	assert_eq!(pmmr.unpruned_size(), 11);
@@ -344,14 +344,14 @@ fn pmmr_push_root() {
 	let pos_13 = (pos_9, pos_12).hash_with_index(13);
 	let pos_14 = (pos_6, pos_13).hash_with_index(14);
 	assert_eq!(pmmr.peaks(), vec![pos_14]);
-	assert_eq!(pmmr.root(), pos_14);
+	assert_eq!(pmmr.root().unwrap(), pos_14);
 	assert_eq!(pmmr.unpruned_size(), 15);
 
 	// nine elements
 	pmmr.push(&elems[8]).unwrap();
 	let pos_15 = elems[8].hash_with_index(15);
 	assert_eq!(pmmr.peaks(), vec![pos_14, pos_15]);
-	assert_eq!(pmmr.root(), (pos_14, pos_15).hash_with_index(16));
+	assert_eq!(pmmr.root().unwrap(), (pos_14, pos_15).hash_with_index(16));
 	assert_eq!(pmmr.unpruned_size(), 16);
 }
 
@@ -427,7 +427,7 @@ fn pmmr_prune() {
 		for elem in &elems[..] {
 			pmmr.push(elem).unwrap();
 		}
-		orig_root = pmmr.root();
+		orig_root = pmmr.root().unwrap();
 		sz = pmmr.unpruned_size();
 	}
 
@@ -439,7 +439,7 @@ fn pmmr_prune() {
 	{
 		let mut pmmr: PMMR<'_, TestElem, _> = PMMR::at(&mut ba, sz);
 		pmmr.prune(16).unwrap();
-		assert_eq!(orig_root, pmmr.root());
+		assert_eq!(orig_root, pmmr.root().unwrap());
 	}
 	assert_eq!(ba.hashes.len(), 16);
 	assert_eq!(ba.remove_list.len(), 1);
@@ -448,15 +448,15 @@ fn pmmr_prune() {
 	{
 		let mut pmmr: PMMR<'_, TestElem, _> = PMMR::at(&mut ba, sz);
 		pmmr.prune(2).unwrap();
-		assert_eq!(orig_root, pmmr.root());
+		assert_eq!(orig_root, pmmr.root().unwrap());
 	}
 	assert_eq!(ba.hashes.len(), 16);
 	assert_eq!(ba.remove_list.len(), 2);
 
 	{
 		let mut pmmr: PMMR<'_, TestElem, _> = PMMR::at(&mut ba, sz);
 		pmmr.prune(4).unwrap();
-		assert_eq!(orig_root, pmmr.root());
+		assert_eq!(orig_root, pmmr.root().unwrap());
 	}
 	assert_eq!(ba.hashes.len(), 16);
 	assert_eq!(ba.remove_list.len(), 3);
@@ -465,7 +465,7 @@ fn pmmr_prune() {
 	{
 		let mut pmmr: PMMR<'_, TestElem, _> = PMMR::at(&mut ba, sz);
 		pmmr.prune(3).unwrap_err();
-		assert_eq!(orig_root, pmmr.root());
+		assert_eq!(orig_root, pmmr.root().unwrap());
 	}
 	assert_eq!(ba.hashes.len(), 16);
 	assert_eq!(ba.remove_list.len(), 3);
@@ -474,7 +474,7 @@ fn pmmr_prune() {
 	{
 		let mut pmmr: PMMR<'_, TestElem, _> = PMMR::at(&mut ba, sz);
 		pmmr.prune(5).unwrap();
-		assert_eq!(orig_root, pmmr.root());
+		assert_eq!(orig_root, pmmr.root().unwrap());
 	}
 	assert_eq!(ba.hashes.len(), 16);
 	assert_eq!(ba.remove_list.len(), 4);
@@ -484,7 +484,7 @@ fn pmmr_prune() {
 	{
 		let mut pmmr: PMMR<'_, TestElem, _> = PMMR::at(&mut ba, sz);
 		pmmr.prune(1).unwrap();
-		assert_eq!(orig_root, pmmr.root());
+		assert_eq!(orig_root, pmmr.root().unwrap());
 	}
 	assert_eq!(ba.hashes.len(), 16);
 	assert_eq!(ba.remove_list.len(), 5);
@@ -495,7 +495,7 @@ fn pmmr_prune() {
 		for n in 1..16 {
 			let _ = pmmr.prune(n);
 		}
-		assert_eq!(orig_root, pmmr.root());
+		assert_eq!(orig_root, pmmr.root().unwrap());
 	}
 	assert_eq!(ba.hashes.len(), 16);
 	assert_eq!(ba.remove_list.len(), 9);