Handle OOM in Linker::instantiate_pre (#12683)

Author: fitzgen

crates/fuzzing/tests/oom.rs

@@ -508,6 +508,51 @@ fn linker_func_wrap() -> Result<()> {
     })
 }
 
+#[test]
+#[cfg(arc_try_new)]
+fn linker_instantiate_pre() -> 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
+                    (import "module" "func" (func (param i32) (result i32)))
+
+                    (memory (export "memory") 1)
+                    (data (i32.const 0) "a")
+
+                    (table (export "table") 1 funcref)
+                    (elem (i32.const 0) func 1)
+
+                    (func (export "func") (param i32) (result i32)
+                        (call 0 (local.get 0))
+                    )
+                )
+            "#,
+        )?;
+        module.serialize()?
+    };
+
+    let mut config = Config::new();
+    config.enable_compiler(false);
+    config.concurrency_support(false);
+
+    let engine = Engine::new(&config)?;
+
+    let mut linker = Linker::<()>::new(&engine);
+    linker.func_wrap("module", "func", |x: i32| x * 2)?;
+
+    let module = unsafe { Module::deserialize(&engine, &module_bytes)? };
+
+    OomTest::new().test(|| {
+        let _ = linker.instantiate_pre(&module)?;
+        Ok(())
+    })
+}
+
 #[test]
 #[cfg(arc_try_new)]
 fn module_deserialize() -> Result<()> {

crates/wasmtime/src/runtime/instance.rs

@@ -17,7 +17,7 @@ use core::ptr::NonNull;
 use wasmparser::WasmFeatures;
 use wasmtime_environ::{
     EntityIndex, EntityType, FuncIndex, GlobalIndex, MemoryIndex, PrimaryMap, TableIndex, TagIndex,
-    TypeTrace,
+    TypeTrace, collections,
 };
 
 /// An instantiated WebAssembly module.
@@ -787,10 +787,10 @@ pub struct InstancePre<T> {
     /// provided, passed to `Instance::new_started` after inserting them into a
     /// `Store`.
     ///
-    /// Note that this is stored as an `Arc<[T]>` to quickly move a strong
-    /// reference to everything internally into a `Store<T>` without having to
-    /// clone each individual item.
-    items: Arc<[Definition]>,
+    /// Note that this is stored as an `Arc` to quickly move a strong reference
+    /// to everything internally into a `Store<T>` without having to clone each
+    /// individual item.
+    items: Arc<collections::Vec<Definition>>,
 
     /// A count of `Definition::HostFunc` entries in `items` above to
     /// preallocate space in a `Store` up front for all entries to be inserted.
@@ -801,8 +801,8 @@ pub struct InstancePre<T> {
     /// `VMFuncRef`s so that we don't have to do it at
     /// instantiation time.
     ///
-    /// This is an `Arc<[T]>` for the same reason as `items`.
-    func_refs: Arc<[VMFuncRef]>,
+    /// This is an `Arc` for the same reason as `items`.
+    func_refs: Arc<collections::Vec<VMFuncRef>>,
 
     /// Whether or not any import in `items` is flagged as needing async.
     ///
@@ -836,10 +836,13 @@ impl<T: 'static> InstancePre<T> {
     /// This method is unsafe as the `T` of the `InstancePre<T>` is not
     /// guaranteed to be the same as the `T` within the `Store`, the caller must
     /// verify that.
-    pub(crate) unsafe fn new(module: &Module, items: Vec<Definition>) -> Result<InstancePre<T>> {
+    pub(crate) unsafe fn new(
+        module: &Module,
+        items: collections::Vec<Definition>,
+    ) -> Result<InstancePre<T>> {
         typecheck(module, &items, |cx, ty, item| cx.definition(ty, &item.ty()))?;
 
-        let mut func_refs = vec![];
+        let mut func_refs = collections::Vec::with_capacity(items.len())?;
         let mut host_funcs = 0;
         let mut asyncness = Asyncness::No;
         for item in &items {
@@ -853,7 +856,7 @@ impl<T: 'static> InstancePre<T> {
                                 .wasm_to_array_trampoline(f.sig_index())
                                 .map(|f| f.into()),
                             ..*f.func_ref()
-                        });
+                        })?;
                     }
                     asyncness = asyncness | f.asyncness();
                 }
@@ -862,9 +865,9 @@ impl<T: 'static> InstancePre<T> {
 
         Ok(InstancePre {
             module: module.clone(),
-            items: items.into(),
+            items: try_new::<Arc<_>>(items)?,
             host_funcs,
-            func_refs: func_refs.into(),
+            func_refs: try_new::<Arc<_>>(func_refs)?,
             asyncness,
             _marker: core::marker::PhantomData,
         })
