Handle OOM in {Func,Memory,Table,Global}::new and when calling an instance's exported function (bytecodealliance#12855)

* Use `try_new` for `Box<dyn RuntimeLinearMemory>` in `DefaultMemoryCreator`

* Use `TryPrimaryMap` for `host_globals` in `Store`

* Add `Func::try_wrap` and use `try_new` for `Box<HostFunc>`

Add `Func::try_wrap` as a fallible version of `Func::wrap` that returns an error
on out-of-memory instead of panicking. `Func::wrap` now delegates to `try_wrap`.

Also use `try_new::<Box<_>>` instead of `Box::new` for `HostFunc`.

* Use `bumpalo`'s `try_alloc` for `FuncRefs`

* Use `try_new` for `Arc<Module>` in "trampoline" code

* Test that we handle OOM in `{Func,Memory,Table,Global}::new` and when calling an instance's exported function

* cargo fmt

Author: fitzgen

crates/fuzzing/tests/oom/func.rs

@@ -0,0 +1,18 @@
+#![cfg(arc_try_new)]
+
+use wasmtime::{Config, Engine, Func, Result, Store};
+use wasmtime_fuzzing::oom::OomTest;
+
+#[test]
+fn func_new() -> Result<()> {
+    let mut config = Config::new();
+    config.enable_compiler(false);
+    config.concurrency_support(false);
+    let engine = Engine::new(&config)?;
+
+    OomTest::new().test(|| {
+        let mut store = Store::try_new(&engine, ())?;
+        let _func = Func::try_wrap(&mut store, |x: i32| x * 2)?;
+        Ok(())
+    })
+}

crates/fuzzing/tests/oom/global.rs

@@ -0,0 +1,19 @@
+#![cfg(arc_try_new)]
+
+use wasmtime::{Config, Engine, GlobalType, Mutability, Result, Store, Val, ValType};
+use wasmtime_fuzzing::oom::OomTest;
+
+#[test]
+fn global_new() -> Result<()> {
+    let mut config = Config::new();
+    config.enable_compiler(false);
+    config.concurrency_support(false);
+    let engine = Engine::new(&config)?;
+
+    OomTest::new().test(|| {
+        let mut store = Store::try_new(&engine, ())?;
+        let ty = GlobalType::new(ValType::I32, Mutability::Var);
+        let _global = wasmtime::Global::new(&mut store, ty, Val::I32(42))?;
+        Ok(())
+    })
+}

crates/fuzzing/tests/oom/instance.rs

@@ -0,0 +1,42 @@
+#![cfg(arc_try_new)]
+
+use wasmtime::{Config, Engine, Linker, Module, Result, Store};
+use wasmtime_fuzzing::oom::OomTest;
+
+#[test]
+fn call_exported_func() -> Result<()> {
+    let module_bytes = {
+        let mut config = Config::new();
+        config.concurrency_support(false);
+        let engine = Engine::new(&config)?;
+        let module = Module::new(
+            &engine,
+            r#"
+                (module
+                    (func (export "add") (param i32 i32) (result i32)
+                        (i32.add (local.get 0) (local.get 1))
+                    )
+                )
+            "#,
+        )?;
+        module.serialize()?
+    };
+
+    let mut config = Config::new();
+    config.enable_compiler(false);
+    config.concurrency_support(false);
+    let engine = Engine::new(&config)?;
+
+    let module = unsafe { Module::deserialize(&engine, &module_bytes)? };
+    let linker = Linker::<()>::new(&engine);
+    let instance_pre = linker.instantiate_pre(&module)?;
+
+    OomTest::new().test(|| {
+        let mut store = Store::try_new(&engine, ())?;
+        let instance = instance_pre.instantiate(&mut store)?;
+        let add = instance.get_typed_func::<(i32, i32), i32>(&mut store, "add")?;
+        let result = add.call(&mut store, (1, 2))?;
+        assert_eq!(result, 3);
+        Ok(())
+    })
+}

crates/fuzzing/tests/oom/main.rs

