diff --git a/merk/src/merk/restore.rs b/merk/src/merk/restore.rs index 0f5d6882d..1a169d2b3 100644 --- a/merk/src/merk/restore.rs +++ b/merk/src/merk/restore.rs @@ -599,7 +599,10 @@ impl<'db, S: StorageContext<'db>> Restorer { .ok_or(Error::ChunkRestoringError(ChunkError::InternalError( "tree is None in rewrite_heights", )))?; - let walker = RefWalker::new(&mut tree, self.merk.source()); + let walker = RefWalker::new( + &mut tree, + self.merk.source_without_child_height_validation(), + ); rewrite_child_heights(walker, &mut batch, grove_version)?; diff --git a/merk/src/merk/source.rs b/merk/src/merk/source.rs index 47d0247c1..6d8a368dc 100644 --- a/merk/src/merk/source.rs +++ b/merk/src/merk/source.rs @@ -16,6 +16,15 @@ where MerkSource { storage: &self.storage, tree_type: self.tree_type, + validate_child_heights: true, + } + } + + pub(in crate::merk) fn source_without_child_height_validation(&self) -> MerkSource<'_, S> { + MerkSource { + storage: &self.storage, + tree_type: self.tree_type, + validate_child_heights: false, } } } @@ -25,6 +34,7 @@ where pub struct MerkSource<'s, S> { storage: &'s S, tree_type: TreeType, + validate_child_heights: bool, } impl Clone for MerkSource<'_, S> { @@ -32,6 +42,7 @@ impl Clone for MerkSource<'_, S> { MerkSource { storage: self.storage, tree_type: self.tree_type, + validate_child_heights: self.validate_child_heights, } } } @@ -40,6 +51,14 @@ impl<'db, S> Fetch for MerkSource<'_, S> where S: StorageContext<'db>, { + fn tree_type(&self) -> TreeType { + self.tree_type + } + + fn validate_fetched_link_child_heights(&self) -> bool { + self.validate_child_heights + } + fn fetch( &self, link: &Link, diff --git a/merk/src/tree/mod.rs b/merk/src/tree/mod.rs index c4ebdb6e6..31d92f7f3 100644 --- a/merk/src/tree/mod.rs +++ b/merk/src/tree/mod.rs @@ -1455,6 +1455,66 @@ impl TreeNode { Ok(()).wrap_with_cost(cost) } + /// Validates that a fetched child matches the pruned link metadata. + pub(crate) fn validate_fetched_link( + link: &Link, + tree: &Self, + tree_type: TreeType, + validate_child_heights: bool, + ) -> Result<(), Error> { + if tree.key() != link.key() { + return Err(Error::CorruptedState("fetched link key mismatch")); + } + let expected_child_heights = match link { + Link::Reference { child_heights, .. } + | Link::Uncommitted { child_heights, .. } + | Link::Loaded { child_heights, .. } => *child_heights, + Link::Modified { .. } => { + return Err(Error::CorruptedState( + "cannot validate fetched link against modified link", + )); + } + }; + if validate_child_heights && tree.child_heights() != expected_child_heights { + return Err(Error::CorruptedState("fetched link child heights mismatch")); + } + let aggregate_data = tree.aggregate_data()?; + if aggregate_data != link.aggregate_data() { + return Err(Error::CorruptedState("fetched link aggregate mismatch")); + } + match (tree_type, aggregate_data) { + (TreeType::ProvableCountTree, AggregateData::ProvableCount(_)) + | (TreeType::ProvableCountSumTree, AggregateData::ProvableCountAndSum(..)) + | (TreeType::ProvableSumTree, AggregateData::ProvableSum(_)) + | ( + TreeType::ProvableCountProvableSumTree, + AggregateData::ProvableCountAndProvableSum(..), + ) + | ( + TreeType::NormalTree + | TreeType::SumTree + | TreeType::BigSumTree + | TreeType::CountTree + | TreeType::CountSumTree + | TreeType::CommitmentTree(_) + | TreeType::MmrTree + | TreeType::BulkAppendTree(_) + | TreeType::DenseAppendOnlyFixedSizeTree(_), + _, + ) => {} + _ => { + return Err(Error::CorruptedState( + "fetched link aggregate incompatible with tree type", + )); + } + } + let hash = tree.hash_for_link(tree_type).unwrap(); + if &hash != link.hash() { + return Err(Error::CorruptedState("fetched link hash mismatch")); + } + Ok(()) + } + /// Fetches the child on the given side using the given data source, and /// places it in the child slot (upgrading the link from `Link::Reference` /// to `Link::Loaded`). @@ -1494,7 +1554,15 @@ impl TreeNode { &mut cost, source.fetch(link, value_defined_cost_fn, grove_version) ); - debug_assert_eq!(tree.key(), link.key()); + cost_return_on_error_no_add!( + cost, + Self::validate_fetched_link( + link, + &tree, + source.tree_type(), + source.validate_fetched_link_child_heights(), + ) + ); *self.slot_mut(left) = Some(Link::Loaded { tree, hash: *hash, @@ -1525,10 +1593,15 @@ mod test_provable_count_edge_cases; #[cfg(test)] mod test { - use super::{commit::NoopCommit, hash::NULL_HASH, AggregateData, TreeNode}; + use super::{commit::NoopCommit, hash::NULL_HASH, AggregateData, Link, TreeNode}; use crate::tree::{ tree_feature_type::TreeFeatureType::SummedMerkNode, TreeFeatureType::BasicMerkNode, }; + use crate::tree_type::TreeType; + + fn fetched_child() -> TreeNode { + TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerkNode).unwrap() + } #[test] fn build_tree() { @@ -1670,6 +1743,85 @@ mod test { assert_eq!(tree.child_hash(false), &NULL_HASH); } + #[test] + fn validate_fetched_link_accepts_loaded_and_uncommitted_links() { + let tree = fetched_child(); + let hash = tree.hash_for_link(TreeType::NormalTree).unwrap(); + let child_heights = tree.child_heights(); + let aggregate_data = tree.aggregate_data().unwrap(); + + let uncommitted = Link::Uncommitted { + hash, + child_heights, + tree: fetched_child(), + aggregate_data, + }; + TreeNode::validate_fetched_link(&uncommitted, &tree, TreeType::NormalTree, true) + .expect("uncommitted link should validate"); + + let loaded = Link::Loaded { + hash, + child_heights, + tree: fetched_child(), + aggregate_data, + }; + TreeNode::validate_fetched_link(&loaded, &tree, TreeType::NormalTree, true) + .expect("loaded link should validate"); + } + + #[test] + fn validate_fetched_link_rejects_modified_link() { + let tree = fetched_child(); + let modified = Link::Modified { + pending_writes: 0, + child_heights: tree.child_heights(), + tree: fetched_child(), + }; + + assert!(matches!( + TreeNode::validate_fetched_link(&modified, &tree, TreeType::NormalTree, true), + Err(crate::Error::CorruptedState( + "cannot validate fetched link against modified link" + )) + )); + } + + #[test] + fn validate_fetched_link_rejects_aggregate_mismatch() { + let tree = fetched_child(); + let link = Link::Reference { + hash: tree.hash_for_link(TreeType::NormalTree).unwrap(), + key: tree.key().to_vec(), + child_heights: tree.child_heights(), + aggregate_data: AggregateData::Sum(1), + }; + + assert!(matches!( + TreeNode::validate_fetched_link(&link, &tree, TreeType::NormalTree, true), + Err(crate::Error::CorruptedState( + "fetched link aggregate mismatch" + )) + )); + } + + #[test] + fn validate_fetched_link_rejects_tree_type_aggregate_mismatch() { + let tree = fetched_child(); + let link = Link::Reference { + hash: tree.hash_for_link(TreeType::NormalTree).unwrap(), + key: tree.key().to_vec(), + child_heights: tree.child_heights(), + aggregate_data: tree.aggregate_data().unwrap(), + }; + + assert!(matches!( + TreeNode::validate_fetched_link(&link, &tree, TreeType::ProvableSumTree, true), + Err(crate::Error::CorruptedState( + "fetched link aggregate incompatible with tree type" + )) + )); + } + #[test] fn hash() { let tree = TreeNode::new(vec![0], vec![1], None, BasicMerkNode).unwrap(); diff --git a/merk/src/tree/ops.rs b/merk/src/tree/ops.rs index 7959afd50..7c8ddfb6d 100644 --- a/merk/src/tree/ops.rs +++ b/merk/src/tree/ops.rs @@ -25,7 +25,10 @@ use Op::*; #[cfg(feature = "minimal")] use super::{Fetch, Link, TreeNode, Walker}; #[cfg(feature = "minimal")] -use crate::{error::Error, tree::tree_feature_type::TreeFeatureType, CryptoHash, HASH_LENGTH_U32}; +use crate::{ + error::Error, tree::tree_feature_type::TreeFeatureType, tree_type::TreeType, CryptoHash, + HASH_LENGTH_U32, +}; use crate::{ merk::KeyUpdates, tree::kv::{ValueDefinedCostType, ValueDefinedCostType::SpecializedValueDefinedCost}, @@ -123,6 +126,10 @@ pub struct PanicSource {} #[cfg(feature = "minimal")] impl Fetch for PanicSource { + fn tree_type(&self) -> TreeType { + TreeType::NormalTree + } + fn fetch( &self, _link: &Link, @@ -1059,6 +1066,11 @@ mod test { tree::{tree_feature_type::TreeFeatureType::BasicMerkNode, *}, }; + #[test] + fn panic_source_reports_normal_tree_type() { + assert_eq!((PanicSource {}).tree_type(), TreeType::NormalTree); + } + #[test] fn simple_insert() { let grove_version = GroveVersion::latest(); diff --git a/merk/src/tree/walk/fetch.rs b/merk/src/tree/walk/fetch.rs index 5b43d9fe4..9b5442c80 100644 --- a/merk/src/tree/walk/fetch.rs +++ b/merk/src/tree/walk/fetch.rs @@ -10,12 +10,23 @@ use super::super::{Link, TreeNode}; use crate::error::Error; #[cfg(feature = "minimal")] use crate::tree::kv::ValueDefinedCostType; +#[cfg(feature = "minimal")] +use crate::tree_type::TreeType; #[cfg(feature = "minimal")] /// A source of data to be used by the tree when encountering a pruned node. /// This typically means fetching the tree node from a backing store by its key, /// but could also implement an in-memory cache for example. pub trait Fetch { + /// The tree type used to validate fetched links. + fn tree_type(&self) -> TreeType; + + /// Whether fetched children must match the child-height metadata stored on + /// the pruned link. + fn validate_fetched_link_child_heights(&self) -> bool { + true + } + /// Called when the tree needs to fetch a node with the given `Link`. The /// `link` value will always be a `Link::Reference` variant. fn fetch( diff --git a/merk/src/tree/walk/mod.rs b/merk/src/tree/walk/mod.rs index b7f2788b4..539833b32 100644 --- a/merk/src/tree/walk/mod.rs +++ b/merk/src/tree/walk/mod.rs @@ -18,7 +18,7 @@ use grovedb_version::version::GroveVersion; pub use ref_walker::RefWalker; #[cfg(feature = "minimal")] -use super::{Link, TreeNode}; +use super::TreeNode; use crate::tree::kv::ValueDefinedCostType; #[cfg(feature = "minimal")] use crate::{owner::Owner, tree::tree_feature_type::TreeFeatureType, CryptoHash, Error}; @@ -77,21 +77,27 @@ where } } } else { - let link = self.tree.slot_mut(left).take(); - match link { - Some(Link::Reference { .. }) => (), - _ => { - return Err(Error::CorruptedState( - "expected Link::Reference but found other variant or None", - )) - .wrap_with_cost(cost); - } - } - cost_return_on_error!( + let link = cost_return_on_error_no_add!( + cost, + self.tree.slot_mut(left).take().ok_or(Error::CorruptedState( + "expected Link::Reference but found other variant or None" + )) + ); + let child = cost_return_on_error!( &mut cost, self.source - .fetch(&link.unwrap(), value_defined_cost_fn, grove_version) - ) + .fetch(&link, value_defined_cost_fn, grove_version) + ); + cost_return_on_error_no_add!( + cost, + TreeNode::validate_fetched_link( + &link, + &child, + self.source.tree_type(), + self.source.validate_fetched_link_child_heights(), + ) + ); + child }; let child = self.wrap(child); @@ -421,12 +427,20 @@ mod test { use grovedb_version::version::GroveVersion; use super::{super::NoopCommit, *}; - use crate::tree::{AggregateData, TreeFeatureType::BasicMerkNode, TreeNode}; + use crate::{ + tree::{TreeFeatureType::BasicMerkNode, TreeNode}, + tree_type::TreeType, + Link, + }; #[derive(Clone)] struct MockSource {} impl Fetch for MockSource { + fn tree_type(&self) -> TreeType { + TreeType::NormalTree + } + fn fetch( &self, link: &Link, @@ -435,7 +449,7 @@ mod test { >, _grove_version: &GroveVersion, ) -> CostResult { - TreeNode::new(link.key().to_vec(), b"foo".to_vec(), None, BasicMerkNode).map(Ok) + TreeNode::new(link.key().to_vec(), b"bar".to_vec(), None, BasicMerkNode).map(Ok) } } @@ -503,15 +517,16 @@ mod test { #[test] fn walk_pruned() { let grove_version = GroveVersion::latest(); + let child = TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerkNode).unwrap(); let tree = TreeNode::from_fields( b"test".to_vec(), b"abc".to_vec(), Default::default(), Some(Link::Reference { - hash: Default::default(), + hash: child.hash_for_link(TreeType::NormalTree).unwrap(), key: b"foo".to_vec(), - child_heights: (0, 0), - aggregate_data: AggregateData::NoAggregateData, + child_heights: child.child_heights(), + aggregate_data: child.aggregate_data().unwrap(), }), None, BasicMerkNode, @@ -536,6 +551,141 @@ mod test { assert!(walker.into_inner().child(true).is_none()); } + #[derive(Clone)] + struct BadKeySource; + + impl Fetch for BadKeySource { + fn tree_type(&self) -> TreeType { + TreeType::NormalTree + } + + fn fetch( + &self, + _link: &Link, + _value_defined_cost_fn: Option< + &impl Fn(&[u8], &GroveVersion) -> Option, + >, + _grove_version: &GroveVersion, + ) -> CostResult { + TreeNode::new(b"wrong".to_vec(), b"bar".to_vec(), None, BasicMerkNode).map(Ok) + } + } + + #[derive(Clone)] + struct BadHashSource; + + impl Fetch for BadHashSource { + fn tree_type(&self) -> TreeType { + TreeType::NormalTree + } + + fn fetch( + &self, + link: &Link, + _value_defined_cost_fn: Option< + &impl Fn(&[u8], &GroveVersion) -> Option, + >, + _grove_version: &GroveVersion, + ) -> CostResult { + TreeNode::new(link.key().to_vec(), b"wrong".to_vec(), None, BasicMerkNode).map(Ok) + } + } + + fn tree_with_pruned_child() -> TreeNode { + tree_with_pruned_child_heights(None) + } + + fn tree_with_pruned_child_heights(child_heights: Option<(u8, u8)>) -> TreeNode { + let child = TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerkNode).unwrap(); + TreeNode::from_fields( + b"test".to_vec(), + b"abc".to_vec(), + Default::default(), + Some(Link::Reference { + hash: child.hash_for_link(TreeType::NormalTree).unwrap(), + key: b"foo".to_vec(), + child_heights: child_heights.unwrap_or_else(|| child.child_heights()), + aggregate_data: child.aggregate_data().unwrap(), + }), + None, + BasicMerkNode, + ) + .unwrap() + } + + #[test] + fn walk_pruned_rejects_wrong_fetched_child_heights() { + let grove_version = GroveVersion::latest(); + let walker = Walker::new(tree_with_pruned_child_heights(Some((1, 0))), MockSource {}); + + let result = walker + .detach_expect( + true, + None::<&fn(&[u8], &GroveVersion) -> Option>, + grove_version, + ) + .unwrap(); + assert!(matches!( + result, + Err(Error::CorruptedState("fetched link child heights mismatch")) + )); + } + + #[test] + fn walk_pruned_rejects_wrong_fetched_key() { + let grove_version = GroveVersion::latest(); + let walker = Walker::new(tree_with_pruned_child(), BadKeySource); + + let result = walker + .detach_expect( + true, + None::<&fn(&[u8], &GroveVersion) -> Option>, + grove_version, + ) + .unwrap(); + assert!(matches!( + result, + Err(Error::CorruptedState("fetched link key mismatch")) + )); + } + + #[test] + fn walk_pruned_rejects_wrong_fetched_hash() { + let grove_version = GroveVersion::latest(); + let walker = Walker::new(tree_with_pruned_child(), BadHashSource); + + let result = walker + .detach_expect( + true, + None::<&fn(&[u8], &GroveVersion) -> Option>, + grove_version, + ) + .unwrap(); + assert!(matches!( + result, + Err(Error::CorruptedState("fetched link hash mismatch")) + )); + } + + #[test] + fn ref_walker_rejects_wrong_fetched_hash() { + let grove_version = GroveVersion::latest(); + let mut tree = tree_with_pruned_child(); + let mut walker = RefWalker::new(&mut tree, BadHashSource); + + let result = walker + .walk( + true, + None::<&fn(&[u8], &GroveVersion) -> Option>, + grove_version, + ) + .unwrap(); + assert!(matches!( + result, + Err(Error::CorruptedState("fetched link hash mismatch")) + )); + } + #[test] fn walk_none() { let grove_version = GroveVersion::latest();