implement fuzzing for component types (#4537)

This addresses #4307.

For the static API we generate 100 arbitrary test cases at build time, each of
which includes 0-5 parameter types, a result type, and a WAT fragment containing
an imported function and an exported function.  The exported function calls the
imported function, which is implemented by the host.  At runtime, the fuzz test
selects a test case at random and feeds it zero or more sets of arbitrary
parameters and results, checking that values which flow host-to-guest and
guest-to-host make the transition unchanged.

The fuzz test for the dynamic API follows a similar pattern, the only difference
being that test cases are generated at runtime.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

crates/wasmtime/src/component/func.rs

@@ -257,6 +257,12 @@ impl Func {
             .collect()
     }
 
+    /// Get the result type for this function.
+    pub fn result(&self, store: impl AsContext) -> Type {
+        let data = &store.as_context()[self.0];
+        Type::from(&data.types[data.ty].result, &data.types)
+    }
+
     /// Invokes this function with the `params` given and returns the result.
     ///
     /// The `params` here must match the type signature of this `Func`, or this will return an error. If a trap
@@ -307,11 +313,14 @@ impl Func {
                     self.store_args(store, &options, &params, args, dst)
                 } else {
                     dst.write([ValRaw::u64(0); MAX_FLAT_PARAMS]);
-                    let dst = unsafe {
+
+                    let dst = &mut unsafe {
                         mem::transmute::<_, &mut [MaybeUninit<ValRaw>; MAX_FLAT_PARAMS]>(dst)
-                    };
+                    }
+                    .iter_mut();
+
                     args.iter()
-                        .try_for_each(|arg| arg.lower(store, &options, &mut dst.iter_mut()))
+                        .try_for_each(|arg| arg.lower(store, &options, dst))
                 }
             },
             |store, options, src: &[ValRaw; MAX_FLAT_RESULTS]| {

crates/wasmtime/src/component/func/host.rs

@@ -1,9 +1,10 @@
 use crate::component::func::{Memory, MemoryMut, Options};
-use crate::component::{ComponentParams, ComponentType, Lift, Lower};
+use crate::component::types::SizeAndAlignment;
+use crate::component::{ComponentParams, ComponentType, Lift, Lower, Type, Val};
 use crate::{AsContextMut, StoreContextMut, ValRaw};
-use anyhow::{bail, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use std::any::Any;
-use std::mem::MaybeUninit;
+use std::mem::{self, MaybeUninit};
 use std::panic::{self, AssertUnwindSafe};
 use std::ptr::NonNull;
 use std::sync::Arc;
@@ -43,7 +44,7 @@ pub trait IntoComponentFunc<T, Params, Return> {
 
 pub struct HostFunc {
     entrypoint: VMLoweringCallee,
-    typecheck: fn(TypeFuncIndex, &ComponentTypes) -> Result<()>,
+    typecheck: Box<dyn (Fn(TypeFuncIndex, &Arc<ComponentTypes>) -> Result<()>) + Send + Sync>,
     func: Box<dyn Any + Send + Sync>,
 }
 
@@ -51,17 +52,54 @@ impl HostFunc {
     fn new<F, P, R>(func: F, entrypoint: VMLoweringCallee) -> Arc<HostFunc>
     where
         F: Send + Sync + 'static,
-        P: ComponentParams + Lift,
-        R: Lower,
+        P: ComponentParams + Lift + 'static,
+        R: Lower + 'static,
     {
         Arc::new(HostFunc {
             entrypoint,
-            typecheck: typecheck::<P, R>,
+            typecheck: Box::new(typecheck::<P, R>),
             func: Box::new(func),
         })
     }
 
-    pub fn typecheck(&self, ty: TypeFuncIndex, types: &ComponentTypes) -> Result<()> {
+    pub(crate) fn new_dynamic<
+        T,
+        F: Fn(StoreContextMut<'_, T>, &[Val]) -> Result<Val> + Send + Sync + 'static,
+    >(
+        func: F,
+        index: TypeFuncIndex,
+        types: &Arc<ComponentTypes>,
+    ) -> Arc<HostFunc> {
+        let ty = &types[index];
+
+        Arc::new(HostFunc {
+            entrypoint: dynamic_entrypoint::<T, F>,
+            typecheck: Box::new({
+                let types = types.clone();
+
+                move |expected_index, expected_types| {
+                    if index == expected_index && Arc::ptr_eq(&types, expected_types) {
+                        Ok(())
+                    } else {
+                        Err(anyhow!("function type mismatch"))
+                    }
+                }
+            }),
+            func: Box::new(DynamicContext {
+                func,
+                types: Types {
+                    params: ty
+                        .params
+                        .iter()
+                        .map(|(_, ty)| Type::from(ty, types))
+                        .collect(),
+                    result: Type::from(&ty.result, types),
+                },
+            }),
+        })
+    }
+
+    pub fn typecheck(&self, ty: TypeFuncIndex, types: &Arc<ComponentTypes>) -> Result<()> {
         (self.typecheck)(ty, types)
     }
 
@@ -74,7 +112,7 @@ impl HostFunc {
     }
 }
 
-fn typecheck<P, R>(ty: TypeFuncIndex, types: &ComponentTypes) -> Result<()>
+fn typecheck<P, R>(ty: TypeFuncIndex, types: &Arc<ComponentTypes>) -> Result<()>
 where
     P: ComponentParams + Lift,
     R: Lower,
@@ -256,8 +294,8 @@ macro_rules! impl_into_component_func {
         impl<T, F, $($args,)* R> IntoComponentFunc<T, ($($args,)*), R> for F
         where
             F: Fn($($args),*) -> Result<R> + Send + Sync + 'static,
-            ($($args,)*): ComponentParams + Lift,
-            R: Lower,
+            ($($args,)*): ComponentParams + Lift + 'static,
+            R: Lower + 'static,
         {
             extern "C" fn entrypoint(
                 cx: *mut VMOpaqueContext,
@@ -294,8 +332,8 @@ macro_rules! impl_into_component_func {
         impl<T, F, $($args,)* R> IntoComponentFunc<T, (StoreContextMut<'_, T>, $($args,)*), R> for F
         where
             F: Fn(StoreContextMut<'_, T>, $($args),*) -> Result<R> + Send + Sync + 'static,
-            ($($args,)*): ComponentParams + Lift,
-            R: Lower,
+            ($($args,)*): ComponentParams + Lift + 'static,
+            R: Lower + 'static,
         {
             extern "C" fn entrypoint(
                 cx: *mut VMOpaqueContext,
@@ -330,3 +368,154 @@ macro_rules! impl_into_component_func {
 }
 
 for_each_function_signature!(impl_into_component_func);
+
+unsafe fn call_host_dynamic<T, F>(
+    Types { params, result }: &Types,
+    cx: *mut VMOpaqueContext,
+    mut flags: InstanceFlags,
+    memory: *mut VMMemoryDefinition,
+    realloc: *mut VMCallerCheckedAnyfunc,
+    string_encoding: StringEncoding,
+    storage: &mut [ValRaw],
+    closure: F,
+) -> Result<()>
+where
+    F: FnOnce(StoreContextMut<'_, T>, &[Val]) -> Result<Val>,
+{
+    let cx = VMComponentContext::from_opaque(cx);
+    let instance = (*cx).instance();
+    let mut cx = StoreContextMut::from_raw((*instance).store());
+
+    let options = Options::new(
+        cx.0.id(),
+        NonNull::new(memory),
+        NonNull::new(realloc),
+        string_encoding,
+    );
+
+    // Perform a dynamic check that this instance can indeed be left. Exiting
+    // the component is disallowed, for example, when the `realloc` function
+    // calls a canonical import.
+    if !flags.may_leave() {
+        bail!("cannot leave component instance");
+    }
+
+    let param_count = params.iter().map(|ty| ty.flatten_count()).sum::<usize>();
+
+    let args;
+    let ret_index;
+
+    if param_count <= MAX_FLAT_PARAMS {
+        let iter = &mut storage.iter();
+        args = params
+            .iter()
+            .map(|ty| Val::lift(ty, cx.0, &options, iter))
+            .collect::<Result<Box<[_]>>>()?;
+        ret_index = param_count;
+    } else {
+        let param_layout = {
+            let mut size = 0;
+            let mut alignment = 1;
+            for ty in params.iter() {
+                alignment = alignment.max(ty.size_and_alignment().alignment);
+                ty.next_field(&mut size);
+            }
+            SizeAndAlignment { size, alignment }
+        };
+
+        let memory = Memory::new(cx.0, &options);
+        let mut offset = validate_inbounds_dynamic(param_layout, memory.as_slice(), &storage[0])?;
+        args = params
+            .iter()
+            .map(|ty| {
+                Val::load(
+                    ty,
+                    &memory,
+                    &memory.as_slice()[ty.next_field(&mut offset)..]
+                        [..ty.size_and_alignment().size],
+                )
+            })
+            .collect::<Result<Box<[_]>>>()?;
+        ret_index = 1;
+    };
+
+    let ret = closure(cx.as_context_mut(), &args)?;
+    flags.set_may_leave(false);
+    result.check(&ret)?;
+
+    let result_count = result.flatten_count();
+    if result_count <= MAX_FLAT_RESULTS {
+        let dst = mem::transmute::<&mut [ValRaw], &mut [MaybeUninit<ValRaw>]>(storage);
+        ret.lower(&mut cx, &options, &mut dst.iter_mut())?;
+    } else {
+        let ret_ptr = &storage[ret_index];
+        let mut memory = MemoryMut::new(cx.as_context_mut(), &options);
+        let ptr =
+            validate_inbounds_dynamic(result.size_and_alignment(), memory.as_slice_mut(), ret_ptr)?;
+        ret.store(&mut memory, ptr)?;
+    }
+
+    flags.set_may_leave(true);
+
+    return Ok(());
+}
+
+fn validate_inbounds_dynamic(
+    SizeAndAlignment { size, alignment }: SizeAndAlignment,
+    memory: &[u8],
+    ptr: &ValRaw,
+) -> Result<usize> {
+    // FIXME: needs memory64 support
+    let ptr = usize::try_from(ptr.get_u32())?;
+    if ptr % usize::try_from(alignment)? != 0 {
+        bail!("pointer not aligned");
+    }
+    let end = match ptr.checked_add(size) {
+        Some(n) => n,
+        None => bail!("pointer size overflow"),
+    };
+    if end > memory.len() {
+        bail!("pointer out of bounds")
+    }
+    Ok(ptr)
+}
+
+struct Types {
+    params: Box<[Type]>,
+    result: Type,
+}
+
+struct DynamicContext<F> {
+    func: F,
+    types: Types,
+}
+
+extern "C" fn dynamic_entrypoint<
+    T,
+    F: Fn(StoreContextMut<'_, T>, &[Val]) -> Result<Val> + Send + Sync + 'static,
+>(
+    cx: *mut VMOpaqueContext,
+    data: *mut u8,
+    flags: InstanceFlags,
+    memory: *mut VMMemoryDefinition,
+    realloc: *mut VMCallerCheckedAnyfunc,
+    string_encoding: StringEncoding,
+    storage: *mut ValRaw,
+    storage_len: usize,
+) {
+    let data = data as *const DynamicContext<F>;
+    unsafe {
+        handle_result(|| {
+            call_host_dynamic::<T, _>(
+                &(*data).types,
+                cx,
+                flags,
+                memory,
+                realloc,
+                string_encoding,
+                std::slice::from_raw_parts_mut(storage, storage_len),
+                |store, values| ((*data).func)(store, values),
+            )
+        })
+    }
+}

crates/wasmtime/src/component/func/typed.rs

@@ -1731,7 +1731,7 @@ macro_rules! impl_component_ty_for_tuples {
                     _size = align_to(_size, $t::ALIGN32);
                     _size += $t::SIZE32;
                 )*
-                _size
+                align_to(_size, Self::ALIGN32)
             };
 
             const ALIGN32: u32 = {

crates/wasmtime/src/component/instance.rs

@@ -64,7 +64,7 @@ impl Instance {
     /// Looks up a function by name within this [`Instance`].
     ///
     /// This is a convenience method for calling [`Instance::exports`] followed
-    /// by [`ExportInstance::get_func`].
+    /// by [`ExportInstance::func`].
     ///
     /// # Panics
     ///

crates/wasmtime/src/component/linker.rs

@@ -1,12 +1,13 @@
 use crate::component::func::HostFunc;
 use crate::component::instance::RuntimeImport;
 use crate::component::matching::TypeChecker;
-use crate::component::{Component, Instance, InstancePre, IntoComponentFunc};
-use crate::{AsContextMut, Engine, Module};
+use crate::component::{Component, Instance, InstancePre, IntoComponentFunc, Val};
+use crate::{AsContextMut, Engine, Module, StoreContextMut};
 use anyhow::{anyhow, bail, Context, Result};
 use std::collections::hash_map::{Entry, HashMap};
 use std::marker;
 use std::sync::Arc;
+use wasmtime_environ::component::TypeDef;
 use wasmtime_environ::PrimaryMap;
 
 /// A type used to instantiate [`Component`]s.
@@ -230,6 +231,37 @@ impl<T> LinkerInstance<'_, T> {
         self.insert(name, Definition::Func(func.into_host_func()))
     }
 
+    /// Define a new host-provided function using dynamic types.
+    ///
+    /// `name` must refer to a function type import in `component`.  If and when
+    /// that import is invoked by the component, the specified `func` will be
+    /// called, which must return a `Val` which is an instance of the result
+    /// type of the import.
+    pub fn func_new<
+        F: Fn(StoreContextMut<'_, T>, &[Val]) -> Result<Val> + Send + Sync + 'static,
+    >(
+        &mut self,
+        component: &Component,
+        name: &str,
+        func: F,
+    ) -> Result<()> {
+        for (import_name, ty) in component.env_component().import_types.values() {
+            if name == import_name {
+                if let TypeDef::ComponentFunc(index) = ty {
+                    let name = self.strings.intern(name);
+                    return self.insert(
+                        name,
+                        Definition::Func(HostFunc::new_dynamic(func, *index, component.types())),
+                    );
+                } else {
+                    bail!("import `{name}` has the wrong type (expected a function)");
+                }
+            }
+        }
+
+        Err(anyhow!("import `{name}` not found"))
+    }
+
     /// Defines a [`Module`] within this instance.
     ///
     /// This can be used to provide a core wasm [`Module`] as an import to a

crates/wasmtime/src/component/matching.rs

@@ -3,12 +3,13 @@ use crate::component::linker::{Definition, NameMap, Strings};
 use crate::types::matching;
 use crate::Module;
 use anyhow::{anyhow, bail, Context, Result};
+use std::sync::Arc;
 use wasmtime_environ::component::{
     ComponentTypes, TypeComponentInstance, TypeDef, TypeFuncIndex, TypeModule,
 };
 
 pub struct TypeChecker<'a> {
-    pub types: &'a ComponentTypes,
+    pub types: &'a Arc<ComponentTypes>,
     pub strings: &'a Strings,
 }
 

crates/wasmtime/src/component/types.rs

@@ -222,6 +222,7 @@ impl Flags {
 }
 
 /// Represents the size and alignment requirements of the heap-serialized form of a type
+#[derive(Debug)]
 pub(crate) struct SizeAndAlignment {
     pub(crate) size: usize,
     pub(crate) alignment: u32,
@@ -662,7 +663,10 @@ fn variant_size_and_alignment(types: impl ExactSizeIterator<Item = Type>) -> Siz
     }
 
     SizeAndAlignment {
-        size: func::align_to(usize::from(discriminant_size), alignment) + size,
+        size: func::align_to(
+            func::align_to(usize::from(discriminant_size), alignment) + size,
+            alignment,
+        ),
         alignment,
     }
 }

crates/wasmtime/src/component/values.rs

@@ -604,8 +604,9 @@ impl Val {
             }
             Type::Flags(handle) => {
                 let count = u32::try_from(handle.names().len()).unwrap();
-                assert!(count <= 32);
-                let value = iter::once(u32::lift(store, options, next(src))?).collect();
+                let value = iter::repeat_with(|| u32::lift(store, options, next(src)))
+                    .take(u32_count_for_flag_count(count.try_into()?))
+                    .collect::<Result<_>>()?;
 
                 Val::Flags(Flags {
                     ty: handle.clone(),