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/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();
+}