Contracts: return arbitrary sized buffers (#711)

* WIP

* Direct return of output data.

* Docs and renamings.

* Add get_storage test.

* Overwrite the scratch buffer.

substrate/srml/contract/src/exec.rs

@@ -24,13 +24,13 @@ use runtime_primitives::traits::{Zero, CheckedAdd, CheckedSub};
 use runtime_support::{StorageMap, StorageValue};
 use balances::{self, EnsureAccountLiquid};
 
+// TODO: Add logs
 pub struct CreateReceipt<T: Trait> {
 	pub address: T::AccountId,
 }
 
-pub struct CallReceipt {
-	pub return_data: Vec<u8>,
-}
+// TODO: Add logs.
+pub struct CallReceipt;
 
 pub struct ExecutionContext<'a, T: Trait + 'a> {
 	// typically should be dest
@@ -48,6 +48,7 @@ impl<'a, T: Trait> ExecutionContext<'a, T> {
 		value: T::Balance,
 		gas_meter: &mut GasMeter<T>,
 		data: &[u8],
+		output_data: &mut Vec<u8>,
 	) -> Result<CallReceipt, &'static str> {
 		if self.depth == <MaxDepth<T>>::get() as usize {
 			return Err("reached maximum depth, cannot make a call");
@@ -60,7 +61,7 @@ impl<'a, T: Trait> ExecutionContext<'a, T> {
 
 		let dest_code = <CodeOf<T>>::get(&dest);
 
-		let (exec_result, change_set) = {
+		let change_set = {
 			let mut overlay = OverlayAccountDb::new(&self.overlay);
 
 			if value > T::Balance::zero() {
@@ -79,31 +80,26 @@ impl<'a, T: Trait> ExecutionContext<'a, T> {
 				self_account: dest.clone(),
 				depth: self.depth + 1,
 			};
-			let exec_result = if !dest_code.is_empty() {
+
+			if !dest_code.is_empty() {
 				vm::execute(
 					&dest_code,
 					data,
+					output_data,
 					&mut CallContext {
 						ctx: &mut nested,
 						_caller: caller,
 					},
 					gas_meter,
-				).map_err(|_| "vm execute returned error while call")?
-			} else {
-				// that was a plain transfer
-				vm::ExecutionResult {
-					return_data: Vec::new(),
-				}
-			};
+				).map_err(|_| "vm execute returned error while call")?;
+			}
 
-			(exec_result, nested.overlay.into_change_set())
+			nested.overlay.into_change_set()
 		};
 
 		self.overlay.commit(change_set);
 
-		Ok(CallReceipt {
-			return_data: exec_result.return_data,
-		})
+		Ok(CallReceipt)
 	}
 
 	pub fn create(
@@ -111,7 +107,7 @@ impl<'a, T: Trait> ExecutionContext<'a, T> {
 		caller: T::AccountId,
 		endowment: T::Balance,
 		gas_meter: &mut GasMeter<T>,
-		ctor: &[u8],
+		init_code: &[u8],
 		data: &[u8],
 	) -> Result<CreateReceipt<T>, &'static str> {
 		if self.depth == <MaxDepth<T>>::get() as usize {
@@ -123,7 +119,7 @@ impl<'a, T: Trait> ExecutionContext<'a, T> {
 			return Err("not enough gas to pay base create fee");
 		}
 
-		let dest = T::DetermineContractAddress::contract_address_for(ctor, data, &self.self_account);
+		let dest = T::DetermineContractAddress::contract_address_for(init_code, data, &self.self_account);
 		if <CodeOf<T>>::exists(&dest) {
 			// TODO: Is it enough?
 			return Err("contract already exists");
@@ -148,19 +144,20 @@ impl<'a, T: Trait> ExecutionContext<'a, T> {
 				self_account: dest.clone(),
 				depth: self.depth + 1,
 			};
-			let exec_result = {
-				vm::execute(
-					ctor,
-					data,
-					&mut CallContext {
-						ctx: &mut nested,
-						_caller: caller,
-					},
-					gas_meter,
-				).map_err(|_| "vm execute returned error while create")?
-			};
 
-			nested.overlay.set_code(&dest, exec_result.return_data);
+			let mut contract_code = Vec::new();
+			vm::execute(
+				init_code,
+				data,
+				&mut contract_code,
+				&mut CallContext {
+					ctx: &mut nested,
+					_caller: caller,
+				},
+				gas_meter,
+			).map_err(|_| "vm execute returned error while create")?;
+
+			nested.overlay.set_code(&dest, contract_code);
 			nested.overlay.into_change_set()
 		};
 
@@ -269,10 +266,12 @@ impl<'a, 'b: 'a, T: Trait + 'b> vm::Ext for CallContext<'a, 'b, T> {
 		value: T::Balance,
 		gas_meter: &mut GasMeter<T>,
 		data: &[u8],
-	) -> Result<CallReceipt, ()> {
+		output_data: &mut Vec<u8>,
+	) -> Result<(), ()> {
 		let caller = self.ctx.self_account.clone();
 		self.ctx
-			.call(caller, to.clone(), value, gas_meter, data)
+			.call(caller, to.clone(), value, gas_meter, data, output_data)
 			.map_err(|_| ())
+			.map(|_| ())
 	}
 }

substrate/srml/contract/src/lib.rs

@@ -163,7 +163,7 @@ decl_module! {
 			origin,
 			value: T::Balance,
 			gas_limit: T::Gas,
-			ctor: Vec<u8>,
+			init_code: Vec<u8>,
 			data: Vec<u8>
 		) -> Result;
 	}
@@ -227,7 +227,9 @@ impl<T: Trait> Module<T> {
 			depth: 0,
 			overlay: OverlayAccountDb::<T>::new(&account_db::DirectAccountDb),
 		};
-		let result = ctx.call(origin.clone(), dest, value, &mut gas_meter, &data);
+
+		let mut output_data = Vec::new();
+		let result = ctx.call(origin.clone(), dest, value, &mut gas_meter, &data, &mut output_data);
 
 		if let Ok(_) = result {
 			// Commit all changes that made it thus far into the persistant storage.

substrate/srml/contract/src/tests.rs

@@ -505,7 +505,7 @@ fn top_level_create() {
 		));
 
 		// 11 - value sent with the transaction
-		// (3 * 129) - gas spent by the ctor
+		// (3 * 129) - gas spent by the init_code.
 		// (3 * 175) - base gas fee for create (175) (top level) multipled by gas price (3)
 		// ((21 / 3) * 3) - price for contract creation
 		assert_eq!(

substrate/srml/contract/src/vm/env_def/mod.rs

@@ -14,13 +14,13 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate. If not, see <http://www.gnu.org/licenses/>.
 
-use super::{BalanceOf, CallReceipt, CreateReceipt, Ext, GasMeterResult, Runtime};
-use codec::Decode;
+use super::{SpecialTrap, BalanceOf, CreateReceipt, Ext, GasMeterResult, Runtime};
+use codec::{Encode, Decode};
 use parity_wasm::elements::{FunctionType, ValueType};
 use rstd::prelude::*;
 use rstd::string::String;
 use rstd::collections::btree_map::BTreeMap;
-use runtime_primitives::traits::As;
+use runtime_primitives::traits::{As, CheckedMul};
 use sandbox::{self, TypedValue};
 use system;
 use Trait;
@@ -120,8 +120,6 @@ impl<E: Ext> HostFunction<E> {
 // a function set which can be imported by an executed contract.
 define_env!(init_env, <E: Ext>,
 
-	// gas(amount: u32)
-	//
 	// Account for used gas. Traps if gas used is greater than gas limit.
 	//
 	// - amount: How much gas is used.
@@ -134,8 +132,6 @@ define_env!(init_env, <E: Ext>,
 		}
 	},
 
-	// ext_put_storage(location_ptr: u32, value_non_null: u32, value_ptr: u32);
-	//
 	// Change the value at the given location in storage or remove it.
 	//
 	// - location_ptr: pointer into the linear
@@ -144,48 +140,68 @@ define_env!(init_env, <E: Ext>,
 	//   at the given location will be removed.
 	// - value_ptr: pointer into the linear memory
 	//   where the value to set is placed. If `value_non_null` is set to 0, then this parameter is ignored.
-	ext_set_storage(ctx, location_ptr: u32, value_non_null: u32, value_ptr: u32) => {
-		let mut location = [0; 32];
-
-		ctx.memory().get(location_ptr, &mut location)?;
+	// - value_len: the length of the value. If `value_non_null` is set to 0, then this parameter is ignored.
+	ext_set_storage(ctx, key_ptr: u32, value_non_null: u32, value_ptr: u32, value_len: u32) => {
+		let mut key = [0; 32];
+		ctx.memory().get(key_ptr, &mut key)?;
 
 		let value = if value_non_null != 0 {
-			let mut value = [0; 32];
-			ctx.memory().get(value_ptr, &mut value)?;
-			Some(value.to_vec())
+			let mut value_buf = Vec::new();
+			value_buf.resize(value_len as usize, 0);
+			ctx.memory().get(value_ptr, &mut value_buf)?;
+			Some(value_buf)
 		} else {
 			None
 		};
-		ctx.ext.set_storage(&location, value);
+		ctx.ext.set_storage(&key, value);
 
 		Ok(())
 	},
 
-	// ext_get_storage(location_ptr: u32, dest_ptr: u32);
+	// Retrieve the value at the given location from the strorage and return 0.
+	// If there is no entry at the given location then this function will return 1 and
+	// clear the scratch buffer.
 	//
-	// Retrieve the value at the given location from the strorage.
-	// If there is no entry at the given location then all-zero-value
-	// will be returned.
-	//
-	// - location_ptr: pointer into the linear
-	//   memory where the location of the requested value is placed.
-	// - dest_ptr: pointer where contents of the specified storage location
-	//   should be placed.
-	ext_get_storage(ctx, location_ptr: u32, dest_ptr: u32) => {
-		let mut location = [0; 32];
-		ctx.memory().get(location_ptr, &mut location)?;
+	// - key_ptr: pointer into the linear memory where the key
+	//   of the requested value is placed.
+	ext_get_storage(ctx, key_ptr: u32) -> u32 => {
+		let mut key = [0; 32];
+		ctx.memory().get(key_ptr, &mut key)?;
 
-		if let Some(value) = ctx.ext.get_storage(&location) {
-			ctx.memory().set(dest_ptr, &value)?;
+		if let Some(value) = ctx.ext.get_storage(&key) {
+			ctx.scratch_buf = value;
+			Ok(0)
 		} else {
-			ctx.memory().set(dest_ptr, &[0u8; 32])?;
+			ctx.scratch_buf.clear();
+			Ok(1)
 		}
-
-		Ok(())
 	},
 
-	// ext_call(transfer_to_ptr: u32, transfer_to_len: u32, gas: u64, value_ptr: u32, value_len: u32, input_data_ptr: u32, input_data_len: u32)
-	ext_call(ctx, callee_ptr: u32, callee_len: u32, gas: u64, value_ptr: u32, value_len: u32, input_data_ptr: u32, input_data_len: u32) -> u32 => {
+	// Make a call to another contract.
+	//
+	// Returns 0 on the successful execution and puts the result data returned
+	// by the callee into the scratch buffer. Otherwise, returns 1 and clears the scratch
+	// buffer.
+	//
+	// - callee_ptr: a pointer to the address of the callee contract.
+	//   Should be decodable as an `T::AccountId`. Traps otherwise.
+	// - callee_len: length of the address buffer.
+	// - gas: how much gas to devote to the execution.
+	// - value_ptr: a pointer to the buffer with value, how much value to send.
+	//   Should be decodable as a `T::Balance`. Traps otherwise.
+	// - value_len: length of the value buffer.
+	// - input_data_ptr: a pointer to a buffer to be used as input data to the callee.
+	// - input_data_len: length of the input data buffer.
+	ext_call(
+		ctx,
+		callee_ptr: u32,
+		callee_len: u32,
+		gas: u64,
+		value_ptr: u32,
+		value_len: u32,
+		input_data_ptr: u32,
+		input_data_len: u32
+	) -> u32 => {
 		let mut callee = Vec::new();
 		callee.resize(callee_len as usize, 0);
 		ctx.memory().get(callee_ptr, &mut callee)?;
@@ -203,31 +219,51 @@ define_env!(init_env, <E: Ext>,
 		input_data.resize(input_data_len as usize, 0u8);
 		ctx.memory().get(input_data_ptr, &mut input_data)?;
 
+		// Clear the scratch buffer in any case.
+		ctx.scratch_buf.clear();
+
 		let nested_gas_limit = if gas == 0 {
 			ctx.gas_meter.gas_left()
 		} else {
 			<<<E as Ext>::T as Trait>::Gas as As<u64>>::sa(gas)
 		};
 		let ext = &mut ctx.ext;
+		let scratch_buf = &mut ctx.scratch_buf;
 		let call_outcome = ctx.gas_meter.with_nested(nested_gas_limit, |nested_meter| {
 			match nested_meter {
-				Some(nested_meter) => ext.call(&callee, value, nested_meter, &input_data),
+				Some(nested_meter) => ext.call(&callee, value, nested_meter, &input_data, scratch_buf),
 				// there is not enough gas to allocate for the nested call.
 				None => Err(()),
 			}
 		});
 
 		match call_outcome {
-			// TODO: Find a way how to pass return_data back to the this sandbox.
-			Ok(CallReceipt { .. }) => Ok(0),
+			Ok(()) => Ok(0),
 			Err(_) => Ok(1),
 		}
 	},
 
-	// ext_create(code_ptr: u32, code_len: u32, gas: u64, value_ptr: u32, value_len: u32, input_data_ptr: u32, input_data_len: u32) -> u32
+	// Create a contract with code returned by the specified initializer code.
+	//
+	// This function creates an account and executes initializer code. After the execution,
+	// the returned buffer is saved as the code of the created account.
+	//
+	// Returns 0 on the successful contract creation and puts the address
+	// of the created contract into the scratch buffer.
+	// Otherwise, returns 1 and clears the scratch buffer.
+	//
+	// - init_code_ptr: a pointer to the buffer that contains the initializer code.
+	// - init_code_len: length of the initializer code buffer.
+	// - gas: how much gas to devote to the execution of the initializer code.
+	// - value_ptr: a pointer to the buffer with value, how much value to send.
+	//   Should be decodable as a `T::Balance`. Traps otherwise.
+	// - value_len: length of the value buffer.
+	// - input_data_ptr: a pointer to a buffer to be used as input data to the initializer code.
+	// - input_data_len: length of the input data buffer.
 	ext_create(
-		ctx, code_ptr: u32,
-		code_len: u32,
+		ctx,
+		init_code_ptr: u32,
+		init_code_len: u32,
 		gas: u64,
 		value_ptr: u32,
 		value_len: u32,
@@ -240,14 +276,17 @@ define_env!(init_env, <E: Ext>,
 		let value = BalanceOf::<<E as Ext>::T>::decode(&mut &value_buf[..])
 			.ok_or_else(|| sandbox::HostError)?;
 
-		let mut code = Vec::new();
-		code.resize(code_len as usize, 0u8);
-		ctx.memory().get(code_ptr, &mut code)?;
+		let mut init_code = Vec::new();
+		init_code.resize(init_code_len as usize, 0u8);
+		ctx.memory().get(init_code_ptr, &mut init_code)?;
 
 		let mut input_data = Vec::new();
 		input_data.resize(input_data_len as usize, 0u8);
 		ctx.memory().get(input_data_ptr, &mut input_data)?;
 
+		// Clear the scratch buffer in any case.
+		ctx.scratch_buf.clear();
+
 		let nested_gas_limit = if gas == 0 {
 			ctx.gas_meter.gas_left()
 		} else {
@@ -256,27 +295,37 @@ define_env!(init_env, <E: Ext>,
 		let ext = &mut ctx.ext;
 		let create_outcome = ctx.gas_meter.with_nested(nested_gas_limit, |nested_meter| {
 			match nested_meter {
-				Some(nested_meter) => ext.create(&code, value, nested_meter, &input_data),
+				Some(nested_meter) => ext.create(&init_code, value, nested_meter, &input_data),
 				// there is not enough gas to allocate for the nested call.
 				None => Err(()),
 			}
 		});
-
 		match create_outcome {
-			// TODO: Copy an address of the created contract in the sandbox.
-			Ok(CreateReceipt { .. }) => Ok(0),
+			Ok(CreateReceipt { address }) => {
+				// Write the address to the scratch buffer.
+				address.encode_to(&mut ctx.scratch_buf);
+				Ok(0)
+			},
 			Err(_) => Ok(1),
 		}
 	},
 
-	// ext_return(data_ptr: u32, data_len: u32) -> !
+	// Save a data buffer as a result of the execution.
 	ext_return(ctx, data_ptr: u32, data_len: u32) => {
-		let mut data_buf = Vec::new();
-		data_buf.resize(data_len as usize, 0);
-		ctx.memory().get(data_ptr, &mut data_buf)?;
+		let data_len_in_gas = <<<E as Ext>::T as Trait>::Gas as As<u64>>::sa(data_len as u64);
+		let price = (ctx.config.return_data_per_byte_cost)
+			.checked_mul(&data_len_in_gas)
+			.ok_or_else(|| sandbox::HostError)?;
 
-		ctx.store_return_data(data_buf)
-			.map_err(|_| sandbox::HostError)?;
+		match ctx.gas_meter.charge(price) {
+			GasMeterResult::Proceed => (),
+			GasMeterResult::OutOfGas => return Err(sandbox::HostError),
+		}
+
+		ctx.output_data.resize(data_len as usize, 0);
+		ctx.memory.get(data_ptr, &mut ctx.output_data)?;
+
+		ctx.special_trap = Some(SpecialTrap::Return);
 
 		// The trap mechanism is used to immediately terminate the execution.
 		// This trap should be handled appropriately before returning the result
@@ -284,16 +333,12 @@ define_env!(init_env, <E: Ext>,
 		Err(sandbox::HostError)
 	},
 
-	// ext_input_size() -> u32
-	//
-	// Returns size of an input buffer.
+	// Returns the size of the input buffer.
 	ext_input_size(ctx) -> u32 => {
 		Ok(ctx.input_data.len() as u32)
 	},
 
-	// ext_input_copy(dest_ptr: u32, offset: u32, len: u32)
-	//
-	// Copy data from an input buffer starting from `offset` with length `len` into the contract memory.
+	// Copy data from the input buffer starting from `offset` with length `len` into the contract memory.
 	// The region at which the data should be put is specified by `dest_ptr`.
 	ext_input_copy(ctx, dest_ptr: u32, offset: u32, len: u32) => {
 		let offset = offset as usize;
@@ -312,4 +357,29 @@ define_env!(init_env, <E: Ext>,
 
 		Ok(())
 	},
+
+	// Returns the size of the scratch buffer.
+	ext_scratch_size(ctx) -> u32 => {
+		Ok(ctx.scratch_buf.len() as u32)
+	},
+
+	// Copy data from the scratch buffer starting from `offset` with length `len` into the contract memory.
+	// The region at which the data should be put is specified by `dest_ptr`.
+	ext_scratch_copy(ctx, dest_ptr: u32, offset: u32, len: u32) => {
+		let offset = offset as usize;
+		if offset > ctx.scratch_buf.len() {
+			// Offset can't be larger than scratch buffer length.
+			return Err(sandbox::HostError);
+		}
+
+		// This can't panic since `offset <= ctx.scratch_buf.len()`.
+		let src = &ctx.scratch_buf[offset..];
+		if src.len() != len as usize {
+			return Err(sandbox::HostError);
+		}
+
+		ctx.memory().set(dest_ptr, src)?;
+
+		Ok(())
+	},
 );

substrate/srml/contract/src/vm/mod.rs

@@ -17,10 +17,10 @@
 //! This module provides a means for executing contracts
 //! represented in wasm.
 
-use exec::{CallReceipt, CreateReceipt};
+use exec::{CreateReceipt};
 use gas::{GasMeter, GasMeterResult};
 use rstd::prelude::*;
-use runtime_primitives::traits::{As, CheckedMul};
+use runtime_primitives::traits::As;
 use {sandbox, balances, system};
 use Trait;
 
@@ -48,7 +48,6 @@ pub trait Ext {
 	/// Sets the storage entry by the given key to the specified value.
 	fn set_storage(&mut self, key: &[u8], value: Option<Vec<u8>>);
 
-	// TODO: Return the address of the created contract.
 	/// Create a new account for a contract.
 	///
 	/// The newly created account will be associated with the `code`. `value` specifies the amount of value
@@ -68,7 +67,8 @@ pub trait Ext {
 		value: BalanceOf<Self::T>,
 		gas_meter: &mut GasMeter<Self::T>,
 		data: &[u8],
-	) -> Result<CallReceipt, ()>;
+		output_data: &mut Vec<u8>,
+	) -> Result<(), ()>;
 }
 
 /// Error that can occur while preparing or executing wasm smart-contract.
@@ -116,14 +116,15 @@ pub enum Error {
 /// In this runtime traps used not only for signaling about errors but also
 /// to just terminate quickly in some cases.
 enum SpecialTrap {
-	// TODO: Can we pass wrapped memory instance instead of copying?
 	/// Signals that trap was generated in response to call `ext_return` host function.
-	Return(Vec<u8>),
+	Return,
 }
 
 pub(crate) struct Runtime<'a, 'data, E: Ext + 'a> {
 	ext: &'a mut E,
 	input_data: &'data [u8],
+	output_data: &'data mut Vec<u8>,
+	scratch_buf: Vec<u8>,
 	config: &'a Config<E::T>,
 	memory: sandbox::Memory,
 	gas_meter: &'a mut GasMeter<E::T>,
@@ -133,68 +134,35 @@ impl<'a, 'data, E: Ext + 'a> Runtime<'a, 'data, E> {
 	fn memory(&self) -> &sandbox::Memory {
 		&self.memory
 	}
-	/// Save a data buffer as a result of the execution.
-	///
-	/// This function also charges gas for the returning.
-	///
-	/// Returns `Err` if there is not enough gas.
-	fn store_return_data(&mut self, data: Vec<u8>) -> Result<(), ()> {
-		let data_len = <<<E as Ext>::T as Trait>::Gas as As<u64>>::sa(data.len() as u64);
-		let price = (self.config.return_data_per_byte_cost)
-			.checked_mul(&data_len)
-			.ok_or_else(|| ())?;
-
-		match self.gas_meter.charge(price) {
-			GasMeterResult::Proceed => {
-				self.special_trap = Some(SpecialTrap::Return(data));
-				Ok(())
-			}
-			GasMeterResult::OutOfGas => Err(()),
-		}
-	}
 }
 
 fn to_execution_result<E: Ext>(
 	runtime: Runtime<E>,
 	run_err: Option<sandbox::Error>,
-) -> Result<ExecutionResult, Error> {
+) -> Result<(), Error> {
 	// Check the exact type of the error. It could be plain trap or
 	// special runtime trap the we must recognize.
-	let return_data = match (run_err, runtime.special_trap) {
+	match (run_err, runtime.special_trap) {
 		// No traps were generated. Proceed normally.
-		(None, None) => Vec::new(),
+		(None, None) => Ok(()),
 		// Special case. The trap was the result of the execution `return` host function.
-		(Some(sandbox::Error::Execution), Some(SpecialTrap::Return(rd))) => rd,
+		(Some(sandbox::Error::Execution), Some(SpecialTrap::Return)) => Ok(()),
 		// Any other kind of a trap should result in a failure.
-		(Some(_), _) => return Err(Error::Invoke),
+		(Some(_), _) => Err(Error::Invoke),
 		// Any other case (such as special trap flag without actual trap) signifies
 		// a logic error.
 		_ => unreachable!(),
-	};
-
-	Ok(ExecutionResult { return_data })
-}
-
-/// The result of execution of a smart-contract.
-#[derive(PartialEq, Eq)]
-#[cfg_attr(feature = "std", derive(Debug))]
-pub struct ExecutionResult {
-	/// The result produced by the execution of the contract.
-	///
-	/// The contract can designate some buffer at the execution time via a special function.
-	/// If contract called this function with non-empty buffer it will be copied here.
-	///
-	/// Note that gas is already charged for returning the data.
-	pub return_data: Vec<u8>,
+	}
 }
 
 /// Execute the given code as a contract.
 pub fn execute<'a, E: Ext>(
 	code: &[u8],
 	input_data: &[u8],
+	output_data: &mut Vec<u8>,
 	ext: &'a mut E,
 	gas_meter: &mut GasMeter<E::T>,
-) -> Result<ExecutionResult, Error> {
+) -> Result<(), Error> {
 	let config = Config::default();
 	let env = env_def::init_env();
 
@@ -212,7 +180,9 @@ pub fn execute<'a, E: Ext>(
 	let mut runtime = Runtime {
 		ext,
 		input_data,
+		output_data,
 		config: &config,
+		scratch_buf: Vec::new(),
 		memory,
 		gas_meter,
 		special_trap: None,
@@ -323,7 +293,8 @@ mod tests {
 			value: u64,
 			gas_meter: &mut GasMeter<Test>,
 			data: &[u8],
-		) -> Result<CallReceipt, ()> {
+			_output_data: &mut Vec<u8>,
+		) -> Result<(), ()> {
 			self.transfers.push(TransferEntry {
 				to: *to,
 				value,
@@ -332,9 +303,7 @@ mod tests {
 			});
 			// Assume for now that it was just a plain transfer.
 			// TODO: Add tests for different call outcomes.
-			Ok(CallReceipt {
-				return_data: Vec::new(),
-			})
+			Ok(())
 		}
 	}
 
@@ -383,6 +352,7 @@ mod tests {
 		execute(
 			&code_transfer,
 			&[],
+			&mut Vec::new(),
 			&mut mock_ext,
 			&mut GasMeter::with_limit(50_000, 1),
 		).unwrap();
@@ -444,6 +414,7 @@ mod tests {
 		execute(
 			&code_create,
 			&[],
+			&mut Vec::new(),
 			&mut mock_ext,
 			&mut GasMeter::with_limit(50_000, 1),
 		).unwrap();
@@ -482,6 +453,7 @@ mod tests {
 			execute(
 				&code_mem,
 				&[],
+				&mut Vec::new(),
 				&mut mock_ext,
 				&mut GasMeter::with_limit(100_000, 1)
 			),
@@ -534,6 +506,7 @@ mod tests {
 		execute(
 			&code_transfer,
 			&[],
+			&mut Vec::new(),
 			&mut mock_ext,
 			&mut GasMeter::with_limit(50_000, 1),
 		).unwrap();
@@ -550,4 +523,89 @@ mod tests {
 			}]
 		);
 	}
+
+	const CODE_GET_STORAGE: &str = r#"
+(module
+	(import "env" "ext_get_storage" (func $ext_get_storage (param i32) (result i32)))
+	(import "env" "ext_scratch_size" (func $ext_scratch_size (result i32)))
+	(import "env" "ext_scratch_copy" (func $ext_scratch_copy (param i32 i32 i32)))
+	(import "env" "ext_return" (func $ext_return (param i32 i32)))
+	(import "env" "memory" (memory 1 1))
+
+	(func $assert (param i32)
+		(block $ok
+			(br_if $ok
+				(get_local 0)
+			)
+			(unreachable)
+		)
+	)
+
+	(func (export "call")
+		(local $buf_size i32)
+
+
+		;; Load a storage value into the scratch buf.
+		(call $assert
+			(i32.eq
+				(call $ext_get_storage
+					(i32.const 4)		;; The pointer to the storage key to fetch
+				)
+
+				;; Return value 0 means that the value is found and there were
+				;; no errors.
+				(i32.const 0)
+			)
+		)
+
+		;; Find out the size of the scratch buffer
+		(set_local $buf_size
+			(call $ext_scratch_size)
+		)
+
+		;; Copy scratch buffer into this contract memory.
+		(call $ext_scratch_copy
+			(i32.const 36)		;; The pointer where to store the scratch buffer contents,
+								;; 36 = 4 + 32
+			(i32.const 0)		;; Offset from the start of the scratch buffer.
+			(get_local			;; Count of bytes to copy.
+				$buf_size
+			)
+		)
+
+		;; Return the contents of the buffer
+		(call $ext_return
+			(i32.const 36)
+			(get_local $buf_size)
+		)
+
+		;; env:ext_return doesn't return, so this is effectively unreachable.
+		(unreachable)
+	)
+
+	(data (i32.const 4) "\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11")
+)
+"#;
+
+	#[test]
+	fn get_storage_puts_data_into_scratch_buf() {
+		let code_get_storage = wabt::wat2wasm(CODE_GET_STORAGE).unwrap();
+
+		let mut mock_ext = MockExt::default();
+		mock_ext.storage.insert([0x11; 32].to_vec(), [0x22; 32].to_vec());
+
+		let mut return_buf = Vec::new();
+		execute(
+			&code_get_storage,
+			&[],
+			&mut return_buf,
+			&mut mock_ext,
+			&mut GasMeter::with_limit(50_000, 1),
+		).unwrap();
+
+		assert_eq!(
+			return_buf,
+			[0x22; 32].to_vec(),
+		);
+	}
 }