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/client/executor/src/integration_tests/mod.rs

@@ -49,6 +49,7 @@ fn call_in_wasm<E: Externalities>(
 	);
 	executor.call_in_wasm(
 		&WASM_BINARY[..],
+		None,
 		function,
 		call_data,
 		ext,
@@ -513,6 +514,7 @@ fn should_trap_when_heap_exhausted(wasm_method: WasmExecutionMethod) {
 	);
 	executor.call_in_wasm(
 		&WASM_BINARY[..],
+		None,
 		"test_exhaust_heap",
 		&[0],
 		&mut ext.ext(),

substrate/client/executor/src/lib.rs

@@ -52,8 +52,12 @@ pub trait RuntimeInfo {
 	/// Native runtime information.
 	fn native_version(&self) -> &NativeVersion;
 
-	/// Extract RuntimeVersion of given :code block
-	fn runtime_version(&self, ext: &mut dyn Externalities) -> error::Result<RuntimeVersion>;
+	/// Extract [`RuntimeVersion`](sp_version::RuntimeVersion) of the given `runtime_code`.
+	fn runtime_version(
+		&self,
+		ext: &mut dyn Externalities,
+		runtime_code: &sp_core::traits::RuntimeCode,
+	) -> error::Result<RuntimeVersion>;
 }
 
 #[cfg(test)]
@@ -77,6 +81,7 @@ mod tests {
 		);
 		let res = executor.call_in_wasm(
 			&WASM_BINARY[..],
+			None,
 			"test_empty_return",
 			&[],
 			&mut ext,

substrate/client/executor/src/native_executor.rs

@@ -16,11 +16,11 @@
 
 use crate::{
 	RuntimeInfo, error::{Error, Result},
-	wasm_runtime::{RuntimeCache, WasmExecutionMethod, CodeSource},
+	wasm_runtime::{RuntimeCache, WasmExecutionMethod},
 };
 use sp_version::{NativeVersion, RuntimeVersion};
 use codec::{Decode, Encode};
-use sp_core::{NativeOrEncoded, traits::{CodeExecutor, Externalities}};
+use sp_core::{NativeOrEncoded, traits::{CodeExecutor, Externalities, RuntimeCode}};
 use log::trace;
 use std::{result, panic::{UnwindSafe, AssertUnwindSafe}, sync::Arc};
 use sp_wasm_interface::{HostFunctions, Function};
@@ -111,7 +111,7 @@ impl WasmExecutor {
 		}
 	}
 
-	/// Execute the given closure `f` with the latest runtime (based on the `CODE` key in `ext`).
+	/// Execute the given closure `f` with the latest runtime (based on `runtime_code`).
 	///
 	/// The closure `f` is expected to return `Err(_)` when there happened a `panic!` in native code
 	/// while executing the runtime in Wasm. If a `panic!` occurred, the runtime is invalidated to
@@ -124,9 +124,9 @@ impl WasmExecutor {
 	/// runtime is invalidated on any `panic!` to prevent a poisoned state. `ext` is already
 	/// implicitly handled as unwind safe, as we store it in a global variable while executing the
 	/// native runtime.
-	fn with_instance<'c, R, F>(
+	fn with_instance<R, F>(
 		&self,
-		code: CodeSource<'c>,
+		runtime_code: &RuntimeCode,
 		ext: &mut dyn Externalities,
 		f: F,
 	) -> Result<R>
@@ -137,7 +137,7 @@ impl WasmExecutor {
 		) -> Result<Result<R>>,
 	{
 		match self.cache.with_instance(
-			code,
+			runtime_code,
 			ext,
 			self.method,
 			self.default_heap_pages,
@@ -158,17 +158,48 @@ impl WasmExecutor {
 impl sp_core::traits::CallInWasm for WasmExecutor {
 	fn call_in_wasm(
 		&self,
-		wasm_blob: &[u8],
+		wasm_code: &[u8],
+		code_hash: Option<Vec<u8>>,
 		method: &str,
 		call_data: &[u8],
 		ext: &mut dyn Externalities,
 	) -> std::result::Result<Vec<u8>, String> {
-		self.with_instance(CodeSource::Custom(wasm_blob), ext, |instance, _, mut ext| {
+		if let Some(hash) = code_hash {
+			let code = RuntimeCode {
+				code_fetcher: &sp_core::traits::WrappedRuntimeCode(wasm_code.into()),
+				hash,
+				heap_pages: None,
+			};
+
+			self.with_instance(&code, ext, |instance, _, mut ext| {
+				with_externalities_safe(
+					&mut **ext,
+					move || instance.call(method, call_data),
+				)
+			}).map_err(|e| e.to_string())
+		} else {
+			let module = crate::wasm_runtime::create_wasm_runtime_with_code(
+				self.method,
+				self.default_heap_pages,
+				&wasm_code,
+				self.host_functions.to_vec(),
+				self.allow_missing_func_imports,
+			)
+				.map_err(|e| format!("Failed to create module: {:?}", e))?;
+
+			let instance = module.new_instance()
+				.map_err(|e| format!("Failed to create instance: {:?}", e))?;
+
+			let instance = AssertUnwindSafe(instance);
+			let mut ext = AssertUnwindSafe(ext);
+
 			with_externalities_safe(
 				&mut **ext,
 				move || instance.call(method, call_data),
 			)
-		}).map_err(|e| e.to_string())
+			.and_then(|r| r)
+			.map_err(|e| e.to_string())
+		}
 	}
 }
 
@@ -220,8 +251,11 @@ impl<D: NativeExecutionDispatch> RuntimeInfo for NativeExecutor<D> {
 	fn runtime_version(
 		&self,
 		ext: &mut dyn Externalities,
+		runtime_code: &RuntimeCode,
 	) -> Result<RuntimeVersion> {
-		self.wasm.with_instance(CodeSource::Externalities, ext,
+		self.wasm.with_instance(
+			runtime_code,
+			ext,
 			|_instance, version, _ext|
 				Ok(version.cloned().ok_or_else(|| Error::ApiError("Unknown version".into())))
 		)
@@ -237,6 +271,7 @@ impl<D: NativeExecutionDispatch + 'static> CodeExecutor for NativeExecutor<D> {
 	>(
 		&self,
 		ext: &mut dyn Externalities,
+		runtime_code: &RuntimeCode,
 		method: &str,
 		data: &[u8],
 		use_native: bool,
@@ -244,7 +279,7 @@ impl<D: NativeExecutionDispatch + 'static> CodeExecutor for NativeExecutor<D> {
 	) -> (Result<NativeOrEncoded<R>>, bool) {
 		let mut used_native = false;
 		let result = self.wasm.with_instance(
-			CodeSource::Externalities,
+			runtime_code,
 			ext,
 			|instance, onchain_version, mut ext| {
 				let onchain_version = onchain_version.ok_or_else(
@@ -324,11 +359,12 @@ impl<D: NativeExecutionDispatch> sp_core::traits::CallInWasm for NativeExecutor<
 	fn call_in_wasm(
 		&self,
 		wasm_blob: &[u8],
+		code_hash: Option<Vec<u8>>,
 		method: &str,
 		call_data: &[u8],
 		ext: &mut dyn Externalities,
 	) -> std::result::Result<Vec<u8>, String> {
-		sp_core::traits::CallInWasm::call_in_wasm(&self.wasm, wasm_blob, method, call_data, ext)
+		self.wasm.call_in_wasm(wasm_blob, code_hash, method, call_data, ext)
 	}
 }
 

substrate/client/executor/src/wasm_runtime.rs

@@ -20,11 +20,10 @@
 //! components of the runtime that are expensive to initialize.
 
 use std::sync::Arc;
-use std::borrow::Cow;
 use crate::error::{Error, WasmError};
 use parking_lot::{Mutex, RwLock};
 use codec::Decode;
-use sp_core::{storage::well_known_keys, traits::Externalities};
+use sp_core::traits::{Externalities, RuntimeCode, FetchRuntimeCode};
 use sp_version::RuntimeVersion;
 use std::panic::AssertUnwindSafe;
 use sc_executor_common::wasm_runtime::{WasmModule, WasmInstance};
@@ -41,14 +40,6 @@ pub enum WasmExecutionMethod {
 	Compiled,
 }
 
-/// Executoed code origin.
-pub enum CodeSource<'a> {
-	/// Take code from storage,
-	Externalities,
-	/// Use provided code,
-	Custom(&'a [u8]),
-}
-
 /// A Wasm runtime object along with its cached runtime version.
 struct VersionedRuntime {
 	/// Runtime code hash.
@@ -102,8 +93,7 @@ impl RuntimeCache {
 	///
 	/// `code` - Provides external code or tells the executor to fetch it from storage.
 	///
-	/// `ext` - Externalities to use for the runtime. This is used for setting
-	/// up an initial runtime instance.
+	/// `runtime_code` - The runtime wasm code used setup the runtime.
 	///
 	/// `default_heap_pages` - Number of 64KB pages to allocate for Wasm execution.
 	///
@@ -124,7 +114,7 @@ impl RuntimeCache {
 	/// identifier `memory` can be found in the runtime.
 	pub fn with_instance<'c, R, F>(
 		&self,
-		code: CodeSource<'c>,
+		runtime_code: &'c RuntimeCode<'c>,
 		ext: &mut dyn Externalities,
 		wasm_method: WasmExecutionMethod,
 		default_heap_pages: u64,
@@ -138,28 +128,14 @@ impl RuntimeCache {
 			&mut dyn Externalities)
 		-> Result<R, Error>,
 	{
-		let (code_hash, heap_pages) = match &code {
-			CodeSource::Externalities => {
-				(
-					ext
-						.original_storage_hash(well_known_keys::CODE)
-						.ok_or(Error::InvalidCode("`CODE` not found in storage.".into()))?,
-					ext
-						.storage(well_known_keys::HEAP_PAGES)
-						.and_then(|pages| u64::decode(&mut &pages[..]).ok())
-						.unwrap_or(default_heap_pages),
-				)
-			},
-			CodeSource::Custom(code) => {
-				(sp_core::blake2_256(code).to_vec(), default_heap_pages)
-			}
-		};
+		let code_hash = &runtime_code.hash;
+		let heap_pages = runtime_code.heap_pages.unwrap_or(default_heap_pages);
 
 		let mut runtimes = self.runtimes.lock(); // this must be released prior to calling f
 		let pos = runtimes.iter().position(|r| r.as_ref().map_or(
 			false,
 			|r| r.wasm_method == wasm_method &&
-				r.code_hash == code_hash &&
+				r.code_hash == *code_hash &&
 				r.heap_pages == heap_pages
 		));
 
@@ -168,19 +144,11 @@ impl RuntimeCache {
 				.clone()
 				.expect("`position` only returns `Some` for entries that are `Some`"),
 			None =>  {
-				let code = match code {
-					CodeSource::Externalities => {
-						Cow::Owned(ext.original_storage(well_known_keys::CODE)
-							.ok_or(WasmError::CodeNotFound)?)
-					}
-					CodeSource::Custom(code) => {
-						Cow::Borrowed(code)
-					}
-				};
+				let code = runtime_code.fetch_runtime_code().ok_or(WasmError::CodeNotFound)?;
 
 				let result = create_versioned_wasm_runtime(
 					&code,
-					code_hash,
+					code_hash.clone(),
 					ext,
 					wasm_method,
 					heap_pages,