Replace unwrap with expect (#1684)

* Replace unwrap with expect

* Move expect to call sites

* Bubble errors up and trap

* Update wasm

* Update invalid pointer

* Remove test which makes problems in CI

* Check for underflow

substrate/core/executor/src/allocator.rs

@@ -17,7 +17,9 @@
 //! This module implements a freeing-bump allocator.
 //! See more details at https://github.com/paritytech/substrate/issues/1615.
 
+use crate::wasm_utils::UserError;
 use log::trace;
+use wasmi::Error;
 use wasmi::MemoryRef;
 use wasmi::memory_units::Bytes;
 
@@ -32,6 +34,9 @@ const ALIGNMENT: u32 = 8;
 const N: usize = 22;
 const MAX_POSSIBLE_ALLOCATION: u32 = 16777216; // 2^24 bytes
 
+pub const OUT_OF_SPACE: &str = "Requested allocation size does not fit into remaining heap space";
+pub const REQUESTED_SIZE_TOO_LARGE: &str = "Requested size to allocate is too large";
+
 pub struct FreeingBumpHeapAllocator {
 	bumper: u32,
 	heads: [u32; N],
@@ -82,54 +87,73 @@ impl FreeingBumpHeapAllocator {
 
 	/// Gets requested number of bytes to allocate and returns a pointer.
 	/// The maximum size which can be allocated at once is 16 MiB.
-	pub fn allocate(&mut self, size: u32) -> u32 {
+	pub fn allocate(&mut self, size: u32) -> Result<u32, UserError> {
 		if size > MAX_POSSIBLE_ALLOCATION {
-			return 0;
+			return Err(UserError(REQUESTED_SIZE_TOO_LARGE));
 		}
 
 		let size = size.max(8);
 		let item_size = size.next_power_of_two();
 		if item_size + 8 + self.total_size > self.max_heap_size {
-			return 0;
+			return Err(UserError(OUT_OF_SPACE));
 		}
 
 		let list_index = (item_size.trailing_zeros() - 3) as usize;
 		let ptr: u32 = if self.heads[list_index] != 0 {
 			// Something from the free list
 			let item = self.heads[list_index];
-			self.heads[list_index] = FreeingBumpHeapAllocator::le_bytes_to_u32(self.get_heap_4bytes(item));
+			let four_bytes = self.get_heap_4bytes(item)
+				.map_err(|_| UserError("Unable to get bytes at pointer taken from list of free items"))?;
+			self.heads[list_index] = FreeingBumpHeapAllocator::le_bytes_to_u32(four_bytes);
 			item + 8
 		} else {
 			// Nothing to be freed. Bump.
 			self.bump(item_size + 8) + 8
 		};
 
-		for i in 1..8 { self.set_heap(ptr - i, 255); }
+		for i in 1..8 {
+			self.set_heap(ptr - i, 255)
+				.map_err(|_| UserError("Unable to successively write bytes into heap at pointer prefix"))?;
+		}
 
-		self.set_heap(ptr - 8, list_index as u8);
+		self.set_heap(ptr - 8, list_index as u8)
+			.map_err(|_| UserError("Unable to write byte into heap at pointer prefix"))?;
 
 		self.total_size = self.total_size + item_size + 8;
 		trace!(target: "wasm-heap", "Heap size is {} bytes after allocation", self.total_size);
 
-		self.ptr_offset + ptr
+		Ok(self.ptr_offset + ptr)
 	}
 
 	/// Deallocates the space which was allocated for a pointer.
-	pub fn deallocate(&mut self, ptr: u32) {
+	pub fn deallocate(&mut self, ptr: u32) -> Result<(), UserError> {
 		let ptr = ptr - self.ptr_offset;
+		if ptr < 8 {
+			return Err(UserError("Invalid pointer for deallocation"));
+		}
 
-		let list_index = self.get_heap_byte(ptr - 8) as usize;
-		for i in 1..8 { debug_assert!(self.get_heap_byte(ptr - i) == 255); }
+		let list_index = self.get_heap_byte(ptr - 8)
+			.map_err(|_| UserError("Unable to access pointer prefix"))? as usize;
+		for i in 1..8 {
+			let heap_byte = self.get_heap_byte(ptr - i)
+				.map_err(|_| UserError("Unable to write single bytes into heap at pointer"))?;
+			debug_assert!(heap_byte == 255)
+		}
 		let tail = self.heads[list_index];
 		self.heads[list_index] = ptr - 8;
 
-		let mut slice = self.get_heap_4bytes(ptr - 8);
+		let mut slice = self.get_heap_4bytes(ptr - 8)
+			.map_err(|_| UserError("Unable to get 4 bytes from heap at pointer prefix"))?;
 		FreeingBumpHeapAllocator::write_u32_into_le_bytes(tail, &mut slice);
-		self.set_heap_4bytes(ptr - 8, slice);
+		self.set_heap_4bytes(ptr - 8, slice)
+			.map_err(|_| UserError("Unable to write 4 bytes into heap at pointer prefix"))?;
 
 		let item_size = FreeingBumpHeapAllocator::get_item_size_from_index(list_index);
-		self.total_size = self.total_size.checked_sub(item_size as u32 + 8).unwrap_or(0);
+		self.total_size = self.total_size.checked_sub(item_size as u32 + 8)
+			.ok_or_else(|| UserError("Unable to subtract from total heap size without overflow"))?;
 		trace!(target: "wasm-heap", "Heap size is {} bytes after deallocation", self.total_size);
+
+		Ok(())
 	}
 
 	fn bump(&mut self, n: u32) -> u32 {
@@ -153,24 +177,24 @@ impl FreeingBumpHeapAllocator {
 		1 << 3 << index
 	}
 
-	fn get_heap_4bytes(&mut self, ptr: u32) -> [u8; 4] {
+	fn get_heap_4bytes(&mut self, ptr: u32) -> Result<[u8; 4], Error> {
 		let mut arr = [0u8; 4];
-		self.heap.get_into(self.ptr_offset + ptr, &mut arr).unwrap();
-		arr
+		self.heap.get_into(self.ptr_offset + ptr, &mut arr)?;
+		Ok(arr)
 	}
 
-	fn get_heap_byte(&mut self, ptr: u32) -> u8 {
+	fn get_heap_byte(&mut self, ptr: u32) -> Result<u8, Error> {
 		let mut arr = [0u8; 1];
-		self.heap.get_into(self.ptr_offset + ptr, &mut arr).unwrap();
-		arr[0]
+		self.heap.get_into(self.ptr_offset + ptr, &mut arr)?;
+		Ok(arr[0])
 	}
 
-	fn set_heap(&mut self, ptr: u32, value: u8) {
-		self.heap.set(self.ptr_offset + ptr, &[value]).unwrap()
+	fn set_heap(&mut self, ptr: u32, value: u8) -> Result<(), Error> {
+		self.heap.set(self.ptr_offset + ptr, &[value])
 	}
 
-	fn set_heap_4bytes(&mut self, ptr: u32, value: [u8; 4]) {
-		self.heap.set(self.ptr_offset + ptr, &value).unwrap()
+	fn set_heap_4bytes(&mut self, ptr: u32, value: [u8; 4]) -> Result<(), Error> {
+		self.heap.set(self.ptr_offset + ptr, &value)
 	}
 
 }
@@ -195,7 +219,7 @@ mod tests {
 		let mut heap = FreeingBumpHeapAllocator::new(mem);
 
 		// when
-		let ptr = heap.allocate(1);
+		let ptr = heap.allocate(1).unwrap();
 
 		// then
 		assert_eq!(ptr, 8);
@@ -209,7 +233,7 @@ mod tests {
 		let mut heap = FreeingBumpHeapAllocator::new(mem);
 
 		// when
-		let ptr = heap.allocate(1);
+		let ptr = heap.allocate(1).unwrap();
 
 		// then
 		// the pointer must start at the next multiple of 8 from 13
@@ -224,9 +248,9 @@ mod tests {
 		let mut heap = FreeingBumpHeapAllocator::new(mem);
 
 		// when
-		let ptr1 = heap.allocate(1);
-		let ptr2 = heap.allocate(9);
-		let ptr3 = heap.allocate(1);
+		let ptr1 = heap.allocate(1).unwrap();
+		let ptr2 = heap.allocate(9).unwrap();
+		let ptr3 = heap.allocate(1).unwrap();
 
 		// then
 		// a prefix of 8 bytes is prepended to each pointer
@@ -245,16 +269,16 @@ mod tests {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
 		let mut heap = FreeingBumpHeapAllocator::new(mem);
-		let ptr1 = heap.allocate(1);
+		let ptr1 = heap.allocate(1).unwrap();
 		// the prefix of 8 bytes is prepended to the pointer
 		assert_eq!(ptr1, 8);
 
-		let ptr2 = heap.allocate(1);
+		let ptr2 = heap.allocate(1).unwrap();
 		// the prefix of 8 bytes + the content of ptr 1 is prepended to the pointer
 		assert_eq!(ptr2, 24);
 
 		// when
-		heap.deallocate(ptr2);
+		heap.deallocate(ptr2).unwrap();
 
 		// then
 		// then the heads table should contain a pointer to the
@@ -270,19 +294,19 @@ mod tests {
 		let padded_offset = 16;
 		let mut heap = FreeingBumpHeapAllocator::new(mem);
 
-		let ptr1 = heap.allocate(1);
+		let ptr1 = heap.allocate(1).unwrap();
 		// the prefix of 8 bytes is prepended to the pointer
 		assert_eq!(ptr1, padded_offset + 8);
 
-		let ptr2 = heap.allocate(9);
+		let ptr2 = heap.allocate(9).unwrap();
 		// the padded_offset + the previously allocated ptr (8 bytes prefix +
 		// 8 bytes content) + the prefix of 8 bytes which is prepended to the
 		// current pointer
 		assert_eq!(ptr2, padded_offset + 16 + 8);
 
 		// when
-		heap.deallocate(ptr2);
-		let ptr3 = heap.allocate(9);
+		heap.deallocate(ptr2).unwrap();
+		let ptr3 = heap.allocate(9).unwrap();
 
 		// then
 		// should have re-allocated
@@ -296,21 +320,21 @@ mod tests {
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
 		let mut heap = FreeingBumpHeapAllocator::new(mem);
 
-		let ptr1 = heap.allocate(8);
-		let ptr2 = heap.allocate(8);
-		let ptr3 = heap.allocate(8);
+		let ptr1 = heap.allocate(8).unwrap();
+		let ptr2 = heap.allocate(8).unwrap();
+		let ptr3 = heap.allocate(8).unwrap();
 
 		// when
-		heap.deallocate(ptr1);
-		heap.deallocate(ptr2);
-		heap.deallocate(ptr3);
+		heap.deallocate(ptr1).unwrap();
+		heap.deallocate(ptr2).unwrap();
+		heap.deallocate(ptr3).unwrap();
 
 		// then
 		let mut expected = [0; N];
 		expected[0] = ptr3 - 8;
 		assert_eq!(heap.heads, expected);
 
-		let ptr4 = heap.allocate(8);
+		let ptr4 = heap.allocate(8).unwrap();
 		assert_eq!(ptr4, ptr3);
 
 		expected[0] = ptr2 - 8;
@@ -328,7 +352,10 @@ mod tests {
 		let ptr = heap.allocate(PAGE_SIZE - 13);
 
 		// then
-		assert_eq!(ptr, 0);
+		assert_eq!(ptr.is_err(), true);
+		if let Err(err) = ptr {
+			assert_eq!(err, UserError(OUT_OF_SPACE));
+		}
 	}
 
 	#[test]
@@ -336,7 +363,7 @@ mod tests {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), Some(Pages(1))).unwrap();
 		let mut heap = FreeingBumpHeapAllocator::new(mem);
-		let ptr1 = heap.allocate((PAGE_SIZE / 2) - 8);
+		let ptr1 = heap.allocate((PAGE_SIZE / 2) - 8).unwrap();
 		assert_eq!(ptr1, 8);
 
 		// when
@@ -344,7 +371,10 @@ mod tests {
 
 		// then
 		// there is no room for another half page incl. its 8 byte prefix
-		assert_eq!(ptr2, 0);
+		assert_eq!(ptr2.is_err(), true);
+		if let Err(err) = ptr2 {
+			assert_eq!(err, UserError(OUT_OF_SPACE));
+		}
 	}
 
 	#[test]
@@ -355,7 +385,7 @@ mod tests {
 		let mut heap = FreeingBumpHeapAllocator::new(mem);
 
 		// when
-		let ptr = heap.allocate(MAX_POSSIBLE_ALLOCATION);
+		let ptr = heap.allocate(MAX_POSSIBLE_ALLOCATION).unwrap();
 
 		// then
 		assert_eq!(ptr, 8);
@@ -371,7 +401,10 @@ mod tests {
 		let ptr = heap.allocate(MAX_POSSIBLE_ALLOCATION + 1);
 
 		// then
-		assert_eq!(ptr, 0);
+		assert_eq!(ptr.is_err(), true);
+		if let Err(err) = ptr {
+			assert_eq!(err, UserError(REQUESTED_SIZE_TOO_LARGE));
+		}
 	}
 
 	#[test]
@@ -383,7 +416,7 @@ mod tests {
 
 		// when
 		// an item size of 16 must be used then
-		heap.allocate(9);
+		heap.allocate(9).unwrap();
 
 		// then
 		assert_eq!(heap.total_size, 8 + 16);
@@ -397,9 +430,9 @@ mod tests {
 		let mut heap = FreeingBumpHeapAllocator::new(mem);
 
 		// when
-		let ptr = heap.allocate(42);
+		let ptr = heap.allocate(42).unwrap();
 		assert_eq!(ptr, 16 + 8);
-		heap.deallocate(ptr);
+		heap.deallocate(ptr).unwrap();
 
 		// then
 		assert_eq!(heap.total_size, 0);
@@ -414,8 +447,8 @@ mod tests {
 
 		// when
 		for _ in 1..10 {
-			let ptr = heap.allocate(42);
-			heap.deallocate(ptr);
+			let ptr = heap.allocate(42).unwrap();
+			heap.deallocate(ptr).unwrap();
 		}
 
 		// then

substrate/core/executor/src/sandbox.rs

@@ -148,13 +148,21 @@ pub trait SandboxCapabilities {
 
 	/// Allocate space of the specified length in the supervisor memory.
 	///
+	/// # Errors
+	///
+	/// Returns `Err` if allocation not possible or errors during heap management.
+	///
 	/// Returns pointer to the allocated block.
-	fn allocate(&mut self, len: u32) -> u32;
+	fn allocate(&mut self, len: u32) -> Result<u32, UserError>;
 
 	/// Deallocate space specified by the pointer that was previously returned by [`allocate`].
 	///
+	/// # Errors
+	///
+	/// Returns `Err` if deallocation not possible or because of errors in heap management.
+	///
 	/// [`allocate`]: #tymethod.allocate
-	fn deallocate(&mut self, ptr: u32);
+	fn deallocate(&mut self, ptr: u32) -> Result<(), UserError>;
 
 	/// Write `data` into the supervisor memory at offset specified by `ptr`.
 	///
@@ -232,7 +240,7 @@ impl<'a, FE: SandboxCapabilities + Externals + 'a> Externals for GuestExternals<
 		// Move serialized arguments inside the memory and invoke dispatch thunk and
 		// then free allocated memory.
 		let invoke_args_ptr = self.supervisor_externals
-			.allocate(invoke_args_data.len() as u32);
+			.allocate(invoke_args_data.len() as u32)?;
 		self.supervisor_externals
 			.write_memory(invoke_args_ptr, &invoke_args_data)?;
 		let result = ::wasmi::FuncInstance::invoke(
@@ -245,7 +253,7 @@ impl<'a, FE: SandboxCapabilities + Externals + 'a> Externals for GuestExternals<
 			],
 			self.supervisor_externals,
 		);
-		self.supervisor_externals.deallocate(invoke_args_ptr);
+		self.supervisor_externals.deallocate(invoke_args_ptr)?;
 
 		// dispatch_thunk returns pointer to serialized arguments.
 		let (serialized_result_val_ptr, serialized_result_val_len) = match result {
@@ -264,7 +272,7 @@ impl<'a, FE: SandboxCapabilities + Externals + 'a> Externals for GuestExternals<
 		let serialized_result_val = self.supervisor_externals
 			.read_memory(serialized_result_val_ptr, serialized_result_val_len)?;
 		self.supervisor_externals
-			.deallocate(serialized_result_val_ptr);
+			.deallocate(serialized_result_val_ptr)?;
 
 		// We do not have to check the signature here, because it's automatically
 		// checked by wasmi.
@@ -558,6 +566,8 @@ impl Store {
 #[cfg(test)]
 mod tests {
 	use primitives::{Blake2Hasher};
+	use crate::allocator;
+	use crate::sandbox::trap;
 	use crate::wasm_executor::WasmExecutor;
 	use state_machine::TestExternalities;
 	use wabt;
@@ -615,6 +625,32 @@ mod tests {
 		);
 	}
 
+	#[test]
+	fn sandbox_should_trap_when_heap_exhausted() {
+		let mut ext = TestExternalities::<Blake2Hasher>::default();
+		let test_code = include_bytes!("../wasm/target/wasm32-unknown-unknown/release/runtime_test.compact.wasm");
+
+		let code = wabt::wat2wasm(r#"
+		(module
+			(import "env" "assert" (func $assert (param i32)))
+			(func (export "call")
+				i32.const 0
+				call $assert
+			)
+		)
+		"#).unwrap();
+
+		let res = WasmExecutor::new().call(&mut ext, 8, &test_code[..], "test_exhaust_heap", &code);
+		assert_eq!(res.is_err(), true);
+		if let Err(err) = res {
+			let inner_err = err.iter().next().unwrap();
+			assert_eq!(
+				format!("{}", inner_err),
+				format!("{}", wasmi::Error::Trap(trap(allocator::OUT_OF_SPACE)))
+			);
+		}
+	}
+
 	#[test]
 	fn start_called() {
 		let mut ext = TestExternalities::<Blake2Hasher>::default();

substrate/core/executor/src/wasm_executor.rs

@@ -75,10 +75,10 @@ impl<'e, E: Externalities<Blake2Hasher>> sandbox::SandboxCapabilities for Functi
 	fn store_mut(&mut self) -> &mut sandbox::Store {
 		&mut self.sandbox_store
 	}
-	fn allocate(&mut self, len: u32) -> u32 {
+	fn allocate(&mut self, len: u32) -> ::std::result::Result<u32, UserError> {
 		self.heap.allocate(len)
 	}
-	fn deallocate(&mut self, ptr: u32) {
+	fn deallocate(&mut self, ptr: u32) -> ::std::result::Result<(), UserError> {
 		self.heap.deallocate(ptr)
 	}
 	fn write_memory(&mut self, ptr: u32, data: &[u8]) -> ::std::result::Result<(), UserError> {
@@ -133,12 +133,12 @@ impl_function_executor!(this: FunctionExecutor<'e, E>,
 		Ok(())
 	},
 	ext_malloc(size: usize) -> *mut u8 => {
-		let r = this.heap.allocate(size);
+		let r = this.heap.allocate(size)?;
 		debug_trace!(target: "sr-io", "malloc {} bytes at {}", size, r);
 		Ok(r)
 	},
 	ext_free(addr: *mut u8) => {
-		this.heap.deallocate(addr);
+		this.heap.deallocate(addr)?;
 		debug_trace!(target: "sr-io", "free {}", addr);
 		Ok(())
 	},
@@ -252,7 +252,7 @@ impl_function_executor!(this: FunctionExecutor<'e, E>,
 		);
 
 		if let Some(value) = maybe_value {
-			let offset = this.heap.allocate(value.len() as u32) as u32;
+			let offset = this.heap.allocate(value.len() as u32)? as u32;
 			this.memory.set(offset, &value).map_err(|_| UserError("Invalid attempt to set memory in ext_get_allocated_storage"))?;
 			this.memory.write_primitive(written_out, value.len() as u32)
 				.map_err(|_| UserError("Invalid attempt to write written_out in ext_get_allocated_storage"))?;
@@ -291,7 +291,7 @@ impl_function_executor!(this: FunctionExecutor<'e, E>,
 		);
 
 		if let Some(value) = maybe_value {
-			let offset = this.heap.allocate(value.len() as u32) as u32;
+			let offset = this.heap.allocate(value.len() as u32)? as u32;
 			this.memory.set(offset, &value).map_err(|_| UserError("Invalid attempt to set memory in ext_get_allocated_child_storage"))?;
 			this.memory.write_primitive(written_out, value.len() as u32)
 				.map_err(|_| UserError("Invalid attempt to write written_out in ext_get_allocated_child_storage"))?;
@@ -373,7 +373,7 @@ impl_function_executor!(this: FunctionExecutor<'e, E>,
 		let storage_key = this.memory.get(storage_key_data, storage_key_len as usize).map_err(|_| UserError("Invalid attempt to determine storage_key in ext_child_storage_root"))?;
 		let r = this.ext.child_storage_root(&storage_key);
 		if let Some(value) = r {
-			let offset = this.heap.allocate(value.len() as u32) as u32;
+			let offset = this.heap.allocate(value.len() as u32)? as u32;
 			this.memory.set(offset, &value).map_err(|_| UserError("Invalid attempt to set memory in ext_child_storage_root"))?;
 			this.memory.write_primitive(written_out, value.len() as u32)
 				.map_err(|_| UserError("Invalid attempt to write written_out in ext_child_storage_root"))?;
@@ -677,7 +677,7 @@ impl WasmExecutor {
 		let used_mem = memory.used_size();
 		let mut fec = FunctionExecutor::new(memory.clone(), table, ext)?;
 		let size = data.len() as u32;
-		let offset = fec.heap.allocate(size);
+		let offset = fec.heap.allocate(size).map_err(|_| ErrorKind::Runtime)?;
 		memory.set(offset, &data)?;
 
 		let result = module_instance.invoke_export(

substrate/core/executor/src/wasm_utils.rs

@@ -20,7 +20,7 @@ use wasmi::{ValueType, RuntimeValue, HostError};
 use wasmi::nan_preserving_float::{F32, F64};
 use std::fmt;
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub struct UserError(pub &'static str);
 impl fmt::Display for UserError {
 	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {