From a67dcfb281ee54f3bd9a4fc217134d7f7da0e9ee Mon Sep 17 00:00:00 2001
From: Dmitry Kashitsyn <korvin@deeptown.org>
Date: Wed, 20 May 2020 17:28:13 +0700
Subject: [PATCH] Refactor resource and error handling in wasm (#6074)

* Refactor resource and error handling in wasm

* Fixes based on review
---
 client/executor/common/src/sandbox.rs | 35 ++++++++++++++++++---------
 1 file changed, 23 insertions(+), 12 deletions(-)

diff --git a/client/executor/common/src/sandbox.rs b/client/executor/common/src/sandbox.rs
index d897ef583cac9..37860474cb07c 100644
--- a/client/executor/common/src/sandbox.rs
+++ b/client/executor/common/src/sandbox.rs
@@ -237,21 +237,33 @@ impl<'a, FE: SandboxCapabilities + 'a> Externals for GuestExternals<'a, FE> {
 			.supervisor_externals
 			.allocate_memory(invoke_args_len)
 			.map_err(|_| trap("Can't allocate memory in supervisor for the arguments"))?;
-		self
+
+		let deallocate = |this: &mut GuestExternals<FE>, ptr, fail_msg| {
+			this
+				.supervisor_externals
+				.deallocate_memory(ptr)
+				.map_err(|_| trap(fail_msg))
+		};
+
+		if self
 			.supervisor_externals
 			.write_memory(invoke_args_ptr, &invoke_args_data)
-			.map_err(|_| trap("Can't write invoke args into memory"))?;
+			.is_err()
+		{
+			deallocate(self, invoke_args_ptr, "Failed dealloction after failed write of invoke arguments")?;
+			return Err(trap("Can't write invoke args into memory"));
+		}
+
 		let result = self.supervisor_externals.invoke(
 			&self.sandbox_instance.dispatch_thunk,
 			invoke_args_ptr,
 			invoke_args_len,
 			state,
 			func_idx,
-		)?;
-		self
-			.supervisor_externals
-			.deallocate_memory(invoke_args_ptr)
-			.map_err(|_| trap("Can't deallocate memory for dispatch thunk's invoke arguments"))?;
+		);
+
+		deallocate(self, invoke_args_ptr, "Can't deallocate memory for dispatch thunk's invoke arguments")?;
+		let result = result?;
 
 		// dispatch_thunk returns pointer to serialized arguments.
 		// Unpack pointer and len of the serialized result data.
@@ -265,12 +277,11 @@ impl<'a, FE: SandboxCapabilities + 'a> Externals for GuestExternals<'a, FE> {
 
 		let serialized_result_val = self.supervisor_externals
 			.read_memory(serialized_result_val_ptr, serialized_result_val_len)
-			.map_err(|_| trap("Can't read the serialized result from dispatch thunk"))?;
-		self.supervisor_externals
-			.deallocate_memory(serialized_result_val_ptr)
-			.map_err(|_| trap("Can't deallocate memory for dispatch thunk's result"))?;
+			.map_err(|_| trap("Can't read the serialized result from dispatch thunk"));
 
-		deserialize_result(&serialized_result_val)
+		deallocate(self, serialized_result_val_ptr, "Can't deallocate memory for dispatch thunk's result")
+			.and_then(|_| serialized_result_val)
+			.and_then(|serialized_result_val| deserialize_result(&serialized_result_val))
 	}
 }