From 70e27144673795a388878a9d5065f595cd214cac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Thei=C3=9Fen?= <alex.theissen@me.com>
Date: Tue, 18 Oct 2022 12:42:21 +0200
Subject: [PATCH] Only use a minimal stack when never executing module

---
 frame/contracts/src/benchmarking/sandbox.rs | 12 ++++++++----
 frame/contracts/src/wasm/mod.rs             |  7 ++++++-
 frame/contracts/src/wasm/prepare.rs         | 19 ++++++++++++++-----
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/frame/contracts/src/benchmarking/sandbox.rs b/frame/contracts/src/benchmarking/sandbox.rs
index 1abc6b2d2f7c7..ff3988d8a8fb7 100644
--- a/frame/contracts/src/benchmarking/sandbox.rs
+++ b/frame/contracts/src/benchmarking/sandbox.rs
@@ -21,7 +21,7 @@ use super::{code::WasmModule, Config};
 /// ! environment that provides the seal interface as imported functions.
 use crate::wasm::{Environment, PrefabWasmModule};
 use sp_core::crypto::UncheckedFrom;
-use wasmi::{errors::LinkerError, Func, Linker, Store};
+use wasmi::{errors::LinkerError, Func, Linker, StackLimits, Store};
 
 /// Minimal execution environment without any imported functions.
 pub struct Sandbox {
@@ -49,9 +49,13 @@ where
 			.as_ref()
 			.map(|mem| (mem.min_pages, mem.max_pages))
 			.unwrap_or((0, 0));
-		let (store, _memory, instance) =
-			PrefabWasmModule::<T>::instantiate::<EmptyEnv, _>(&module.code, (), memory)
-				.expect("Failed to create benchmarking Sandbox instance");
+		let (store, _memory, instance) = PrefabWasmModule::<T>::instantiate::<EmptyEnv, _>(
+			&module.code,
+			(),
+			memory,
+			StackLimits::default(),
+		)
+		.expect("Failed to create benchmarking Sandbox instance");
 		let entry_point = instance.get_export(&store, "call").unwrap().into_func().unwrap();
 		Self { entry_point, store }
 	}
diff --git a/frame/contracts/src/wasm/mod.rs b/frame/contracts/src/wasm/mod.rs
index c5d721dbfaa79..f71a2c7375038 100644
--- a/frame/contracts/src/wasm/mod.rs
+++ b/frame/contracts/src/wasm/mod.rs
@@ -43,7 +43,9 @@ use sp_core::crypto::UncheckedFrom;
 use sp_std::prelude::*;
 #[cfg(test)]
 pub use tests::MockExt;
-use wasmi::{Config as WasmiConfig, Engine, Instance, Linker, Memory, MemoryType, Module, Store};
+use wasmi::{
+	Config as WasmiConfig, Engine, Instance, Linker, Memory, MemoryType, Module, StackLimits, Store,
+};
 
 /// A prepared wasm module ready for execution.
 ///
@@ -170,12 +172,14 @@ where
 		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)
@@ -260,6 +264,7 @@ where
 			self.code.as_slice(),
 			runtime,
 			(self.initial, self.maximum),
+			StackLimits::default(),
 		)
 		.map_err(|msg| {
 			log::debug!(target: "runtime::contracts", "failed to instantiate code: {}", msg);
diff --git a/frame/contracts/src/wasm/prepare.rs b/frame/contracts/src/wasm/prepare.rs
index b35d5449eb2e0..588f7b909ec2f 100644
--- a/frame/contracts/src/wasm/prepare.rs
+++ b/frame/contracts/src/wasm/prepare.rs
@@ -32,6 +32,7 @@ use sp_std::prelude::*;
 use wasm_instrument::parity_wasm::elements::{
 	self, External, Internal, MemoryType, Type, ValueType,
 };
+use wasmi::StackLimits;
 
 /// 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.
@@ -447,11 +448,19 @@ where
 	// - Doesn't use any unknown imports.
 	// - Doesn't explode the wasmi bytecode generation.
 	if matches!(try_instantiate, TryInstantiate::Instantiate) {
-		PrefabWasmModule::<T>::instantiate::<E, _>(original_code.as_ref(), (), (initial, maximum))
-			.map_err(|err| {
-				log::debug!(target: "runtime::contracts", "{}", err);
-				(Error::<T>::CodeRejected.into(), "new code rejected after instrumentation")
-			})?;
+		// 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, _>(
+			original_code.as_ref(),
+			(),
+			(initial, maximum),
+			stack_limits,
+		)
+		.map_err(|err| {
+			log::debug!(target: "runtime::contracts", "{}", err);
+			(Error::<T>::CodeRejected.into(), "new code rejected after instrumentation")
+		})?;
 	}
 
 	let original_code_len = original_code.len();