Merge branch 'alex/CRP-2001-add-multi-leaf-witnesses-to-witness-equality-fuzzer' into 'master'

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

Currently, the witness equality fuzzer checks the witness equality only for paths to leaves. This MR adds 1) support to the fuzzer and tests in `canonical_state` for empty subtrees in addition to leaves (also fixes an inconsistency for handling empty subtrees in `canonical_state`) and 2) support for multi-leaf/empty-subtree witnesses to the fuzzer.

For multi-leaf fuzzing, the fuzzer generates random partial trees, which include `[2,min(10, nun_leaves_and_empty_subtrees(full_tree) - 1]` leaves/empty subtrees each.

The fuzzer now also tests pruning of the full tree.

Closes CRP-2001 

Closes CRP-2001

See merge request dfinity-lab/public/ic!12410

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)]