@@ -8,18 +8,23 @@ mod config;
 mod engine;
 mod entity_set;
 mod error;
+mod func;
 mod func_type;
+mod global;
 mod hash_map;
 mod hash_set;
 mod index_map;
+mod instance;
 mod instance_pre;
 mod linker;
+mod memory;
 mod module;
 mod primary_map;
 mod secondary_map;
 mod smoke;
 mod store;
 mod string;
+mod table;
 mod vec;
 
 use wasmtime_fuzzing::oom::OomTestAllocator;

crates/fuzzing/tests/oom/memory.rs

@@ -0,0 +1,22 @@
+#![cfg(arc_try_new)]
+
+use wasmtime::{Config, Engine, MemoryType, Result, Store};
+use wasmtime_fuzzing::oom::OomTest;
+
+#[test]
+fn memory_new() -> Result<()> {
+    let mut config = Config::new();
+    config.enable_compiler(false);
+    config.concurrency_support(false);
+    let engine = Engine::new(&config)?;
+
+    OomTest::new()
+        // `IndexMap::reserve` will try to allocate double space, but if that
+        // fails, will attempt to allocate the minimal space necessary.
+        .allow_alloc_after_oom(true)
+        .test(|| {
+            let mut store = Store::try_new(&engine, ())?;
+            let _memory = wasmtime::Memory::new(&mut store, MemoryType::new(1, None))?;
+            Ok(())
+        })
+}

crates/fuzzing/tests/oom/table.rs

@@ -0,0 +1,23 @@
+#![cfg(arc_try_new)]
+
+use wasmtime::{Config, Engine, Ref, RefType, Result, Store, TableType};
+use wasmtime_fuzzing::oom::OomTest;
+
+#[test]
+fn table_new() -> Result<()> {
+    let mut config = Config::new();
+    config.enable_compiler(false);
+    config.concurrency_support(false);
+    let engine = Engine::new(&config)?;
+
+    OomTest::new()
+        // `IndexMap::reserve` will try to allocate double space, but if that
+        // fails, will attempt to allocate the minimal space necessary.
+        .allow_alloc_after_oom(true)
+        .test(|| {
+            let mut store = Store::try_new(&engine, ())?;
+            let ty = TableType::new(RefType::FUNCREF, 1, None);
+            let _table = wasmtime::Table::new(&mut store, ty, Ref::Func(None))?;
+            Ok(())
+        })
+}

crates/wasmtime/src/runtime/func.rs

