diff --git a/primitives/state-machine/Cargo.toml b/primitives/state-machine/Cargo.toml index e5c9ea9890687..dc54486e2078c 100644 --- a/primitives/state-machine/Cargo.toml +++ b/primitives/state-machine/Cargo.toml @@ -35,6 +35,7 @@ tracing = { version = "0.1.22", optional = true } hex-literal = "0.3.1" sp-runtime = { version = "4.0.0-dev", path = "../runtime" } pretty_assertions = "0.6.1" +rand = { version = "0.7.2", feature = ["small_rng"] } [features] default = ["std"] diff --git a/primitives/state-machine/src/lib.rs b/primitives/state-machine/src/lib.rs index e12be0c586b7d..032899faeb523 100644 --- a/primitives/state-machine/src/lib.rs +++ b/primitives/state-machine/src/lib.rs @@ -1515,7 +1515,9 @@ mod tests { #[test] fn prove_read_and_proof_check_works() { let child_info = ChildInfo::new_default(b"sub1"); + let missing_child_info = ChildInfo::new_default(b"sub1sub2"); // key will include other child root to proof. let child_info = &child_info; + let missing_child_info = &missing_child_info; // fetch read proof from 'remote' full node let remote_backend = trie_backend::tests::test_trie(); let remote_root = remote_backend.storage_root(std::iter::empty()).0; @@ -1553,11 +1555,111 @@ mod tests { &[b"value2"], ) .unwrap(); + let local_result3 = read_child_proof_check::( + remote_root, + remote_proof.clone(), + missing_child_info, + &[b"dummy"], + ) + .unwrap(); + assert_eq!( local_result1.into_iter().collect::>(), vec![(b"value3".to_vec(), Some(vec![142]))], ); assert_eq!(local_result2.into_iter().collect::>(), vec![(b"value2".to_vec(), None)]); + assert_eq!(local_result3.into_iter().collect::>(), vec![(b"dummy".to_vec(), None)]); + } + + #[test] + fn child_read_compact_stress_test() { + use rand::{rngs::SmallRng, RngCore, SeedableRng}; + let mut storage: HashMap, BTreeMap> = + Default::default(); + let mut seed = [0; 16]; + for i in 0..50u32 { + let mut child_infos = Vec::new(); + &seed[0..4].copy_from_slice(&i.to_be_bytes()[..]); + let mut rand = SmallRng::from_seed(seed); + + let nb_child_trie = rand.next_u32() as usize % 25; + for _ in 0..nb_child_trie { + let key_len = 1 + (rand.next_u32() % 10); + let mut key = vec![0; key_len as usize]; + rand.fill_bytes(&mut key[..]); + let child_info = ChildInfo::new_default(key.as_slice()); + let nb_item = 1 + rand.next_u32() % 25; + let mut items = BTreeMap::new(); + for item in 0..nb_item { + let key_len = 1 + (rand.next_u32() % 10); + let mut key = vec![0; key_len as usize]; + rand.fill_bytes(&mut key[..]); + let value = vec![item as u8; item as usize + 28]; + items.insert(key, value); + } + child_infos.push(child_info.clone()); + storage.insert(Some(child_info), items); + } + + let trie: InMemoryBackend = storage.clone().into(); + let trie_root = trie.root().clone(); + let backend = crate::ProvingBackend::new(&trie); + let mut queries = Vec::new(); + for c in 0..(5 + nb_child_trie / 2) { + // random existing query + let child_info = if c < 5 { + // 4 missing child trie + let key_len = 1 + (rand.next_u32() % 10); + let mut key = vec![0; key_len as usize]; + rand.fill_bytes(&mut key[..]); + ChildInfo::new_default(key.as_slice()) + } else { + child_infos[rand.next_u32() as usize % nb_child_trie].clone() + }; + + if let Some(values) = storage.get(&Some(child_info.clone())) { + for _ in 0..(1 + values.len() / 2) { + let ix = rand.next_u32() as usize % values.len(); + for (i, (key, value)) in values.iter().enumerate() { + if i == ix { + assert_eq!( + &backend + .child_storage(&child_info, key.as_slice()) + .unwrap() + .unwrap(), + value + ); + queries.push(( + child_info.clone(), + key.clone(), + Some(value.clone()), + )); + break + } + } + } + } + for _ in 0..4 { + let key_len = 1 + (rand.next_u32() % 10); + let mut key = vec![0; key_len as usize]; + rand.fill_bytes(&mut key[..]); + let result = backend.child_storage(&child_info, key.as_slice()).unwrap(); + queries.push((child_info.clone(), key, result)); + } + } + + let storage_proof = backend.extract_proof(); + let remote_proof = test_compact(storage_proof, &trie_root); + let proof_check = + create_proof_check_backend::(trie_root, remote_proof).unwrap(); + + for (child_info, key, expected) in queries { + assert_eq!( + proof_check.child_storage(&child_info, key.as_slice()).unwrap(), + expected, + ); + } + } } #[test] diff --git a/primitives/trie/src/trie_codec.rs b/primitives/trie/src/trie_codec.rs index 8f2f44317649b..1596229f2b5de 100644 --- a/primitives/trie/src/trie_codec.rs +++ b/primitives/trie/src/trie_codec.rs @@ -161,8 +161,9 @@ where } let mut previous_extracted_child_trie = None; + let mut nodes_iter = nodes_iter.peekable(); for child_root in child_tries.into_iter() { - if previous_extracted_child_trie.is_none() { + if previous_extracted_child_trie.is_none() && nodes_iter.peek().is_some() { let (top_root, _) = trie_db::decode_compact_from_iter::(db, &mut nodes_iter)?; previous_extracted_child_trie = Some(top_root);