fix(merk): validate fetched links during lazy load

merk/src/merk/restore.rs

@@ -599,7 +599,10 @@ impl<'db, S: StorageContext<'db>> Restorer<S> {
                 .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)?;
 

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,13 +34,15 @@ where
 pub struct MerkSource<'s, S> {
     storage: &'s S,
     tree_type: TreeType,
+    validate_child_heights: bool,
 }
 
 impl<S> Clone for MerkSource<'_, S> {
     fn clone(&self) -> Self {
         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,

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();

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();

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(