From e92a3dae0fd335ff57ed886dccad08b61d8519ce Mon Sep 17 00:00:00 2001 From: Boyu Yang Date: Wed, 28 Jun 2023 17:59:25 +0800 Subject: [PATCH] feat: verify a old root and all incremental items after it --- src/mmr.rs | 198 +++++++++++++++++++++++++++++++++- src/tests/mod.rs | 1 + src/tests/test_incremental.rs | 54 ++++++++++ 3 files changed, 251 insertions(+), 2 deletions(-) create mode 100644 src/tests/test_incremental.rs diff --git a/src/mmr.rs b/src/mmr.rs index 91efb85..123a25e 100644 --- a/src/mmr.rs +++ b/src/mmr.rs @@ -5,8 +5,11 @@ //! https://github.com/mimblewimble/grin/blob/0ff6763ee64e5a14e70ddd4642b99789a1648a32/core/src/core/pmmr.rs#L606 use crate::borrow::Cow; -use crate::collections::VecDeque; -use crate::helper::{get_peak_map, get_peaks, parent_offset, pos_height_in_tree, sibling_offset}; +use crate::collections::{BTreeMap, VecDeque}; +use crate::helper::{ + get_peak_map, get_peaks, leaf_index_to_mmr_size, leaf_index_to_pos, parent_offset, + pos_height_in_tree, sibling_offset, +}; use crate::mmr_store::{MMRBatch, MMRStoreReadOps, MMRStoreWriteOps}; use crate::vec; use crate::vec::Vec; @@ -291,6 +294,75 @@ impl> MerkleProof { self.calculate_root(leaves) .map(|calculated_root| calculated_root == root) } + + /// Verifies a old root and all incremental leaves. + /// + /// If this method returns `true`, it means the following assertion are true: + /// - The old root could be generated in the history of the current MMR. + /// - All incremental leaves are on the current MMR. + /// - The MMR, which could generate the old root, appends all incremental leaves, becomes the + /// current MMR. + // TODO refactor: Use leaves count to instead of mmr size in MerkleProof. + // + // Reason: + // - Calculate mmr size with leaves count is easy; but the reverse is difficult and it may + // cause errors. + // - The `prev_leaves_count` in this method could be removed, it could be calculated via + // `current_leaves_count - incremental.len()`. + pub fn verify_incremental( + &self, + root: T, + prev_leaves_count: u64, + prev_root: T, + incremental: Vec, + ) -> Result { + let (positions, leaves) = incremental.into_iter().enumerate().fold( + (Vec::new(), Vec::new()), + |(mut positions, mut leaves), (index, leaf)| { + let pos = leaf_index_to_pos(prev_leaves_count + index as u64); + positions.push(pos); + leaves.push((pos, leaf)); + (positions, leaves) + }, + ); + // Test if previous root is correct. + let prev_peaks_positions = { + let prev_index = prev_leaves_count - 1; + let prev_mmr_size = leaf_index_to_mmr_size(prev_index); + let prev_peaks_positions = get_peaks(prev_mmr_size); + if prev_peaks_positions.len() != self.proof.len() { + return Err(Error::CorruptedProof); + } + prev_peaks_positions + }; + let proof_items_positions = calculate_positions_of_proof_items_via_peaks_hashes( + positions, + self.mmr_size, + self.proof.len(), + )?; + if proof_items_positions.len() != self.proof.len() { + return Err(Error::CorruptedProof); + } + let map = proof_items_positions + .into_iter() + .enumerate() + .map(|(index, pos)| (pos, index)) + .collect::>(); + let prev_peaks_res = prev_peaks_positions + .iter() + .map(|pos| map.get(pos).map(|index| self.proof[*index].clone())) + .collect::>(); + if let Some(prev_peaks) = prev_peaks_res { + let calculated_prev_root = bagging_peaks_hashes::(prev_peaks)?; + if calculated_prev_root != prev_root { + return Ok(false); + } + } else { + return Ok(false); + } + // Test if incremental leaves are correct. + self.verify(root, leaves) + } } fn calculate_peak_root<'a, T: 'a, M: Merge, I: Iterator>( @@ -356,6 +428,70 @@ fn calculate_peak_root<'a, T: 'a, M: Merge, I: Iterator> Err(Error::CorruptedProof) } +// Simplified version of `calculate_peak_root`. +fn calculate_positions_of_proof_items_via_peak_root( + leaves: Vec, + peak_pos: u64, + proof_items_count: &mut usize, + return_data: &mut Vec, +) -> Result<()> { + debug_assert!(!leaves.is_empty(), "can't be empty"); + + let mut queue: VecDeque<_> = leaves.into_iter().map(|pos| (pos, 0)).collect(); + + // calculate tree root from each items + while let Some((pos, height)) = queue.pop_front() { + if pos == peak_pos { + if queue.is_empty() { + return Ok(()); + } else { + return Err(Error::CorruptedProof); + } + } + // calculate sibling + let next_height = pos_height_in_tree(pos + 1); + let parent_pos = { + let sibling_offset = sibling_offset(height); + if next_height > height { + // implies pos is right sibling + let sib_pos = pos - sibling_offset; + let parent_pos = pos + 1; + if Some(&sib_pos) == queue.front().map(|(pos, _)| pos) { + let _ = queue.pop_front(); + } else { + if *proof_items_count == 0 { + return Err(Error::CorruptedProof); + } + *proof_items_count -= 1; + return_data.push(sib_pos); + } + parent_pos + } else { + // pos is left sibling + let sib_pos = pos + sibling_offset; + let parent_pos = pos + parent_offset(height); + if Some(&sib_pos) == queue.front().map(|(pos, _)| pos) { + let _ = queue.pop_front(); + } else { + if *proof_items_count == 0 { + return Err(Error::CorruptedProof); + } + *proof_items_count -= 1; + return_data.push(sib_pos); + } + parent_pos + } + }; + + if parent_pos <= peak_pos { + queue.push_back((parent_pos, height + 1)) + } else { + return Err(Error::CorruptedProof); + } + } + Err(Error::CorruptedProof) +} + fn calculate_peaks_hashes<'a, T: 'a + Clone, M: Merge, I: Iterator>( mut leaves: Vec<(u64, T)>, mmr_size: u64, @@ -411,6 +547,64 @@ fn calculate_peaks_hashes<'a, T: 'a + Clone, M: Merge, I: Iterator, + mmr_size: u64, + mut proof_items_count: usize, +) -> Result> { + if leaves.iter().any(|pos| pos_height_in_tree(*pos) > 0) { + return Err(Error::NodeProofsNotSupported); + } + + // special handle the only 1 leaf MMR + if mmr_size == 1 && leaves.len() == 1 && leaves[0] == 0 { + return Ok(Default::default()); + } + // ensure leaves are sorted and unique + leaves.sort_by_key(|pos| *pos); + leaves.dedup_by(|a, b| a == b); + let peaks = get_peaks(mmr_size); + + let mut peaks_hashes: Vec = Vec::with_capacity(peaks.len() + 1); + for peak_pos in peaks { + let mut leaves: Vec<_> = take_while_vec(&mut leaves, |pos| *pos <= peak_pos); + if leaves.len() == 1 && leaves[0] == peak_pos { + // leaf is the peak + let _ = leaves.remove(0); + } else if leaves.is_empty() { + // if empty, means the next proof is a peak root or rhs bagged root + if proof_items_count > 0 { + proof_items_count -= 1; + peaks_hashes.push(peak_pos); + } else { + // means that either all right peaks are bagged, or proof is corrupted + // so we break loop and check no items left + break; + } + } else { + calculate_positions_of_proof_items_via_peak_root( + leaves, + peak_pos, + &mut proof_items_count, + &mut peaks_hashes, + )?; + }; + } + + // ensure nothing left in leaves + if !leaves.is_empty() { + return Err(Error::CorruptedProof); + } + + // check rhs peaks + if proof_items_count > 0 { + return Err(Error::CorruptedProof); + } + + Ok(peaks_hashes) +} + fn bagging_peaks_hashes>(mut peaks_hashes: Vec) -> Result { // bagging peaks // bagging from right to left via hash(right, left). diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 5248702..81fb485 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -1,5 +1,6 @@ mod test_accumulate_headers; mod test_helper; +mod test_incremental; mod test_mmr; mod test_sequence; diff --git a/src/tests/test_incremental.rs b/src/tests/test_incremental.rs new file mode 100644 index 0000000..a014b13 --- /dev/null +++ b/src/tests/test_incremental.rs @@ -0,0 +1,54 @@ +use proptest::proptest; + +use super::{MergeNumberHash, NumberHash}; +use crate::util::{MemMMR, MemStore}; + +proptest! { + #[test] + fn test_incremental(start in 1u32..500, steps in 1usize..50, turns in 10usize..20) { + test_incremental_with_params(start, steps, turns); + } +} + +fn test_incremental_with_params(start: u32, steps: usize, turns: usize) { + let store = MemStore::default(); + let mut mmr = MemMMR::<_, MergeNumberHash>::new(0, &store); + + let mut curr = 0; + + let _positions: Vec = (0u32..start) + .map(|_| { + let pos = mmr.push(NumberHash::from(curr)).unwrap(); + curr += 1; + pos + }) + .collect(); + mmr.commit().expect("commit changes"); + + for turn in 0..turns { + let prev_root = mmr.get_root().expect("get root"); + let prev_leaves_count = u64::from(curr); + let (positions, leaves) = (0..steps).fold( + (Vec::new(), Vec::new()), + |(mut positions, mut leaves), _| { + let leaf = NumberHash::from(curr); + let pos = mmr.push(leaf.clone()).unwrap(); + curr += 1; + positions.push(pos); + leaves.push(leaf); + (positions, leaves) + }, + ); + mmr.commit().expect("commit changes"); + let proof = mmr.gen_proof(positions).expect("gen proof"); + let root = mmr.get_root().expect("get root"); + let result = proof + .verify_incremental(root, prev_leaves_count, prev_root, leaves) + .unwrap(); + assert!( + result, + "start: {}, steps: {}, turn: {}, curr: {}", + start, steps, turn, curr + ); + } +}