wip: (wasm-smith): allow specifying the set of valid imports

crates/wasm-smith/src/config.rs

@@ -2,6 +2,8 @@
 
 use arbitrary::{Arbitrary, Result, Unstructured};
 
+use crate::FuncType;
+
 /// Configuration for a generated module.
 ///
 /// Don't care to configure your generated modules? Just use
@@ -48,6 +50,16 @@ pub trait Config: Clone {
         100
     }
 
+    /// Return the available set of imports.
+    ///
+    /// By default, returns `None` which means that any arbitrary import can be generated.
+    ///
+    /// To only allow imports from a specific set, override this to return a vec of
+    /// `(module name, field name, entity type)` describing each available import.
+    fn available_imports(&self) -> Option<Vec<Import>> {
+        None
+    }
+
     /// The minimum number of functions to generate. Defaults to 0.  This
     /// includes imported functions.
     fn min_funcs(&self) -> usize {
@@ -275,6 +287,31 @@ pub trait Config: Clone {
     }
 }
 
+/// TODO
+#[derive(Debug)]
+pub struct Import {
+    /// TODO
+    pub module: String,
+    /// TODO
+    pub name: Option<String>,
+    /// TODO
+    pub desc: ImportDesc,
+}
+
+/// TODO
+#[non_exhaustive]
+#[derive(Debug)]
+pub enum ImportDesc {
+    /// TODO
+    Func(FuncType),
+    /// TODO
+    Global(()),
+    /// TODO
+    Table(()),
+    /// TODO
+    Memory(()),
+}
+
 /// The default configuration.
 #[derive(Arbitrary, Debug, Default, Copy, Clone)]
 pub struct DefaultConfig;

crates/wasm-smith/src/lib.rs

@@ -59,11 +59,12 @@ mod terminate;
 
 use crate::code_builder::CodeBuilderAllocations;
 use arbitrary::{Arbitrary, Result, Unstructured};
+use std::cell::Cell;
 use std::collections::{HashMap, HashSet};
 use std::rc::Rc;
 use std::str;
 
-pub use config::{Config, DefaultConfig, SwarmConfig};
+pub use config::{Config, DefaultConfig, Import, ImportDesc, SwarmConfig};
 
 /// A pseudo-random WebAssembly module.
 ///
@@ -293,10 +294,13 @@ enum Type {
     Instance(Rc<InstanceType>),
 }
 
