Add custom resource limiting for instances

lucet-runtime/lucet-runtime-internals/src/instance.rs

@@ -76,6 +76,7 @@ pub fn new_instance_handle(
     module: Arc<dyn Module>,
     alloc: Alloc,
     embed_ctx: CtxMap,
+    limiter: Option<Box<dyn ResourceLimiter>>,
 ) -> Result<InstanceHandle, Error> {
     let inst = NonNull::new(instance)
         .ok_or_else(|| lucet_format_err!("instance pointer is null; this is a bug"))?;
@@ -90,7 +91,7 @@ pub fn new_instance_handle(
         needs_inst_drop: false,
     };
 
-    let inst = Instance::new(alloc, module, embed_ctx);
+    let inst = Instance::new(alloc, module, embed_ctx, limiter);
 
     unsafe {
         // this is wildly unsafe! you must be very careful to not let the drop impls run on the
@@ -276,6 +277,8 @@ pub struct Instance {
     /// `panic!()` is called, etc. Terminating allows the error to be more directly identifiable.
     terminate_on_heap_oom: bool,
 
+    resource_limiter: Option<Box<dyn ResourceLimiter>>,
+
     /// `_padding` must be the last member of the structure.
     /// This marks where the padding starts to make the structure exactly 4096 bytes long.
     /// It is also used to compute the size of the structure up to that point, i.e. without padding.
@@ -710,6 +713,16 @@ impl Instance {
         let additional_bytes = additional_pages
             .checked_mul(WASM_PAGE_SIZE)
             .ok_or_else(|| lucet_format_err!("additional pages larger than wasm address space",))?;
+
+        if let Some(limiter) = &self.resource_limiter {
+            let current = self.alloc.heap_accessible_size as u32;
+            let desired = current + additional_bytes;
+            let limit = self.alloc.heap_memory_size_limit as u32;
+            limiter
+                .memory_growing(self, current, desired, limit)
+                .map_err(|err| Error::LimitsExceeded(err))?;
+        }
+
         let orig_len = self
             .alloc
             .expand_heap(additional_bytes, self.module.as_ref())?;
@@ -992,7 +1005,12 @@ impl Instance {
 
 // Private API
 impl Instance {
-    fn new(alloc: Alloc, module: Arc<dyn Module>, embed_ctx: CtxMap) -> Self {
+    fn new(
+        alloc: Alloc,
+        module: Arc<dyn Module>,
+        embed_ctx: CtxMap,
+        resource_limiter: Option<Box<dyn ResourceLimiter>>,
+    ) -> Self {
         let globals_ptr = alloc.slot().globals as *mut i64;
 
         #[cfg(feature = "concurrent_testpoints")]
@@ -1021,6 +1039,7 @@ impl Instance {
             entrypoint: None,
             resumed_val: None,
             terminate_on_heap_oom: false,
+            resource_limiter,
             _padding: (),
         };
         inst.set_globals_ptr(globals_ptr);
@@ -1611,3 +1630,25 @@ pub(crate) struct EmptyYieldVal;
 fn default_fatal_handler(inst: &Instance) -> ! {
     panic!("> instance {:p} had fatal error: {}", inst, inst.state);
 }
+
+/// Used by hosts to limit resource consumption of instances.
+///
+/// An instance can be configured with a resource limiter so that hosts can take into account non-WebAssembly
+/// resource usage to determine if linear memory should be permitted to grow.
+pub trait ResourceLimiter {
+    /// Called when the the instance's linear memory has been requested to grow.
+    ///
+    /// * `current` gives the current size of the linear memory, in bytes.
+    /// * `desired` gives the desired new size of the linear memory, in bytes.
+    /// * `limit` gives the configured maximum size of the linear memory, in bytes.
+    ///
+    /// Returning `Ok` signals that the memory should be permitted to grow; returning an `Err(err)` prevents
+    /// the growth from occurring, yielding an `Error::LimitsExceeded(err)` from the call to grow the memory.
+    fn memory_growing(
+        &self,
+        instance: &Instance,
+        current: u32,
+        desired: u32,
+        limit: u32,
+    ) -> Result<(), String>;
+}

lucet-runtime/lucet-runtime-internals/src/region.rs

@@ -6,7 +6,7 @@ pub mod uffd;
 use crate::alloc::{Alloc, AllocStrategy, Limits, Slot};
 use crate::embed_ctx::CtxMap;
 use crate::error::Error;
-use crate::instance::InstanceHandle;
+use crate::instance::{InstanceHandle, ResourceLimiter};
 use crate::module::Module;
 use std::any::Any;
 use std::sync::Arc;
@@ -98,6 +98,7 @@ pub struct NewInstanceArgs {
     pub heap_memory_size_limit: usize,
     pub alloc_strategy: AllocStrategy,
     pub terminate_on_heap_oom: bool,
+    pub resource_limiter: Option<Box<dyn ResourceLimiter>>,
 }
 
 impl<'a> InstanceBuilder<'a> {
@@ -110,6 +111,7 @@ impl<'a> InstanceBuilder<'a> {
                 heap_memory_size_limit: region.get_limits().heap_memory_size,
                 alloc_strategy: AllocStrategy::Linear,
                 terminate_on_heap_oom: false,
+                resource_limiter: None,
             },
         }
     }
@@ -157,6 +159,16 @@ impl<'a> InstanceBuilder<'a> {
         self
     }
 
+    /// Add a resource limiter to the built instance.
+    ///
+    /// This call is optional. It can be used to add additional checks when the instance requests
+    /// additional resources, e.g. when growing memory. These checks are useful to take dynamic,
+    /// non-WebAssembly-related concerns into account.
+    pub fn with_resource_limiter(mut self, limiter: impl ResourceLimiter + 'static) -> Self {
+        self.args.resource_limiter = Some(Box::new(limiter));
+        self
+    }
+
     /// Build the instance.
     pub fn build(self) -> Result<InstanceHandle, Error> {
         self.region.new_instance_with(self.args)

lucet-runtime/lucet-runtime-internals/src/region/mmap.rs

@@ -83,6 +83,7 @@ impl RegionInternal for MmapRegion {
             heap_memory_size_limit,
             mut alloc_strategy,
             terminate_on_heap_oom,
+            resource_limiter,
             ..
         }: NewInstanceArgs,
     ) -> Result<InstanceHandle, Error> {
@@ -146,7 +147,7 @@ impl RegionInternal for MmapRegion {
 
         // Though this is a potential early return from the function, the Drop impl
         // on the Alloc will put the slot back on the freelist.
-        let mut inst = new_instance_handle(inst_ptr, module, alloc, embed_ctx)?;
+        let mut inst = new_instance_handle(inst_ptr, module, alloc, embed_ctx, resource_limiter)?;
         inst.set_terminate_on_heap_oom(terminate_on_heap_oom);
 
         Ok(inst)

lucet-runtime/lucet-runtime-internals/src/region/uffd.rs

@@ -317,6 +317,7 @@ impl RegionInternal for UffdRegion {
             heap_memory_size_limit,
             mut alloc_strategy,
             terminate_on_heap_oom,
+            resource_limiter,
             ..
         }: NewInstanceArgs,
     ) -> Result<InstanceHandle, Error> {
@@ -416,7 +417,7 @@ impl RegionInternal for UffdRegion {
             invalid_pages: Vec::new(),
         };
 
-        let mut inst = new_instance_handle(inst_ptr, module, alloc, embed_ctx)?;
+        let mut inst = new_instance_handle(inst_ptr, module, alloc, embed_ctx, resource_limiter)?;
         inst.set_terminate_on_heap_oom(terminate_on_heap_oom);
 
         Ok(inst)