Support for ResourceLimiter API

crates/wasmi/src/engine/executor.rs

@@ -24,6 +24,7 @@ use crate::{
         ValueStack,
     },
     func::FuncEntity,
+    store::ResourceLimiterRef,
     table::TableEntity,
     FuelConsumptionMode,
     Func,
@@ -92,15 +93,25 @@ pub enum ReturnOutcome {
 ///
 /// If the Wasm execution traps.
 #[inline(never)]
-pub fn execute_wasm<'engine>(
-    ctx: &mut StoreInner,
+pub fn execute_wasm<'ctx, 'engine>(
+    ctx: &'ctx mut StoreInner,
     cache: &'engine mut InstanceCache,
     value_stack: &'engine mut ValueStack,
     call_stack: &'engine mut CallStack,
     code_map: &'engine CodeMap,
     const_pool: ConstPoolView<'engine>,
+    resource_limiter: ResourceLimiterRef<'ctx>,
 ) -> Result<WasmOutcome, TrapCode> {
-    Executor::new(ctx, cache, value_stack, call_stack, code_map, const_pool).execute()
+    Executor::new(
+        ctx,
+        cache,
+        value_stack,
+        call_stack,
+        code_map,
+        const_pool,
+        resource_limiter,
+    )
+    .execute()
 }
 
 /// The function signature of Wasm load operations.
@@ -168,6 +179,7 @@ struct Executor<'ctx, 'engine> {
     code_map: &'engine CodeMap,
     /// A read-only view to a pool of constant values.
     const_pool: ConstPoolView<'engine>,
+    resource_limiter: ResourceLimiterRef<'ctx>,
 }
 
 macro_rules! forward_call {
@@ -195,6 +207,7 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> {
         call_stack: &'engine mut CallStack,
         code_map: &'engine CodeMap,
         const_pool: ConstPoolView<'engine>,
+        resource_limiter: ResourceLimiterRef<'ctx>,
     ) -> Self {
         let frame = call_stack.pop().expect("must have frame on the call stack");
         let sp = value_stack.stack_ptr();
@@ -208,6 +221,7 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> {
             call_stack,
             code_map,
             const_pool,
+            resource_limiter,
         }
     }
 
@@ -1097,7 +1111,7 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> {
                 let new_pages = this
                     .ctx
                     .resolve_memory_mut(memory)
-                    .grow(delta)
+                    .grow(delta, &mut this.resource_limiter)
                     .map(u32::from)
                     .map_err(|_| EntityGrowError::InvalidGrow)?;
                 // The `memory.grow` operation might have invalidated the cached
@@ -1219,7 +1233,7 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> {
                 let table = this.cache.get_table(this.ctx, table_index);
                 this.ctx
                     .resolve_table_mut(&table)
-                    .grow_untyped(delta, init)
+                    .grow_untyped(delta, init, &mut this.resource_limiter)
                     .map_err(|_| EntityGrowError::InvalidGrow)
             },
         );

crates/wasmi/src/engine/mod.rs

@@ -48,6 +48,7 @@ pub(crate) use self::{
 use crate::{
     core::{Trap, TrapCode},
     func::FuncEntity,
+    store::ResourceLimiterRef,
     AsContext,
     AsContextMut,
     Func,
@@ -798,13 +799,19 @@ impl<'engine> EngineExecutor<'engine> {
         let call_stack = &mut self.stack.frames;
         let code_map = &self.res.code_map;
         let const_pool = self.res.const_pool.view();
+        let resource_limiter = ResourceLimiterRef(match &mut ctx.store.limiter {
+            Some(q) => Some(q.0(&mut ctx.store.data)),
+            None => None,
+        });
+
         execute_wasm(
             store_inner,
             cache,
             value_stack,
             call_stack,
             code_map,
             const_pool,
+            resource_limiter,
         )
         .map_err(make_trap)
     }

crates/wasmi/src/lib.rs

@@ -93,6 +93,7 @@ mod externref;
 mod func;
 mod global;
 mod instance;
+mod limits;
 mod linker;
 mod memory;
 mod module;
@@ -145,6 +146,7 @@ pub use self::{
     },
     global::{Global, GlobalType, Mutability},
     instance::{Export, ExportsIter, Extern, ExternType, Instance},
+    limits::ResourceLimiter,
     linker::Linker,
     memory::{Memory, MemoryType},
     module::{

crates/wasmi/src/limits.rs

@@ -0,0 +1,39 @@
+use crate::{memory::MemoryError, table::TableError};
+
+/// Value returned by [`ResourceLimiter::instances`] default method
+pub const DEFAULT_INSTANCE_LIMIT: usize = 10000;
+
+/// Value returned by [`ResourceLimiter::tables`] default method
+pub const DEFAULT_TABLE_LIMIT: usize = 10000;
+
+/// Value returned by [`ResourceLimiter::memories`] default method
+pub const DEFAULT_MEMORY_LIMIT: usize = 10000;
+
+pub trait ResourceLimiter {
+    fn memory_growing(
+        &mut self,
+        current: usize,
+        desired: usize,
+        maximum: Option<usize>,
+    ) -> Result<bool, MemoryError>;
+
+    fn table_growing(
+        &mut self,
+        current: u32,
+        desired: u32,
+        maximum: Option<u32>,
+    ) -> Result<bool, TableError>;
+
+    // Provided methods
+    fn memory_grow_failed(&mut self, _error: &MemoryError) {}
+    fn table_grow_failed(&mut self, _error: &TableError) {}
+    fn instances(&self) -> usize {
+        DEFAULT_INSTANCE_LIMIT
+    }
+    fn tables(&self) -> usize {
+        DEFAULT_TABLE_LIMIT
+    }
+    fn memories(&self) -> usize {
+        DEFAULT_MEMORY_LIMIT
+    }
+}

crates/wasmi/src/linker.rs

@@ -649,6 +649,8 @@ impl<T> Linker<T> {
         module: &Module,
     ) -> Result<InstancePre, Error> {
         assert!(Engine::same(self.engine(), context.as_context().engine()));
+        // TODO: possibly add further resource limtation here on number of externals.
+        // Not clear that user can't import the same external lots of times to inflate this.
         let externals = module
             .imports()
             .map(|import| self.process_import(&mut context, import))

crates/wasmi/src/memory/error.rs

@@ -20,6 +20,8 @@ pub enum MemoryError {
         /// The [`MemoryType`] which is supposed to be a supertype of `ty`.
         other: MemoryType,
     },
+    /// Tried to create too many memories
+    TooManyMemories,
 }
 
 impl Display for MemoryError {
@@ -40,6 +42,9 @@ impl Display for MemoryError {
             Self::InvalidSubtype { ty, other } => {
                 write!(f, "memory type {ty:?} is not a subtype of {other:?}",)
             }
+            Self::TooManyMemories => {
+                write!(f, "too many memories")
+            }
         }
     }
 }

crates/wasmi/src/memory/mod.rs

@@ -5,6 +5,8 @@ mod error;
 #[cfg(test)]
 mod tests;
 
+use crate::store::ResourceLimiterRef;
+
 use self::buffer::ByteBuffer;
 pub use self::{
     data::{DataSegment, DataSegmentEntity, DataSegmentIdx},
@@ -127,17 +129,35 @@ pub struct MemoryEntity {
 
 impl MemoryEntity {
     /// Creates a new memory entity with the given memory type.
-    pub fn new(memory_type: MemoryType) -> Result<Self, MemoryError> {
+    pub fn new(
+        memory_type: MemoryType,
+        limiter: &mut ResourceLimiterRef<'_>,
+    ) -> Result<Self, MemoryError> {
         let initial_pages = memory_type.initial_pages();
-        let initial_len = initial_pages
-            .to_bytes()
-            .ok_or(MemoryError::OutOfBoundsAllocation)?;
-        let memory = Self {
-            bytes: ByteBuffer::new(initial_len),
-            memory_type,
-            current_pages: initial_pages,
-        };
-        Ok(memory)
+        let initial_len = initial_pages.to_bytes();
+        let maximum_pages = memory_type.maximum_pages().unwrap_or_else(Pages::max);
+        let maximum_len = maximum_pages.to_bytes();
+
+        if let Some(limiter) = &mut limiter.0 {
+            if !limiter.memory_growing(0, initial_len.unwrap_or(usize::MAX), maximum_len)? {
+                return Err(MemoryError::OutOfBoundsAllocation);
+            }
+        }
+
+        if let Some(initial_len) = initial_len {
+            let memory = Self {
+                bytes: ByteBuffer::new(initial_len),
+                memory_type,
+                current_pages: initial_pages,
+            };
+            Ok(memory)
+        } else {
+            let err = MemoryError::OutOfBoundsAllocation;
+            if let Some(limiter) = &mut limiter.0 {
+                limiter.memory_grow_failed(&err)
+            }
+            Err(err)
+        }
     }
 
     /// Returns the memory type of the linear memory.
@@ -171,25 +191,61 @@ impl MemoryEntity {
     ///
     /// If the linear memory would grow beyond its maximum limit after
     /// the grow operation.
-    pub fn grow(&mut self, additional: Pages) -> Result<Pages, MemoryError> {
+    pub fn grow(
+        &mut self,
+        additional: Pages,
+        limiter: &mut ResourceLimiterRef<'_>,
+    ) -> Result<Pages, MemoryError> {
         let current_pages = self.current_pages();
         if additional == Pages::from(0) {
             // Nothing to do in this case. Bail out early.
             return Ok(current_pages);
         }
+
         let maximum_pages = self.ty().maximum_pages().unwrap_or_else(Pages::max);
-        let new_pages = current_pages
-            .checked_add(additional)
-            .filter(|&new_pages| new_pages <= maximum_pages)
-            .ok_or(MemoryError::OutOfBoundsGrowth)?;
-        let new_size = new_pages
-            .to_bytes()
-            .ok_or(MemoryError::OutOfBoundsAllocation)?;
-        // At this point it is okay to grow the underlying virtual memory
-        // by the given amount of additional pages.
-        self.bytes.grow(new_size);
-        self.current_pages = new_pages;
-        Ok(current_pages)
+        let desired_pages = current_pages.checked_add(additional);
+
+        // ResourceLimiter gets first look at the request.
+        if let Some(limiter) = &mut limiter.0 {
+            let current_size = current_pages.to_bytes().unwrap_or(usize::MAX);
+            let desired_size = desired_pages
+                .unwrap_or_else(Pages::max)
+                .to_bytes()
+                .unwrap_or(usize::MAX);
+            let maximum_size = maximum_pages.to_bytes();
+            if !limiter.memory_growing(current_size, desired_size, maximum_size)? {
+                return Err(MemoryError::OutOfBoundsGrowth);
+            }
+        }
+
+        let ret: Result<Pages, MemoryError>;
+
+        if let Some(new_pages) = desired_pages {
+            if new_pages <= maximum_pages {
+                if let Some(new_size) = new_pages.to_bytes() {
+                    // At this point it is okay to grow the underlying virtual memory
+                    // by the given amount of additional pages.
+                    self.bytes.grow(new_size);
+                    self.current_pages = new_pages;
+                    ret = Ok(current_pages)
+                } else {
+                    ret = Err(MemoryError::OutOfBoundsAllocation);
+                }
+            } else {
+                ret = Err(MemoryError::OutOfBoundsGrowth);
+            }
+        } else {
+            ret = Err(MemoryError::OutOfBoundsGrowth);
+        }
+
+        // If there was an error, ResourceLimiter gets to see.
+        if let Err(e) = &ret {
+            if let Some(limiter) = &mut limiter.0 {
+                limiter.memory_grow_failed(e)
+            }
+        }
+
+        ret
     }
 
     /// Returns a shared slice to the bytes underlying to the byte buffer.
@@ -257,8 +313,13 @@ impl Memory {
     ///
     /// If more than [`u32::MAX`] much linear memory is allocated.
     pub fn new(mut ctx: impl AsContextMut, ty: MemoryType) -> Result<Self, MemoryError> {
-        let entity = MemoryEntity::new(ty)?;
-        let memory = ctx.as_context_mut().store.inner.alloc_memory(entity);
+        let store = ctx.as_context_mut().store;
+        let mut resource_limiter = ResourceLimiterRef(match &mut store.limiter {
+            Some(q) => Some(q.0(&mut store.data)),
+            None => None,
+        });
+        let entity = MemoryEntity::new(ty, &mut resource_limiter)?;
+        let memory = store.inner.alloc_memory(entity);
         Ok(memory)
     }
 
@@ -318,12 +379,13 @@ impl Memory {
         &self,
         mut ctx: impl AsContextMut,
         additional: Pages,
+        limiter: &mut ResourceLimiterRef<'_>,
     ) -> Result<Pages, MemoryError> {
         ctx.as_context_mut()
             .store
             .inner
             .resolve_memory_mut(self)
-            .grow(additional)
+            .grow(additional, limiter)
     }
 
     /// Returns a shared slice to the bytes underlying the [`Memory`].

crates/wasmi/src/module/instantiate/error.rs

@@ -49,6 +49,7 @@ pub enum InstantiationError {
         /// The index of the found `start` function.
         index: u32,
     },
+    TooManyInstances,
 }
 
 #[cfg(feature = "std")]
@@ -85,6 +86,7 @@ impl Display for InstantiationError {
             Self::Table(error) => Display::fmt(error, f),
             Self::Memory(error) => Display::fmt(error, f),
             Self::Global(error) => Display::fmt(error, f),
+            Self::TooManyInstances => write!(f, "too many instances")
         }
     }
 }

crates/wasmi/src/module/instantiate/mod.rs

@@ -57,6 +57,8 @@ impl Module {
         let handle = context.as_context_mut().store.inner.alloc_instance();
         let mut builder = InstanceEntity::build(self);
 
+        context.as_context_mut().store.bump_resource_counts(self)?;
+
         self.extract_imports(&mut context, &mut builder, externals)?;
         self.extract_functions(&mut context, &mut builder, handle);
         self.extract_tables(&mut context, &mut builder)?;