@@ -959,9 +962,9 @@ impl<T: 'static> InstancePre<T> {
 fn pre_instantiate_raw(
     store: &mut StoreOpaque,
     module: &Module,
-    items: &Arc<[Definition]>,
+    items: &Arc<collections::Vec<Definition>>,
     host_funcs: usize,
-    func_refs: &Arc<[VMFuncRef]>,
+    func_refs: &Arc<collections::Vec<VMFuncRef>>,
     asyncness: Asyncness,
 ) -> Result<OwnedImports> {
     // Register this module and use it to fill out any funcref wasm_call holes

crates/wasmtime/src/runtime/linker.rs

@@ -13,7 +13,10 @@ use core::future::Future;
 use core::marker;
 use core::mem::MaybeUninit;
 use log::warn;
-use wasmtime_environ::{Atom, StringPool, collections::HashMap};
+use wasmtime_environ::{
+    Atom, StringPool,
+    collections::{HashMap, Vec},
+};
 
 /// Structure used to link wasm modules/instances together.
 ///
@@ -633,15 +636,15 @@ impl<T> Linker<T> {
         T: 'static,
     {
         let mut store = store.as_context_mut();
-        let exports = instance
+        let exports: Vec<_> = instance
             .exports(&mut store)
             .map(|e| {
                 Ok((
                     self.import_key(module_name, Some(e.name()))?,
                     e.into_extern(),
                 ))
             })
-            .collect::<Result<Vec<_>>>()?;
+            .try_collect::<_, Error>()?;
         for (key, export) in exports {
             self.insert(key, Definition::new(store.0, export))?;
         }
@@ -988,12 +991,12 @@ impl<T> Linker<T> {
     pub fn alias_module(&mut self, module: &str, as_module: &str) -> Result<()> {
         let module = self.pool.insert(module)?;
         let as_module = self.pool.insert(as_module)?;
-        let items = self
+        let items: Vec<_> = self
             .map
             .iter()
             .filter(|(key, _def)| key.module == module)
-            .map(|(key, def)| (key.name, def.clone()))
-            .collect::<Vec<_>>();
+            .map(|(key, def)| Ok((key.name, def.clone())))
+            .try_collect::<_, Error>()?;
         for (name, item) in items {
             self.insert(
                 ImportKey {
@@ -1177,10 +1180,10 @@ impl<T> Linker<T> {
     where
         T: 'static,
     {
-        let mut imports = module
+        let mut imports: Vec<_> = module
             .imports()
-            .map(|import| self._get_by_import(&import))
-            .collect::<Result<Vec<_>, _>>()?;
+            .map(|import| Ok(self._get_by_import(&import)?))
+            .try_collect::<_, Error>()?;
         if let Some(store) = store {
             for import in imports.iter_mut() {
                 import.update_size(store);

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

@@ -8,6 +8,7 @@ use crate::runtime::HostFunc;
 use crate::runtime::vm::{AlwaysMut, SendSyncPtr, VMArrayCallHostFuncContext, VMFuncRef};
 use alloc::sync::Arc;
 use core::ptr::NonNull;
+use wasmtime_environ::collections;
 
 /// An arena of `VMFuncRef`s.
 ///
@@ -39,15 +40,15 @@ enum Storage {
     /// inside them, etc.
     InstancePreDefinitions {
         #[expect(dead_code, reason = "only here to keep the original value alive")]
-        defs: Arc<[Definition]>,
+        defs: Arc<collections::Vec<Definition>>,
     },
 
     /// Pinned `VMFuncRef`s that had their `wasm_call` field
     /// pre-patched when constructing an `InstancePre`, and which we need to
     /// keep alive for our owning store's lifetime.
     InstancePreFuncRefs {
         #[expect(dead_code, reason = "only here to keep the original value alive")]
-        funcs: Arc<[VMFuncRef]>,
+        funcs: Arc<collections::Vec<VMFuncRef>>,
     },
 
     /// A uniquely-owned host function within a `Store`. This comes about with
@@ -118,7 +119,7 @@ impl FuncRefs {
     ///
     /// This is used to ensure that the store itself persists the entire list of
     /// `funcs` for the entire lifetime of the store.
-    pub fn push_instance_pre_func_refs(&mut self, funcs: Arc<[VMFuncRef]>) {
+    pub fn push_instance_pre_func_refs(&mut self, funcs: Arc<collections::Vec<VMFuncRef>>) {
         self.storage.push(Storage::InstancePreFuncRefs { funcs });
     }
 
@@ -127,7 +128,7 @@ impl FuncRefs {
     ///
     /// This is used to keep linker-defined functions' vmctx values alive, for
     /// example.
-    pub fn push_instance_pre_definitions(&mut self, defs: Arc<[Definition]>) {
+    pub fn push_instance_pre_definitions(&mut self, defs: Arc<collections::Vec<Definition>>) {
         self.storage.push(Storage::InstancePreDefinitions { defs });
     }