contracts: Replace sp-sandbox and wasmi-validation by newest wasmi (#12501)

* Replace sp-sandbox and wasmi-validation by just wasmi

* ".git/.scripts/bench-bot.sh" pallet dev pallet_contracts

* Re-check original code on re-instrumentation

* Fix clippy

* ".git/.scripts/bench-bot.sh" pallet dev pallet_contracts

* Apply suggestions from code review

Co-authored-by: Robin Freyler <robin.freyler@gmail.com>

* Replace wasmi by ::wasmi

* Bump wasmi to 0.20

* Add explanation for `unreachable`

* Change proof

* Fixup master merge

* ".git/.scripts/bench-bot.sh" pallet dev pallet_contracts

* Fixup naming inconsistencies introduced by reentrancy PR

* Fix `scan_imports` docs

* Apply suggestions from code review

Co-authored-by: Sasha Gryaznov <hi@agryaznov.com>

* Fixup suggestions

* Remove unnecessary &mut

* Fix test

* ".git/.scripts/bench-bot.sh" pallet dev pallet_contracts

* Fix benchmark merge fail

* ".git/.scripts/bench-bot.sh" pallet dev pallet_contracts

* Fix docs as suggested by code review

* Improve docs for `CodeRejected`

* Apply suggestions from code review

Co-authored-by: Sasha Gryaznov <hi@agryaznov.com>

* Fix logic bug when setting `deterministic_only`

* Don't panic when module fails to compile

* Apply suggestions from code review

Co-authored-by: Robin Freyler <robin.freyler@gmail.com>

Co-authored-by: command-bot <>
Co-authored-by: Robin Freyler <robin.freyler@gmail.com>
Co-authored-by: Sasha Gryaznov <hi@agryaznov.com>

substrate/frame/contracts/src/wasm/code_cache.rs

@@ -192,7 +192,10 @@ where
 pub fn reinstrument<T: Config>(
 	prefab_module: &mut PrefabWasmModule<T>,
 	schedule: &Schedule<T>,
-) -> Result<u32, DispatchError> {
+) -> Result<u32, DispatchError>
+where
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>,
+{
 	let original_code =
 		<PristineCode<T>>::get(&prefab_module.code_hash).ok_or(Error::<T>::CodeNotFound)?;
 	let original_code_len = original_code.len();
@@ -201,9 +204,12 @@ pub fn reinstrument<T: Config>(
 	// as the contract is already deployed and every change in size would be the result
 	// of changes in the instrumentation algorithm controlled by the chain authors.
 	prefab_module.code = WeakBoundedVec::force_from(
-		prepare::reinstrument_contract::<T>(&original_code, schedule, prefab_module.determinism)
-			.map_err(|_| <Error<T>>::CodeRejected)?,
-		Some("Contract exceeds limit after re-instrumentation."),
+		prepare::reinstrument::<super::runtime::Env, T>(
+			&original_code,
+			schedule,
+			prefab_module.determinism,
+		)?,
+		Some("Contract exceeds size limit after re-instrumentation."),
 	);
 	prefab_module.instruction_weights_version = schedule.instruction_weights.version;
 	<CodeStorage<T>>::insert(&prefab_module.code_hash, &*prefab_module);

substrate/frame/contracts/src/wasm/env_def/mod.rs

@@ -1,83 +0,0 @@
-// This file is part of Substrate.
-
-// Copyright (C) 2018-2022 Parity Technologies (UK) Ltd.
-// SPDX-License-Identifier: Apache-2.0
-
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// 	http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use super::Runtime;
-use crate::exec::Ext;
-
-use sp_sandbox::Value;
-use wasm_instrument::parity_wasm::elements::{FunctionType, ValueType};
-
-pub trait ConvertibleToWasm: Sized {
-	const VALUE_TYPE: ValueType;
-	type NativeType;
-	fn to_typed_value(self) -> Value;
-	fn from_typed_value(_: Value) -> Option<Self>;
-}
-impl ConvertibleToWasm for i32 {
-	const VALUE_TYPE: ValueType = ValueType::I32;
-	type NativeType = i32;
-	fn to_typed_value(self) -> Value {
-		Value::I32(self)
-	}
-	fn from_typed_value(v: Value) -> Option<Self> {
-		v.as_i32()
-	}
-}
-impl ConvertibleToWasm for u32 {
-	const VALUE_TYPE: ValueType = ValueType::I32;
-	type NativeType = u32;
-	fn to_typed_value(self) -> Value {
-		Value::I32(self as i32)
-	}
-	fn from_typed_value(v: Value) -> Option<Self> {
-		match v {
-			Value::I32(v) => Some(v as u32),
-			_ => None,
-		}
-	}
-}
-impl ConvertibleToWasm for u64 {
-	const VALUE_TYPE: ValueType = ValueType::I64;
-	type NativeType = u64;
-	fn to_typed_value(self) -> Value {
-		Value::I64(self as i64)
-	}
-	fn from_typed_value(v: Value) -> Option<Self> {
-		match v {
-			Value::I64(v) => Some(v as u64),
-			_ => None,
-		}
-	}
-}
-
-pub type HostFunc<E> = fn(
-	&mut Runtime<E>,
-	&[sp_sandbox::Value],
-) -> Result<sp_sandbox::ReturnValue, sp_sandbox::HostError>;
-
-pub trait FunctionImplProvider<E: Ext> {
-	fn impls<F: FnMut(&[u8], &[u8], HostFunc<E>)>(f: &mut F);
-}
-
-/// This trait can be used to check whether the host environment can satisfy
-/// a requested function import.
-pub trait ImportSatisfyCheck {
-	/// Returns `true` if the host environment contains a function with
-	/// the specified name and its type matches to the given type, or `false`
-	/// otherwise.
-	fn can_satisfy(module: &[u8], name: &[u8], func_type: &FunctionType) -> bool;
-}

substrate/frame/contracts/src/wasm/mod.rs

@@ -18,19 +18,19 @@
 //! This module provides a means for executing contracts
 //! represented in wasm.
 
-#[macro_use]
-mod env_def;
 mod code_cache;
 mod prepare;
 mod runtime;
 
 #[cfg(feature = "runtime-benchmarks")]
 pub use crate::wasm::code_cache::reinstrument;
-pub use crate::wasm::runtime::{CallFlags, ReturnCode, Runtime, RuntimeCosts};
+pub use crate::wasm::{
+	prepare::TryInstantiate,
+	runtime::{CallFlags, Environment, ReturnCode, Runtime, RuntimeCosts},
+};
 use crate::{
 	exec::{ExecResult, Executable, ExportedFunction, Ext},
 	gas::GasMeter,
-	wasm::env_def::FunctionImplProvider,
 	AccountIdOf, BalanceOf, CodeHash, CodeStorage, CodeVec, Config, Error, RelaxedCodeVec,
 	Schedule,
 };
@@ -38,10 +38,12 @@ use codec::{Decode, Encode, MaxEncodedLen};
 use frame_support::dispatch::{DispatchError, DispatchResult};
 use sp_core::crypto::UncheckedFrom;
 use sp_runtime::RuntimeDebug;
-use sp_sandbox::{SandboxEnvironmentBuilder, SandboxInstance, SandboxMemory};
 use sp_std::prelude::*;
 #[cfg(test)]
 pub use tests::MockExt;
+use wasmi::{
+	Config as WasmiConfig, Engine, Instance, Linker, Memory, MemoryType, Module, StackLimits, Store,
+};
 
 /// A prepared wasm module ready for execution.
 ///
@@ -151,12 +153,14 @@ where
 		schedule: &Schedule<T>,
 		owner: AccountIdOf<T>,
 		determinism: Determinism,
+		try_instantiate: TryInstantiate,
 	) -> Result<Self, (DispatchError, &'static str)> {
-		let module = prepare::prepare_contract(
+		let module = prepare::prepare::<runtime::Env, T>(
 			original_code.try_into().map_err(|_| (<Error<T>>::CodeTooLarge.into(), ""))?,
 			schedule,
 			owner,
 			determinism,
+			try_instantiate,
 		)?;
 		Ok(module)
 	}
@@ -189,6 +193,44 @@ where
 		}
 	}
 
+	/// Creates and returns an instance of the supplied code.
+	///
+	/// This is either used for later executing a contract or for validation of a contract.
+	/// When validating we pass `()` as `host_state`. Please note that such a dummy instance must
+	/// **never** be called/executed since it will panic the executor.
+	pub fn instantiate<E, H>(
+		code: &[u8],
+		host_state: H,
+		memory: (u32, u32),
+		stack_limits: StackLimits,
+	) -> Result<(Store<H>, Memory, Instance), wasmi::Error>
+	where
+		E: Environment<H>,
+	{
+		let mut config = WasmiConfig::default();
+		config
+			.set_stack_limits(stack_limits)
+			.wasm_multi_value(false)
+			.wasm_mutable_global(false)
+			.wasm_sign_extension(false)
+			.wasm_saturating_float_to_int(false);
+		let engine = Engine::new(&config);
+		let module = Module::new(&engine, code)?;
+		let mut store = Store::new(&engine, host_state);
+		let mut linker = Linker::new();
+		E::define(&mut store, &mut linker)?;
+		let memory = Memory::new(&mut store, MemoryType::new(memory.0, Some(memory.1))?).expect(
+			"The limits defined in our `Schedule` limit the amount of memory well below u32::MAX; qed",
+		);
+		linker
+			.define("env", "memory", memory)
+			.expect("We just created the linker. It has no define with this name attached; qed");
+
+		let instance = linker.instantiate(&mut store, &module)?.ensure_no_start(&mut store)?;
+
+		Ok((store, memory, instance))
+	}
+
 	/// Create and store the module without checking nor instrumenting the passed code.
 	///
 	/// # Note
@@ -201,7 +243,7 @@ where
 		schedule: &Schedule<T>,
 		owner: T::AccountId,
 	) -> DispatchResult {
-		let executable = prepare::benchmarking::prepare_contract(original_code, schedule, owner)
+		let executable = prepare::benchmarking::prepare(original_code, schedule, owner)
 			.map_err::<DispatchError, _>(Into::into)?;
 		code_cache::store(executable, false)
 	}
