Don't include :code by default in storage proofs (#5179)

* Don't include `:code` by default in storage proofs (#5060)

* Adds test to verify that the runtime currently is always contained in
the proof

* Start passing the runtime wasm code from the outside

* Fix compilation

* More build fixes

* Make the test work as expected now :)

* Last fixes

* Fixes benchmarks

* Review feedback

* Apply suggestions from code review

Co-Authored-By: Sergei Pepyakin <sergei@parity.io>

* Review feedback

* Fix compilation

Co-authored-by: Sergei Pepyakin <s.pepyakin@gmail.com>

* Fix compilation and change the way `RuntimeCode` works

* Fix tests

* Switch to `Cow`

Co-authored-by: Benjamin Kampmann <ben@gnunicorn.org>
Co-authored-by: Sergei Pepyakin <s.pepyakin@gmail.com>

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

@@ -18,9 +18,9 @@
 
 use log::warn;
 use hash_db::Hasher;
-use codec::Encode;
+use codec::{Decode, Encode};
 
-use sp_core::storage::{ChildInfo, OwnedChildInfo};
+use sp_core::{traits::RuntimeCode, storage::{ChildInfo, OwnedChildInfo, well_known_keys}};
 use sp_trie::{TrieMut, MemoryDB, trie_types::TrieDBMut};
 
 use crate::{
@@ -359,3 +359,42 @@ pub(crate) fn insert_into_memory_db<H, I>(mdb: &mut MemoryDB<H>, input: I) -> Op
 
 	Some(root)
 }
+
+/// Wrapper to create a [`RuntimeCode`] from a type that implements [`Backend`].
+pub struct BackendRuntimeCode<'a, B, H> {
+	backend: &'a B,
+	_marker: std::marker::PhantomData<H>,
+}
+
+impl<'a, B: Backend<H>, H: Hasher> sp_core::traits::FetchRuntimeCode for
+	BackendRuntimeCode<'a, B, H>
+{
+	fn fetch_runtime_code<'b>(&'b self) -> Option<std::borrow::Cow<'b, [u8]>> {
+		self.backend.storage(well_known_keys::CODE).ok().flatten().map(Into::into)
+	}
+}
+
+impl<'a, B: Backend<H>, H: Hasher> BackendRuntimeCode<'a, B, H> where H::Out: Encode {
+	/// Create a new instance.
+	pub fn new(backend: &'a B) -> Self {
+		Self {
+			backend,
+			_marker: std::marker::PhantomData,
+		}
+	}
+
+	/// Return the [`RuntimeCode`] build from the wrapped `backend`.
+	pub fn runtime_code(&self) -> Result<RuntimeCode, &'static str> {
+		let hash = self.backend.storage_hash(well_known_keys::CODE)
+			.ok()
+			.flatten()
+			.ok_or("`:code` hash not found")?
+			.encode();
+		let heap_pages = self.backend.storage(well_known_keys::HEAP_PAGES)
+			.ok()
+			.flatten()
+			.and_then(|d| Decode::decode(&mut &d[..]).ok());
+
+		Ok(RuntimeCode { code_fetcher: self, hash, heap_pages })
+	}
+}

substrate/primitives/state-machine/src/basic.rs

@@ -119,14 +119,6 @@ impl Externalities for BasicExternalities {
 		self.storage(key).map(|v| Blake2Hasher::hash(&v).encode())
 	}
 