@@ -783,19 +783,33 @@ impl Func {
     /// # }
     /// ```
     pub fn wrap<T, Params, Results>(
-        mut store: impl AsContextMut<Data = T>,
+        store: impl AsContextMut<Data = T>,
         func: impl IntoFunc<T, Params, Results>,
     ) -> Func
+    where
+        T: 'static,
+    {
+        Self::try_wrap(store, func).expect(
+            "allocation failure during `Func::wrap` (use `Func::try_wrap` to handle such errors)",
+        )
+    }
+
+    /// Fallible version of [`Func::wrap`] that returns an error on
+    /// out-of-memory instead of panicking.
+    pub fn try_wrap<T, Params, Results>(
+        mut store: impl AsContextMut<Data = T>,
+        func: impl IntoFunc<T, Params, Results>,
+    ) -> Result<Func>
     where
         T: 'static,
     {
         let store = store.as_context_mut().0;
         let engine = store.engine();
-        let host = func.into_func(engine).panic_on_oom();
+        let host = func.into_func(engine)?;
 
         // SAFETY: The `T` the closure takes is the same as the `T` of the store
         // we're inserting into via the type signature above.
-        unsafe { host.into_func(store).panic_on_oom() }
+        Ok(unsafe { host.into_func(store)? })
     }
 
     /// Same as [`Func::wrap`], except the closure asynchronously produces the
@@ -2693,7 +2707,7 @@ impl HostFunc {
         store.set_async_required(self.asyncness);
 
         let (funcrefs, modules) = store.func_refs_and_modules();
-        let funcref = funcrefs.push_box_host(Box::new(self), modules)?;
+        let funcref = funcrefs.push_box_host(try_new::<Box<_>>(self)?, modules)?;
         // SAFETY: this funcref was just pushed within `store`, so it's safe to
         // say it's owned by the store's id.
         Ok(unsafe { Func::from_vm_func_ref(store.id(), funcref) })

crates/wasmtime/src/runtime/store.rs

@@ -112,7 +112,7 @@ use core::num::NonZeroU64;
 use core::ops::{Deref, DerefMut};
 use core::pin::Pin;
 use core::ptr::NonNull;
-use wasmtime_environ::{DefinedGlobalIndex, DefinedTableIndex, EntityRef, PrimaryMap, TripleExt};
+use wasmtime_environ::{DefinedGlobalIndex, DefinedTableIndex, EntityRef, TripleExt};
 
 mod context;
 pub use self::context::*;
@@ -478,7 +478,7 @@ pub struct StoreOpaque {
     signal_handler: Option<SignalHandler>,
     modules: ModuleRegistry,
     func_refs: FuncRefs,
-    host_globals: PrimaryMap<DefinedGlobalIndex, StoreBox<VMHostGlobalContext>>,
+    host_globals: TryPrimaryMap<DefinedGlobalIndex, StoreBox<VMHostGlobalContext>>,
     // GC-related fields.
     gc_store: Option<GcStore>,
     gc_roots: RootSet,
@@ -756,7 +756,7 @@ impl<T> Store<T> {
             pending_exception: None,
             modules: ModuleRegistry::default(),
             func_refs: FuncRefs::default(),
-            host_globals: PrimaryMap::new(),
+            host_globals: TryPrimaryMap::new(),
             instance_count: 0,
             instance_limit: crate::DEFAULT_INSTANCE_LIMIT,
             memory_count: 0,
@@ -1687,13 +1687,13 @@ impl StoreOpaque {
 
     pub(crate) fn host_globals(
         &self,
-    ) -> &PrimaryMap<DefinedGlobalIndex, StoreBox<VMHostGlobalContext>> {
+    ) -> &TryPrimaryMap<DefinedGlobalIndex, StoreBox<VMHostGlobalContext>> {
         &self.host_globals
     }
 
     pub(crate) fn host_globals_mut(
         &mut self,
-    ) -> &mut PrimaryMap<DefinedGlobalIndex, StoreBox<VMHostGlobalContext>> {
+    ) -> &mut TryPrimaryMap<DefinedGlobalIndex, StoreBox<VMHostGlobalContext>> {
         &mut self.host_globals
     }
 

crates/wasmtime/src/runtime/store/func_refs.rs

@@ -7,6 +7,7 @@ use crate::prelude::*;
 use crate::runtime::HostFunc;
 use crate::runtime::vm::{AlwaysMut, SendSyncPtr, VMArrayCallHostFuncContext, VMFuncRef};
 use alloc::sync::Arc;
+use core::mem::size_of;
 use core::ptr::NonNull;
 
 /// An arena of `VMFuncRef`s.
@@ -92,7 +93,11 @@ impl FuncRefs {
         modules: &ModuleRegistry,
     ) -> Result<NonNull<VMFuncRef>, OutOfMemory> {
         debug_assert!(func_ref.wasm_call.is_none());
-        let func_ref = self.bump.get_mut().alloc(func_ref);
+        let func_ref = self
+            .bump
+            .get_mut()
+            .try_alloc(func_ref)
+            .map_err(|_| OutOfMemory::new(size_of::<VMFuncRef>()))?;
         // SAFETY: it's a contract of this function itself that `func_ref` has a
         // valid vmctx field to read.
         let has_hole = unsafe { !try_fill(func_ref, modules) };

crates/wasmtime/src/runtime/trampoline/global.rs

@@ -83,6 +83,6 @@ pub fn generate_global_export(
         }
     }
 
-    let index = store.host_globals_mut().push(ctx);
+    let index = store.host_globals_mut().push(ctx)?;
     Ok(crate::Global::from_host(store.id(), index))
 }