Add OOM test for component Func::call_async (#13017)

* Add OOM test for component Func::call_async

Add NameMap::get_by_str to avoid infallible String allocation in
component export name lookups. The NameMapNoIntern implementation
allocated a String on every lookup; the new method uses borrowed
&str keys directly on the underlying IndexMap instead.

* Address review feedback

* Remove some unnnecessary Result<_, OutOfMemory>

* fix clippy warning

* address review feedback

Author: fitzgen

crates/core/src/alloc/string.rs

@@ -60,13 +60,31 @@ impl Borrow<str> for TryString {
     }
 }
 
+impl TryFrom<&str> for TryString {
+    type Error = OutOfMemory;
+
+    #[inline]
+    fn try_from(value: &str) -> Result<Self, Self::Error> {
+        let mut s = TryString::new();
+        s.push_str(value)?;
+        Ok(s)
+    }
+}
+
 impl From<inner::String> for TryString {
     #[inline]
     fn from(inner: inner::String) -> Self {
         Self { inner }
     }
 }
 
+impl From<TryString> for inner::String {
+    #[inline]
+    fn from(s: TryString) -> Self {
+        s.inner
+    }
+}
+
 #[cfg(feature = "serde")]
 impl serde::ser::Serialize for TryString {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>

crates/environ/src/component/info.rs

@@ -107,7 +107,7 @@ pub struct Component {
     pub imports: PrimaryMap<RuntimeImportIndex, (ImportIndex, Vec<String>)>,
 
     /// This component's own root exports from the component itself.
-    pub exports: NameMap<String, ExportIndex>,
+    pub exports: NameMap<TryString, ExportIndex>,
 
     /// All exports of this component and exported instances of this component.
     ///
@@ -459,7 +459,7 @@ pub enum ExportItem<T> {
 }
 
 /// Possible exports from a component.
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize)]
 pub enum Export {
     /// A lifted function being exported which is an adaptation of a core wasm
     /// function.
@@ -491,7 +491,7 @@ pub enum Export {
         /// Instance type index, if such is assigned
         ty: TypeComponentInstanceIndex,
         /// Instance export map
-        exports: NameMap<String, ExportIndex>,
+        exports: NameMap<TryString, ExportIndex>,
     },
     /// An exported type from a component or instance, currently only
     /// informational.

crates/environ/src/component/names.rs

@@ -1,5 +1,8 @@
+use crate::collections::TryCow;
 use crate::error::{Result, bail};
 use crate::prelude::*;
+use alloc::sync::Arc;
+use core::borrow::Borrow;
 use core::hash::Hash;
 use semver::Version;
 use serde_derive::{Deserialize, Serialize};
@@ -12,16 +15,19 @@ use serde_derive::{Deserialize, Serialize};
 /// which is currently considered a key feature of WASI's compatibility story.
 ///
 /// On the outside this looks like a map of `K` to `V`.
-#[derive(Clone, Serialize, Deserialize, Debug)]
-pub struct NameMap<K: Clone + Hash + Eq + Ord, V> {
+#[derive(Serialize, Deserialize, Debug)]
+pub struct NameMap<K, V>
+where
+    K: TryClone + Hash + Eq + Ord,
+{
     /// A map of keys to the value that they define.
     ///
     /// Note that this map is "exact" where the name here is the exact name that
     /// was specified when the `insert` was called. This doesn't have any
     /// semver-mangling or anything like that.
     ///
     /// This map is always consulted first during lookups.
-    definitions: IndexMap<K, V>,
+    definitions: TryIndexMap<K, V>,
 
     /// An auxiliary map tracking semver-compatible names. This is a map from
     /// "semver compatible alternate name" to a name present in `definitions`
@@ -43,12 +49,81 @@ pub struct NameMap<K: Clone + Hash + Eq + Ord, V> {
     ///
     /// The `Version` here is tracked to ensure that when multiple versions on
     /// one track are defined that only the maximal version here is retained.
-    alternate_lookups: IndexMap<K, (K, Version)>,
+    alternate_lookups: TryIndexMap<K, (K, TryVersion)>,
+}
+
+/// A wrapper around `Version` that implements `TryClone`.
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+struct TryVersion(Arc<Version>);
+
+impl TryFrom<Version> for TryVersion {
+    type Error = OutOfMemory;
+
+    fn try_from(value: Version) -> Result<Self, Self::Error> {
+        Ok(Self(try_new::<Arc<_>>(value)?))
+    }
+}
+
+impl TryClone for TryVersion {
+    #[inline]
+    fn try_clone(&self) -> Result<Self, OutOfMemory> {
+        Ok(Self(self.0.clone()))
+    }
+}
+
+impl core::ops::Deref for TryVersion {
+    type Target = Version;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl Borrow<Version> for TryVersion {
+    #[inline]
+    fn borrow(&self) -> &Version {
+        &self.0
+    }
+}
+
+impl serde::Serialize for TryVersion {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        self.0.serialize(serializer)
+    }
+}
+
+impl<'de> serde::Deserialize<'de> for TryVersion {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        use serde::de::Error;
+        let v = Version::deserialize(deserializer)?;
+        let v = try_new::<Arc<_>>(v).map_err(|oom| D::Error::custom(oom))?;
+        Ok(Self(v))
+    }
+}
+
+impl<K, V> TryClone for NameMap<K, V>
+where
+    K: TryClone + Hash + Eq + Ord,
+    V: TryClone,
+{
+    fn try_clone(&self) -> Result<Self, OutOfMemory> {
+        Ok(Self {
+            definitions: self.definitions.try_clone()?,
+            alternate_lookups: self.alternate_lookups.try_clone()?,
+        })
+    }
 }
 
 impl<K, V> NameMap<K, V>
 where
-    K: Clone + Hash + Eq + Ord,
+    K: TryClone + Hash + Eq + Ord,
 {
     /// Inserts the `name` specified into this map.
     ///
@@ -64,32 +139,32 @@ where
     pub fn insert<I>(&mut self, name: &str, cx: &mut I, allow_shadowing: bool, item: V) -> Result<K>
     where
         I: NameMapIntern<Key = K>,
+        I::BorrowedKey: Hash + Eq,
     {
         // Always insert `name` and `item` as an exact definition.
-        let key = cx.intern(name);
-        if let Some(prev) = self.definitions.insert(key.clone(), item) {
-            if !allow_shadowing {
-                self.definitions.insert(key, prev);
-                bail!("map entry `{name}` defined twice")
-            }
+        let key = cx.intern(name)?;
+        if !allow_shadowing && self.definitions.contains_key(&key) {
+            bail!("map entry `{name}` defined twice")
         }
+        self.definitions.insert(key.try_to_owned()?, item)?;
 
         // If `name` is a semver-looking thing, like `a:b/c@1.0.0`, then also
         // insert an entry in the semver-compatible map under a key such as
         // `a:b/c@1`.
         //
         // This key is used during `get` later on.
         if let Some((alternate_key, version)) = alternate_lookup_key(name) {
-            let alternate_key = cx.intern(alternate_key);
-            if let Some((prev_key, prev_version)) = self
-                .alternate_lookups
-                .insert(alternate_key.clone(), (key.clone(), version.clone()))
-            {
+            let alternate_key = cx.intern(alternate_key)?;
+            let version = TryVersion::try_from(version)?;
+            if let Some((prev_key, prev_version)) = self.alternate_lookups.insert(
+                alternate_key.try_clone()?,
+                (key.try_clone()?, version.clone()),
+            )? {
                 // Prefer the latest version, so only do this if we're
                 // greater than the prior version.
                 if version < prev_version {
                     self.alternate_lookups
-                        .insert(alternate_key, (prev_key, prev_version));
+                        .insert(alternate_key, (prev_key, prev_version))?;
                 }
             }
         }
@@ -105,11 +180,13 @@ where
     pub fn get<I>(&self, name: &str, cx: &I) -> Option<&V>
     where
         I: NameMapIntern<Key = K>,
+        I::Key: Borrow<I::BorrowedKey>,
+        I::BorrowedKey: Hash + Eq,
     {
         // First look up an exact match and if that's found return that. This
         // enables defining multiple versions in the map and the requested
         // version is returned if it matches exactly.
-        let candidate = cx.lookup(name).and_then(|k| self.definitions.get(&k));
+        let candidate = cx.lookup(name).and_then(|k| self.definitions.get(&*k));
         if let Some(def) = candidate {
             return Some(def);
         }
@@ -122,7 +199,7 @@ where
         let (alternate_name, _version) = alternate_lookup_key(name)?;
         let alternate_key = cx.lookup(alternate_name)?;
         let (exact_key, _version) = self.alternate_lookups.get(&alternate_key)?;
-        self.definitions.get(exact_key)
+        self.definitions.get(exact_key.borrow())
     }
 
     /// Returns an iterator over inserted values in this map.
@@ -142,7 +219,7 @@ where
 
 impl<K, V> Default for NameMap<K, V>
 where
-    K: Clone + Hash + Eq + Ord,
+    K: TryClone + Hash + Eq + Ord,
 {
     fn default() -> NameMap<K, V> {
         NameMap {
@@ -156,29 +233,33 @@ where
 /// keys to non-strings.
 pub trait NameMapIntern {
     /// The key that this interning context generates.
-    type Key;
+    type Key: Borrow<Self::BorrowedKey>;
+
+    /// The borrowed version of the key type.
+    type BorrowedKey: ?Sized + TryToOwned<Owned = Self::Key>;
 
     /// Inserts `s` into `self` and returns the intern'd key `Self::Key`.
-    fn intern(&mut self, s: &str) -> Self::Key;
+    fn intern(&mut self, s: &str) -> Result<Self::Key, OutOfMemory>;
 
     /// Looks up `s` in `self` returning `Some` if it was found or `None` if
     /// it's not present.
-    fn lookup(&self, s: &str) -> Option<Self::Key>;
+    fn lookup<'a>(&'a self, s: &'a str) -> Option<TryCow<'a, Self::BorrowedKey>>;
 }
 
 /// For use with [`NameMap`] when no interning should happen and instead string
 /// keys are copied as-is.
 pub struct NameMapNoIntern;
 
 impl NameMapIntern for NameMapNoIntern {
-    type Key = String;
+    type Key = TryString;
+    type BorrowedKey = str;
 
-    fn intern(&mut self, s: &str) -> String {
-        s.to_string()
+    fn intern(&mut self, s: &str) -> Result<Self::Key, OutOfMemory> {
+        TryString::try_from(s)
     }
 
-    fn lookup(&self, s: &str) -> Option<String> {
-        Some(s.to_string())
+    fn lookup<'a>(&'a self, s: &'a str) -> Option<TryCow<'a, Self::BorrowedKey>> {
+        Some(TryCow::Borrowed(s))
     }
 }
 

crates/environ/src/component/types.rs

@@ -22,6 +22,12 @@ macro_rules! indices {
         #[repr(transparent)]
         pub struct $name(u32);
         cranelift_entity::entity_impl!($name);
+        impl TryClone for $name {
+            #[inline]
+            fn try_clone(&self) -> Result<Self, OutOfMemory> {
+                Ok(*self)
+            }
+        }
     )*);
 }
 

crates/environ/src/component/types_builder.rs

@@ -150,7 +150,7 @@ impl ComponentTypesBuilder {
         }
         for (name, ty) in component.exports.raw_iter() {
             component_ty.exports.insert(
-                name.clone(),
+                name.clone_panic_on_oom().into(),
                 self.export_type_def(&component.export_items, *ty),
             );
         }

crates/fuzzing/tests/oom/component_func.rs

@@ -1,6 +1,6 @@
 #![cfg(arc_try_new)]
 
-use wasmtime::component::{Component, Linker};
+use wasmtime::component::{Component, Linker, Val};
 use wasmtime::{Config, Engine, Result, Store};
 use wasmtime_fuzzing::oom::OomTest;
 
@@ -42,3 +42,47 @@ async fn component_instance_pre_instantiate_async() -> Result<()> {
         })
         .await
 }
+
+#[tokio::test]
+async fn component_func_call_async() -> Result<()> {
+    let component_bytes = {
+        let mut config = Config::new();
+        config.concurrency_support(false);
+        let engine = Engine::new(&config)?;
+        Component::new(
+            &engine,
+            r#"
+                (component
+                    (core module $m
+                        (func (export "id") (param i32) (result i32) (local.get 0))
+                    )
+                    (core instance $i (instantiate $m))
+                    (func (export "id") (param "x" s32) (result s32)
+                        (canon lift (core func $i "id"))
+                    )
+                )
+            "#,
+        )?
+        .serialize()?
+    };
+    let mut config = Config::new();
+    config.enable_compiler(false);
+    config.concurrency_support(false);
+    let engine = Engine::new(&config)?;
+    let component = unsafe { Component::deserialize(&engine, &component_bytes)? };
+    let linker = Linker::<()>::new(&engine);
+    let instance_pre = linker.instantiate_pre(&component)?;
+
+    OomTest::new()
+        .test_async(|| async {
+            let mut store = Store::try_new(&engine, ())?;
+            let instance = instance_pre.instantiate_async(&mut store).await?;
+            let func = instance.get_func(&mut store, "id").unwrap();
+            let mut results = [Val::S32(0)];
+            func.call_async(&mut store, &[Val::S32(42)], &mut results)
+                .await?;
+            assert_eq!(results[0], Val::S32(42));
+            Ok(())
+        })
+        .await
+}