-	fn original_storage(&self, key: &[u8]) -> Option<StorageValue> {
-		self.storage(key)
-	}
-
-	fn original_storage_hash(&self, key: &[u8]) -> Option<Vec<u8>> {
-		self.storage_hash(key)
-	}
-
 	fn child_storage(
 		&self,
 		storage_key: ChildStorageKey,
@@ -145,24 +137,6 @@ impl Externalities for BasicExternalities {
 		self.child_storage(storage_key, child_info, key).map(|v| Blake2Hasher::hash(&v).encode())
 	}
 
-	fn original_child_storage_hash(
-		&self,
-		storage_key: ChildStorageKey,
-		child_info: ChildInfo,
-		key: &[u8],
-	) -> Option<Vec<u8>> {
-		self.child_storage_hash(storage_key, child_info, key)
-	}
-
-	fn original_child_storage(
-		&self,
-		storage_key: ChildStorageKey,
-		child_info: ChildInfo,
-		key: &[u8],
-	) -> Option<StorageValue> {
-		Externalities::child_storage(self, storage_key, child_info, key)
-	}
-
 	fn next_storage_key(&self, key: &[u8]) -> Option<StorageKey> {
 		let range = (Bound::Excluded(key), Bound::Unbounded);
 		self.inner.top.range::<[u8], _>(range).next().map(|(k, _)| k).cloned()

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

@@ -179,30 +179,6 @@ where
 		result.map(|r| r.encode())
 	}
 
-	fn original_storage(&self, key: &[u8]) -> Option<StorageValue> {
-		let _guard = sp_panic_handler::AbortGuard::force_abort();
-		let result = self.backend.storage(key).expect(EXT_NOT_ALLOWED_TO_FAIL);
-
-		trace!(target: "state-trace", "{:04x}: GetOriginal {}={:?}",
-			self.id,
-			HexDisplay::from(&key),
-			result.as_ref().map(HexDisplay::from)
-		);
-		result
-	}
-
-	fn original_storage_hash(&self, key: &[u8]) -> Option<Vec<u8>> {
-		let _guard = sp_panic_handler::AbortGuard::force_abort();
-		let result = self.backend.storage_hash(key).expect(EXT_NOT_ALLOWED_TO_FAIL);
-
-		trace!(target: "state-trace", "{:04x}: GetOriginalHash {}={:?}",
-			self.id,
-			HexDisplay::from(&key),
-			result,
-		);
-		result.map(|r| r.encode())
-	}
-
 	fn child_storage(
 		&self,
 		storage_key: ChildStorageKey,
@@ -253,47 +229,6 @@ where
 		result.map(|r| r.encode())
 	}
 
-	fn original_child_storage(
-		&self,
-		storage_key: ChildStorageKey,
-		child_info: ChildInfo,
-		key: &[u8],
-	) -> Option<StorageValue> {
-		let _guard = sp_panic_handler::AbortGuard::force_abort();
-		let result = self.backend
-			.child_storage(storage_key.as_ref(), child_info, key)
-			.expect(EXT_NOT_ALLOWED_TO_FAIL);
-
-		trace!(target: "state-trace", "{:04x}: ChildOriginal({}) {}={:?}",
-			self.id,
-			HexDisplay::from(&storage_key.as_ref()),
-			HexDisplay::from(&key),
-			result.as_ref().map(HexDisplay::from),
-		);
-
-		result
-	}
-
-	fn original_child_storage_hash(
-		&self,
-		storage_key: ChildStorageKey,
-		child_info: ChildInfo,
-		key: &[u8],
-	) -> Option<Vec<u8>> {
-		let _guard = sp_panic_handler::AbortGuard::force_abort();
-		let result = self.backend
-			.child_storage_hash(storage_key.as_ref(), child_info, key)
-			.expect(EXT_NOT_ALLOWED_TO_FAIL);
-
-		trace!(target: "state-trace", "{}: ChildHashOriginal({}) {}={:?}",
-			self.id,
-			HexDisplay::from(&storage_key.as_ref()),
-			HexDisplay::from(&key),
-			result,
-		);
-		result.map(|r| r.encode())
-	}
-
 	fn exists_storage(&self, key: &[u8]) -> bool {
 		let _guard = sp_panic_handler::AbortGuard::force_abort();
 		let result = match self.overlay.storage(key) {
@@ -820,21 +755,18 @@ mod tests {
 		let ext = TestExt::new(&mut overlay, &mut cache, &backend, None, None);
 
 		assert_eq!(ext.child_storage(child(), CHILD_INFO_1, &[10]), Some(vec![10]));
-		assert_eq!(ext.original_child_storage(child(), CHILD_INFO_1, &[10]), Some(vec![10]));
 		assert_eq!(
 			ext.child_storage_hash(child(), CHILD_INFO_1, &[10]),
 			Some(Blake2Hasher::hash(&[10]).as_ref().to_vec()),
 		);
 
 		assert_eq!(ext.child_storage(child(), CHILD_INFO_1, &[20]), None);
-		assert_eq!(ext.original_child_storage(child(), CHILD_INFO_1, &[20]), Some(vec![20]));
 		assert_eq!(
 			ext.child_storage_hash(child(), CHILD_INFO_1, &[20]),
 			None,
 		);
 
 		assert_eq!(ext.child_storage(child(), CHILD_INFO_1, &[30]), Some(vec![31]));
-		assert_eq!(ext.original_child_storage(child(), CHILD_INFO_1, &[30]), Some(vec![40]));
 		assert_eq!(
 			ext.child_storage_hash(child(), CHILD_INFO_1, &[30]),
 			Some(Blake2Hasher::hash(&[31]).as_ref().to_vec()),

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

@@ -23,8 +23,8 @@ use log::{warn, trace};
 use hash_db::Hasher;
 use codec::{Decode, Encode, Codec};
 use sp_core::{
-	storage::ChildInfo, NativeOrEncoded, NeverNativeValue,
-	traits::{CodeExecutor, CallInWasmExt}, hexdisplay::HexDisplay,
+	storage::ChildInfo, NativeOrEncoded, NeverNativeValue, hexdisplay::HexDisplay,
+	traits::{CodeExecutor, CallInWasmExt, RuntimeCode},
 };
 use overlayed_changes::OverlayedChangeSet;
 use sp_externalities::Extensions;
@@ -42,7 +42,7 @@ mod trie_backend;
 mod trie_backend_essence;
 mod stats;
 
-pub use sp_trie::{trie_types::{Layout, TrieDBMut}, TrieMut, DBValue, MemoryDB};
+pub use sp_trie::{trie_types::{Layout, TrieDBMut}, StorageProof, TrieMut, DBValue, MemoryDB};
 pub use testing::TestExternalities;
 pub use basic::BasicExternalities;
 pub use ext::Ext;
@@ -67,8 +67,7 @@ pub use overlayed_changes::{
 	StorageCollection, ChildStorageCollection,
 };
 pub use proving_backend::{
-	create_proof_check_backend, create_proof_check_backend_storage, merge_storage_proofs,
-	ProofRecorder, ProvingBackend, ProvingBackendRecorder, StorageProof,
+	create_proof_check_backend, ProofRecorder, ProvingBackend, ProvingBackendRecorder,
 };
 pub use trie_backend_essence::{TrieBackendStorage, Storage};
 pub use trie_backend::TrieBackend;
@@ -191,6 +190,7 @@ pub struct StateMachine<'a, B, H, N, Exec>
 	changes_trie_state: Option<ChangesTrieState<'a, H, N>>,
 	_marker: PhantomData<(H, N)>,
 	storage_transaction_cache: Option<&'a mut StorageTransactionCache<B::Transaction, H, N>>,
+	runtime_code: &'a RuntimeCode<'a>,
 }
 
 impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
@@ -209,6 +209,7 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 		method: &'a str,
 		call_data: &'a [u8],
 		mut extensions: Extensions,
+		runtime_code: &'a RuntimeCode,
 	) -> Self {
 		extensions.register(CallInWasmExt::new(exec.clone()));
 
@@ -222,6 +223,7 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 			changes_trie_state,
 			_marker: PhantomData,
 			storage_transaction_cache: None,
+			runtime_code,
 		}
 	}
 
@@ -292,6 +294,7 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 
 		let (result, was_native) = self.exec.call(
 			&mut ext,
+			self.runtime_code,
 			self.method,
 			self.call_data,
 			use_native,
@@ -436,6 +439,7 @@ pub fn prove_execution<B, H, N, Exec>(
 	exec: &Exec,
 	method: &str,
 	call_data: &[u8],
+	runtime_code: &RuntimeCode,
 ) -> Result<(Vec<u8>, StorageProof), Box<dyn Error>>
 where
 	B: Backend<H>,
@@ -446,7 +450,14 @@ where
 {
 	let trie_backend = backend.as_trie_backend()
 		.ok_or_else(|| Box::new(ExecutionError::UnableToGenerateProof) as Box<dyn Error>)?;
-	prove_execution_on_trie_backend::<_, _, N, _>(trie_backend, overlay, exec, method, call_data)
+	prove_execution_on_trie_backend::<_, _, N, _>(
+		trie_backend,
+		overlay,
+		exec,
+		method,
+		call_data,
+		runtime_code,
+	)
 }
 
 /// Prove execution using the given trie backend, overlayed changes, and call executor.
@@ -464,6 +475,7 @@ pub fn prove_execution_on_trie_backend<S, H, N, Exec>(
 	exec: &Exec,
 	method: &str,
 	call_data: &[u8],
+	runtime_code: &RuntimeCode,
 ) -> Result<(Vec<u8>, StorageProof), Box<dyn Error>>
 where
 	S: trie_backend_essence::TrieBackendStorage<H>,
@@ -474,7 +486,14 @@ where
 {
 	let proving_backend = proving_backend::ProvingBackend::new(trie_backend);
 	let mut sm = StateMachine::<_, H, N, Exec>::new(
-		&proving_backend, None, overlay, exec, method, call_data, Extensions::default(),
+		&proving_backend,
+		None,
+		overlay,
+		exec,
+		method,
+		call_data,
+		Extensions::default(),
+		runtime_code,
 	);
 
 	let result = sm.execute_using_consensus_failure_handler::<_, NeverNativeValue, fn() -> _>(
@@ -493,6 +512,7 @@ pub fn execution_proof_check<H, N, Exec>(
 	exec: &Exec,
 	method: &str,
 	call_data: &[u8],
+	runtime_code: &RuntimeCode,
 ) -> Result<Vec<u8>, Box<dyn Error>>
 where
 	H: Hasher,
@@ -501,7 +521,14 @@ where
 	N: crate::changes_trie::BlockNumber,
 {
 	let trie_backend = create_proof_check_backend::<H>(root.into(), proof)?;
-	execution_proof_check_on_trie_backend::<_, N, _>(&trie_backend, overlay, exec, method, call_data)
+	execution_proof_check_on_trie_backend::<_, N, _>(
+		&trie_backend,
+		overlay,
+		exec,
+		method,
+		call_data,
+		runtime_code,
+	)
 }
 
 /// Check execution proof on proving backend, generated by `prove_execution` call.
@@ -511,6 +538,7 @@ pub fn execution_proof_check_on_trie_backend<H, N, Exec>(
 	exec: &Exec,
 	method: &str,
 	call_data: &[u8],
+	runtime_code: &RuntimeCode,
 ) -> Result<Vec<u8>, Box<dyn Error>>
 where
 	H: Hasher,
@@ -519,7 +547,14 @@ where
 	N: crate::changes_trie::BlockNumber,
 {
 	let mut sm = StateMachine::<_, H, N, Exec>::new(
-		trie_backend, None, overlay, exec, method, call_data, Extensions::default(),
+		trie_backend,
+		None,
+		overlay,
+		exec,
+		method,
+		call_data,
+		Extensions::default(),
+		runtime_code,
 	);
 
 	sm.execute_using_consensus_failure_handler::<_, NeverNativeValue, fn() -> _>(
@@ -692,7 +727,7 @@ mod tests {
 	use super::*;
 	use super::ext::Ext;
 	use super::changes_trie::Configuration as ChangesTrieConfig;
-	use sp_core::{map, traits::Externalities, storage::ChildStorageKey};
+	use sp_core::{map, traits::{Externalities, RuntimeCode}, storage::ChildStorageKey};
 	use sp_runtime::traits::BlakeTwo256;
 
 	#[derive(Clone)]
@@ -714,6 +749,7 @@ mod tests {
 		>(
 			&self,
 			ext: &mut dyn Externalities,
+			_: &RuntimeCode,
 			_method: &str,
 			_data: &[u8],
 			use_native: bool,
@@ -755,6 +791,7 @@ mod tests {
 		fn call_in_wasm(
 			&self,
 			_: &[u8],
+			_: Option<Vec<u8>>,
 			_: &str,
 			_: &[u8],
 			_: &mut dyn Externalities,
@@ -767,6 +804,7 @@ mod tests {
 	fn execute_works() {
 		let backend = trie_backend::tests::test_trie();
 		let mut overlayed_changes = Default::default();
+		let wasm_code = RuntimeCode::empty();
 
 		let mut state_machine = StateMachine::new(
 			&backend,
@@ -781,6 +819,7 @@ mod tests {
 			"test",
 			&[],
 			Default::default(),
+			&wasm_code,
 		);
 
 		assert_eq!(
@@ -794,6 +833,7 @@ mod tests {
 	fn execute_works_with_native_else_wasm() {
 		let backend = trie_backend::tests::test_trie();
 		let mut overlayed_changes = Default::default();
+		let wasm_code = RuntimeCode::empty();
 
 		let mut state_machine = StateMachine::new(
 			&backend,
@@ -808,6 +848,7 @@ mod tests {
 			"test",
 			&[],
 			Default::default(),
+			&wasm_code,
 		);
 
 		assert_eq!(state_machine.execute(ExecutionStrategy::NativeElseWasm).unwrap(), vec![66]);
@@ -818,6 +859,7 @@ mod tests {
 		let mut consensus_failed = false;
 		let backend = trie_backend::tests::test_trie();
 		let mut overlayed_changes = Default::default();
+		let wasm_code = RuntimeCode::empty();
 
 		let mut state_machine = StateMachine::new(
 			&backend,
@@ -832,6 +874,7 @@ mod tests {
 			"test",
 			&[],
 			Default::default(),
+			&wasm_code,
 		);
 
 		assert!(
@@ -864,6 +907,7 @@ mod tests {
 			&executor,
 			"test",
 			&[],
+			&RuntimeCode::empty(),
 		).unwrap();
 
 		// check proof locally
@@ -874,6 +918,7 @@ mod tests {
 			&executor,
 			"test",
 			&[],
+			&RuntimeCode::empty(),
 		).unwrap();
 
 		// check that both results are correct

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

@@ -18,19 +18,19 @@
 
 use std::sync::Arc;
 use parking_lot::RwLock;
-use codec::{Decode, Encode, Codec};
+use codec::{Decode, Codec};
 use log::debug;
 use hash_db::{Hasher, HashDB, EMPTY_PREFIX, Prefix};
 use sp_trie::{
 	MemoryDB, default_child_trie_root, read_trie_value_with, read_child_trie_value_with,
-	record_all_keys
+	record_all_keys, StorageProof,
 };
 pub use sp_trie::Recorder;
 pub use sp_trie::trie_types::{Layout, TrieError};
 use crate::trie_backend::TrieBackend;
 use crate::trie_backend_essence::{Ephemeral, TrieBackendEssence, TrieBackendStorage};
 use crate::{Error, ExecutionError, Backend};
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use crate::DBValue;
 use sp_core::storage::ChildInfo;
 
@@ -40,82 +40,6 @@ pub struct ProvingBackendRecorder<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Has
 	pub(crate) proof_recorder: &'a mut Recorder<H::Out>,
 }
 
-/// A proof that some set of key-value pairs are included in the storage trie. The proof contains
-/// the storage values so that the partial storage backend can be reconstructed by a verifier that
-/// does not already have access to the key-value pairs.
-///
-/// The proof consists of the set of serialized nodes in the storage trie accessed when looking up
-/// the keys covered by the proof. Verifying the proof requires constructing the partial trie from
-/// the serialized nodes and performing the key lookups.
-#[derive(Debug, PartialEq, Eq, Clone, Encode, Decode)]
-pub struct StorageProof {
-	trie_nodes: Vec<Vec<u8>>,
-}
-
-impl StorageProof {
-	/// Constructs a storage proof from a subset of encoded trie nodes in a storage backend.
-	pub fn new(trie_nodes: Vec<Vec<u8>>) -> Self {
-		StorageProof { trie_nodes }
-	}
-
-	/// Returns a new empty proof.
-	///
-	/// An empty proof is capable of only proving trivial statements (ie. that an empty set of
-	/// key-value pairs exist in storage).
-	pub fn empty() -> Self {
-		StorageProof {
-			trie_nodes: Vec::new(),
-		}
-	}
-
-	/// Returns whether this is an empty proof.
-	pub fn is_empty(&self) -> bool {
-		self.trie_nodes.is_empty()
-	}
-
-	/// Create an iterator over trie nodes constructed from the proof. The nodes are not guaranteed
-	/// to be traversed in any particular order.
-	pub fn iter_nodes(self) -> StorageProofNodeIterator {
-		StorageProofNodeIterator::new(self)
-	}
-}
-
-/// An iterator over trie nodes constructed from a storage proof. The nodes are not guaranteed to
-/// be traversed in any particular order.
-pub struct StorageProofNodeIterator {
-	inner: <Vec<Vec<u8>> as IntoIterator>::IntoIter,
-}
-
-impl StorageProofNodeIterator {
-	fn new(proof: StorageProof) -> Self {
-		StorageProofNodeIterator {
-			inner: proof.trie_nodes.into_iter(),
-		}
-	}
-}
-
-impl Iterator for StorageProofNodeIterator {
-	type Item = Vec<u8>;
-
-	fn next(&mut self) -> Option<Self::Item> {
-		self.inner.next()
-	}
-}
-
-/// Merges multiple storage proofs covering potentially different sets of keys into one proof
-/// covering all keys. The merged proof output may be smaller than the aggregate size of the input
-/// proofs due to deduplication of trie nodes.
-pub fn merge_storage_proofs<I>(proofs: I) -> StorageProof
-	where I: IntoIterator<Item=StorageProof>
-{
-	let trie_nodes = proofs.into_iter()
-		.flat_map(|proof| proof.iter_nodes())
-		.collect::<HashSet<_>>()
-		.into_iter()
-		.collect();
-	StorageProof { trie_nodes }
-}
-
 impl<'a, S, H> ProvingBackendRecorder<'a, S, H>
 	where
 		S: TrieBackendStorage<H>,
@@ -222,7 +146,7 @@ impl<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> ProvingBackend<'a, S, H>
 		let root = essence.root().clone();
 		let recorder = ProofRecorderBackend {
 			backend: essence.backend_storage(),
-			proof_recorder: proof_recorder,
+			proof_recorder,
 		};
 		ProvingBackend(TrieBackend::new(recorder, root))
 	}
@@ -370,7 +294,7 @@ where
 	H: Hasher,
 	H::Out: Codec,
 {
-	let db = create_proof_check_backend_storage(proof);
+	let db = proof.into_memory_db();
 
 	if db.contains(&root, EMPTY_PREFIX) {
 		Ok(TrieBackend::new(db, root))
@@ -379,20 +303,6 @@ where
 	}
 }
 
-/// Create in-memory storage of proof check backend.
-pub fn create_proof_check_backend_storage<H>(
-	proof: StorageProof,
-) -> MemoryDB<H>
-where
-	H: Hasher,
-{
-	let mut db = MemoryDB::default();
-	for item in proof.iter_nodes() {
-		db.insert(EMPTY_PREFIX, &item);
-	}
-	db
-}
-
 #[cfg(test)]
 mod tests {
 	use crate::InMemoryBackend;