Contracts: Reuse module when validating (#3789)

Reuse wasmi Module when validating.
Prepare the code for 0.32 and the addition of Module::new_unchecked

prdoc/pr_3789.prdoc

@@ -0,0 +1,12 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: "[pallet-contracts] Benchmarks improvements"
+
+doc:
+  - audience: Runtime Dev
+    description: Reuse wasmi module when validating the wasm code.
+crates:
+- name: pallet-contracts
+  bump: patch
+

substrate/frame/contracts/src/benchmarking/sandbox.rs

@@ -20,7 +20,8 @@
 /// ! environment that provides the seal interface as imported functions.
 use super::{code::WasmModule, Config};
 use crate::wasm::{
-	AllowDeprecatedInterface, AllowUnstableInterface, Determinism, Environment, WasmBlob,
+	AllowDeprecatedInterface, AllowUnstableInterface, Determinism, Environment, LoadedModule,
+	LoadingMode, WasmBlob,
 };
 use sp_core::Get;
 use wasmi::{errors::LinkerError, Func, Linker, StackLimits, Store};
@@ -42,12 +43,18 @@ impl<T: Config> From<&WasmModule<T>> for Sandbox {
 	/// Creates an instance from the supplied module.
 	/// Sets the execution engine fuel level to `u64::MAX`.
 	fn from(module: &WasmModule<T>) -> Self {
-		let (mut store, _memory, instance) = WasmBlob::<T>::instantiate::<EmptyEnv, _>(
+		let contract = LoadedModule::new::<T>(
 			&module.code,
+			Determinism::Relaxed,
+			Some(StackLimits::default()),
+			LoadingMode::Checked,
+		)
+		.expect("Failed to load Wasm module");
+
+		let (mut store, _memory, instance) = WasmBlob::<T>::instantiate::<EmptyEnv, _>(
+			contract,
 			(),
 			&<T>::Schedule::get(),
-			Determinism::Relaxed,
-			StackLimits::default(),
 			// We are testing with an empty environment anyways
 			AllowDeprecatedInterface::No,
 		)

substrate/frame/contracts/src/tests.rs

@@ -34,7 +34,7 @@ use crate::{
 	primitives::CodeUploadReturnValue,
 	storage::DeletionQueueManager,
 	tests::test_utils::{get_contract, get_contract_checked},
-	wasm::{Determinism, ReturnErrorCode as RuntimeReturnCode},
+	wasm::{Determinism, LoadingMode, ReturnErrorCode as RuntimeReturnCode},
 	weights::WeightInfo,
 	Array, BalanceOf, Code, CodeHash, CodeInfoOf, CollectEvents, Config, ContractInfo,
 	ContractInfoOf, DebugInfo, DefaultAddressGenerator, DeletionQueueCounter, Error, HoldReason,
@@ -1102,6 +1102,20 @@ fn delegate_call() {
 	});
 }
 
+#[test]
+fn track_check_uncheck_module_call() {
+	let (wasm, code_hash) = compile_module::<Test>("dummy").unwrap();
+	ExtBuilder::default().build().execute_with(|| {
+		let _ = <Test as Config>::Currency::set_balance(&ALICE, 1_000_000);
+		Contracts::bare_upload_code(ALICE, wasm, None, Determinism::Enforced).unwrap();
+		builder::bare_instantiate(Code::Existing(code_hash)).build_and_unwrap_result();
+	});
+
+	// It should have recorded 1 `Checked` for the upload and 1 `Unchecked` for the instantiate.
+	let record = crate::wasm::tracker::LOADED_MODULE.with(|stack| stack.borrow().clone());
+	assert_eq!(record, vec![LoadingMode::Checked, LoadingMode::Unchecked]);
+}
+
 #[test]
 fn transfer_expendable_cannot_kill_account() {
 	let (wasm, _code_hash) = compile_module::<Test>("dummy").unwrap();

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

@@ -21,26 +21,26 @@
 mod prepare;
 mod runtime;
 
-#[cfg(test)]
-pub use runtime::STABLE_API_COUNT;
-
 #[cfg(doc)]
 pub use crate::wasm::runtime::api_doc;
 
-pub use crate::wasm::runtime::{
-	AllowDeprecatedInterface, AllowUnstableInterface, Environment, Runtime, RuntimeCosts,
+#[cfg(test)]
+pub use {
+	crate::wasm::{prepare::tracker, runtime::ReturnErrorCode},
+	runtime::STABLE_API_COUNT,
+	tests::MockExt,
 };
 
-#[cfg(test)]
-pub use tests::MockExt;
-
-#[cfg(test)]
-pub use crate::wasm::runtime::ReturnErrorCode;
+pub use crate::wasm::{
+	prepare::{LoadedModule, LoadingMode},
+	runtime::{
+		AllowDeprecatedInterface, AllowUnstableInterface, Environment, Runtime, RuntimeCosts,
+	},
+};
 
 use crate::{
 	exec::{ExecResult, Executable, ExportedFunction, Ext},
 	gas::{GasMeter, Token},
-	wasm::prepare::LoadedModule,
 	weights::WeightInfo,
 	AccountIdOf, BadOrigin, BalanceOf, CodeHash, CodeInfoOf, CodeVec, Config, Error, Event,
 	HoldReason, Pallet, PristineCode, Schedule, Weight, LOG_TARGET,
@@ -201,17 +201,14 @@ impl<T: Config> WasmBlob<T> {
 	/// 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],
+		contract: LoadedModule,
 		host_state: H,
 		schedule: &Schedule<T>,
-		determinism: Determinism,
-		stack_limits: StackLimits,
 		allow_deprecated: AllowDeprecatedInterface,
 	) -> Result<(Store<H>, Memory, InstancePre), &'static str>
 	where
 		E: Environment<H>,
 	{
-		let contract = LoadedModule::new::<T>(&code, determinism, Some(stack_limits))?;
 		let mut store = Store::new(&contract.engine, host_state);
 		let mut linker = Linker::new(&contract.engine);
 		E::define(
@@ -382,12 +379,22 @@ impl<T: Config> WasmBlob<T> {
 		let code = self.code.as_slice();
 		// Instantiate the Wasm module to the engine.
 		let schedule = <T>::Schedule::get();
+
+		let contract = LoadedModule::new::<T>(
+			&code,
+			self.code_info.determinism,
+			Some(StackLimits::default()),
+			LoadingMode::Unchecked,
+		)
+		.map_err(|err| {
+			log::debug!(target: LOG_TARGET, "failed to create wasmi module: {err:?}");
+			Error::<T>::CodeRejected
+		})?;
+
 		let (mut store, memory, instance) = Self::instantiate::<crate::wasm::runtime::Env, _>(
-			code,
+			contract,
 			runtime,
 			&schedule,
-			self.code_info.determinism,
-			StackLimits::default(),
 			match function {
 				ExportedFunction::Call => AllowDeprecatedInterface::Yes,
 				ExportedFunction::Constructor => AllowDeprecatedInterface::No,

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

@@ -48,6 +48,20 @@ pub struct LoadedModule {
 	pub engine: Engine,
 }
 
+#[derive(PartialEq, Debug, Clone)]
+pub enum LoadingMode {
+	Checked,
+	Unchecked,
+}
+
+#[cfg(test)]
+pub mod tracker {
+	use sp_std::cell::RefCell;
+	thread_local! {
+		pub static LOADED_MODULE: RefCell<Vec<super::LoadingMode>> = RefCell::new(Vec::new());
+	}
+}
+
 impl LoadedModule {
 	/// Creates a new instance of `LoadedModule`.
 	///
@@ -57,6 +71,7 @@ impl LoadedModule {
 		code: &[u8],
 		determinism: Determinism,
 		stack_limits: Option<StackLimits>,
+		_mode: LoadingMode,
 	) -> Result<Self, &'static str> {
 		// NOTE: wasmi does not support unstable WebAssembly features. The module is implicitly
 		// checked for not having those ones when creating `wasmi::Module` below.
@@ -79,11 +94,16 @@ impl LoadedModule {
 		}
 
 		let engine = Engine::new(&config);
+
+		// TODO use Module::new_unchecked when validate_module is true once we are on wasmi@0.32
 		let module = Module::new(&engine, code).map_err(|err| {
 			log::debug!(target: LOG_TARGET, "Module creation failed: {:?}", err);
 			"Can't load the module into wasmi!"
 		})?;
 
+		#[cfg(test)]
+		tracker::LOADED_MODULE.with(|t| t.borrow_mut().push(_mode));
+
 		// Return a `LoadedModule` instance with
 		// __valid__ module.
 		Ok(LoadedModule { module, engine })
@@ -229,24 +249,38 @@ where
 	E: Environment<()>,
 	T: Config,
 {
-	(|| {
+	let module = (|| {
+		// We don't actually ever execute this instance so we can get away with a minimal stack
+		// which reduces the amount of memory that needs to be zeroed.
+		let stack_limits = Some(StackLimits::new(1, 1, 0).expect("initial <= max; qed"));
+
 		// We check that the module is generally valid,
 		// and does not have restricted WebAssembly features, here.
 		let contract_module = match *determinism {
 			Determinism::Relaxed =>
-				if let Ok(module) = LoadedModule::new::<T>(code, Determinism::Enforced, None) {
+				if let Ok(module) = LoadedModule::new::<T>(
+					code,
+					Determinism::Enforced,
+					stack_limits,
+					LoadingMode::Checked,
+				) {
 					*determinism = Determinism::Enforced;
 					module
 				} else {
-					LoadedModule::new::<T>(code, Determinism::Relaxed, None)?
+					LoadedModule::new::<T>(code, Determinism::Relaxed, None, LoadingMode::Checked)?
 				},
-			Determinism::Enforced => LoadedModule::new::<T>(code, Determinism::Enforced, None)?,
+			Determinism::Enforced => LoadedModule::new::<T>(
+				code,
+				Determinism::Enforced,
+				stack_limits,
+				LoadingMode::Checked,
+			)?,
 		};
 
 		// The we check that module satisfies constraints the pallet puts on contracts.
 		contract_module.scan_exports()?;
 		contract_module.scan_imports::<T>(schedule)?;
-		Ok(())
+		Ok(contract_module)
 	})()
 	.map_err(|msg: &str| {
 		log::debug!(target: LOG_TARGET, "New code rejected on validation: {}", msg);
@@ -257,22 +291,11 @@ where
 	//
 	// - It doesn't use any unknown imports.
 	// - It doesn't explode the wasmi bytecode generation.
-	//
-	// We don't actually ever execute this instance 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");
-	WasmBlob::<T>::instantiate::<E, _>(
-		&code,
-		(),
-		schedule,
-		*determinism,
-		stack_limits,
-		AllowDeprecatedInterface::No,
-	)
-	.map_err(|err| {
-		log::debug!(target: LOG_TARGET, "{}", err);
-		(Error::<T>::CodeRejected.into(), "New code rejected on wasmi instantiation!")
-	})?;
+	WasmBlob::<T>::instantiate::<E, _>(module, (), schedule, AllowDeprecatedInterface::No)
+		.map_err(|err| {
+			log::debug!(target: LOG_TARGET, "{err}");
+			(Error::<T>::CodeRejected.into(), "New code rejected on wasmi instantiation!")
+		})?;
 
 	Ok(())
 }
@@ -325,7 +348,8 @@ pub mod benchmarking {
 		owner: AccountIdOf<T>,
 	) -> Result<WasmBlob<T>, DispatchError> {
 		let determinism = Determinism::Enforced;
-		let contract_module = LoadedModule::new::<T>(&code, determinism, None)?;
+		let contract_module =
+			LoadedModule::new::<T>(&code, determinism, None, LoadingMode::Checked)?;
 		let _ = contract_module.scan_imports::<T>(schedule)?;
 		let code: CodeVec<T> = code.try_into().map_err(|_| <Error<T>>::CodeTooLarge)?;
 		let code_info = CodeInfo {