From d0ef63a2b800b11811ca9c6ed6f5c68ef871b667 Mon Sep 17 00:00:00 2001
From: Vadim Smirnov <78953304+ukint-vs@users.noreply.github.com>
Date: Mon, 18 Oct 2021 15:27:59 +0300
Subject: [PATCH] memory::grow(), memory::size() (#1)

---
 primitives/sandbox/src/embedded_executor.rs | 11 +++++++++++
 primitives/sandbox/src/lib.rs               | 13 +++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/primitives/sandbox/src/embedded_executor.rs b/primitives/sandbox/src/embedded_executor.rs
index 115c3192f3d89..2739674a76a0d 100644
--- a/primitives/sandbox/src/embedded_executor.rs
+++ b/primitives/sandbox/src/embedded_executor.rs
@@ -57,6 +57,17 @@ impl super::SandboxMemory for Memory {
 		self.memref.set(ptr, value).map_err(|_| Error::OutOfBounds)?;
 		Ok(())
 	}
+
+	pub fn grow(&self, pages: u32) -> Result<u32, Error> {
+		self.memref
+			.grow(Pages(pages as usize))
+			.map(|prev| (prev.0 as u32))
+			.map_err(|_| Error::MemoryGrow)
+	}
+
+	pub fn size(&self) -> u32 {
+		self.memref.current_size().0 as u32
+	}
 }
 
 struct HostFuncIndex(usize);
diff --git a/primitives/sandbox/src/lib.rs b/primitives/sandbox/src/lib.rs
index b6b4a5a97da8c..82cadc1c9a710 100644
--- a/primitives/sandbox/src/lib.rs
+++ b/primitives/sandbox/src/lib.rs
@@ -70,6 +70,9 @@ pub enum Error {
 	/// Note that if wasm module makes an out-of-bounds access then trap will occur.
 	OutOfBounds,
 
+	/// Trying to grow memory by more than maximum limit.
+	MemoryGrow,
+
 	/// Failed to invoke the start function or an exported function for some reason.
 	Execution,
 }
@@ -113,6 +116,16 @@ pub trait SandboxMemory: Sized + Clone {
 	///
 	/// Returns `Err` if the range is out-of-bounds.
 	fn set(&self, ptr: u32, value: &[u8]) -> Result<(), Error>;
+
+	/// Grow memory with provided number of pages.
+	///
+	/// Returns `Err` if attempted to allocate more memory than permited by the limit.
+	pub fn grow(&self, pages: u32) -> Result<u32, Error>;
+
+	/// Returns current memory size.
+	///
+	/// Maximum memory size cannot exceed 65536 pages or 4GiB.
+	pub fn size(&self) -> u32;
 }
 
 /// Struct that can be used for defining an environment for a sandboxed module.