Add a limits and trap-on-OOM options to the CLI (#6149)

* Add a limits and trap-on-OOM options to the CLI

This commit adds new options to the `wasmtime` CLI to control the
`Store::limiter` behavior at runtime. This enables artificially
restriction the memory usage of the wasm instance, for example.
Additionally a new option is added to `StoreLimits` to force a trap on
growth failure. This is intended to help quickly debug modules with
backtraces if OOM is happening, or even diagnosing if OOM is happening
in the first place.

* Fix compile of fuzzing oracle

crates/fuzzing/src/oracles.rs

@@ -97,13 +97,18 @@ impl StoreLimits {
 }
 
 impl ResourceLimiter for StoreLimits {
-    fn memory_growing(&mut self, current: usize, desired: usize, _maximum: Option<usize>) -> bool {
-        self.alloc(desired - current)
+    fn memory_growing(
+        &mut self,
+        current: usize,
+        desired: usize,
+        _maximum: Option<usize>,
+    ) -> Result<bool> {
+        Ok(self.alloc(desired - current))
     }
 
-    fn table_growing(&mut self, current: u32, desired: u32, _maximum: Option<u32>) -> bool {
+    fn table_growing(&mut self, current: u32, desired: u32, _maximum: Option<u32>) -> Result<bool> {
         let delta = (desired - current) as usize * std::mem::size_of::<usize>();
-        self.alloc(delta)
+        Ok(self.alloc(delta))
     }
 }
 

crates/wasmtime/src/limits.rs

@@ -1,3 +1,5 @@
+use anyhow::{bail, Result};
+
 /// Value returned by [`ResourceLimiter::instances`] default method
 pub const DEFAULT_INSTANCE_LIMIT: usize = 10000;
 /// Value returned by [`ResourceLimiter::tables`] default method
@@ -40,22 +42,36 @@ pub trait ResourceLimiter {
     /// The `current` and `desired` amounts are guaranteed to always be
     /// multiples of the WebAssembly page size, 64KiB.
     ///
-    /// This function should return `true` to indicate that the growing
-    /// operation is permitted or `false` if not permitted. Returning `true`
-    /// when a maximum has been exceeded will have no effect as the linear
-    /// memory will not grow.
+    /// This function is not invoked when the requested size doesn't fit in
+    /// `usize`. Additionally this function is not invoked for shared memories
+    /// at this time. Otherwise even when `desired` exceeds `maximum` this
+    /// function will still be called.
+    ///
+    /// ## Return Value
     ///
-    /// This function is not guaranteed to be invoked for all requests to
-    /// `memory.grow`. Requests where the allocation requested size doesn't fit
-    /// in `usize` or exceeds the memory's listed maximum size may not invoke
-    /// this method.
+    /// If `Ok(true)` is returned from this function then the growth operation
+    /// is allowed. This means that the wasm `memory.grow` instruction will
+    /// return with the `desired` size, in wasm pages. Note that even if
+    /// `Ok(true)` is returned, though, if `desired` exceeds `maximum` then the
+    /// growth operation will still fail.
     ///
-    /// Returning `false` from this method will cause the `memory.grow`
+    /// If `Ok(false)` is returned then this will cause the `memory.grow`
     /// instruction in a module to return -1 (failure), or in the case of an
     /// embedder API calling [`Memory::new`](crate::Memory::new) or
     /// [`Memory::grow`](crate::Memory::grow) an error will be returned from
     /// those methods.
-    fn memory_growing(&mut self, current: usize, desired: usize, maximum: Option<usize>) -> bool;
+    ///
+    /// If `Err(e)` is returned then the `memory.grow` function will behave
+    /// as if a trap has been raised. Note that this is not necessarily
+    /// compliant with the WebAssembly specification but it can be a handy and
+    /// useful tool to get a precise backtrace at "what requested so much memory
+    /// to cause a growth failure?".
+    fn memory_growing(
+        &mut self,
+        current: usize,
+        desired: usize,
+        maximum: Option<usize>,
+    ) -> Result<bool>;
 
     /// Notifies the resource limiter that growing a linear memory, permitted by
     /// the `memory_growing` method, has failed.
@@ -73,17 +89,12 @@ pub trait ResourceLimiter {
     /// * `maximum` is either the table's maximum or a maximum from an instance
     ///   allocator.  A value of `None` indicates that the table is unbounded.
     ///
-    /// This function should return `true` to indicate that the growing
-    /// operation is permitted or `false` if not permitted. Returning `true`
-    /// when a maximum has been exceeded will have no effect as the table will
-    /// not grow.
-    ///
     /// Currently in Wasmtime each table element requires a pointer's worth of
     /// space (e.g. `mem::size_of::<usize>()`).
     ///
-    /// Like `memory_growing` returning `false` from this function will cause
-    /// `table.grow` to return -1 or embedder APIs will return an error.
-    fn table_growing(&mut self, current: u32, desired: u32, maximum: Option<u32>) -> bool;
+    /// See the details on the return values for `memory_growing` for what the
+    /// return value of this function indicates.
+    fn table_growing(&mut self, current: u32, desired: u32, maximum: Option<u32>) -> Result<bool>;
 
     /// Notifies the resource limiter that growing a linear memory, permitted by
     /// the `table_growing` method, has failed.
@@ -146,13 +157,18 @@ pub trait ResourceLimiterAsync {
         current: usize,
         desired: usize,
         maximum: Option<usize>,
-    ) -> bool;
+    ) -> Result<bool>;
 
     /// Identical to [`ResourceLimiter::memory_grow_failed`]
     fn memory_grow_failed(&mut self, _error: &anyhow::Error) {}
 
     /// Asynchronous version of [`ResourceLimiter::table_growing`]
-    async fn table_growing(&mut self, current: u32, desired: u32, maximum: Option<u32>) -> bool;
+    async fn table_growing(
+        &mut self,
+        current: u32,
+        desired: u32,
+        maximum: Option<u32>,
+    ) -> Result<bool>;
 
     /// Identical to [`ResourceLimiter::table_grow_failed`]
     fn table_grow_failed(&mut self, _error: &anyhow::Error) {}
@@ -187,7 +203,10 @@ impl StoreLimitsBuilder {
 
     /// The maximum number of bytes a linear memory can grow to.
     ///
-    /// Growing a linear memory beyond this limit will fail.
+    /// Growing a linear memory beyond this limit will fail. This limit is
+    /// applied to each linear memory individually, so if a wasm module has
+    /// multiple linear memories then they're all allowed to reach up to the
+    /// `limit` specified.
     ///
     /// By default, linear memory will not be limited.
     pub fn memory_size(mut self, limit: usize) -> Self {
@@ -197,7 +216,9 @@ impl StoreLimitsBuilder {
 
     /// The maximum number of elements in a table.
     ///
-    /// Growing a table beyond this limit will fail.
+    /// Growing a table beyond this limit will fail. This limit is applied to
+    /// each table individually, so if a wasm module has multiple tables then
+    /// they're all allowed to reach up to the `limit` specified.
     ///
     /// By default, table elements will not be limited.
     pub fn table_elements(mut self, limit: u32) -> Self {
@@ -235,6 +256,20 @@ impl StoreLimitsBuilder {
         self
     }
 
+    /// Indicates that a trap should be raised whenever a growth operation
+    /// would fail.
+    ///
+    /// This operation will force `memory.grow` and `table.grow` instructions
+    /// to raise a trap on failure instead of returning -1. This is not
+    /// necessarily spec-compliant, but it can be quite handy when debugging a
+    /// module that fails to allocate memory and might behave oddly as a result.
+    ///
+    /// This value defaults to `false`.
+    pub fn trap_on_grow_failure(mut self, trap: bool) -> Self {
+        self.0.trap_on_grow_failure = trap;
+        self
+    }
+
     /// Consumes this builder and returns the [`StoreLimits`].
     pub fn build(self) -> StoreLimits {
         self.0
@@ -249,12 +284,14 @@ impl StoreLimitsBuilder {
 /// This is a convenience type included to avoid needing to implement the
 /// [`ResourceLimiter`] trait if your use case fits in the static configuration
 /// that this [`StoreLimits`] provides.
+#[derive(Clone, Debug)]
 pub struct StoreLimits {
     memory_size: Option<usize>,
     table_elements: Option<u32>,
     instances: usize,
     tables: usize,
     memories: usize,
+    trap_on_grow_failure: bool,
 }
 
 impl Default for StoreLimits {
@@ -265,22 +302,44 @@ impl Default for StoreLimits {
             instances: DEFAULT_INSTANCE_LIMIT,
             tables: DEFAULT_TABLE_LIMIT,
             memories: DEFAULT_MEMORY_LIMIT,
+            trap_on_grow_failure: false,
         }
     }
 }
 
 impl ResourceLimiter for StoreLimits {
-    fn memory_growing(&mut self, _current: usize, desired: usize, _maximum: Option<usize>) -> bool {
-        match self.memory_size {
+    fn memory_growing(
+        &mut self,
+        _current: usize,
+        desired: usize,
+        maximum: Option<usize>,
+    ) -> Result<bool> {
+        let allow = match self.memory_size {
             Some(limit) if desired > limit => false,
-            _ => true,
+            _ => match maximum {
+                Some(max) if desired > max => false,
+                _ => true,
+            },
+        };
+        if !allow && self.trap_on_grow_failure {
+            bail!("forcing trap when growing memory to {desired} bytes")
+        } else {
+            Ok(allow)
         }
     }
 
-    fn table_growing(&mut self, _current: u32, desired: u32, _maximum: Option<u32>) -> bool {
-        match self.table_elements {
+    fn table_growing(&mut self, _current: u32, desired: u32, maximum: Option<u32>) -> Result<bool> {
+        let allow = match self.table_elements {
             Some(limit) if desired > limit => false,
-            _ => true,
+            _ => match maximum {
+                Some(max) if desired > max => false,
+                _ => true,
+            },
+        };
+        if !allow && self.trap_on_grow_failure {
+            bail!("forcing trap when growing table to {desired} elements")
+        } else {
+            Ok(allow)
         }
     }
 

crates/wasmtime/src/store.rs

@@ -1874,19 +1874,18 @@ unsafe impl<T> wasmtime_runtime::Store for StoreInner<T> {
     ) -> Result<bool, anyhow::Error> {
         match self.limiter {
             Some(ResourceLimiterInner::Sync(ref mut limiter)) => {
-                Ok(limiter(&mut self.data).memory_growing(current, desired, maximum))
+                limiter(&mut self.data).memory_growing(current, desired, maximum)
             }
             #[cfg(feature = "async")]
             Some(ResourceLimiterInner::Async(ref mut limiter)) => unsafe {
-                Ok(self
-                    .inner
+                self.inner
                     .async_cx()
                     .expect("ResourceLimiterAsync requires async Store")
                     .block_on(
                         limiter(&mut self.data)
                             .memory_growing(current, desired, maximum)
                             .as_mut(),
-                    )?)
+                    )?
             },
             None => Ok(true),
         }
@@ -1923,17 +1922,17 @@ unsafe impl<T> wasmtime_runtime::Store for StoreInner<T> {
 
         match self.limiter {
             Some(ResourceLimiterInner::Sync(ref mut limiter)) => {
-                Ok(limiter(&mut self.data).table_growing(current, desired, maximum))
+                limiter(&mut self.data).table_growing(current, desired, maximum)
             }
             #[cfg(feature = "async")]
             Some(ResourceLimiterInner::Async(ref mut limiter)) => unsafe {
-                Ok(async_cx
+                async_cx
                     .expect("ResourceLimiterAsync requires async Store")
                     .block_on(
                         limiter(&mut self.data)
                             .table_growing(current, desired, maximum)
                             .as_mut(),
-                    )?)
+                    )?
             },
             None => Ok(true),
         }

src/commands/run.rs

@@ -9,7 +9,9 @@ use std::io::Write;
 use std::path::{Component, Path, PathBuf};
 use std::thread;
 use std::time::Duration;
-use wasmtime::{Engine, Func, Linker, Module, Store, Val, ValType};
+use wasmtime::{
+    Engine, Func, Linker, Module, Store, StoreLimits, StoreLimitsBuilder, Val, ValType,
+};
 use wasmtime_cli_flags::{CommonOptions, WasiModules};
 use wasmtime_wasi::maybe_exit_on_error;
 use wasmtime_wasi::sync::{ambient_authority, Dir, TcpListener, WasiCtxBuilder};
@@ -166,6 +168,38 @@ pub struct RunCommand {
     /// The arguments to pass to the module
     #[clap(value_name = "ARGS")]
     module_args: Vec<String>,
+
+    /// Maximum size, in bytes, that a linear memory is allowed to reach.
+    ///
+    /// Growth beyond this limit will cause `memory.grow` instructions in
+    /// WebAssembly modules to return -1 and fail.
+    #[clap(long, value_name = "BYTES")]
+    max_memory_size: Option<usize>,
+
+    /// Maximum size, in table elements, that a table is allowed to reach.
+    #[clap(long)]
+    max_table_elements: Option<u32>,
+
+    /// Maximum number of WebAssembly instances allowed to be created.
+    #[clap(long)]
+    max_instances: Option<usize>,
+
+    /// Maximum number of WebAssembly tables allowed to be created.
+    #[clap(long)]
+    max_tables: Option<usize>,
+
+    /// Maximum number of WebAssembly linear memories allowed to be created.
+    #[clap(long)]
+    max_memories: Option<usize>,
+
+    /// Force a trap to be raised on `memory.grow` and `table.grow` failure
+    /// instead of returning -1 from these instructions.
+    ///
+    /// This is not necessarily a spec-compliant option to enable but can be
+    /// useful for tracking down a backtrace of what is requesting so much
+    /// memory, for example.
+    #[clap(long)]
+    trap_on_grow_failure: bool,
 }
 
 impl RunCommand {
@@ -212,6 +246,27 @@ impl RunCommand {
             preopen_sockets,
         )?;
 
+        let mut limits = StoreLimitsBuilder::new();
+        if let Some(max) = self.max_memory_size {
+            limits = limits.memory_size(max);
+        }
+        if let Some(max) = self.max_table_elements {
+            limits = limits.table_elements(max);
+        }
+        if let Some(max) = self.max_instances {
+            limits = limits.instances(max);
+        }
+        if let Some(max) = self.max_tables {
+            limits = limits.tables(max);
+        }
+        if let Some(max) = self.max_memories {
+            limits = limits.memories(max);
+        }
+        store.data_mut().limits = limits
+            .trap_on_grow_failure(self.trap_on_grow_failure)
+            .build();
+        store.limiter(|t| &mut t.limits);
+
         // If fuel has been configured, we want to add the configured
         // fuel amount to this store.
         if let Some(fuel) = self.common.fuel {
@@ -470,6 +525,7 @@ struct Host {
     wasi_nn: Option<Arc<WasiNnCtx>>,
     #[cfg(feature = "wasi-threads")]
     wasi_threads: Option<Arc<WasiThreadsCtx<Host>>>,
+    limits: StoreLimits,
 }
 
 /// Populates the given `Linker` with WASI APIs.