Fix compact proof decoding unaccessed last child trie. (#9715)

* fix no child proof attached but root included.

* small stress test for proof of child tries.

* rust fmt

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

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::<BlakeTwo256, _>(
+			remote_root,
+			remote_proof.clone(),
+			missing_child_info,
+			&[b"dummy"],
+		)
+		.unwrap();
+
 		assert_eq!(
 			local_result1.into_iter().collect::<Vec<_>>(),
 			vec![(b"value3".to_vec(), Some(vec![142]))],
 		);
 		assert_eq!(local_result2.into_iter().collect::<Vec<_>>(), vec![(b"value2".to_vec(), None)]);
+		assert_eq!(local_result3.into_iter().collect::<Vec<_>>(), vec![(b"dummy".to_vec(), None)]);
+	}
+
+	#[test]
+	fn child_read_compact_stress_test() {
+		use rand::{rngs::SmallRng, RngCore, SeedableRng};
+		let mut storage: HashMap<Option<ChildInfo>, BTreeMap<StorageKey, StorageValue>> =
+			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<BlakeTwo256> = 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::<BlakeTwo256>(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]

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::<L, _, _, _>(db, &mut nodes_iter)?;
 			previous_extracted_child_trie = Some(top_root);