From b9fb71fb8afcc13368c5742247a9e4aa87b8d890 Mon Sep 17 00:00:00 2001 From: Martin Stefcek Date: Mon, 12 Feb 2024 19:19:00 +0400 Subject: [PATCH 1/2] fix: balanced binary merkle tree merged proof --- .../mmr/src/balanced_binary_merkle_proof.rs | 59 +++++++++++++++---- 1 file changed, 49 insertions(+), 10 deletions(-) diff --git a/base_layer/mmr/src/balanced_binary_merkle_proof.rs b/base_layer/mmr/src/balanced_binary_merkle_proof.rs index 23fe5d8522..b9b91cbd12 100644 --- a/base_layer/mmr/src/balanced_binary_merkle_proof.rs +++ b/base_layer/mmr/src/balanced_binary_merkle_proof.rs @@ -20,7 +20,12 @@ // WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE // USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -use std::{collections::HashMap, convert::TryFrom, marker::PhantomData}; +use std::{ + collections::{HashMap, HashSet}, + convert::TryFrom, + f32::consts::E, + marker::PhantomData, +}; use borsh::{BorshDeserialize, BorshSerialize}; use digest::Digest; @@ -198,22 +203,35 @@ where D: Digest + DomainDigest .iter() .max() .ok_or(BalancedBinaryMerkleProofError::CantMergeZeroProofs)?; - + let mut consumed = HashSet::new(); // We need to compute the hashes row by row to be sure they are processed correctly. for height in (0..max_height).rev() { let hashes = computed_hashes.clone(); - for (index, leaf) in computed_hashes.iter_mut().enumerate() { + let mut dangling_paths = HashSet::new(); + for (index, leaf) in computed_hashes.iter_mut().enumerate().rev() { if self.heights[index] <= height { continue; } let Some(hash_or_index) = self.paths[index].pop() else { + // Check if we already joined with other path. + if !consumed.contains(&index) { + // If the path ended, it's going to be merged to some other path. + if !dangling_paths.insert(index) { + return Err(BalancedBinaryMerkleProofError::BadProofSemantics); + } + } // Path at this index already completely processed continue; }; let hash = match hash_or_index { MergedBalancedBinaryMerkleIndexOrHash::Index(index) => { + if !dangling_paths.remove(&(index as usize)) { + // If some path is joining our path, that path should have ended. + return Err(BalancedBinaryMerkleProofError::BadProofSemantics); + } + consumed.insert(index as usize); let index = usize::try_from(index).map_err(|_| BalancedBinaryMerkleProofError::MathOverflow)?; // The index must also point to one of the proofs @@ -235,6 +253,14 @@ where D: Digest + DomainDigest // Parent self.node_indices[index] = (self.node_indices[index] - 1) >> 1; } + if !dangling_paths.is_empty() { + // Something path ended, but it's not joined with any other path. + return Err(BalancedBinaryMerkleProofError::BadProofSemantics); + } + } + if consumed.len() + 1 < self.paths.len() { + // If the proof is valid then all but one paths will be consumed by other paths. + return Err(BalancedBinaryMerkleProofError::BadProofSemantics); } Ok(computed_hashes[0] == *root) } @@ -292,7 +318,8 @@ mod test { heights: vec![1], _phantom: PhantomData, }; - assert!(!proof.verify_consume(&vec![0u8; 32], vec![vec![]]).unwrap()); + // This will fail because the node height is 1 and it's empty, so it's not going to compute the root hash. + proof.verify_consume(&vec![0u8; 32], vec![vec![]]).unwrap_err(); } #[test] @@ -334,10 +361,10 @@ mod test { #[test] fn test_merge_proof_full_tree() { - let leaves = (0..255).map(|i| vec![i; 32]).collect::>(); + let leaves = (0..=255).map(|i| vec![i; 32]).collect::>(); let bmt = BalancedBinaryMerkleTree::::create(leaves.clone()); let root = bmt.get_merkle_root(); - let proofs = (0..255) + let proofs = (0..=255) .map(|i| BalancedBinaryMerkleProof::generate_proof(&bmt, i)) .collect::, _>>() .unwrap(); @@ -382,11 +409,23 @@ mod test { heights: vec![0, 0], _phantom: PhantomData, }; - // This should fail but does not - // proof .verify_consume(&vec![5u8; 32], vec![vec![5u8; 32], vec![2u8; 32]]) .unwrap_err(); - assert!(proof + // This will fail because there are more hashes on the same level as there can be. + proof .verify_consume(&vec![5u8; 32], vec![vec![5u8; 32], vec![2u8; 32]]) - .unwrap()); + .unwrap_err(); + + let proof = MergedBalancedBinaryMerkleProof:: { + paths: vec![vec![MergedBalancedBinaryMerkleIndexOrHash::Hash(vec![5u8; 32])], vec![ + MergedBalancedBinaryMerkleIndexOrHash::Index(1), + ]], + node_indices: vec![1, 1], + heights: vec![0, 1], + _phantom: PhantomData, + }; + // This will fail because we can't have any more nodes if we have leaf at the root. + proof + .verify_consume(&vec![5u8; 32], vec![vec![5u8; 32], vec![2u8; 32]]) + .unwrap_err(); } #[test] From 9283e74e97bf761a162f7491a6f7fd5db63f7255 Mon Sep 17 00:00:00 2001 From: Martin Stefcek Date: Tue, 13 Feb 2024 11:13:23 +0400 Subject: [PATCH 2/2] fix clippy --- base_layer/mmr/src/balanced_binary_merkle_proof.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/base_layer/mmr/src/balanced_binary_merkle_proof.rs b/base_layer/mmr/src/balanced_binary_merkle_proof.rs index b9b91cbd12..eb8ee789ce 100644 --- a/base_layer/mmr/src/balanced_binary_merkle_proof.rs +++ b/base_layer/mmr/src/balanced_binary_merkle_proof.rs @@ -23,7 +23,6 @@ use std::{ collections::{HashMap, HashSet}, convert::TryFrom, - f32::consts::E, marker::PhantomData, }; @@ -227,11 +226,14 @@ where D: Digest + DomainDigest let hash = match hash_or_index { MergedBalancedBinaryMerkleIndexOrHash::Index(index) => { - if !dangling_paths.remove(&(index as usize)) { + if !dangling_paths + .remove(&usize::try_from(index).map_err(|_| BalancedBinaryMerkleProofError::MathOverflow)?) + { // If some path is joining our path, that path should have ended. return Err(BalancedBinaryMerkleProofError::BadProofSemantics); } - consumed.insert(index as usize); + consumed + .insert(usize::try_from(index).map_err(|_| BalancedBinaryMerkleProofError::MathOverflow)?); let index = usize::try_from(index).map_err(|_| BalancedBinaryMerkleProofError::MathOverflow)?; // The index must also point to one of the proofs