+/// A function type.
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
-struct FuncType {
-    params: Vec<ValType>,
-    results: Vec<ValType>,
+pub struct FuncType {
+    /// Types of parameters of the function.
+    pub params: Vec<ValType>,
+    /// Types of results of the function.
+    pub results: Vec<ValType>,
 }
 
 #[derive(Clone, Debug, Default)]
@@ -326,14 +330,23 @@ enum EntityType {
     Module(u32, Rc<ModuleType>),
 }
 
+/// Primitive WASM type.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
-enum ValType {
+#[non_exhaustive]
+pub enum ValType {
+    /// Signed 32-bit integer.
     I32,
+    /// Signed 64-bit integer.
     I64,
+    /// 32-bit floating point number.
     F32,
+    /// 64-bit floating point number.
     F64,
+    /// TODO:
     V128,
+    /// TODO:
     FuncRef,
+    /// TODO:
     ExternRef,
 }
 
@@ -1251,9 +1264,13 @@ where
             return Ok(());
         }
 
+        if let Some(imports) = self.config.available_imports() {
+            return self.arbitrary_available_imports(imports, min, u);
+        }
+
         let mut choices: Vec<
             fn(&mut Unstructured, &mut ConfiguredModule<C>) -> Result<EntityType>,
-        > = Vec::with_capacity(4);
+        > = Vec::with_capacity(6);
 
         let mut imports = Vec::new();
         arbitrary_loop(u, min, self.config.max_imports() - self.num_imports, |u| {
@@ -1373,6 +1390,111 @@ where
         Ok(())
     }
 
+    fn arbitrary_available_imports(
+        &mut self,
+        available_imports: Vec<Import>,
+        min: usize,
+        u: &mut Unstructured,
+    ) -> Result<()> {
+        let mut available_funcs = Vec::new();
+        for import in available_imports {
+            match import.desc {
+                ImportDesc::Func(_) => &mut available_funcs,
+                ImportDesc::Global(_) => todo!(),
+                ImportDesc::Table(_) => todo!(),
+                ImportDesc::Memory(_) => todo!(),
+            }
+            .push((import, Cell::new(None)))
+        }
+
+        let mut choices: Vec<
+            for<'i> fn(
+                &mut Unstructured,
+                &'i [(Import, Cell<Option<u32>>)],
+            ) -> Result<&'i (Import, Cell<Option<u32>>)>,
+        > = Vec::with_capacity(1);
+
+        let mut imports = Vec::new();
+        arbitrary_loop(u, min, self.config.max_imports() - self.num_imports, |u| {
+            choices.clear();
+            if self.can_add_local_or_import_func() && !available_funcs.is_empty() {
+                choices.push(|u, f| u.choose(f))
+            }
+
+            if choices.is_empty() {
+                return Ok(false);
+            }
+
+            let f = u.choose(&choices)?;
+            let (import, idx) = f(u, &available_funcs)?;
+
+            let module = import.module.clone();
+            let name = import.name.clone();
+            let ty: EntityType = match &import.desc {
+                ImportDesc::Func(func) => {
+                    if idx.get().is_none() && self.types.len() < self.config.max_types() {
+                        // Try defining the type in the first section.
+                        // TODO: pick section arbitrary
+                        for (section_idx, section) in self.initial_sections.iter_mut().enumerate() {
+                            if let InitialSection::Type(types) = section {
+                                idx.set(Some(self.types.len() as u32));
+                                self.func_types.push(self.types.len() as u32);
+                                self.types.push(LocalType::Defined {
+                                    section: section_idx,
+                                    nth: types.len(),
+                                });
+                                types.push(Type::Func(Rc::new(func.clone())));
+                                break;
+                            }
+                        }
+                    }
+                    match idx.get() {
+                        Some(idx) => {
+                            let ty = self.func_type(idx).clone();
+                            EntityType::Func(idx, ty)
+                        }
+                        None => return Ok(false),
+                    }
+                }
+                ImportDesc::Global(_) => todo!(),
+                ImportDesc::Table(_) => todo!(),
+                ImportDesc::Memory(_) => todo!(),
+            };
+
+            let budget = self.config.max_type_size() - self.type_size;
+            if ty.size() + 1 > budget {
+                return Ok(false);
+            }
+            self.type_size += ty.size() + 1;
+
+            match &ty {
+                EntityType::Func(idx, ty) => self.funcs.push((Some(*idx), ty.clone())),
+                EntityType::Global(ty) => self.globals.push(ty.clone()),
+                EntityType::Table(ty) => self.tables.push(ty.clone()),
+                EntityType::Memory(ty) => self.memories.push(ty.clone()),
+                EntityType::Module(_idx, ty) => self.modules.push(ty.clone()),
+                EntityType::Instance(_idx, ty) => self.instances.push(ty.clone()),
+            }
+
+            self.num_imports += 1;
+            imports.push((module, name, ty));
+            Ok(true)
+        })?;
+
+        if !imports.is_empty() || u.arbitrary()? {
+            self.initial_sections.push(InitialSection::Import(imports));
+        }
+
+        // After an import section we can no longer update previously-defined
+        // pseudo-instance imports, so set them all to `None` indicating that
+        // the bare name is imported and finalized.
+        for val in self.import_names.values_mut() {
+            *val = None;
+        }
+
+        Ok(())
+    }
+
     fn arbitrary_aliases(
         &mut self,
         available: &mut AvailableAliases,

crates/wasm-smith/tests/tests.rs

@@ -1,7 +1,9 @@
 use arbitrary::{Arbitrary, Unstructured};
 use rand::{rngs::SmallRng, RngCore, SeedableRng};
-use wasm_smith::{Config, ConfiguredModule, Module, SwarmConfig};
-use wasmparser::{Validator, WasmFeatures};
+use wasm_smith::{
+    Config, ConfiguredModule, FuncType, Import, ImportDesc, Module, SwarmConfig, ValType,
+};
+use wasmparser::{ImportSectionEntryType, Parser, Validator, WasmFeatures};
 
 fn wasm_features() -> WasmFeatures {
     WasmFeatures {
@@ -68,6 +70,79 @@ fn smoke_test_swarm_config() {
     }
 }
 
+#[test]
+fn smoke_test_imports_config() {
+    #[derive(Clone, Copy)]
+    struct ImportConfig;
+
+    impl Config for ImportConfig {
+        fn available_imports(&self) -> Option<Vec<Import>> {
+            Some(vec![
+                Import {
+                    module: "env".to_string(),
+                    name: Some("ping".to_string()),
+                    desc: ImportDesc::Func(FuncType {
+                        params: vec![ValType::I32],
+                        results: vec![],
+                    }),
+                },
+                Import {
+                    module: "env".to_string(),
+                    name: Some("pong".to_string()),
+                    desc: ImportDesc::Func(FuncType {
+                        params: vec![],
+                        results: vec![ValType::I32],
+                    }),
+                },
+            ])
+        }
+    }
+
+    impl<'a> Arbitrary<'a> for ImportConfig {
+        fn arbitrary(_: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
+            Ok(ImportConfig)
+        }
+    }
+
+    let mut n_pings = 0;
+    let mut n_pongs = 0;
+
+    let mut rng = SmallRng::seed_from_u64(0);
+    let mut buf = vec![0; 1024];
+    for _ in 0..1024 {
+        rng.fill_bytes(&mut buf);
+
+        let u = Unstructured::new(&buf);
+        if let Ok(module) = ConfiguredModule::<ImportConfig>::arbitrary_take_rest(u) {
+            let wasm_bytes = module.to_bytes();
+
+            let mut validator = Validator::new();
+            let mut features = wasm_features();
+            features.module_linking = module.config().module_linking_enabled();
+            validator.wasm_features(features);
+            validate(&mut validator, &wasm_bytes);
+
+            for payload in Parser::new(0).parse_all(&wasm_bytes) {
+                let payload = payload.unwrap();
+                if let wasmparser::Payload::ImportSection(mut rdr) = payload {
+                    while let Ok(import) = rdr.read() {
+                        if let ImportSectionEntryType::Function(_) = import.ty {
+                            assert_eq!(import.module, "env");
+                            match import.field {
+                                Some("ping") => n_pings += 1,
+                                Some("pong") => n_pongs += 1,
+                                other => panic!("unexpected import: {:?}", other),
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    assert!(n_pings > 0 && n_pongs > 0);
+}
+
 fn validate(validator: &mut Validator, bytes: &[u8]) {
     let err = match validator.validate_all(bytes) {
         Ok(()) => return,