Child trie storage proof (#2433)

* proof on child trie

* higher level api for child storage proof

* boilerplate for proof from light fetch

* actually check proof on light fetch

* Do not break former encoding

* tabify

* tabify2

* Add child trie root tx to full_storage_root transaction.

* Shorten long lines.

* Temp rename for audit

* Make full_storage a trait method

* Name back and replace some code with full_storage where it looks fine.

* fix indentations, remove unused import

* flush child root to top when calculated

* impl +1

substrate/core/state-machine/src/backend.rs

@@ -70,6 +70,7 @@ pub trait Backend<H: Hasher> {
 
 	/// Calculate the storage root, with given delta over what is already stored in
 	/// the backend, and produce a "transaction" that can be used to commit.
+	/// Does not include child storage updates.
 	fn storage_root<I>(&self, delta: I) -> (H::Out, Self::Transaction)
 	where
 		I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>,
@@ -91,6 +92,41 @@ pub trait Backend<H: Hasher> {
 
 	/// Try convert into trie backend.
 	fn try_into_trie_backend(self) -> Option<TrieBackend<Self::TrieBackendStorage, H>>;
+
+	/// Calculate the storage root, with given delta over what is already stored
+	/// in the backend, and produce a "transaction" that can be used to commit.
+	/// Does include child storage updates.
+	fn full_storage_root<I1, I2i, I2>(
+		&self,
+		delta: I1,
+		child_deltas: I2)
+	-> (H::Out, Self::Transaction)
+	where
+		I1: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>,
+		I2i: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>,
+		I2: IntoIterator<Item=(Vec<u8>, I2i)>,
+		<H as Hasher>::Out: Ord,
+	{
+		let mut txs: Self::Transaction = Default::default();
+		let mut child_roots: Vec<_> = Default::default();
+		// child first
+		for (storage_key, child_delta) in child_deltas {
+			let (child_root, empty, child_txs) =
+				self.child_storage_root(&storage_key[..], child_delta);
+			txs.consolidate(child_txs);
+			if empty {
+				child_roots.push((storage_key, None));
+			} else {
+				child_roots.push((storage_key, Some(child_root)));
+			}
+		}
+		let (root, parent_txs) = self.storage_root(
+			delta.into_iter().chain(child_roots.into_iter())
+		);
+		txs.consolidate(parent_txs);
+		(root, txs)
+	}
+
 }
 
 /// Trait that allows consolidate two transactions together.
@@ -213,6 +249,13 @@ impl<H> From<Vec<(Option<Vec<u8>>, Vec<u8>, Option<Vec<u8>>)>> for InMemory<H> {
 
 impl super::Error for Void {}
 
+impl<H: Hasher> InMemory<H> {
+	/// child storage key iterator
+	pub fn child_storage_keys(&self) -> impl Iterator<Item=&[u8]> {
+		self.inner.iter().filter_map(|item| item.0.as_ref().map(|v|&v[..]))
+	}
+}
+
 impl<H: Hasher> Backend<H> for InMemory<H> {
 	type Error = Void;
 	type Transaction = Vec<(Option<Vec<u8>>, Vec<u8>, Option<Vec<u8>>)>;
@@ -290,16 +333,28 @@ impl<H: Hasher> Backend<H> for InMemory<H> {
 		self.inner.get(&None).into_iter().flat_map(|map| map.keys().filter(|k| k.starts_with(prefix)).cloned()).collect()
 	}
 
-	fn try_into_trie_backend(self) -> Option<TrieBackend<Self::TrieBackendStorage, H>> {
+	fn try_into_trie_backend(
+		self
+	)-> Option<TrieBackend<Self::TrieBackendStorage, H>> {
 		let mut mdb = MemoryDB::default();
 		let mut root = None;
+		let mut new_child_roots = Vec::new();
+		let mut root_map = None;
 		for (storage_key, map) in self.inner {
-			if storage_key != None {
-				let _ = insert_into_memory_db::<H, _>(&mut mdb, map.into_iter())?;
+			if let Some(storage_key) = storage_key.as_ref() {
+				let ch = insert_into_memory_db::<H, _>(&mut mdb, map.into_iter())?;
+				new_child_roots.push((storage_key.clone(), ch.as_ref().into()));
 			} else {
-				root = Some(insert_into_memory_db::<H, _>(&mut mdb, map.into_iter())?);
+				root_map = Some(map);
 			}
 		}
+		// root handling
+		if let Some(map) = root_map.take() {
+			root = Some(insert_into_memory_db::<H, _>(
+				&mut mdb,
+				map.into_iter().chain(new_child_roots.into_iter())
+			)?);
+		}
 		let root = match root {
 			Some(root) => root,
 			None => insert_into_memory_db::<H, _>(&mut mdb, ::std::iter::empty())?,

substrate/core/state-machine/src/ext.rs

@@ -18,7 +18,7 @@
 
 use std::{error, fmt, cmp::Ord};
 use log::warn;
-use crate::backend::{Backend, Consolidate};
+use crate::backend::Backend;
 use crate::changes_trie::{AnchorBlockId, Storage as ChangesTrieStorage, compute_changes_trie_root};
 use crate::{Externalities, OverlayedChanges, OffchainExt, ChildStorageKey};
 use hash_db::Hasher;
@@ -134,30 +134,6 @@ where
 		self.storage_transaction = None;
 	}
 
-	/// Fetch child storage root together with its transaction.
-	fn child_storage_root_transaction(&mut self, storage_key: &[u8]) -> (Vec<u8>, B::Transaction) {
-		self.mark_dirty();
-
-		let (root, is_default, transaction) = {
-			let delta = self.overlay.committed.children.get(storage_key)
-				.into_iter()
-				.flat_map(|map| map.1.iter().map(|(k, v)| (k.clone(), v.clone())))
-				.chain(self.overlay.prospective.children.get(storage_key)
-						.into_iter()
-						.flat_map(|map| map.1.iter().map(|(k, v)| (k.clone(), v.clone()))));
-
-			self.backend.child_storage_root(storage_key, delta)
-		};
-
-		let root_val = if is_default {
-			None
-		} else {
-			Some(root.clone())
-		};
-		self.overlay.sync_child_storage_root(storage_key, root_val);
-
-		(root, transaction)
-	}
 }
 
 #[cfg(test)]
@@ -288,21 +264,24 @@ where
 			return root.clone();
 		}
 
-		let mut transaction = B::Transaction::default();
-		let child_storage_keys: std::collections::BTreeSet<_> = self.overlay.prospective.children.keys().cloned()
-			.chain(self.overlay.committed.children.keys().cloned()).collect();
+		let child_storage_keys =
+			self.overlay.prospective.children.keys()
+				.chain(self.overlay.committed.children.keys());
+
+		let child_delta_iter = child_storage_keys.map(|storage_key|
+			(storage_key.clone(), self.overlay.committed.children.get(storage_key)
+				.into_iter()
+				.flat_map(|map| map.1.iter().map(|(k, v)| (k.clone(), v.clone())))
+				.chain(self.overlay.prospective.children.get(storage_key)
+					.into_iter()
+					.flat_map(|map| map.1.iter().map(|(k, v)| (k.clone(), v.clone()))))));
 
-		for key in child_storage_keys {
-			let (_, t) = self.child_storage_root_transaction(&key);
-			transaction.consolidate(t);
-		}
 
 		// compute and memoize
 		let delta = self.overlay.committed.top.iter().map(|(k, v)| (k.clone(), v.value.clone()))
 			.chain(self.overlay.prospective.top.iter().map(|(k, v)| (k.clone(), v.value.clone())));
 
-		let (root, t) = self.backend.storage_root(delta);
-		transaction.consolidate(t);
+		let (root, transaction) = self.backend.full_storage_root(delta, child_delta_iter);
 		self.storage_transaction = Some((transaction, root));
 		root
 	}
@@ -316,7 +295,21 @@ where
 					default_child_trie_root::<H>(storage_key.as_ref())
 				)
 		} else {
-			self.child_storage_root_transaction(storage_key.as_ref()).0
+			let storage_key = storage_key.as_ref();
+
+			let delta = self.overlay.committed.children.get(storage_key)
+				.into_iter()
+				.flat_map(|map| map.1.iter().map(|(k, v)| (k.clone(), v.clone())))
+				.chain(self.overlay.prospective.children.get(storage_key)
+						.into_iter()
+						.flat_map(|map| map.1.iter().map(|(k, v)| (k.clone(), v.clone()))));
+
+			let root = self.backend.child_storage_root(storage_key, delta).0;
+
+			self.overlay.set_storage(storage_key.to_vec(), Some(root.to_vec()));
+
+			root
+
 		}
 	}
 

substrate/core/state-machine/src/lib.rs

@@ -631,7 +631,7 @@ where
 	Exec: CodeExecutor<H>,
 	H::Out: Ord,
 {
-	let trie_backend = proving_backend::create_proof_check_backend::<H>(root.into(), proof)?;
+	let trie_backend = create_proof_check_backend::<H>(root.into(), proof)?;
 	execution_proof_check_on_trie_backend(&trie_backend, overlay, exec, method, call_data)
 }
 
@@ -676,10 +676,29 @@ where
 	H::Out: Ord
 {
 	let trie_backend = backend.try_into_trie_backend()
-		.ok_or_else(|| Box::new(ExecutionError::UnableToGenerateProof) as Box<Error>)?;
+		.ok_or_else(
+			||Box::new(ExecutionError::UnableToGenerateProof) as Box<Error>
+		)?;
 	prove_read_on_trie_backend(&trie_backend, key)
 }
 
+/// Generate child storage read proof.
+pub fn prove_child_read<B, H>(
+	backend: B,
+	storage_key: &[u8],
+	key: &[u8],
+) -> Result<(Option<Vec<u8>>, Vec<Vec<u8>>), Box<Error>>
+where
+	B: Backend<H>,
+	H: Hasher,
+	H::Out: Ord
+{
+	let trie_backend = backend.try_into_trie_backend()
+		.ok_or_else(|| Box::new(ExecutionError::UnableToGenerateProof) as Box<Error>)?;
+	prove_child_read_on_trie_backend(&trie_backend, storage_key, key)
+}
+
+
 /// Generate storage read proof on pre-created trie backend.
 pub fn prove_read_on_trie_backend<S, H>(
 	trie_backend: &TrieBackend<S, H>,
@@ -695,6 +714,23 @@ where
 	Ok((result, proving_backend.extract_proof()))
 }
 
+/// Generate storage read proof on pre-created trie backend.
+pub fn prove_child_read_on_trie_backend<S, H>(
+	trie_backend: &TrieBackend<S, H>,
+	storage_key: &[u8],
+	key: &[u8]
+) -> Result<(Option<Vec<u8>>, Vec<Vec<u8>>), Box<Error>>
+where
+	S: trie_backend_essence::TrieBackendStorage<H>,
+	H: Hasher,
+	H::Out: Ord
+{
+	let proving_backend = proving_backend::ProvingBackend::<_, H>::new(trie_backend);
+	let result = proving_backend.child_storage(storage_key, key)
+		.map_err(|e| Box::new(e) as Box<Error>)?;
+	Ok((result, proving_backend.extract_proof()))
+}
+
 /// Check storage read proof, generated by `prove_read` call.
 pub fn read_proof_check<H>(
 	root: H::Out,
@@ -705,10 +741,26 @@ where
 	H: Hasher,
 	H::Out: Ord
 {
-	let proving_backend = proving_backend::create_proof_check_backend::<H>(root, proof)?;
+	let proving_backend = create_proof_check_backend::<H>(root, proof)?;
 	read_proof_check_on_proving_backend(&proving_backend, key)
 }
 
+/// Check child storage read proof, generated by `prove_child_read` call.
+pub fn read_child_proof_check<H>(
+	root: H::Out,
+	proof: Vec<Vec<u8>>,
+	storage_key: &[u8],
+	key: &[u8],
+) -> Result<Option<Vec<u8>>, Box<Error>>
+where
+	H: Hasher,
+	H::Out: Ord
+{
+	let proving_backend = create_proof_check_backend::<H>(root, proof)?;
+	read_child_proof_check_on_proving_backend(&proving_backend, storage_key, key)
+}
+
+
 /// Check storage read proof on pre-created proving backend.
 pub fn read_proof_check_on_proving_backend<H>(
 	proving_backend: &TrieBackend<MemoryDB<H>, H>,
@@ -721,6 +773,19 @@ where
 	proving_backend.storage(key).map_err(|e| Box::new(e) as Box<Error>)
 }
 
+/// Check child storage read proof on pre-created proving backend.
+pub fn read_child_proof_check_on_proving_backend<H>(
+	proving_backend: &TrieBackend<MemoryDB<H>, H>,
+	storage_key: &[u8],
+	key: &[u8],
+) -> Result<Option<Vec<u8>>, Box<Error>>
+where
+	H: Hasher,
+	H::Out: Ord
+{
+	proving_backend.child_storage(storage_key, key).map_err(|e| Box::new(e) as Box<Error>)
+}
+
 /// Sets overlayed changes' changes trie configuration. Returns error if configuration
 /// differs from previous OR config decode has failed.
 pub(crate) fn set_changes_trie_config(overlay: &mut OverlayedChanges, config: Option<Vec<u8>>, final_check: bool) -> Result<(), Box<Error>> {
@@ -1001,11 +1066,40 @@ mod tests {
 		let remote_root = remote_backend.storage_root(::std::iter::empty()).0;
 		let remote_proof = prove_read(remote_backend, b"value2").unwrap().1;
  		// check proof locally
-		let local_result1 = read_proof_check::<Blake2Hasher>(remote_root, remote_proof.clone(), b"value2").unwrap();
-		let local_result2 = read_proof_check::<Blake2Hasher>(remote_root, remote_proof.clone(), &[0xff]).is_ok();
+		let local_result1 = read_proof_check::<Blake2Hasher>(
+			remote_root,
+			remote_proof.clone(),
+			b"value2"
+		).unwrap();
+		let local_result2 = read_proof_check::<Blake2Hasher>(
+			remote_root,
+			remote_proof.clone(),
+			&[0xff]
+		).is_ok();
  		// check that results are correct
 		assert_eq!(local_result1, Some(vec![24]));
 		assert_eq!(local_result2, false);
+		// on child trie
+		let remote_backend = trie_backend::tests::test_trie();
+		let remote_root = remote_backend.storage_root(::std::iter::empty()).0;
+		let remote_proof = prove_child_read(
+			remote_backend,
+			b":child_storage:default:sub1",
+			b"value3"
+		).unwrap().1;
+		let local_result1 = read_child_proof_check::<Blake2Hasher>(
+			remote_root,
+			remote_proof.clone(),
+			b":child_storage:default:sub1",b"value3"
+		).unwrap();
+		let local_result2 = read_child_proof_check::<Blake2Hasher>(
+			remote_root,
+			remote_proof.clone(),
+			b":child_storage:default:sub1",
+			b"value2"
+		).unwrap();
+		assert_eq!(local_result1, Some(vec![142]));
+		assert_eq!(local_result2, None);
 	}
 
 	#[test]

substrate/core/state-machine/src/overlayed_changes.rs

@@ -159,19 +159,6 @@ impl OverlayedChanges {
 		}
 	}
 
-	/// Sync the child storage root.
-	pub(crate) fn sync_child_storage_root(&mut self, storage_key: &[u8], root: Option<Vec<u8>>) {
-		let entry = self.prospective.top.entry(storage_key.to_vec()).or_default();
-		entry.value = root;
-
-		if let Some((Some(extrinsics), _)) = self.prospective.children.get(storage_key) {
-			for extrinsic in extrinsics {
-				entry.extrinsics.get_or_insert_with(Default::default)
-					.insert(*extrinsic);
-			}
-		}
-	}
-
 	/// Clear child storage of given storage key.
 	///
 	/// NOTE that this doesn't take place immediately but written into the prospective

substrate/core/state-machine/src/proving_backend.rs

@@ -222,6 +222,7 @@ mod tests {
 	use crate::trie_backend::tests::test_trie;
 	use super::*;
 	use primitives::{Blake2Hasher};
+	use crate::ChildStorageKey;
 
 	fn test_proving<'a>(trie_backend: &'a TrieBackend<PrefixedMemoryDB<Blake2Hasher>, Blake2Hasher>) -> ProvingBackend<'a, PrefixedMemoryDB<Blake2Hasher>, Blake2Hasher> {
 		ProvingBackend::new(trie_backend)
@@ -281,4 +282,75 @@ mod tests {
 		let proof_check = create_proof_check_backend::<Blake2Hasher>(in_memory_root.into(), proof).unwrap();
 		assert_eq!(proof_check.storage(&[42]).unwrap().unwrap(), vec![42]);
 	}
+
+	#[test]
+	fn proof_recorded_and_checked_with_child() {
+		let subtrie1 = ChildStorageKey::<Blake2Hasher>::from_slice(
+			b":child_storage:default:sub1"
+		).unwrap();
+		let subtrie2 = ChildStorageKey::<Blake2Hasher>::from_slice(
+			b":child_storage:default:sub2"
+		).unwrap();
+		let own1 = subtrie1.into_owned();
+		let own2 = subtrie2.into_owned();
+		let contents = (0..64).map(|i| (None, vec![i], Some(vec![i])))
+			.chain((28..65).map(|i| (Some(own1.clone()), vec![i], Some(vec![i]))))
+			.chain((10..15).map(|i| (Some(own2.clone()), vec![i], Some(vec![i]))))
+			.collect::<Vec<_>>();
+		let in_memory = InMemory::<Blake2Hasher>::default();
+		let in_memory = in_memory.update(contents);
+		let in_memory_root = in_memory.full_storage_root::<_, Vec<_>, _>(
+			::std::iter::empty(),
+      in_memory.child_storage_keys().map(|k|(k.to_vec(), Vec::new()))
+		).0;
+		(0..64).for_each(|i| assert_eq!(
+			in_memory.storage(&[i]).unwrap().unwrap(),
+			vec![i]
+		));
+		(28..65).for_each(|i| assert_eq!(
+			in_memory.child_storage(&own1[..], &[i]).unwrap().unwrap(),
+			vec![i]
+		));
+		(10..15).for_each(|i| assert_eq!(
+			in_memory.child_storage(&own2[..], &[i]).unwrap().unwrap(),
+			vec![i]
+		));
+
+		let trie = in_memory.try_into_trie_backend().unwrap();
+		let trie_root = trie.storage_root(::std::iter::empty()).0;
+		assert_eq!(in_memory_root, trie_root);
+		(0..64).for_each(|i| assert_eq!(
+			trie.storage(&[i]).unwrap().unwrap(),
+			vec![i]
+		));
+
+		let proving = ProvingBackend::new(&trie);
+		assert_eq!(proving.storage(&[42]).unwrap().unwrap(), vec![42]);
+
+		let proof = proving.extract_proof();
+
+		let proof_check = create_proof_check_backend::<Blake2Hasher>(
+			in_memory_root.into(),
+			proof
+		).unwrap();
+		assert!(proof_check.storage(&[0]).is_err());
+		assert_eq!(proof_check.storage(&[42]).unwrap().unwrap(), vec![42]);
+		// note that it is include in root because proof close
+		assert_eq!(proof_check.storage(&[41]).unwrap().unwrap(), vec![41]);
+		assert_eq!(proof_check.storage(&[64]).unwrap(), None);
+
+		let proving = ProvingBackend::new(&trie);
+		assert_eq!(proving.child_storage(&own1[..], &[64]), Ok(Some(vec![64])));
+
+		let proof = proving.extract_proof();
+		let proof_check = create_proof_check_backend::<Blake2Hasher>(
+			in_memory_root.into(),
+			proof
+		).unwrap();
+		assert_eq!(
+			proof_check.child_storage(&own1[..], &[64]).unwrap().unwrap(),
+			vec![64]
+		);
+	}
+
 }

substrate/core/state-machine/src/trie_backend.rs

@@ -188,6 +188,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> Backend<H> for TrieBackend<S, H> where
 pub mod tests {
 	use std::collections::HashSet;
 	use primitives::{Blake2Hasher, H256};
+	use parity_codec::Encode;
 	use trie::{TrieMut, TrieDBMut, PrefixedMemoryDB};
 	use super::*;
 
@@ -196,6 +197,15 @@ pub mod tests {
 		let mut mdb = PrefixedMemoryDB::<Blake2Hasher>::default();
 		{
 			let mut trie = TrieDBMut::new(&mut mdb, &mut root);
+			trie.insert(b"value3", &[142]).expect("insert failed");
+			trie.insert(b"value4", &[124]).expect("insert failed");
+		};
+
+		{
+			let mut sub_root = Vec::new();
+			root.encode_to(&mut sub_root);
+			let mut trie = TrieDBMut::new(&mut mdb, &mut root);
+			trie.insert(b":child_storage:default:sub1", &sub_root).expect("insert failed");
 			trie.insert(b"key", b"value").expect("insert failed");
 			trie.insert(b"value1", &[42]).expect("insert failed");
 			trie.insert(b"value2", &[24]).expect("insert failed");