Catch panics on the FFI boundary between the runtime and the host for wasmtime (#11189)

* Catch panics on the FFI boundary between the runtime and the host for `wasmtime`

* Use an already existing test runtime function

* Merge the tests together

substrate/client/executor/runtime-test/src/lib.rs

@@ -113,7 +113,9 @@ sp_core::wasm_export_functions! {
 	   }
    }
 
-   fn test_exhaust_heap() -> Vec<u8> { Vec::with_capacity(16777216) }
+	fn test_allocate_vec(size: u32) -> Vec<u8> {
+		Vec::with_capacity(size as usize)
+	}
 
    fn test_fp_f32add(a: [u8; 4], b: [u8; 4]) -> [u8; 4] {
 	   let a = f32::from_le_bytes(a);

substrate/client/executor/src/integration_tests/mod.rs

@@ -466,13 +466,24 @@ fn should_trap_when_heap_exhausted(wasm_method: WasmExecutionMethod) {
 			RuntimeBlob::uncompress_if_needed(wasm_binary_unwrap()).unwrap(),
 			&mut ext.ext(),
 			true,
-			"test_exhaust_heap",
-			&[0],
+			"test_allocate_vec",
+			&16777216_u32.encode(),
 		)
-		.map_err(|e| e.to_string())
 		.unwrap_err();
 
-	assert!(err.contains("Allocator ran out of space"));
+	match err {
+		#[cfg(feature = "wasmtime")]
+		Error::AbortedDueToTrap(error) if wasm_method == WasmExecutionMethod::Compiled => {
+			assert_eq!(
+				error.message,
+				r#"host code panicked while being called by the runtime: Failed to allocate memory: "Allocator ran out of space""#
+			);
+		},
+		Error::RuntimePanicked(error) if wasm_method == WasmExecutionMethod::Interpreted => {
+			assert_eq!(error, r#"Failed to allocate memory: "Allocator ran out of space""#);
+		},
+		error => panic!("unexpected error: {:?}", error),
+	}
 }
 
 fn mk_test_runtime(wasm_method: WasmExecutionMethod, pages: u64) -> Arc<dyn WasmModule> {

substrate/primitives/runtime-interface/proc-macro/src/runtime_interface/host_function_interface.rs

@@ -374,11 +374,27 @@ fn generate_host_function_implementation(
 				-> std::result::Result<#ffi_return_ty, #crate_::sp_wasm_interface::wasmtime::Trap>
 			{
 				T::with_function_context(caller, move |__function_context__| {
-					#struct_name::call(
-						__function_context__,
-						#(#ffi_names,)*
-					)
-				}).map_err(#crate_::sp_wasm_interface::wasmtime::Trap::new)
+					let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+						#struct_name::call(
+							__function_context__,
+							#(#ffi_names,)*
+						).map_err(#crate_::sp_wasm_interface::wasmtime::Trap::new)
+					}));
+					match result {
+						Ok(result) => result,
+						Err(panic) => {
+							let message =
+								if let Some(message) = panic.downcast_ref::<String>() {
+									format!("host code panicked while being called by the runtime: {}", message)
+								} else if let Some(message) = panic.downcast_ref::<&'static str>() {
+									format!("host code panicked while being called by the runtime: {}", message)
+								} else {
+									"host code panicked while being called by the runtime".to_owned()
+								};
+							return Err(#crate_::sp_wasm_interface::wasmtime::Trap::new(message));
+						}
+					}
+				})
 			}
 		)?;
 	};