From 106ac7afdb8bc5e496b7d5aef5f9224c417725c6 Mon Sep 17 00:00:00 2001
From: Sergey Pepyakin <s.pepyakin@gmail.com>
Date: Fri, 2 Feb 2018 00:44:06 +0300
Subject: [PATCH] Fix memory max pages (#33)

* Fix memory max pages

* Get rid of calculate_memory_size.
---
 src/memory.rs | 64 +++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 47 insertions(+), 17 deletions(-)

diff --git a/src/memory.rs b/src/memory.rs
index e6dcd98..d408f44 100644
--- a/src/memory.rs
+++ b/src/memory.rs
@@ -111,14 +111,18 @@ impl MemoryInstance {
 	fn new(limits: ResizableLimits) -> Result<Self, Error> {
 		check_limits(&limits)?;
 
-		let maximum_size = match limits.maximum() {
-			Some(maximum_pages) if maximum_pages > LINEAR_MEMORY_MAX_PAGES =>
-				return Err(Error::Memory(format!("maximum memory size must be at most {} pages", LINEAR_MEMORY_MAX_PAGES))),
-			Some(maximum_pages) => maximum_pages.saturating_mul(LINEAR_MEMORY_PAGE_SIZE),
-			None => u32::MAX,
-		};
-		let initial_size = calculate_memory_size(0, limits.initial(), maximum_size)
-			.ok_or_else(|| Error::Memory(format!("initial memory size must be at most {} pages", LINEAR_MEMORY_MAX_PAGES)))?;
+		let initial_pages = limits.initial();
+		let maximum_pages = limits.maximum().unwrap_or(LINEAR_MEMORY_MAX_PAGES);
+
+		if initial_pages > LINEAR_MEMORY_MAX_PAGES {
+			return Err(Error::Memory(format!("initial memory size must be at most {} pages", LINEAR_MEMORY_MAX_PAGES)));
+		}
+		if maximum_pages > LINEAR_MEMORY_MAX_PAGES {
+			return Err(Error::Memory(format!("maximum memory size must be at most {} pages", LINEAR_MEMORY_MAX_PAGES)));
+		}
+
+		let maximum_size = maximum_pages.saturating_mul(LINEAR_MEMORY_PAGE_SIZE);
+		let initial_size = initial_pages.saturating_mul(LINEAR_MEMORY_PAGE_SIZE);
 
 		let memory = MemoryInstance {
 			limits: limits,
@@ -274,23 +278,49 @@ impl MemoryInstance {
 }
 
 fn calculate_memory_size(old_size: u32, additional_pages: u32, maximum_size: u32) -> Option<u32> {
-	additional_pages
-		.checked_mul(LINEAR_MEMORY_PAGE_SIZE)
-		.and_then(|size| size.checked_add(old_size))
-		.and_then(|size| if size > maximum_size {
-			None
-		} else {
-			Some(size)
-		})
+	let size = additional_pages
+		.saturating_mul(LINEAR_MEMORY_PAGE_SIZE);
+	match size.checked_add(old_size) {
+		Some(size) if size <= maximum_size => Some(size),
+		_ => None,
+	}
 }
 
 #[cfg(test)]
 mod tests {
 
-	use super::MemoryInstance;
+	use super::{MemoryInstance, LINEAR_MEMORY_MAX_PAGES};
 	use Error;
 	use parity_wasm::elements::ResizableLimits;
 
+	#[test]
+	fn alloc() {
+		let fixtures = &[
+			(0, None, true),
+			(0, Some(0), true),
+			(1, None, true),
+			(1, Some(1), true),
+			(0, Some(1), true),
+			(1, Some(0), false),
+			(0, Some(LINEAR_MEMORY_MAX_PAGES), true),
+			(LINEAR_MEMORY_MAX_PAGES, Some(LINEAR_MEMORY_MAX_PAGES), true),
+			(LINEAR_MEMORY_MAX_PAGES, Some(0), false),
+			(LINEAR_MEMORY_MAX_PAGES, None, true),
+		];
+		for &(initial, maybe_max, expected_ok) in fixtures {
+			let result = MemoryInstance::alloc(initial, maybe_max);
+			if result.is_ok() != expected_ok {
+				panic!(
+					"unexpected error, initial={}, max={:?}, expected={}, result={:?}",
+					initial,
+					maybe_max,
+					expected_ok,
+					result,
+				);
+			}
+		}
+	}
+
 	fn create_memory(initial_content: &[u8]) -> MemoryInstance {
 		let mem = MemoryInstance::new(ResizableLimits::new(1, Some(1)))
 			.expect("MemoryInstance created successfuly");