Code review feedback and bug fix.

* `Memory::new` now returns `Result<Self>` so that an error can be returned if
  the initial requested memory exceeds any limits placed on the store.

* Changed an `Arc` to `Rc` as the `Arc` wasn't necessary.

* Removed `Store` from the `ResourceLimiter` callbacks. Custom resource limiter
  implementations are free to capture any context they want, so no need to
  unnecessarily store a weak reference to `Store` from the proxy type.

* Fixed a bug in the pooling instance allocator where an instance would be
  leaked from the pool. Previously, this would only have happened if the OS was
  unable to make the necessary linear memory available for the instance. With
  these changes, however, the instance might not be created due to limits
  placed on the store. We now properly deallocate the instance on error.

* Added more tests, including one that covers the fix mentioned above.

RELEASES.md

@@ -4,6 +4,19 @@
 
 ## Unreleased
 
+### Added
+
+* Added `Store::with_limits`, `StoreLimits`, and `ResourceLimiter` to the
+  Wasmtime API to help with enforcing resource limits at runtime. The
+  `ResourceLimiter` trait can be implemented by custom resource limiters to
+  decide if linear memories or tables can be grown.
+
+### Changed
+
+* Breaking: `Memory::new` has been changed to return `Result` as creating a
+  host memory object is now a fallible operation when the initial size of
+  the memory exceeds the store limits.
+
 ## 0.26.0
 
 Released 2021-04-05.

crates/c-api/src/memory.rs