@@ -247,36 +289,35 @@ where
 		function: &ExportedFunction,
 		input_data: Vec<u8>,
 	) -> ExecResult {
-		let memory = sp_sandbox::default_executor::Memory::new(self.initial, Some(self.maximum))
-			.unwrap_or_else(|_| {
-				// unlike `.expect`, explicit panic preserves the source location.
-				// Needed as we can't use `RUST_BACKTRACE` in here.
-				panic!(
-					"exec.prefab_module.initial can't be greater than exec.prefab_module.maximum;
-						thus Memory::new must not fail;
-						qed"
-				)
-			});
+		let runtime = Runtime::new(ext, input_data);
+		let (mut store, memory, instance) = Self::instantiate::<crate::wasm::runtime::Env, _>(
+			self.code.as_slice(),
+			runtime,
+			(self.initial, self.maximum),
+			StackLimits::default(),
+		)
+		.map_err(|msg| {
+			log::debug!(target: "runtime::contracts", "failed to instantiate code: {}", msg);
+			Error::<T>::CodeRejected
+		})?;
+		store.state_mut().set_memory(memory);
 
-		let mut imports = sp_sandbox::default_executor::EnvironmentDefinitionBuilder::new();
-		imports.add_memory(self::prepare::IMPORT_MODULE_MEMORY, "memory", memory.clone());
-		runtime::Env::impls(&mut |module, name, func_ptr| {
-			imports.add_host_func(module, name, func_ptr);
-		});
+		let exported_func = instance
+			.get_export(&store, function.identifier())
+			.and_then(|export| export.into_func())
+			.ok_or_else(|| {
+				log::error!(target: "runtime::contracts", "failed to find entry point");
+				Error::<T>::CodeRejected
+			})?;
 
 		// We store before executing so that the code hash is available in the constructor.
-		let code = self.code.clone();
 		if let &ExportedFunction::Constructor = function {
 			code_cache::store(self, true)?;
 		}
 
-		// Instantiate the instance from the instrumented module code and invoke the contract
-		// entrypoint.
-		let mut runtime = Runtime::new(ext, input_data, memory);
-		let result = sp_sandbox::default_executor::Instance::new(&code, &imports, &mut runtime)
-			.and_then(|mut instance| instance.invoke(function.identifier(), &[], &mut runtime));
+		let result = exported_func.call(&mut store, &[], &mut []);
 
-		runtime.to_execution_result(result)
+		store.into_state().to_execution_result(result)
 	}
 
 	fn code_hash(&self) -> &CodeHash<T> {
@@ -307,7 +348,7 @@ mod tests {
 	};
 	use assert_matches::assert_matches;
 	use frame_support::{
-		assert_ok,
+		assert_err, assert_ok,
 		dispatch::DispatchResultWithPostInfo,
 		weights::{OldWeight, Weight},
 	};
@@ -578,7 +619,7 @@ mod tests {
 		fn ecdsa_to_eth_address(&self, _pk: &[u8; 33]) -> Result<[u8; 20], ()> {
 			Ok([2u8; 20])
 		}
-		fn reentrant_count(&self) -> u32 {
+		fn reentrance_count(&self) -> u32 {
 			12
 		}
 		fn account_reentrance_count(&self, _account_id: &AccountIdOf<Self::T>) -> u32 {
@@ -594,8 +635,9 @@ mod tests {
 			&schedule,
 			ALICE,
 			Determinism::Deterministic,
+			TryInstantiate::Skip,
 		)
-		.unwrap();
+		.map_err(|err| err.0)?;
 		executable.execute(ext.borrow_mut(), &ExportedFunction::Call, input_data)
 	}
 
@@ -788,10 +830,7 @@ mod tests {
 "#;
 		let mut mock_ext = MockExt::default();
 		let input = vec![0xff, 0x2a, 0x99, 0x88];
-		frame_support::assert_err!(
-			execute(CODE, input.clone(), &mut mock_ext),
-			<Error<Test>>::InputForwarded,
-		);
+		assert_err!(execute(CODE, input.clone(), &mut mock_ext), <Error<Test>>::InputForwarded,);
 
 		assert_eq!(
 			&mock_ext.calls,
@@ -1584,35 +1623,32 @@ mod tests {
 		assert_ok!(execute(CODE_VALUE_TRANSFERRED, vec![], MockExt::default()));
 	}
 
-	const CODE_RETURN_FROM_START_FN: &str = r#"
+	const START_FN_ILLEGAL: &str = r#"
 (module
 	(import "seal0" "seal_return" (func $seal_return (param i32 i32 i32)))
 	(import "env" "memory" (memory 1 1))
 
 	(start $start)
 	(func $start
-		(call $seal_return
-			(i32.const 0)
-			(i32.const 8)
-			(i32.const 4)
-		)
 		(unreachable)
 	)
 
 	(func (export "call")
 		(unreachable)
 	)
-	(func (export "deploy"))
+
+	(func (export "deploy")
+		(unreachable)
+	)
 
 	(data (i32.const 8) "\01\02\03\04")
 )
 "#;
 
 	#[test]
-	fn return_from_start_fn() {
-		let output = execute(CODE_RETURN_FROM_START_FN, vec![], MockExt::default()).unwrap();
-
-		assert_eq!(output, ExecReturnValue { flags: ReturnFlags::empty(), data: vec![1, 2, 3, 4] });
+	fn start_fn_illegal() {
+		let output = execute(START_FN_ILLEGAL, vec![], MockExt::default());
+		assert_err!(output, <Error<Test>>::CodeRejected,);
 	}
 
 	const CODE_TIMESTAMP_NOW: &str = r#"
@@ -2859,10 +2895,10 @@ mod tests {
 
 	#[test]
 	#[cfg(feature = "unstable-interface")]
-	fn reentrant_count_works() {
+	fn reentrance_count_works() {
 		const CODE: &str = r#"
 (module
-	(import "__unstable__" "reentrant_count" (func $reentrant_count (result i32)))
+	(import "__unstable__" "reentrance_count" (func $reentrance_count (result i32)))
 	(import "env" "memory" (memory 1 1))
 	(func $assert (param i32)
 		(block $ok
@@ -2875,7 +2911,7 @@ mod tests {
 	(func (export "call")
 		(local $return_val i32)
 		(set_local $return_val
-			(call $reentrant_count)
+			(call $reentrance_count)
 		)
 		(call $assert
 			(i32.eq (get_local $return_val) (i32.const 12))

substrate/frame/contracts/src/wasm/prepare.rs

@@ -22,20 +22,38 @@
 use crate::{
 	chain_extension::ChainExtension,
 	storage::meter::Diff,
-	wasm::{env_def::ImportSatisfyCheck, Determinism, OwnerInfo, PrefabWasmModule},
+	wasm::{Determinism, Environment, OwnerInfo, PrefabWasmModule},
 	AccountIdOf, CodeVec, Config, Error, Schedule,
 };
 use codec::{Encode, MaxEncodedLen};
+use sp_core::crypto::UncheckedFrom;
 use sp_runtime::{traits::Hash, DispatchError};
 use sp_std::prelude::*;
 use wasm_instrument::parity_wasm::elements::{
 	self, External, Internal, MemoryType, Type, ValueType,
 };
+use wasmi::StackLimits;
+use wasmparser::{Validator, WasmFeatures};
 
 /// Imported memory must be located inside this module. The reason for hardcoding is that current
 /// compiler toolchains might not support specifying other modules than "env" for memory imports.
 pub const IMPORT_MODULE_MEMORY: &str = "env";
 
+/// Determines whether a module should be instantiated during preparation.
+pub enum TryInstantiate {
+	/// Do the instantiation to make sure that the module is valid.
+	///
+	/// This should be used if a module is only uploaded but not executed. We need
+	/// to make sure that it can be actually instantiated.
+	Instantiate,
+	/// Skip the instantiation during preparation.
+	///
+	/// This makes sense when the preparation takes place as part of an instantation. Then
+	/// this instantiation would fail the whole transaction and an extra check is not
+	/// necessary.
+	Skip,
+}
+
 struct ContractModule<'a, T: Config> {
 	/// A deserialized module. The module is valid (this is Guaranteed by `new` method).
 	module: elements::Module,
@@ -48,14 +66,9 @@ impl<'a, T: Config> ContractModule<'a, T> {
 	/// Returns `Err` if the `original_code` couldn't be decoded or
 	/// if it contains an invalid module.
 	fn new(original_code: &[u8], schedule: &'a Schedule<T>) -> Result<Self, &'static str> {
-		use wasmi_validation::{validate_module, PlainValidator};
-
 		let module =
 			elements::deserialize_buffer(original_code).map_err(|_| "Can't decode wasm code")?;
 
-		// Make sure that the module is valid.
-		validate_module::<PlainValidator>(&module, ()).map_err(|_| "Module is not valid")?;
-
 		// Return a `ContractModule` instance with
 		// __valid__ module.
 		Ok(ContractModule { module, schedule })
@@ -279,27 +292,33 @@ impl<'a, T: Config> ContractModule<'a, T> {
 
 	/// Scan an import section if any.
 	///
-	/// This accomplishes two tasks:
+	/// This makes sure that the import section looks as we expect it from a contract
+	/// and enforces and returns the memory type declared by the contract if any.
 	///
-	/// - checks any imported function against defined host functions set, incl. their signatures.
-	/// - if there is a memory import, returns it's descriptor
 	/// `import_fn_banlist`: list of function names that are disallowed to be imported
-	fn scan_imports<C: ImportSatisfyCheck>(
+	fn scan_imports(
 		&self,
 		import_fn_banlist: &[&[u8]],
 	) -> Result<Option<&MemoryType>, &'static str> {
 		let module = &self.module;
-
-		let types = module.type_section().map(|ts| ts.types()).unwrap_or(&[]);
 		let import_entries = module.import_section().map(|is| is.entries()).unwrap_or(&[]);
-
 		let mut imported_mem_type = None;
 
 		for import in import_entries {
-			let type_idx = match *import.external() {
+			match *import.external() {
 				External::Table(_) => return Err("Cannot import tables"),
 				External::Global(_) => return Err("Cannot import globals"),
-				External::Function(ref type_idx) => type_idx,
+				External::Function(_) => {
+					if !T::ChainExtension::enabled() &&
+						import.field().as_bytes() == b"seal_call_chain_extension"
+					{
+						return Err("module uses chain extensions but chain extensions are disabled")
+					}
+
+					if import_fn_banlist.iter().any(|f| import.field().as_bytes() == *f) {
+						return Err("module imports a banned function")
+					}
+				},
 				External::Memory(ref memory_type) => {
 					if import.module() != IMPORT_MODULE_MEMORY {
 						return Err("Invalid module for imported memory")
@@ -313,22 +332,6 @@ impl<'a, T: Config> ContractModule<'a, T> {
 					imported_mem_type = Some(memory_type);
 					continue
 				},
-			};
-
-			let Type::Function(ref func_ty) = types
-				.get(*type_idx as usize)
-				.ok_or("validation: import entry points to a non-existent type")?;
-
-			if !T::ChainExtension::enabled() &&
-				import.field().as_bytes() == b"seal_call_chain_extension"
-			{
-				return Err("module uses chain extensions but chain extensions are disabled")
-			}
-
-			if import_fn_banlist.iter().any(|f| import.field().as_bytes() == *f) ||
-				!C::can_satisfy(import.module().as_bytes(), import.field().as_bytes(), func_ty)
-			{
-				return Err("module imports a non-existent function")
 			}
 		}
 		Ok(imported_mem_type)
@@ -366,12 +369,54 @@ fn get_memory_limits<T: Config>(
 	}
 }
 
-fn check_and_instrument<C: ImportSatisfyCheck, T: Config>(
+/// Check and instrument the given `original_code`.
+///
+/// On success it returns the instrumented versions together with its `(initial, maximum)`
+/// error requirement. The memory requirement was also validated against the `schedule`.
+fn instrument<E, T>(
 	original_code: &[u8],
 	schedule: &Schedule<T>,
 	determinism: Determinism,
-) -> Result<(Vec<u8>, (u32, u32)), &'static str> {
-	let result = (|| {
+	try_instantiate: TryInstantiate,
+) -> Result<(Vec<u8>, (u32, u32)), (DispatchError, &'static str)>
+where
+	E: Environment<()>,
+	T: Config,
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>,
+{
+	// Do not enable any features here. Any additional feature needs to be carefully
+	// checked for potential security issues. For example, enabling multi value could lead
+	// to a DoS vector: It breaks our assumption that branch instructions are of constant time.
+	// Depending on the implementation they can linearly depend on the amount of values returned
+	// from a block.
+	Validator::new_with_features(WasmFeatures {
+		relaxed_simd: false,
+		threads: false,
+		tail_call: false,
+		multi_memory: false,
+		exceptions: false,
+		memory64: false,
+		extended_const: false,
+		component_model: false,
+		// This is not our only defense: We check for float types later in the preparation
+		// process. Additionally, all instructions explictily  need to have weights assigned
+		// or the deployment will fail. We have none assigned for float instructions.
+		deterministic_only: matches!(determinism, Determinism::Deterministic),
+		mutable_global: false,
+		saturating_float_to_int: false,
+		sign_extension: false,
+		bulk_memory: false,
+		multi_value: false,
+		reference_types: false,
+		simd: false,
+	})
+	.validate_all(original_code)
+	.map_err(|err| {
+		log::debug!(target: "runtime::contracts", "{}", err);
+		(Error::<T>::CodeRejected.into(), "validation of new code failed")
+	})?;
+
+	let (code, (initial, maximum)) = (|| {
 		let contract_module = ContractModule::new(original_code, schedule)?;
 		contract_module.scan_exports()?;
 		contract_module.ensure_no_internal_memory()?;
@@ -387,7 +432,7 @@ fn check_and_instrument<C: ImportSatisfyCheck, T: Config>(
 		// We disallow importing `gas` function here since it is treated as implementation detail.
 		let disallowed_imports = [b"gas".as_ref()];
 		let memory_limits =
-			get_memory_limits(contract_module.scan_imports::<C>(&disallowed_imports)?, schedule)?;
+			get_memory_limits(contract_module.scan_imports(&disallowed_imports)?, schedule)?;
 
 		let code = contract_module
 			.inject_gas_metering(determinism)?
@@ -395,24 +440,56 @@ fn check_and_instrument<C: ImportSatisfyCheck, T: Config>(
 			.into_wasm_code()?;
 
 		Ok((code, memory_limits))
-	})();
+	})()
+	.map_err(|msg: &str| {
+		log::debug!(target: "runtime::contracts", "new code rejected: {}", msg);
+		(Error::<T>::CodeRejected.into(), msg)
+	})?;
 
-	if let Err(msg) = &result {
-		log::debug!(target: "runtime::contracts", "CodeRejected: {}", msg);
+	// This will make sure that the module can be actually run within wasmi:
+	//
+	// - Doesn't use any unknown imports.
+	// - Doesn't explode the wasmi bytecode generation.
+	if matches!(try_instantiate, TryInstantiate::Instantiate) {
+		// We don't actually ever run any code so we can get away with a minimal stack which
+		// reduces the amount of memory that needs to be zeroed.
+		let stack_limits = StackLimits::new(1, 1, 0).expect("initial <= max; qed");
+		PrefabWasmModule::<T>::instantiate::<E, _>(&code, (), (initial, maximum), stack_limits)
+			.map_err(|err| {
+				log::debug!(target: "runtime::contracts", "{}", err);
+				(Error::<T>::CodeRejected.into(), "new code rejected after instrumentation")
+			})?;
 	}
 
-	result
+	Ok((code, (initial, maximum)))
 }
 
-fn do_preparation<C: ImportSatisfyCheck, T: Config>(
+/// Loads the given module given in `original_code`, performs some checks on it and
+/// does some preprocessing.
+///
+/// The checks are:
+///
+/// - the provided code is a valid wasm module
+/// - the module doesn't define an internal memory instance
+/// - imported memory (if any) doesn't reserve more memory than permitted by the `schedule`
+/// - all imported functions from the external environment matches defined by `env` module
+///
+/// The preprocessing includes injecting code for gas metering and metering the height of stack.
+pub fn prepare<E, T>(
 	original_code: CodeVec<T>,
 	schedule: &Schedule<T>,
 	owner: AccountIdOf<T>,
 	determinism: Determinism,
-) -> Result<PrefabWasmModule<T>, (DispatchError, &'static str)> {
+	try_instantiate: TryInstantiate,
+) -> Result<PrefabWasmModule<T>, (DispatchError, &'static str)>
+where
+	E: Environment<()>,
+	T: Config,
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>,
+{
 	let (code, (initial, maximum)) =
-		check_and_instrument::<C, T>(original_code.as_ref(), schedule, determinism)
-			.map_err(|msg| (<Error<T>>::CodeRejected.into(), msg))?;
+		instrument::<E, T>(original_code.as_ref(), schedule, determinism, try_instantiate)?;
+
 	let original_code_len = original_code.len();
 
 	let mut module = PrefabWasmModule {
@@ -420,10 +497,10 @@ fn do_preparation<C: ImportSatisfyCheck, T: Config>(
 		initial,
 		maximum,
 		code: code.try_into().map_err(|_| (<Error<T>>::CodeTooLarge.into(), ""))?,
-		determinism,
 		code_hash: T::Hashing::hash(&original_code),
 		original_code: Some(original_code),
 		owner_info: None,
+		determinism,
 	};
 
 	// We need to add the sizes of the `#[codec(skip)]` fields which are stored in different
@@ -441,37 +518,28 @@ fn do_preparation<C: ImportSatisfyCheck, T: Config>(
 	Ok(module)
 }
 
-/// Loads the given module given in `original_code`, performs some checks on it and
-/// does some preprocessing.
+/// Same as [`prepare`] but without constructing a new module.
 ///
-/// The checks are:
-///
-/// - provided code is a valid wasm module.
-/// - the module doesn't define an internal memory instance,
-/// - imported memory (if any) doesn't reserve more memory than permitted by the `schedule`,
-/// - all imported functions from the external environment matches defined by `env` module,
-///
-/// The preprocessing includes injecting code for gas metering and metering the height of stack.
-pub fn prepare_contract<T: Config>(
-	original_code: CodeVec<T>,
-	schedule: &Schedule<T>,
-	owner: AccountIdOf<T>,
-	determinism: Determinism,
-) -> Result<PrefabWasmModule<T>, (DispatchError, &'static str)> {
-	do_preparation::<super::runtime::Env, T>(original_code, schedule, owner, determinism)
-}
-
-/// The same as [`prepare_contract`] but without constructing a new [`PrefabWasmModule`]
-///
-/// # Note
-///
-/// Use this when an existing contract should be re-instrumented with a newer schedule version.
-pub fn reinstrument_contract<T: Config>(
+/// Used to update the code of an existing module to the newest [`Schedule`] version.
+/// Stictly speaking is not necessary to check the existing code before reinstrumenting because
+/// it can't change in the meantime. However, since we recently switched the validation library
+/// we want to re-validate to weed out any bugs that were lurking in the old version.
+pub fn reinstrument<E, T>(
 	original_code: &[u8],
 	schedule: &Schedule<T>,
 	determinism: Determinism,
-) -> Result<Vec<u8>, &'static str> {
-	Ok(check_and_instrument::<super::runtime::Env, T>(original_code, schedule, determinism)?.0)
+) -> Result<Vec<u8>, DispatchError>
+where
+	E: Environment<()>,
+	T: Config,
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>,
+{
+	instrument::<E, T>(original_code, schedule, determinism, TryInstantiate::Skip)
+		.map_err(|(err, msg)| {
+			log::error!(target: "runtime::contracts", "CodeRejected during reinstrument: {}", msg);
+			err
+		})
+		.map(|(code, _)| code)
 }
 
 /// Alternate (possibly unsafe) preparation functions used only for benchmarking.
@@ -482,29 +550,22 @@ pub fn reinstrument_contract<T: Config>(
 /// in production code.
 #[cfg(feature = "runtime-benchmarks")]
 pub mod benchmarking {
-	use super::{elements::FunctionType, *};
-
-	impl ImportSatisfyCheck for () {
-		fn can_satisfy(_module: &[u8], _name: &[u8], _func_type: &FunctionType) -> bool {
-			true
-		}
-	}
+	use super::*;
 
 	/// Prepare function that neither checks nor instruments the passed in code.
-	pub fn prepare_contract<T: Config>(
+	pub fn prepare<T: Config>(
 		original_code: Vec<u8>,
 		schedule: &Schedule<T>,
 		owner: AccountIdOf<T>,
 	) -> Result<PrefabWasmModule<T>, &'static str> {
 		let contract_module = ContractModule::new(&original_code, schedule)?;
-		let memory_limits = get_memory_limits(contract_module.scan_imports::<()>(&[])?, schedule)?;
+		let memory_limits = get_memory_limits(contract_module.scan_imports(&[])?, schedule)?;
 		Ok(PrefabWasmModule {
 			instruction_weights_version: schedule.instruction_weights.version,
 			initial: memory_limits.0,
 			maximum: memory_limits.1,
 			code_hash: T::Hashing::hash(&original_code),
 			original_code: Some(original_code.try_into().map_err(|_| "Original code too large")?),
-			determinism: Determinism::Deterministic,
 			code: contract_module
 				.into_wasm_code()?
 				.try_into()
@@ -515,6 +576,7 @@ pub mod benchmarking {
 				deposit: Default::default(),
 				refcount: 0,
 			}),
+			determinism: Determinism::Deterministic,
 		})
 	}
 }
@@ -540,27 +602,28 @@ mod tests {
 	#[allow(unreachable_code)]
 	mod env {
 		use super::*;
+		use crate::wasm::runtime::TrapReason;
 
 		// Define test environment for tests. We need ImportSatisfyCheck
 		// implementation from it. So actual implementations doesn't matter.
 		#[define_env]
 		pub mod test_env {
-			fn panic(_ctx: crate::wasm::Runtime<E>) -> Result<(), TrapReason> {
+			fn panic(_ctx: _, _memory: _) -> Result<(), TrapReason> {
 				Ok(())
 			}
 
 			// gas is an implementation defined function and a contract can't import it.
-			fn gas(_ctx: crate::wasm::Runtime<E>, _amount: u32) -> Result<(), TrapReason> {
+			fn gas(_ctx: _, _memory: _, _amount: u64) -> Result<(), TrapReason> {
 				Ok(())
 			}
 
-			fn nop(_ctx: crate::wasm::Runtime<E>, _unused: u64) -> Result<(), TrapReason> {
+			fn nop(_ctx: _, _memory: _, _unused: u64) -> Result<(), TrapReason> {
 				Ok(())
 			}
 
-			// new version of nop with other data type for argumebt
+			// new version of nop with other data type for argument
 			#[version(1)]
-			fn nop(_ctx: crate::wasm::Runtime<E>, _unused: i32) -> Result<(), TrapReason> {
+			fn nop(_ctx: _, _memory: _, _unused: i32) -> Result<(), TrapReason> {
 				Ok(())
 			}
 		}
@@ -582,7 +645,13 @@ mod tests {
 					},
 					.. Default::default()
 				};
-				let r = do_preparation::<env::Env, Test>(wasm, &schedule, ALICE, Determinism::Deterministic);
+				let r = prepare::<env::Env, Test>(
+					wasm,
+					&schedule,
+					ALICE,
+					Determinism::Deterministic,
+					TryInstantiate::Instantiate,
+				);
 				assert_matches::assert_matches!(r.map_err(|(_, msg)| msg), $($expected)*);
 			}
 		};
@@ -718,7 +787,7 @@ mod tests {
 				(func (export "deploy"))
 			)
 			"#,
-			Err("Module is not valid")
+			Err("validation of new code failed")
 		);
 
 		prepare_test!(
@@ -784,7 +853,7 @@ mod tests {
 				(func (export "deploy"))
 			)
 			"#,
-			Err("Module is not valid")
+			Err("validation of new code failed")
 		);
 
 		prepare_test!(
@@ -910,7 +979,7 @@ mod tests {
 				(func (export "deploy"))
 			)
 			"#,
-			Err("module imports a non-existent function")
+			Err("module imports a banned function")
 		);
 
 		// memory is in "env" and not in "seal0"
@@ -965,7 +1034,7 @@ mod tests {
 				(func (export "deploy"))
 			)
 			"#,
-			Err("module imports a non-existent function")
+			Err("module imports a banned function")
 		);
 
 		prepare_test!(
@@ -978,7 +1047,7 @@ mod tests {
 				(func (export "deploy"))
 			)
 			"#,
-			Err("module imports a non-existent function")
+			Err("new code rejected after instrumentation")
 		);
 	}