From 8faba4ef83511b09c50229749f586428a45430e6 Mon Sep 17 00:00:00 2001 From: jaspervdm Date: Tue, 17 Nov 2020 19:38:44 +0100 Subject: [PATCH] PMMR segment creation and validation (#3453) * Chunk generation and validation * Rename chunk -> segment * Missed a few * Generate and validate merkle proof * Fix bugs in generation and validation * Add test for unprunable MMR of various sizes * Add missing docs * Remove unused functions * Remove segment error variant on chain error type * Simplify calculation by using a Vec instead of HashMap * Use vectors in segment definition * Compare subtree root during tests * Add test of segments for a prunable mmr * Remove assertion * Only send intermediary hashes for prunable MMRs * Get hash from file directly * Require both leaves if one of them is not pruned * More pruning tests * Add segment (de)serialization * Require sorted vectors in segment deser * Store pos and data separately in segment * Rename log_size -> height * Fix bitmap index in root calculation * Add validation function for output (bitmap) MMRs * Remove left over debug statements * Fix test * Edge case: final segment with uneven number of leaves * Use last_pos instead of segment_last_pos * Simplify pruning in test * Add leaf and hash iterators * Support fully pruned segments * Drop backend before deleting dir in pruned_segment test * Simplify output of first_unpruned_parent --- core/src/core.rs | 1 + core/src/core/pmmr.rs | 1 + core/src/core/pmmr/segment.rs | 680 ++++++++++++++++++++++++++++++++++ core/tests/segment.rs | 61 +++ store/tests/segment.rs | 422 +++++++++++++++++++++ 5 files changed, 1165 insertions(+) create mode 100644 core/src/core/pmmr/segment.rs create mode 100644 core/tests/segment.rs create mode 100644 store/tests/segment.rs diff --git a/core/src/core.rs b/core/src/core.rs index 89c6b67a5..c90aa7969 100644 --- a/core/src/core.rs +++ b/core/src/core.rs @@ -33,6 +33,7 @@ pub use self::block_sums::*; pub use self::committed::Committed; pub use self::compact_block::*; pub use self::id::ShortId; +pub use self::pmmr::segment::*; pub use self::transaction::*; /// Common errors diff --git a/core/src/core/pmmr.rs b/core/src/core/pmmr.rs index b41d89467..a294eba24 100644 --- a/core/src/core/pmmr.rs +++ b/core/src/core/pmmr.rs @@ -40,6 +40,7 @@ mod backend; mod pmmr; mod readonly_pmmr; mod rewindable_pmmr; +pub mod segment; mod vec_backend; pub use self::backend::*; diff --git a/core/src/core/pmmr/segment.rs b/core/src/core/pmmr/segment.rs new file mode 100644 index 000000000..551254c8b --- /dev/null +++ b/core/src/core/pmmr/segment.rs @@ -0,0 +1,680 @@ +// Copyright 2020 The Grin Developers +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Segment of a PMMR. + +use crate::core::hash::Hash; +use crate::core::pmmr::{self, Backend, ReadablePMMR, ReadonlyPMMR}; +use crate::ser::{Error, PMMRIndexHashable, PMMRable, Readable, Reader, Writeable, Writer}; +use croaring::Bitmap; +use std::cmp::min; +use std::fmt::{self, Debug}; + +#[derive(Clone, Debug, PartialEq, Eq)] +/// Error related to segment creation or validation +pub enum SegmentError { + /// An expected leaf was missing + MissingLeaf(u64), + /// An expected hash was missing + MissingHash(u64), + /// The segment does not exist + NonExistent, + /// Mismatch between expected and actual root hash + Mismatch, +} + +impl fmt::Display for SegmentError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + SegmentError::MissingLeaf(idx) => write!(f, "Missing leaf at pos {}", idx), + SegmentError::MissingHash(idx) => write!(f, "Missing hash at pos {}", idx), + SegmentError::NonExistent => write!(f, "Segment does not exist"), + SegmentError::Mismatch => write!(f, "Root hash mismatch"), + } + } +} + +/// Tuple that defines a segment of a given PMMR +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +pub struct SegmentIdentifier { + /// Height of a segment + pub height: u8, + /// Zero-based index of the segment + pub idx: u64, +} + +impl Readable for SegmentIdentifier { + fn read(reader: &mut R) -> Result { + let height = reader.read_u8()?; + let idx = reader.read_u64()?; + Ok(Self { height, idx }) + } +} + +impl Writeable for SegmentIdentifier { + fn write(&self, writer: &mut W) -> Result<(), Error> { + writer.write_u8(self.height)?; + writer.write_u64(self.idx) + } +} + +/// Segment of a PMMR: unpruned leaves and the necessary data to verify +/// segment membership in the original MMR. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct Segment { + identifier: SegmentIdentifier, + hash_pos: Vec, + hashes: Vec, + leaf_pos: Vec, + leaf_data: Vec, + proof: SegmentProof, +} + +impl Segment { + /// Creates an empty segment + fn empty(identifier: SegmentIdentifier) -> Self { + Segment { + identifier, + hash_pos: Vec::new(), + hashes: Vec::new(), + leaf_pos: Vec::new(), + leaf_data: Vec::new(), + proof: SegmentProof::empty(), + } + } + + /// Maximum number of leaves in a segment, given by `2**height` + fn segment_capacity(&self) -> u64 { + 1 << self.identifier.height + } + + /// Offset (in leaf idx) of first leaf in the segment + fn leaf_offset(&self) -> u64 { + self.identifier.idx * self.segment_capacity() + } + + /// Number of leaves in this segment. Equal to capacity except for the final segment, which can be smaller + fn segment_unpruned_size(&self, last_pos: u64) -> u64 { + min( + self.segment_capacity(), + pmmr::n_leaves(last_pos).saturating_sub(self.leaf_offset()), + ) + } + + /// Whether the segment is full (size == capacity) + fn full_segment(&self, last_pos: u64) -> bool { + self.segment_unpruned_size(last_pos) == self.segment_capacity() + } + + /// Inclusive range of MMR positions for this segment + pub fn segment_pos_range(&self, last_pos: u64) -> (u64, u64) { + let segment_size = self.segment_unpruned_size(last_pos); + let leaf_offset = self.leaf_offset(); + let first = pmmr::insertion_to_pmmr_index(leaf_offset + 1); + let last = if self.full_segment(last_pos) { + pmmr::insertion_to_pmmr_index(leaf_offset + segment_size) + + (self.identifier.height as u64) + } else { + last_pos + }; + + (first, last) + } + + fn get_hash(&self, pos: u64) -> Result { + self.hash_pos + .iter() + .zip(&self.hashes) + .find(|&(&p, _)| p == pos) + .map(|(_, &h)| h) + .ok_or_else(|| SegmentError::MissingHash(pos)) + } + + /// Iterator of all the leaves in the segment + pub fn leaf_iter(&self) -> impl Iterator + '_ { + self.leaf_pos.iter().map(|&p| p).zip(&self.leaf_data) + } + + /// Iterator of all the hashes in the segment + pub fn hash_iter(&self) -> impl Iterator + '_ { + self.hash_pos + .iter() + .zip(&self.hashes) + .map(|(&p, &h)| (p, h)) + } +} + +impl Segment +where + T: Readable + Writeable + Debug, +{ + /// Generate a segment from a PMMR + pub fn from_pmmr( + segment_id: SegmentIdentifier, + pmmr: &ReadonlyPMMR<'_, U, B>, + prunable: bool, + ) -> Result + where + U: PMMRable, + B: Backend, + { + let mut segment = Segment::empty(segment_id); + + let last_pos = pmmr.unpruned_size(); + if segment.segment_unpruned_size(last_pos) == 0 { + return Err(SegmentError::NonExistent); + } + + // Fill leaf data and hashes + let (segment_first_pos, segment_last_pos) = segment.segment_pos_range(last_pos); + for pos in segment_first_pos..=segment_last_pos { + if pmmr::is_leaf(pos) { + if let Some(data) = pmmr.get_data_from_file(pos) { + segment.leaf_data.push(data); + segment.leaf_pos.push(pos); + continue; + } else if !prunable { + return Err(SegmentError::MissingLeaf(pos)); + } + } + // TODO: optimize, no need to send every intermediary hash + if prunable { + if let Some(hash) = pmmr.get_from_file(pos) { + segment.hashes.push(hash); + segment.hash_pos.push(pos); + } + } + } + + let mut start_pos = None; + // Fully pruned segment: only include a single hash, the first unpruned parent + if segment.leaf_data.is_empty() && segment.hashes.is_empty() { + let family_branch = pmmr::family_branch(segment_last_pos, last_pos); + for (pos, _) in family_branch { + if let Some(hash) = pmmr.get_from_file(pos) { + segment.hashes.push(hash); + segment.hash_pos.push(pos); + start_pos = Some(pos); + break; + } + } + } + + // Segment merkle proof + segment.proof = SegmentProof::generate( + pmmr, + last_pos, + segment_first_pos, + segment_last_pos, + start_pos, + )?; + + Ok(segment) + } +} + +impl Segment +where + T: PMMRIndexHashable, +{ + /// Calculate root hash of this segment + /// Returns `None` iff the segment is full and completely pruned + pub fn root( + &self, + last_pos: u64, + bitmap: Option<&Bitmap>, + ) -> Result, SegmentError> { + let (segment_first_pos, segment_last_pos) = self.segment_pos_range(last_pos); + let mut hashes = Vec::>::with_capacity(2 * (self.identifier.height as usize)); + let mut leaves = self.leaf_pos.iter().zip(&self.leaf_data); + for pos in segment_first_pos..=segment_last_pos { + let height = pmmr::bintree_postorder_height(pos); + let hash = if height == 0 { + // Leaf + if bitmap + .map(|b| { + let idx_1 = pmmr::n_leaves(pos) - 1; + let idx_2 = if pmmr::is_left_sibling(pos) { + idx_1 + 1 + } else { + idx_1 - 1 + }; + b.contains(idx_1 as u32) || b.contains(idx_2 as u32) || pos == last_pos + }) + .unwrap_or(true) + { + // We require the data of this leaf if either the mmr is not prunable or if + // the bitmap indicates it (or its sibling) should be here. + // Edge case: if the final segment has an uneven number of leaves, we + // require the last leaf to be present regardless of the status in the bitmap. + // TODO: possibly remove requirement on the sibling when we no longer support + // syncing through the txhashset.zip method. + let data = leaves + .find(|&(&p, _)| p == pos) + .map(|(_, l)| l) + .ok_or_else(|| SegmentError::MissingLeaf(pos))?; + Some(data.hash_with_index(pos - 1)) + } else { + None + } + } else { + let left_child_pos = pos - (1 << height); + let right_child_pos = pos - 1; + + let right_child = hashes.pop().unwrap(); + let left_child = hashes.pop().unwrap(); + + if bitmap.is_some() { + // Prunable MMR + match (left_child, right_child) { + (None, None) => None, + (Some(l), Some(r)) => Some((l, r).hash_with_index(pos - 1)), + (None, Some(r)) => { + let l = self.get_hash(left_child_pos)?; + Some((l, r).hash_with_index(pos - 1)) + } + (Some(l), None) => { + let r = self.get_hash(right_child_pos)?; + Some((l, r).hash_with_index(pos - 1)) + } + } + } else { + // Non-prunable MMR: require both children + Some( + ( + left_child.ok_or_else(|| SegmentError::MissingHash(left_child_pos))?, + right_child + .ok_or_else(|| SegmentError::MissingHash(right_child_pos))?, + ) + .hash_with_index(pos - 1), + ) + } + }; + hashes.push(hash); + } + + if self.full_segment(last_pos) { + // Full segment: last position of segment is subtree root + Ok(hashes.pop().unwrap()) + } else { + // Not full (only final segment): peaks in segment, bag them together + let peaks = pmmr::peaks(last_pos) + .into_iter() + .filter(|&pos| pos >= segment_first_pos && pos <= segment_last_pos) + .rev(); + let mut hash = None; + for pos in peaks { + let mut lhash = hashes.pop().ok_or_else(|| SegmentError::MissingHash(pos))?; + if lhash.is_none() && bitmap.is_some() { + // If this entire peak is pruned, load it from the segment hashes + lhash = Some(self.get_hash(pos)?); + } + let lhash = lhash.ok_or_else(|| SegmentError::MissingHash(pos))?; + + hash = match hash { + None => Some(lhash), + Some(rhash) => Some((lhash, rhash).hash_with_index(last_pos)), + }; + } + Ok(Some(hash.unwrap())) + } + } + + /// Get the first unpruned parent hash of this segment + pub fn first_unpruned_parent( + &self, + last_pos: u64, + bitmap: Option<&Bitmap>, + ) -> Result<(Hash, u64), SegmentError> { + let root = self.root(last_pos, bitmap)?; + let (_, last) = self.segment_pos_range(last_pos); + if let Some(root) = root { + return Ok((root, last)); + } + let bitmap = bitmap.unwrap(); + let n_leaves = pmmr::n_leaves(last_pos); + + let mut cardinality = 0; + let mut pos = last; + let mut hash = Err(SegmentError::MissingHash(last)); + let mut family_branch = pmmr::family_branch(last, last_pos).into_iter(); + while cardinality == 0 { + hash = self.get_hash(pos).map(|h| (h, pos)); + if hash.is_ok() { + // Return early in case a lower level hash is already present + // This can occur if both child trees are pruned but compaction hasn't run yet + return hash; + } + + if let Some((p, _)) = family_branch.next() { + pos = p; + let range = (pmmr::n_leaves(pmmr::bintree_leftmost(p)) - 1) + ..min(pmmr::n_leaves(pmmr::bintree_rightmost(p)), n_leaves); + cardinality = bitmap.range_cardinality(range); + } else { + break; + } + } + hash + } + + /// Check validity of the segment by calculating its root and validating the merkle proof + pub fn validate( + &self, + last_pos: u64, + bitmap: Option<&Bitmap>, + mmr_root: Hash, + ) -> Result<(), SegmentError> { + let (first, last) = self.segment_pos_range(last_pos); + let (segment_root, segment_unpruned_pos) = self.first_unpruned_parent(last_pos, bitmap)?; + self.proof.validate( + last_pos, + mmr_root, + first, + last, + segment_root, + segment_unpruned_pos, + ) + } + + /// Check validity of the segment by calculating its root and validating the merkle proof + /// This function assumes a final hashing step together with `other_root` + pub fn validate_with( + &self, + last_pos: u64, + bitmap: Option<&Bitmap>, + mmr_root: Hash, + other_root: Hash, + other_is_left: bool, + ) -> Result<(), SegmentError> { + let (first, last) = self.segment_pos_range(last_pos); + let (segment_root, segment_unpruned_pos) = self.first_unpruned_parent(last_pos, bitmap)?; + self.proof.validate_with( + last_pos, + mmr_root, + first, + last, + segment_root, + segment_unpruned_pos, + other_root, + other_is_left, + ) + } +} + +impl Readable for Segment { + fn read(reader: &mut R) -> Result { + let identifier = Readable::read(reader)?; + + let mut last_pos = 0; + let n_hashes = reader.read_u64()? as usize; + let mut hash_pos = Vec::with_capacity(n_hashes); + for _ in 0..n_hashes { + let pos = reader.read_u64()?; + if pos <= last_pos { + return Err(Error::SortError); + } + last_pos = pos; + hash_pos.push(pos); + } + + let mut hashes = Vec::::with_capacity(n_hashes); + for _ in 0..n_hashes { + hashes.push(Readable::read(reader)?); + } + + let n_leaves = reader.read_u64()? as usize; + let mut leaf_pos = Vec::with_capacity(n_leaves); + last_pos = 0; + for _ in 0..n_leaves { + let pos = reader.read_u64()?; + if pos <= last_pos { + return Err(Error::SortError); + } + last_pos = pos; + leaf_pos.push(pos); + } + + let mut leaf_data = Vec::::with_capacity(n_leaves); + for _ in 0..n_leaves { + leaf_data.push(Readable::read(reader)?); + } + + let proof = Readable::read(reader)?; + + Ok(Self { + identifier, + hash_pos, + hashes, + leaf_pos, + leaf_data, + proof, + }) + } +} + +impl Writeable for Segment { + fn write(&self, writer: &mut W) -> Result<(), Error> { + Writeable::write(&self.identifier, writer)?; + writer.write_u64(self.hashes.len() as u64)?; + for &pos in &self.hash_pos { + writer.write_u64(pos)?; + } + for hash in &self.hashes { + Writeable::write(hash, writer)?; + } + writer.write_u64(self.leaf_data.len() as u64)?; + for &pos in &self.leaf_pos { + writer.write_u64(pos)?; + } + for data in &self.leaf_data { + Writeable::write(data, writer)?; + } + Writeable::write(&self.proof, writer)?; + Ok(()) + } +} + +/// Merkle proof of a segment +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SegmentProof { + hashes: Vec, +} + +impl SegmentProof { + fn empty() -> Self { + Self { hashes: Vec::new() } + } + + fn generate( + pmmr: &ReadonlyPMMR<'_, U, B>, + last_pos: u64, + segment_first_pos: u64, + segment_last_pos: u64, + start_pos: Option, + ) -> Result + where + U: PMMRable, + B: Backend, + { + let family_branch = pmmr::family_branch(segment_last_pos, last_pos); + + // 1. siblings along the path from the subtree root to the peak + let hashes: Result, _> = family_branch + .iter() + .filter(|&&(p, _)| start_pos.map(|s| p > s).unwrap_or(true)) + .map(|&(_, s)| pmmr.get_hash(s).ok_or_else(|| SegmentError::MissingHash(s))) + .collect(); + let mut proof = Self { hashes: hashes? }; + + // 2. bagged peaks to the right + let peak_pos = family_branch + .last() + .map(|&(p, _)| p) + .unwrap_or(segment_last_pos); + if let Some(h) = pmmr.bag_the_rhs(peak_pos) { + proof.hashes.push(h); + } + + // 3. peaks to the left + let peaks: Result, _> = pmmr::peaks(last_pos) + .into_iter() + .filter(|&x| x < segment_first_pos) + .rev() + .map(|p| pmmr.get_hash(p).ok_or_else(|| SegmentError::MissingHash(p))) + .collect(); + proof.hashes.extend(peaks?); + + Ok(proof) + } + + /// Reconstruct PMMR root using this proof + pub fn reconstruct_root( + &self, + last_pos: u64, + segment_first_pos: u64, + segment_last_pos: u64, + segment_root: Hash, + segment_unpruned_pos: u64, + ) -> Result { + let mut iter = self.hashes.iter(); + let family_branch = pmmr::family_branch(segment_last_pos, last_pos); + + // 1. siblings along the path from the subtree root to the peak + let mut root = segment_root; + for &(p, s) in family_branch + .iter() + .filter(|&&(p, _)| p > segment_unpruned_pos) + { + let sibling_hash = iter.next().ok_or_else(|| SegmentError::MissingHash(s))?; + root = if pmmr::is_left_sibling(s) { + (sibling_hash, root).hash_with_index(p - 1) + } else { + (root, sibling_hash).hash_with_index(p - 1) + }; + } + + // 2. bagged peaks to the right + let peak_pos = family_branch + .last() + .map(|&(p, _)| p) + .unwrap_or(segment_last_pos); + + let rhs = pmmr::peaks(last_pos) + .into_iter() + .filter(|&x| x > peak_pos) + .next(); + + if let Some(pos) = rhs { + root = ( + root, + iter.next().ok_or_else(|| SegmentError::MissingHash(pos))?, + ) + .hash_with_index(last_pos) + } + + // 3. peaks to the left + let peaks = pmmr::peaks(last_pos) + .into_iter() + .filter(|&x| x < segment_first_pos) + .rev(); + for pos in peaks { + root = ( + iter.next().ok_or_else(|| SegmentError::MissingHash(pos))?, + root, + ) + .hash_with_index(last_pos); + } + + Ok(root) + } + + /// Check validity of the proof by equating the reconstructed root with the actual root + pub fn validate( + &self, + last_pos: u64, + mmr_root: Hash, + segment_first_pos: u64, + segment_last_pos: u64, + segment_root: Hash, + segment_unpruned_pos: u64, + ) -> Result<(), SegmentError> { + let root = self.reconstruct_root( + last_pos, + segment_first_pos, + segment_last_pos, + segment_root, + segment_unpruned_pos, + )?; + if root == mmr_root { + Ok(()) + } else { + Err(SegmentError::Mismatch) + } + } + + /// Check validity of the proof by equating the reconstructed root with the actual root + /// This function assumes a final hashing step together with `other_root` + pub fn validate_with( + &self, + last_pos: u64, + mmr_root: Hash, + segment_first_pos: u64, + segment_last_pos: u64, + segment_root: Hash, + segment_unpruned_pos: u64, + other_root: Hash, + other_is_left: bool, + ) -> Result<(), SegmentError> { + let root = self.reconstruct_root( + last_pos, + segment_first_pos, + segment_last_pos, + segment_root, + segment_unpruned_pos, + )?; + let root = if other_is_left { + (other_root, root).hash_with_index(last_pos) + } else { + (root, other_root).hash_with_index(last_pos) + }; + if root == mmr_root { + Ok(()) + } else { + Err(SegmentError::Mismatch) + } + } +} + +impl Readable for SegmentProof { + fn read(reader: &mut R) -> Result { + let n_hashes = reader.read_u64()? as usize; + let mut hashes = Vec::with_capacity(n_hashes); + for _ in 0..n_hashes { + let hash: Hash = Readable::read(reader)?; + hashes.push(hash); + } + Ok(Self { hashes }) + } +} + +impl Writeable for SegmentProof { + fn write(&self, writer: &mut W) -> Result<(), Error> { + writer.write_u64(self.hashes.len() as u64)?; + for hash in &self.hashes { + Writeable::write(hash, writer)?; + } + Ok(()) + } +} diff --git a/core/tests/segment.rs b/core/tests/segment.rs new file mode 100644 index 000000000..c70a73e6f --- /dev/null +++ b/core/tests/segment.rs @@ -0,0 +1,61 @@ +// Copyright 2020 The Grin Developers +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +mod common; + +use self::core::core::pmmr; +use self::core::core::{Segment, SegmentIdentifier}; +use common::TestElem; +use grin_core as core; +use grin_core::core::pmmr::ReadablePMMR; + +fn test_unprunable_size(height: u8, n_leaves: u32) { + let size = 1u64 << height; + let n_segments = (n_leaves as u64 + size - 1) / size; + + // Build an MMR with n_leaves leaves + let mut ba = pmmr::VecBackend::new(); + let mut mmr = pmmr::PMMR::new(&mut ba); + for i in 0..n_leaves { + mmr.push(&TestElem([i / 7, i / 5, i / 3, i])).unwrap(); + } + let mmr = mmr.readonly_pmmr(); + let last_pos = mmr.unpruned_size(); + let root = mmr.root().unwrap(); + + for idx in 0..n_segments { + let id = SegmentIdentifier { height, idx }; + let segment = Segment::from_pmmr(id, &mmr, false).unwrap(); + println!( + "\n\n>>>>>>> N_LEAVES = {}, LAST_POS = {}, SEGMENT = {}:\n{:#?}", + n_leaves, last_pos, idx, segment + ); + if idx < n_segments - 1 || (n_leaves as u64) % size == 0 { + // Check if the reconstructed subtree root matches with the hash stored in the mmr + let subtree_root = segment.root(last_pos, None).unwrap().unwrap(); + let last = pmmr::insertion_to_pmmr_index((idx + 1) * size) + (height as u64); + assert_eq!(subtree_root, mmr.get_hash(last).unwrap()); + println!(" ROOT OK"); + } + segment.validate(last_pos, None, root).unwrap(); + println!(" PROOF OK"); + } +} + +#[test] +fn unprunable_mmr() { + for i in 1..=64 { + test_unprunable_size(3, i); + } +} diff --git a/store/tests/segment.rs b/store/tests/segment.rs new file mode 100644 index 000000000..042a5365a --- /dev/null +++ b/store/tests/segment.rs @@ -0,0 +1,422 @@ +// Copyright 2020 The Grin Developers +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::core::core::hash::DefaultHashable; +use crate::core::core::pmmr; +use crate::core::core::pmmr::segment::{Segment, SegmentIdentifier}; +use crate::core::core::pmmr::{Backend, ReadablePMMR, ReadonlyPMMR, PMMR}; +use crate::core::ser::{ + BinReader, BinWriter, Error, PMMRable, ProtocolVersion, Readable, Reader, Writeable, Writer, +}; +use crate::store::pmmr::PMMRBackend; +use chrono::Utc; +use croaring::Bitmap; +use grin_core as core; +use grin_store as store; +use std::fs; +use std::io::Cursor; + +#[test] +fn prunable_mmr() { + let t = Utc::now(); + let data_dir = format!( + "./target/tmp/{}.{}-prunable_mmr", + t.timestamp(), + t.timestamp_subsec_nanos() + ); + fs::create_dir_all(&data_dir).unwrap(); + + let n_leaves = 64 + 8 + 4 + 2 + 1; + let mut ba = PMMRBackend::new(&data_dir, true, ProtocolVersion(1), None).unwrap(); + let mut mmr = PMMR::new(&mut ba); + for i in 0..n_leaves { + mmr.push(&TestElem([i / 7, i / 5, i / 3, i])).unwrap(); + } + let last_pos = mmr.unpruned_size(); + let root = mmr.root().unwrap(); + + let mut bitmap = Bitmap::create(); + bitmap.add_range(0..n_leaves as u64); + + let id = SegmentIdentifier { height: 3, idx: 1 }; + + // Validate a segment before any pruning + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + assert_eq!( + segment.root(last_pos, Some(&bitmap)).unwrap().unwrap(), + mmr.get_hash(30).unwrap() + ); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + // Prune a few leaves + let mut mmr = PMMR::at(&mut ba, last_pos); + prune(&mut mmr, &mut bitmap, &[8, 9, 13]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + // Validate + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + assert_eq!( + segment.root(last_pos, Some(&bitmap)).unwrap().unwrap(), + mmr.get_hash(30).unwrap() + ); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + // Prune more + let mut mmr = PMMR::at(&mut ba, last_pos); + prune(&mut mmr, &mut bitmap, &[10, 11]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + // Validate + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + assert_eq!( + segment.root(last_pos, Some(&bitmap)).unwrap().unwrap(), + mmr.get_hash(30).unwrap() + ); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + // Prune all but 1 + let mut mmr = PMMR::at(&mut ba, last_pos); + prune(&mut mmr, &mut bitmap, &[14, 15]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + // Validate + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + assert_eq!( + segment.root(last_pos, Some(&bitmap)).unwrap().unwrap(), + mmr.get_hash(30).unwrap() + ); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + // Prune all + let mut mmr = PMMR::at(&mut ba, last_pos); + prune(&mut mmr, &mut bitmap, &[12]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + assert!(Segment::from_pmmr(id, &mmr, true).is_ok()); + + // Final segment is not full, test it before pruning + let id = SegmentIdentifier { height: 3, idx: 9 }; + + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + // Prune second and third to last leaves (a full peak in the MMR) + let mut mmr = PMMR::at(&mut ba, last_pos); + prune(&mut mmr, &mut bitmap, &[76, 77]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + // Validate + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + // Prune final element + let mut mmr = PMMR::at(&mut ba, last_pos); + prune(&mut mmr, &mut bitmap, &[78]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + // Validate + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + std::mem::drop(ba); + fs::remove_dir_all(&data_dir).unwrap(); +} + +#[test] +fn pruned_segment() { + let t = Utc::now(); + let data_dir = format!( + "./target/tmp/{}.{}-pruned_segment", + t.timestamp(), + t.timestamp_subsec_nanos() + ); + fs::create_dir_all(&data_dir).unwrap(); + + let n_leaves = 16; + let mut ba = PMMRBackend::new(&data_dir, true, ProtocolVersion(1), None).unwrap(); + let mut mmr = PMMR::new(&mut ba); + for i in 0..n_leaves { + mmr.push(&TestElem([i / 7, i / 5, i / 3, i])).unwrap(); + } + let last_pos = mmr.unpruned_size(); + let root = mmr.root().unwrap(); + + let mut bitmap = Bitmap::create(); + bitmap.add_range(0..n_leaves as u64); + + // Prune all leaves of segment 1 + prune(&mut mmr, &mut bitmap, &[4, 5, 6, 7]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + // Validate the empty segment 1 + let id = SegmentIdentifier { height: 2, idx: 1 }; + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + assert_eq!(segment.leaf_iter().count(), 0); + assert_eq!(segment.hash_iter().count(), 1); + assert_eq!( + segment + .first_unpruned_parent(last_pos, Some(&bitmap)) + .unwrap(), + (ba.get_hash(14).unwrap(), 14) + ); + assert!(segment.root(last_pos, Some(&bitmap)).unwrap().is_none()); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + // Prune all leaves of segment 0 + let mut mmr = PMMR::at(&mut ba, last_pos); + prune(&mut mmr, &mut bitmap, &[0, 1, 2, 3]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + // Validate the empty segment 1 again + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + assert_eq!(segment.leaf_iter().count(), 0); + assert_eq!(segment.hash_iter().count(), 1); + // Since both 7 and 14 are now pruned, the first unpruned hash will be at 15 + assert_eq!( + segment + .first_unpruned_parent(last_pos, Some(&bitmap)) + .unwrap(), + (ba.get_hash(15).unwrap(), 15) + ); + assert!(segment.root(last_pos, Some(&bitmap)).unwrap().is_none()); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + // Prune all leaves of segment 2 & 3 + let mut mmr = PMMR::at(&mut ba, last_pos); + prune(&mut mmr, &mut bitmap, &[8, 9, 10, 11, 12, 13, 14, 15]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + // Validate the empty segment 1 again + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + assert_eq!(segment.leaf_iter().count(), 0); + assert_eq!(segment.hash_iter().count(), 1); + // Since both 15 and 30 are now pruned, the first unpruned hash will be at 31: the mmr root + assert_eq!( + segment + .first_unpruned_parent(last_pos, Some(&bitmap)) + .unwrap(), + (root, 31) + ); + assert!(segment.root(last_pos, Some(&bitmap)).unwrap().is_none()); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + let n_leaves = n_leaves + 4 + 2 + 1; + let mut mmr = PMMR::at(&mut ba, last_pos); + for i in 16..n_leaves { + mmr.push(&TestElem([i / 7, i / 5, i / 3, i])).unwrap(); + } + bitmap.add_range(16..n_leaves as u64); + let last_pos = mmr.unpruned_size(); + let root = mmr.root().unwrap(); + + // Prune all leaves of segment 4 + // The root of this segment is a direct peak of the full MMR + prune(&mut mmr, &mut bitmap, &[16, 17, 18, 19]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + // Validate segment 4 + let id = SegmentIdentifier { height: 2, idx: 4 }; + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + assert_eq!(segment.leaf_iter().count(), 0); + assert_eq!(segment.hash_iter().count(), 1); + assert_eq!( + segment + .first_unpruned_parent(last_pos, Some(&bitmap)) + .unwrap(), + (ba.get_hash(38).unwrap(), 38) + ); + assert!(segment.root(last_pos, Some(&bitmap)).unwrap().is_none()); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + // Segment 5 has 2 peaks + let id = SegmentIdentifier { height: 2, idx: 5 }; + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + assert_eq!(segment.leaf_iter().count(), 3); + assert_eq!( + segment + .first_unpruned_parent(last_pos, Some(&bitmap)) + .unwrap() + .1, + segment.segment_pos_range(last_pos).1 + ); + assert!(segment.root(last_pos, Some(&bitmap)).unwrap().is_some()); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + let prev_segment = segment; + + // Prune final leaf (a peak) + let mut mmr = PMMR::at(&mut ba, last_pos); + prune(&mut mmr, &mut bitmap, &[22]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + // Segment 5 should be unchanged + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + assert_eq!(segment, prev_segment); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + // Prune other peak of segment 5 + let mut mmr = PMMR::at(&mut ba, last_pos); + prune(&mut mmr, &mut bitmap, &[20, 21]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + // Validate segment 5 again + let mmr = ReadonlyPMMR::at(&mut ba, last_pos); + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + assert_eq!(segment.leaf_iter().count(), 1); + assert_eq!(segment.hash_iter().count(), 1); + assert_eq!( + segment + .first_unpruned_parent(last_pos, Some(&bitmap)) + .unwrap() + .1, + segment.segment_pos_range(last_pos).1 + ); + assert!(segment.root(last_pos, Some(&bitmap)).unwrap().is_some()); + segment.validate(last_pos, Some(&bitmap), root).unwrap(); + + std::mem::drop(ba); + fs::remove_dir_all(&data_dir).unwrap(); +} + +#[test] +fn ser_round_trip() { + let t = Utc::now(); + let data_dir = format!( + "./target/tmp/{}.{}-segment_ser_round_trip", + t.timestamp(), + t.timestamp_subsec_nanos() + ); + fs::create_dir_all(&data_dir).unwrap(); + + let n_leaves = 32; + let mut ba = PMMRBackend::new(&data_dir, true, ProtocolVersion(1), None).unwrap(); + let mut mmr = pmmr::PMMR::new(&mut ba); + for i in 0..n_leaves { + mmr.push(&TestElem([i / 7, i / 5, i / 3, i])).unwrap(); + } + let mut bitmap = Bitmap::create(); + bitmap.add_range(0..n_leaves as u64); + let last_pos = mmr.unpruned_size(); + + prune(&mut mmr, &mut bitmap, &[0, 1]); + ba.sync().unwrap(); + ba.check_compact(last_pos, &Bitmap::create()).unwrap(); + ba.sync().unwrap(); + + let mmr = ReadonlyPMMR::at(&ba, last_pos); + let id = SegmentIdentifier { height: 3, idx: 0 }; + let segment = Segment::from_pmmr(id, &mmr, true).unwrap(); + + let mut cursor = Cursor::new(Vec::::new()); + let mut writer = BinWriter::new(&mut cursor, ProtocolVersion(1)); + Writeable::write(&segment, &mut writer).unwrap(); + assert_eq!( + cursor.position(), + (9) + (8 + 7 * (8 + 32)) + (8 + 6 * (8 + 16)) + (8 + 2 * 32) + ); + cursor.set_position(0); + + let mut reader = BinReader::new(&mut cursor, ProtocolVersion(1)); + let segment2: Segment = Readable::read(&mut reader).unwrap(); + assert_eq!(segment, segment2); + + std::mem::drop(ba); + fs::remove_dir_all(&data_dir).unwrap(); +} + +fn prune(mmr: &mut PMMR, bitmap: &mut Bitmap, leaf_idxs: &[u64]) +where + T: PMMRable, + B: Backend, +{ + for &leaf_idx in leaf_idxs { + mmr.prune(pmmr::insertion_to_pmmr_index(leaf_idx + 1)) + .unwrap(); + bitmap.remove(leaf_idx as u32); + } +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub struct TestElem(pub [u32; 4]); + +impl DefaultHashable for TestElem {} + +impl PMMRable for TestElem { + type E = Self; + + fn as_elmt(&self) -> Self::E { + *self + } + + fn elmt_size() -> Option { + Some(16) + } +} + +impl Writeable for TestElem { + fn write(&self, writer: &mut W) -> Result<(), Error> { + writer.write_u32(self.0[0])?; + writer.write_u32(self.0[1])?; + writer.write_u32(self.0[2])?; + writer.write_u32(self.0[3]) + } +} + +impl Readable for TestElem { + fn read(reader: &mut R) -> Result { + Ok(TestElem([ + reader.read_u32()?, + reader.read_u32()?, + reader.read_u32()?, + reader.read_u32()?, + ])) + } +}