Limit the maximum number of wasm memory pages a runtime can have (#9308)

* Limit the maximum number of wasm memory pages a runtime can have

* Switch the argument order

* fmt

substrate/client/executor/src/wasm_runtime.rs

@@ -321,6 +321,7 @@ pub fn create_wasm_runtime_with_code(
 			blob,
 			sc_executor_wasmtime::Config {
 				heap_pages: heap_pages as u32,
+				max_memory_pages: None,
 				allow_missing_func_imports,
 				cache_path: cache_path.map(ToOwned::to_owned),
 				semantics: sc_executor_wasmtime::Semantics {

substrate/client/executor/wasmtime/src/runtime.rs

@@ -79,9 +79,22 @@ pub struct WasmtimeRuntime {
 	engine: Engine,
 }
 
+impl WasmtimeRuntime {
+	/// Creates the store respecting the set limits.
+	fn new_store(&self) -> Store {
+		match self.config.max_memory_pages {
+			Some(max_memory_pages) => Store::new_with_limits(
+				&self.engine,
+				wasmtime::StoreLimitsBuilder::new().memory_pages(max_memory_pages).build(),
+			),
+			None => Store::new(&self.engine),
+		}
+	}
+}
+
 impl WasmModule for WasmtimeRuntime {
 	fn new_instance(&self) -> Result<Box<dyn WasmInstance>> {
-		let store = Store::new(&self.engine);
+		let store = self.new_store();
 
 		// Scan all imports, find the matching host functions, and create stubs that adapt arguments
 		// and results.
@@ -353,6 +366,20 @@ pub struct Config {
 	/// The number of wasm pages to be mounted after instantiation.
 	pub heap_pages: u32,
 
+	/// The total number of wasm pages an instance can request.
+	///
+	/// If specified, the runtime will be able to allocate only that much of wasm memory pages. This
+	/// is the total number and therefore the [`heap_pages`] is accounted for.
+	///
+	/// That means that the initial number of pages of a linear memory plus the [`heap_pages`] should
+	/// be less or equal to `max_memory_pages`, otherwise the instance won't be created.
+	///
+	/// Moreover, `memory.grow` will fail (return -1) if the sum of the number of currently mounted
+	/// pages and the number of additional pages exceeds `max_memory_pages`.
+	///
+	/// The default is `None`.
+	pub max_memory_pages: Option<u32>,
+
 	/// The WebAssembly standard requires all imports of an instantiated module to be resolved,
 	/// othewise, the instantiation fails. If this option is set to `true`, then this behavior is
 	/// overriden and imports that are requested by the module and not provided by the host functions

substrate/client/executor/wasmtime/src/tests.rs

@@ -29,6 +29,7 @@ struct RuntimeBuilder {
 	canonicalize_nans: bool,
 	deterministic_stack: bool,
 	heap_pages: u32,
+	max_memory_pages: Option<u32>,
 }
 
 impl RuntimeBuilder {
@@ -41,6 +42,7 @@ impl RuntimeBuilder {
 			canonicalize_nans: false,
 			deterministic_stack: false,
 			heap_pages: 1024,
+			max_memory_pages: None,
 		}
 	}
 
@@ -56,6 +58,10 @@ impl RuntimeBuilder {
 		self.deterministic_stack = deterministic_stack;
 	}
 
+	fn max_memory_pages(&mut self, max_memory_pages: Option<u32>) {
+		self.max_memory_pages = max_memory_pages;
+	}
+
 	fn build(self) -> Arc<dyn WasmModule> {
 		let blob = {
 			let wasm: Vec<u8>;
@@ -63,7 +69,7 @@ impl RuntimeBuilder {
 			let wasm = match self.code {
 				None => wasm_binary_unwrap(),
 				Some(wat) => {
-					wasm = wat::parse_str(wat).unwrap();
+					wasm = wat::parse_str(wat).expect("wat parsing failed");
 					&wasm
 				},
 			};
@@ -76,6 +82,7 @@ impl RuntimeBuilder {
 			blob,
 			crate::Config {
 				heap_pages: self.heap_pages,
+				max_memory_pages: self.max_memory_pages,
 				allow_missing_func_imports: true,
 				cache_path: None,
 				semantics: crate::Semantics {
@@ -160,3 +167,143 @@ fn test_stack_depth_reaching() {
 		format!("{:?}", err).starts_with("Other(\"Wasm execution trapped: wasm trap: unreachable")
 	);
 }
+
+#[test]
+fn test_max_memory_pages() {
+	fn try_instantiate(
+		max_memory_pages: Option<u32>,
+		wat: &'static str,
+	) -> Result<(), Box<dyn std::error::Error>> {
+		let runtime = {
+			let mut builder = RuntimeBuilder::new_on_demand();
+			builder.use_wat(wat);
+			builder.max_memory_pages(max_memory_pages);
+			builder.build()
+		};
+		let instance = runtime.new_instance()?;
+		let _ = instance.call_export("main", &[])?;
+		Ok(())
+	}
+
+	// check the old behavior if preserved. That is, if no limit is set we allow 4 GiB of memory.
+	try_instantiate(
+		None,
+		r#"
+		(module
+			;; we want to allocate the maximum number of pages supported in wasm for this test.
+			;;
+			;; However, due to a bug in wasmtime (I think wasmi is also affected) it is only possible
+			;; to allocate 65536 - 1 pages.
+			;;
+			;; Then, during creation of the Substrate Runtime instance, 1024 (heap_pages) pages are
+			;; mounted.
+			;;
+			;; Thus 65535 = 64511 + 1024
+			(import "env" "memory" (memory 64511))
+
+			(global (export "__heap_base") i32 (i32.const 0))
+			(func (export "main")
+				(param i32 i32) (result i64)
+				(i64.const 0)
+			)
+		)
+		"#,
+	)
+	.unwrap();
+
+	// max is not specified, therefore it's implied to be 65536 pages (4 GiB).
+	//
+	// max_memory_pages = 1 (initial) + 1024 (heap_pages)
+	try_instantiate(
+		Some(1 + 1024),
+		r#"
+		(module
+
+			(import "env" "memory" (memory 1)) ;; <- 1 initial, max is not specified
+
+			(global (export "__heap_base") i32 (i32.const 0))
+			(func (export "main")
+				(param i32 i32) (result i64)
+				(i64.const 0)
+			)
+		)
+		"#,
+	)
+	.unwrap();
+
+	// max is specified explicitly to 2048 pages.
+	try_instantiate(
+		Some(1 + 1024),
+		r#"
+		(module
+
+			(import "env" "memory" (memory 1 2048)) ;; <- max is 2048
+
+			(global (export "__heap_base") i32 (i32.const 0))
+			(func (export "main")
+				(param i32 i32) (result i64)
+				(i64.const 0)
+			)
+		)
+		"#,
+	)
+	.unwrap();
+
+	// memory grow should work as long as it doesn't exceed 1025 pages in total.
+	try_instantiate(
+		Some(0 + 1024 + 25),
+		r#"
+		(module
+			(import "env" "memory" (memory 0)) ;; <- zero starting pages.
+
+			(global (export "__heap_base") i32 (i32.const 0))
+			(func (export "main")
+				(param i32 i32) (result i64)
+
+				;; assert(memory.grow returns != -1)
+				(if
+					(i32.eq
+						(memory.grow
+							(i32.const 25)
+						)
+						(i32.const -1)
+					)
+					(unreachable)
+				)
+
+				(i64.const 0)
+			)
+		)
+		"#,
+	)
+	.unwrap();
+
+	// We start with 1025 pages and try to grow at least one.
+	try_instantiate(
+		Some(1 + 1024),
+		r#"
+		(module
+			(import "env" "memory" (memory 1))  ;; <- initial=1, meaning after heap pages mount the
+												;; total will be already 1025
+			(global (export "__heap_base") i32 (i32.const 0))
+			(func (export "main")
+				(param i32 i32) (result i64)
+
+				;; assert(memory.grow returns == -1)
+				(if
+					(i32.ne
+						(memory.grow
+							(i32.const 1)
+						)
+						(i32.const -1)
+					)
+					(unreachable)
+				)
+
+				(i64.const 0)
+			)
+		)
+		"#,
+	)
+	.unwrap();
+}