@@ -31,13 +31,13 @@ impl wasm_memory_t {
 pub extern "C" fn wasm_memory_new(
     store: &wasm_store_t,
     mt: &wasm_memorytype_t,
-) -> Box<wasm_memory_t> {
-    let memory = Memory::new(&store.store, mt.ty().ty.clone());
-    Box::new(wasm_memory_t {
+) -> Option<Box<wasm_memory_t>> {
+    let memory = Memory::new(&store.store, mt.ty().ty.clone()).ok()?;
+    Some(Box::new(wasm_memory_t {
         ext: wasm_extern_t {
             which: memory.into(),
         },
-    })
+    }))
 }
 
 #[no_mangle]

crates/fuzzing/src/oracles/dummy.rs

@@ -87,7 +87,7 @@ pub fn dummy_table(store: &Store, ty: TableType) -> Table {
 
 /// Construct a dummy memory for the given memory type.
 pub fn dummy_memory(store: &Store, ty: MemoryType) -> Memory {
-    Memory::new(store, ty)
+    Memory::new(store, ty).unwrap()
 }
 
 /// Construct a dummy instance for the given instance type.

crates/runtime/src/instance/allocator.rs

@@ -15,6 +15,7 @@ use std::any::Any;
 use std::cell::RefCell;
 use std::convert::TryFrom;
 use std::ptr::{self, NonNull};
+use std::rc::Rc;
 use std::slice;
 use std::sync::Arc;
 use thiserror::Error;
@@ -61,7 +62,7 @@ pub struct InstanceAllocationRequest<'a> {
     pub module_info_lookup: Option<*const dyn ModuleInfoLookup>,
 
     /// The resource limiter to use for the instance.
-    pub limiter: Option<&'a Arc<dyn ResourceLimiter>>,
+    pub limiter: Option<&'a Rc<dyn ResourceLimiter>>,
 }
 
 /// An link error while instantiating a module.
@@ -595,21 +596,21 @@ impl OnDemandInstanceAllocator {
 
     fn create_tables(
         module: &Module,
-        limiter: &Option<&Arc<dyn ResourceLimiter>>,
-    ) -> PrimaryMap<DefinedTableIndex, Table> {
+        limiter: Option<&Rc<dyn ResourceLimiter>>,
+    ) -> Result<PrimaryMap<DefinedTableIndex, Table>, InstantiationError> {
         let num_imports = module.num_imported_tables;
         let mut tables: PrimaryMap<DefinedTableIndex, _> =
             PrimaryMap::with_capacity(module.table_plans.len() - num_imports);
         for table in &module.table_plans.values().as_slice()[num_imports..] {
-            tables.push(Table::new_dynamic(table, limiter));
+            tables.push(Table::new_dynamic(table, limiter).map_err(InstantiationError::Resource)?);
         }
-        tables
+        Ok(tables)
     }
 
     fn create_memories(
         &self,
         module: &Module,
-        limiter: &Option<&Arc<dyn ResourceLimiter>>,
+        limiter: Option<&Rc<dyn ResourceLimiter>>,
     ) -> Result<PrimaryMap<DefinedMemoryIndex, Memory>, InstantiationError> {
         let creator = self
             .mem_creator
@@ -642,8 +643,8 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
         &self,
         mut req: InstanceAllocationRequest,
     ) -> Result<InstanceHandle, InstantiationError> {
-        let memories = self.create_memories(&req.module, &req.limiter)?;
-        let tables = Self::create_tables(&req.module, &req.limiter);
+        let memories = self.create_memories(&req.module, req.limiter)?;
+        let tables = Self::create_tables(&req.module, req.limiter)?;
 
         let host_state = std::mem::replace(&mut req.host_state, Box::new(()));
 

crates/runtime/src/instance/allocator/pooling.rs

@@ -18,6 +18,7 @@ use std::cell::RefCell;
 use std::cmp::min;
 use std::convert::TryFrom;
 use std::mem;
+use std::rc::Rc;
 use std::sync::{Arc, Mutex};
 use wasmtime_environ::{
     entity::{EntitySet, PrimaryMap},
@@ -376,10 +377,45 @@ impl InstancePool {
         }
     }
 
+    unsafe fn setup_instance(
+        &self,
+        index: usize,
+        mut req: InstanceAllocationRequest,
+    ) -> Result<InstanceHandle, InstantiationError> {
+        let instance = self.instance(index);
+
+        instance.module = req.module.clone();
+        instance.offsets = VMOffsets::new(
+            std::mem::size_of::<*const u8>() as u8,
+            instance.module.as_ref(),
+        );
+        instance.host_state = std::mem::replace(&mut req.host_state, Box::new(()));
+
+        Self::set_instance_memories(
+            instance,
+            self.memories.get(index),
+            self.memories.max_wasm_pages,
+            req.limiter,
+        )?;
+
+        Self::set_instance_tables(
+            instance,
+            self.tables.get(index),
+            self.tables.max_elements,
+            req.limiter,
+        )?;
+
+        initialize_vmcontext(instance, req);
+
+        Ok(InstanceHandle {
+            instance: instance as _,
+        })
+    }
+
     fn allocate(
         &self,
         strategy: PoolingAllocationStrategy,
-        mut req: InstanceAllocationRequest,
+        req: InstanceAllocationRequest,
     ) -> Result<InstanceHandle, InstantiationError> {
         let index = {
             let mut free_list = self.free_list.lock().unwrap();
@@ -390,36 +426,14 @@ impl InstancePool {
             free_list.swap_remove(free_index)
         };
 
-        let host_state = std::mem::replace(&mut req.host_state, Box::new(()));
-
         unsafe {
-            let instance = self.instance(index);
-
-            instance.module = req.module.clone();
-            instance.offsets = VMOffsets::new(
-                std::mem::size_of::<*const u8>() as u8,
-                instance.module.as_ref(),
-            );
-            instance.host_state = host_state;
-
-            Self::set_instance_memories(
-                instance,
-                self.memories.get(index),
-                self.memories.max_wasm_pages,
-                &req.limiter,
-            )?;
-
-            Self::set_instance_tables(
-                instance,
-                self.tables.get(index),
-                self.tables.max_elements,
-                &req.limiter,
-            )?;
-
-            initialize_vmcontext(instance, req);
-
-            Ok(InstanceHandle {
-                instance: instance as _,
+            self.setup_instance(index, req).or_else(|e| {
+                // Deallocate the allocated instance on error
+                let instance = self.instance(index);
+                self.deallocate(&InstanceHandle {
+                    instance: instance as _,
+                });
+                Err(e)
             })
         }
     }
@@ -482,7 +496,7 @@ impl InstancePool {
         instance: &mut Instance,
         mut memories: impl Iterator<Item = *mut u8>,
         max_pages: u32,
-        limiter: &Option<&Arc<dyn ResourceLimiter>>,
+        limiter: Option<&Rc<dyn ResourceLimiter>>,
     ) -> Result<(), InstantiationError> {
         let module = instance.module.as_ref();
 
@@ -514,7 +528,7 @@ impl InstancePool {
         instance: &mut Instance,
         mut tables: impl Iterator<Item = *mut u8>,
         max_elements: u32,
-        limiter: &Option<&Arc<dyn ResourceLimiter>>,
+        limiter: Option<&Rc<dyn ResourceLimiter>>,
     ) -> Result<(), InstantiationError> {
         let module = instance.module.as_ref();
 
@@ -526,9 +540,10 @@ impl InstancePool {
             commit_table_pages(base, max_elements as usize * mem::size_of::<*mut u8>())
                 .map_err(InstantiationError::Resource)?;
 
-            instance
-                .tables
-                .push(Table::new_static(plan, base as _, max_elements, limiter));
+            instance.tables.push(
+                Table::new_static(plan, base as _, max_elements, limiter)
+                    .map_err(InstantiationError::Resource)?,
+            );
         }
 
         let mut dropped_elements = instance.dropped_elements.borrow_mut();

crates/runtime/src/memory.rs

@@ -5,13 +5,13 @@
 use crate::mmap::Mmap;
 use crate::vmcontext::VMMemoryDefinition;
 use crate::ResourceLimiter;
-use anyhow::Result;
+use anyhow::{bail, Result};
 use more_asserts::{assert_ge, assert_le};
 use std::cell::{Cell, RefCell};
 use std::cmp::min;
 use std::convert::TryFrom;
 use std::ptr;
-use std::sync::Arc;
+use std::rc::Rc;
 use wasmtime_environ::{MemoryPlan, MemoryStyle, WASM_MAX_PAGES, WASM_PAGE_SIZE};
 
 /// A memory allocator
@@ -203,20 +203,21 @@ enum MemoryStorage {
 /// Represents an instantiation of a WebAssembly memory.
 pub struct Memory {
     storage: MemoryStorage,
-    limiter: Option<Arc<dyn ResourceLimiter>>,
+    limiter: Option<Rc<dyn ResourceLimiter>>,
 }
 
 impl Memory {
     /// Create a new dynamic (movable) memory instance for the specified plan.
     pub fn new_dynamic(
         plan: &MemoryPlan,
         creator: &dyn RuntimeMemoryCreator,
-        limiter: &Option<&Arc<dyn ResourceLimiter>>,
+        limiter: Option<&Rc<dyn ResourceLimiter>>,
     ) -> Result<Self> {
-        Ok(Self {
-            storage: MemoryStorage::Dynamic(creator.new_memory(plan)?),
-            limiter: limiter.cloned(),
-        })
+        Self::new(
+            plan,
+            MemoryStorage::Dynamic(creator.new_memory(plan)?),
+            limiter,
+        )
     }
 
     /// Create a new static (immovable) memory instance for the specified plan.
@@ -225,21 +226,50 @@ impl Memory {
         base: *mut u8,
         maximum: u32,
         make_accessible: fn(*mut u8, usize) -> Result<()>,
-        limiter: &Option<&Arc<dyn ResourceLimiter>>,
+        limiter: Option<&Rc<dyn ResourceLimiter>>,
+    ) -> Result<Self> {
+        let storage = MemoryStorage::Static {
+            base,
+            size: Cell::new(plan.memory.minimum),
+            maximum: min(plan.memory.maximum.unwrap_or(maximum), maximum),
+            make_accessible,
+            #[cfg(all(feature = "uffd", target_os = "linux"))]
+            guard_page_faults: RefCell::new(Vec::new()),
+        };
+
+        Self::new(plan, storage, limiter)
+    }
+
+    fn new(
+        plan: &MemoryPlan,
+        storage: MemoryStorage,
+        limiter: Option<&Rc<dyn ResourceLimiter>>,
     ) -> Result<Self> {
-        if plan.memory.minimum > 0 {
-            make_accessible(base, plan.memory.minimum as usize * WASM_PAGE_SIZE as usize)?;
+        if let Some(limiter) = limiter {
+            if !limiter.memory_growing(0, plan.memory.minimum, plan.memory.maximum) {
+                bail!(
+                    "memory minimum size of {} pages exceeds memory limits",
+                    plan.memory.minimum
+                );
+            }
+        }
+
+        if let MemoryStorage::Static {
+            base,
+            make_accessible,
+            ..
+        } = &storage
+        {
+            if plan.memory.minimum > 0 {
+                make_accessible(
+                    *base,
+                    plan.memory.minimum as usize * WASM_PAGE_SIZE as usize,
+                )?;
+            }
         }
 
         Ok(Self {
-            storage: MemoryStorage::Static {
-                base,
-                size: Cell::new(plan.memory.minimum),
-                maximum: min(plan.memory.maximum.unwrap_or(maximum), maximum),
-                make_accessible,
-                #[cfg(all(feature = "uffd", target_os = "linux"))]
-                guard_page_faults: RefCell::new(Vec::new()),
-            },
+            storage,
             limiter: limiter.cloned(),
         })
     }
@@ -288,11 +318,15 @@ impl Memory {
     /// Generally, prefer using `InstanceHandle::memory_grow`, which encapsulates
     /// this unsafety.
     pub unsafe fn grow(&self, delta: u32) -> Option<u32> {
-        if let Some(limiter) = &self.limiter {
-            let current = self.size();
-            let desired = current.checked_add(delta)?;
+        let old_size = self.size();
+        if delta == 0 {
+            return Some(old_size);
+        }
+
+        let new_size = old_size.checked_add(delta)?;
 
-            if !limiter.memory_growing(current, desired, self.maximum()) {
+        if let Some(limiter) = &self.limiter {
+            if !limiter.memory_growing(old_size, new_size, self.maximum()) {
                 return None;
             }
         }
@@ -309,13 +343,6 @@ impl Memory {
                 #[cfg(all(feature = "uffd", target_os = "linux"))]
                 self.reset_guard_pages().ok()?;
 
-                let old_size = size.get();
-                if delta == 0 {
-                    return Some(old_size);
-                }
-
-                let new_size = old_size.checked_add(delta)?;
-
                 if new_size > *maximum || new_size >= WASM_MAX_PAGES {
                     return None;
                 }