From 213dfd4f19104030d64394e6a6c568b3b699ac75 Mon Sep 17 00:00:00 2001
From: ard'O'crat <ardocrat@proton.me>
Date: Tue, 31 Oct 2023 11:23:19 +0000
Subject: [PATCH] Handle invalid MMR root to prevent sync thread panic (#3774)

* fix: handle invalid mmr root to prevent sync thread panic

* test: fix roots check
---
 chain/src/txhashset/desegmenter.rs |  2 +-
 chain/src/txhashset/txhashset.rs   | 12 ++++++------
 chain/tests/test_pibd_copy.rs      |  2 +-
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/chain/src/txhashset/desegmenter.rs b/chain/src/txhashset/desegmenter.rs
index 66f6df6b2..fbbe37ec7 100644
--- a/chain/src/txhashset/desegmenter.rs
+++ b/chain/src/txhashset/desegmenter.rs
@@ -237,7 +237,7 @@ impl Desegmenter {
 		// Quick root check first:
 		{
 			let txhashset = self.txhashset.read();
-			txhashset.roots().validate(&self.archive_header)?;
+			txhashset.roots()?.validate(&self.archive_header)?;
 		}
 
 		// TODO: Possibly Keep track of this in the DB so we can pick up where we left off if needed
diff --git a/chain/src/txhashset/txhashset.rs b/chain/src/txhashset/txhashset.rs
index 19f3f1339..db993bf80 100644
--- a/chain/src/txhashset/txhashset.rs
+++ b/chain/src/txhashset/txhashset.rs
@@ -479,19 +479,19 @@ impl TxHashSet {
 	}
 
 	/// Get MMR roots.
-	pub fn roots(&self) -> TxHashSetRoots {
+	pub fn roots(&self) -> Result<TxHashSetRoots, Error> {
 		let output_pmmr = ReadonlyPMMR::at(&self.output_pmmr_h.backend, self.output_pmmr_h.size);
 		let rproof_pmmr = ReadonlyPMMR::at(&self.rproof_pmmr_h.backend, self.rproof_pmmr_h.size);
 		let kernel_pmmr = ReadonlyPMMR::at(&self.kernel_pmmr_h.backend, self.kernel_pmmr_h.size);
 
-		TxHashSetRoots {
+		Ok(TxHashSetRoots {
 			output_roots: OutputRoots {
-				pmmr_root: output_pmmr.root().expect("no root, invalid tree"),
+				pmmr_root: output_pmmr.root().map_err(|_| Error::InvalidRoot)?,
 				bitmap_root: self.bitmap_accumulator.root(),
 			},
-			rproof_root: rproof_pmmr.root().expect("no root, invalid tree"),
-			kernel_root: kernel_pmmr.root().expect("no root, invalid tree"),
-		}
+			rproof_root: rproof_pmmr.root().map_err(|_| Error::InvalidRoot)?,
+			kernel_root: kernel_pmmr.root().map_err(|_| Error::InvalidRoot)?,
+		})
 	}
 
 	/// Return Commit's MMR position
diff --git a/chain/tests/test_pibd_copy.rs b/chain/tests/test_pibd_copy.rs
index 255f56cdd..9af635833 100644
--- a/chain/tests/test_pibd_copy.rs
+++ b/chain/tests/test_pibd_copy.rs
@@ -238,7 +238,7 @@ impl DesegmenterRequestor {
 	}
 
 	pub fn check_roots(&self) {
-		let roots = self.chain.txhashset().read().roots();
+		let roots = self.chain.txhashset().read().roots().unwrap();
 		let archive_header = self.chain.txhashset_archive_header_header_only().unwrap();
 		debug!("Archive Header is {:?}", archive_header);
 		debug!("TXHashset output root is {:?}", roots);