test(crypto): CRP-2001 add multi leaf witnesses to witness equality fuzzer

Author: altkdf

rs/canonical_state/src/hash_tree.rs

@@ -574,6 +574,11 @@ impl HashTree {
                         );
                     }
                 }
+
+                // Empty subtree.
+                LabeledTree::SubTree(children) if children.is_empty() => B::make_empty(),
+
+                // Non-empty subtree.
                 LabeledTree::SubTree(children) => children
                     .iter()
                     .map(|(l, t)| child_witness::<B>(ht, parent, pos, l, t))

rs/canonical_state/src/hash_tree/test.rs

@@ -89,30 +89,31 @@ fn build_witness_gen(t: &LabeledTree<Vec<u8>>) -> WitnessGeneratorImpl {
     builder.witness_generator().unwrap()
 }
 
-fn enumerate_leaves(t: &LabeledTree<Vec<u8>>, mut f: impl FnMut(LabeledTree<Vec<u8>>)) {
+fn enumerate_leaves_and_empty_subtrees(
+    t: &LabeledTree<Vec<u8>>,
+    mut f: impl FnMut(LabeledTree<Vec<u8>>),
+) {
     fn go<'a>(
         t: &'a LabeledTree<Vec<u8>>,
         path: &mut Vec<&'a Label>,
         f: &mut impl FnMut(LabeledTree<Vec<u8>>),
     ) {
         match t {
-            LabeledTree::Leaf(_) => {
-                let mut subtree = t.clone();
-                #[allow(clippy::unnecessary_to_owned)]
-                for label in path.iter().rev().cloned() {
-                    subtree = LabeledTree::SubTree(flatmap! {
-                        label.clone() => subtree,
-                    });
-                }
-                f(subtree)
-            }
-            LabeledTree::SubTree(children) => {
+            LabeledTree::SubTree(children) if !children.is_empty() => {
                 for (k, v) in children.iter() {
                     path.push(k);
                     go(v, path, f);
                     path.pop();
                 }
             }
+            LabeledTree::Leaf(_) | LabeledTree::SubTree(_) => {
+                let subtree = path.iter().rev().fold(t.clone(), |acc, &label| {
+                    LabeledTree::SubTree(flatmap! {
+                        label.clone() => acc,
+                    })
+                });
+                f(subtree)
+            }
         }
     }
     let mut path = vec![];
@@ -123,32 +124,42 @@ fn assert_same_witness(ht: &HashTree, wg: &WitnessGeneratorImpl, data: &LabeledT
     let ht_witness = ht.witness::<Witness>(data);
     let wg_witness = wg.witness(data).expect("failed to construct a witness");
 
-    assert_eq!(
-        recompute_digest(data, &wg_witness).unwrap(),
-        recompute_digest(data, &ht_witness).unwrap()
-    );
     assert_eq!(
         wg_witness, ht_witness,
         "labeled tree: {:?}, hash_tree: {:?}",
         data, ht
-    )
+    );
+
+    assert_eq!(
+        recompute_digest(data, &wg_witness).unwrap(),
+        recompute_digest(data, &ht_witness).unwrap()
+    );
 }
 
-/// Check that for each leaf, the witness looks the same as with the
-/// old way of generating witnesses
-/// Also check that the new and old way of computing hash trees are equivalent
+/// Check that for each leaf or empty subtree, and the tree as a whole, the
+/// witness looks the same as with the old way of generating witnesses.
+///
+/// Also check that the new and old way of computing hash trees are equivalent.
 fn test_tree(t: &LabeledTree<Vec<u8>>) {
     let hash_tree = hash_lazy_tree(&as_lazy(t));
     let witness_gen = build_witness_gen(t);
-    enumerate_leaves(t, |subtree| {
+    enumerate_leaves_and_empty_subtrees(t, |subtree| {
         assert_same_witness(&hash_tree, &witness_gen, &subtree);
     });
 
-    let crypto_tree = crypto_hash_lazy_tree(&as_lazy(t));
+    assert_same_witness(&hash_tree, &witness_gen, t);
 
+    let crypto_tree = crypto_hash_lazy_tree(&as_lazy(t));
     assert_eq!(hash_tree, crypto_tree);
 }
 
+#[test]
+fn test_empty_subtree() {
+    let t = LabeledTree::SubTree(flatmap! {});
+
+    test_tree(&t);
+}
+
 #[test]
 fn test_one_level_tree() {
     let t = LabeledTree::SubTree(flatmap! {
@@ -272,7 +283,7 @@ fn test_non_existence_proof() {
 
 proptest! {
     #[test]
-    fn same_witness_on_all_leaves(t in arbitrary_labeled_tree()) {
+    fn same_witness(t in arbitrary_labeled_tree()) {
         test_tree(&t);
     }
 }

rs/crypto/test_utils/reproducible_rng/src/lib.rs

@@ -1,6 +1,9 @@
 use rand::{CryptoRng, Error, Rng, RngCore, SeedableRng};
 use rand_chacha::ChaCha20Rng;
 
+/// Byte length of the seed type used in [`ReproducibleRng`].
+pub const SEED_LEN: usize = 32;
+
 /// Provides a seeded RNG, where the randomly chosen seed is printed on standard output.
 pub fn reproducible_rng() -> ReproducibleRng {
     ReproducibleRng::new()
@@ -14,7 +17,7 @@ pub fn reproducible_rng() -> ReproducibleRng {
 /// (See [impl trait type](https://doc.rust-lang.org/reference/types/impl-trait.html)).
 pub struct ReproducibleRng {
     rng: ChaCha20Rng,
-    seed: [u8; 32],
+    seed: [u8; SEED_LEN],
 }
 
 impl ReproducibleRng {

rs/crypto/tree_hash/fuzz/check_witness_equality_utils/src/lib.rs

@@ -1,28 +1,52 @@
 use ic_canonical_state::hash_tree::{crypto_hash_lazy_tree, hash_lazy_tree, HashTree};
 use ic_canonical_state::lazy_tree::{LazyFork, LazyTree};
-use ic_crypto_test_utils_reproducible_rng::ReproducibleRng;
+use ic_crypto_test_utils_reproducible_rng::{ReproducibleRng, SEED_LEN};
+use ic_crypto_tree_hash::test_utils::{
+    merge_path_into_labeled_tree, partial_trees_to_leaves_and_empty_subtrees,
+};
 use ic_crypto_tree_hash::{
     flatmap, FlatMap, HashTreeBuilder, HashTreeBuilderImpl, Label, LabeledTree, Witness,
     WitnessGenerator, WitnessGeneratorImpl,
 };
-use rand::{Rng, SeedableRng};
+use rand::{CryptoRng, Rng, SeedableRng};
 use std::sync::Arc;
 
 #[cfg(test)]
 mod tests;
 
 /// Check that for each leaf, the witness looks the same for both implementations
 /// Also check that the new and old way of computing hash trees are equivalent
-pub fn test_tree(t: &LabeledTree<Vec<u8>>) {
-    let hash_tree = hash_lazy_tree(&as_lazy(t));
-    let witness_gen = build_witness_gen(t);
-    // TODO(CRP-2001): add multi-leaf witnesses
-    enumerate_leaves(t, |subtree| {
-        assert_same_witness(&hash_tree, &witness_gen, &subtree);
-    });
+pub fn test_tree<R: Rng + CryptoRng>(full_tree: &LabeledTree<Vec<u8>>, rng: &mut R) {
+    let hash_tree = hash_lazy_tree(&as_lazy(full_tree));
+    let witness_gen = build_witness_gen(full_tree);
+
+    let paths = partial_trees_to_leaves_and_empty_subtrees(full_tree);
+
+    // prune each path (1 node in each level) from `full_tree`
+    for path in paths.iter() {
+        assert_same_witness(&hash_tree, &witness_gen, path);
+    }
+
+    // prune randomly combined paths
+    const MAX_COMBINED_PATHS: usize = 10;
+    for num_leaves_and_empty_subtrees in 2..MAX_COMBINED_PATHS.min(paths.len()) {
+        let mut indices =
+            rand::seq::index::sample(rng, paths.len(), num_leaves_and_empty_subtrees).into_vec();
+        indices.sort_unstable();
 
-    let crypto_tree = crypto_hash_lazy_tree(&as_lazy(t));
+        let mut partial_tree = paths[indices[0]].clone();
 
+        for index in indices[1..].iter() {
+            merge_path_into_labeled_tree(&mut partial_tree, &paths[*index]);
+        }
+        assert_same_witness(&hash_tree, &witness_gen, &partial_tree);
+    }
+
+    // prune the full tree
+    assert_same_witness(&hash_tree, &witness_gen, full_tree);
+
+    // create a hash tree for the full tree
+    let crypto_tree = crypto_hash_lazy_tree(&as_lazy(full_tree));
     assert_eq!(hash_tree, crypto_tree);
 }
 
@@ -32,38 +56,8 @@ fn assert_same_witness(ht: &HashTree, wg: &WitnessGeneratorImpl, data: &LabeledT
 
     assert_eq!(
         wg_witness, ht_witness,
-        "labeled tree: {data:?}, hash_tree: {ht:?}",
-    )
-}
-
-fn enumerate_leaves(t: &LabeledTree<Vec<u8>>, mut f: impl FnMut(LabeledTree<Vec<u8>>)) {
-    fn go<'a>(
-        t: &'a LabeledTree<Vec<u8>>,
-        path: &mut Vec<&'a Label>,
-        f: &mut impl FnMut(LabeledTree<Vec<u8>>),
-    ) {
-        match t {
-            LabeledTree::Leaf(_) => {
-                let mut subtree = t.clone();
-                #[allow(clippy::unnecessary_to_owned)]
-                for label in path.iter().rev().cloned() {
-                    subtree = LabeledTree::SubTree(flatmap! {
-                        label.clone() => subtree,
-                    });
-                }
-                f(subtree)
-            }
-            LabeledTree::SubTree(children) => {
-                for (k, v) in children.iter() {
-                    path.push(k);
-                    go(v, path, f);
-                    path.pop();
-                }
-            }
-        }
-    }
-    let mut path = vec![];
-    go(t, &mut path, &mut f)
+        "labeled tree: {data:?}, hash_tree: {ht:?}, wg: {wg:?}",
+    );
 }
 
 fn as_lazy(t: &LabeledTree<Vec<u8>>) -> LazyTree<'_> {
@@ -117,7 +111,6 @@ fn build_witness_gen(t: &LabeledTree<Vec<u8>>) -> WitnessGeneratorImpl {
 }
 
 pub fn rng_from_u32(seed: u32) -> ReproducibleRng {
-    const SEED_LEN: usize = std::mem::size_of::<<ReproducibleRng as rand::SeedableRng>::Seed>();
     let seed_bytes: Vec<u8> = seed
         .to_le_bytes()
         .into_iter()

rs/crypto/tree_hash/fuzz/fuzz_targets/check_witness_equality.rs

@@ -5,26 +5,42 @@
 // let $data: &mut [u8] = unsafe { std::slice::from_raw_parts_mut($data, len) };"
 #![allow(clippy::not_unsafe_ptr_arg_deref)]
 
+use ic_crypto_test_utils_reproducible_rng::{ReproducibleRng, SEED_LEN};
 use ic_crypto_tree_hash::{flatmap, LabeledTree};
 use ic_crypto_tree_hash_fuzz_check_witness_equality_utils::*;
 use ic_protobuf::messaging::xnet::v1::LabeledTree as ProtobufLabeledTree;
 use ic_protobuf::proxy::ProtoProxy;
 use libfuzzer_sys::fuzz_target;
-use rand::Rng;
+use rand::{Rng, RngCore, SeedableRng};
 
 fuzz_target!(|data: &[u8]| {
-    if let Ok(tree) = ProtobufLabeledTree::proxy_decode(data) {
-        test_tree(&tree);
+    if data.len() < SEED_LEN {
+        return;
+    }
+    if let Ok(tree) = ProtobufLabeledTree::proxy_decode(&data[SEED_LEN..]) {
+        let seed: [u8; SEED_LEN] = data[..SEED_LEN]
+            .try_into()
+            .expect("failed to copy seed bytes");
+        test_tree(&tree, &mut ReproducibleRng::from_seed(seed));
     };
 });
 
 libfuzzer_sys::fuzz_mutator!(|data: &mut [u8], size: usize, max_size: usize, seed: u32| {
-    let mut tree = match ProtobufLabeledTree::proxy_decode(&data[..size]) {
+    let tree_data = if size < SEED_LEN {
+        // invalid tree encoding if there's not enough bytes to construct a slice
+        &[0u8; 0]
+    } else {
+        &data[SEED_LEN..size]
+    };
+
+    let mut tree = match ProtobufLabeledTree::proxy_decode(tree_data) {
         Ok(tree) if matches!(tree, LabeledTree::SubTree(_)) => tree,
         Err(_) | Ok(_) /*if matches!(tree, LabeledTree::Leaf(_))*/ => {
-            let bytes =
+            let seed = [0u8; SEED_LEN];
+            let encoded_tree =
                 ProtobufLabeledTree::proxy_encode(LabeledTree::<Vec<u8>>::SubTree(flatmap!()))
                     .expect("failed to serialize an empty labeled tree");
+            let bytes: Vec<_> = seed.into_iter().chain(encoded_tree.into_iter()).collect();
             let new_size = bytes.len();
             if new_size <= max_size {
                 data[..new_size].copy_from_slice(&bytes[..new_size]);
@@ -37,7 +53,7 @@ libfuzzer_sys::fuzz_mutator!(|data: &mut [u8], size: usize, max_size: usize, see
 
     let mut rng = rng_from_u32(seed);
 
-    let tree_was_modified = match rng.gen_range(0..8) {
+    let data_size_changed = match rng.gen_range(0..9) {
         0 => try_remove_leaf(&mut tree, &mut rng),
         1 => try_remove_empty_subtree(&mut tree, &mut rng),
         // actions that increase the tree's size have twice the probability of those
@@ -50,15 +66,20 @@ libfuzzer_sys::fuzz_mutator!(|data: &mut [u8], size: usize, max_size: usize, see
         7 => try_randomly_change_bytes_label(&mut tree, &mut rng, &|buffer: &mut Vec<u8>| {
             randomly_modify_buffer(buffer)
         }),
+        // generate new seed
+        8 => {
+            rng.fill_bytes(&mut data[..SEED_LEN]);
+            false
+        }
         _ => unreachable!(),
     };
 
-    if tree_was_modified {
-        let bytes = ProtobufLabeledTree::proxy_encode(tree)
+    if data_size_changed {
+        let encoded_tree = ProtobufLabeledTree::proxy_encode(tree)
             .expect("failed to serialize the labeled tree {tree}");
-        let new_size = bytes.len();
+        let new_size = SEED_LEN + encoded_tree.len();
         if new_size <= max_size {
-            data[..new_size].copy_from_slice(&bytes[..]);
+            data[SEED_LEN..new_size].copy_from_slice(&encoded_tree[..]);
             return new_size;
         } else {
             size

rs/crypto/tree_hash/src/lib.rs

@@ -13,6 +13,7 @@ use std::ops::DerefMut;
 pub mod flat_map;
 pub mod hasher;
 pub mod proto;
+pub mod test_utils;
 pub(crate) mod tree_hash;
 
 #[cfg(test)]