wip: allow to specify module's interface in wasm-smith

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_func_imports(&self) -> Option<Vec<(String, Option<String>, FuncType)>> {
+        None
+    }
+
     /// The minimum number of functions to generate. Defaults to 0.  This
     /// includes imported functions.
     fn min_funcs(&self) -> usize {

crates/wasm-smith/src/lib.rs

@@ -135,6 +135,9 @@ where
     /// Indices within `types` that are instance types.
     instance_types: Vec<u32>,
 
+    /// Subset of `func_types` eligible for importing.
+    import_func_types: Option<Vec<((String, Option<String>), u32)>>,
+
     /// Number of imported items into this module.
     num_imports: usize,
 
@@ -223,6 +226,7 @@ impl<C: Config> ConfiguredModule<C> {
             func_types: Vec::new(),
             module_types: Vec::new(),
             instance_types: Vec::new(),
+            import_func_types: None,
             num_imports: 0,
             num_aliases: 0,
             num_defined_funcs: 0,
@@ -293,10 +297,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 +333,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,
 }
 
@@ -970,6 +986,11 @@ where
             self.valtypes.push(ValType::ExternRef);
             self.valtypes.push(ValType::FuncRef);
         }
+
+        if self.config.available_func_imports().is_some() {
+            self.import_func_types = Some(Vec::new());
+        }
+
         self.arbitrary_initial_sections(u)?;
         self.arbitrary_funcs(u)?;
         self.arbitrary_tables(u)?;
@@ -1046,8 +1067,25 @@ where
         let section_idx = self.initial_sections.len();
         self.initial_sections.push(InitialSection::Type(Vec::new()));
         arbitrary_loop(u, min, self.config.max_types() - self.types.len(), |u| {
-            let ty = self.arbitrary_type(u)?;
+            let mut import_name = None;
+
+            let ty = if self.import_func_types.is_none() || u.arbitrary().unwrap_or(false) {
+                self.arbitrary_type(u)?
+            } else {
+                let candidates = self.config.available_func_imports().unwrap();
+                let (module, name, ty) = u.choose(&candidates)?.clone();
+                import_name = Some((module, name));
+                Type::Func(Rc::new(ty))
+            };
+
             self.record_type(&ty);
+            if let Some(name) = import_name {
+                self.import_func_types
+                    .as_mut()
+                    .unwrap()
+                    .push((name, *self.func_types.last().unwrap()))
+            }
+
             let types = match self.initial_sections.last_mut().unwrap() {
                 InitialSection::Type(list) => list,
                 _ => unreachable!(),
@@ -1252,49 +1290,58 @@ where
         }
 
         let mut choices: Vec<
-            fn(&mut Unstructured, &mut ConfiguredModule<C>) -> Result<EntityType>,
+            fn(
+                &mut Unstructured,
+                &mut ConfiguredModule<C>,
+            ) -> Result<(Option<(String, Option<String>)>, EntityType)>,
         > = Vec::with_capacity(4);
 
         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() {
                 choices.push(|u, m| {
-                    let idx = *u.choose(&m.func_types)?;
+                    let (n, idx) = match &m.import_func_types {
+                        Some(import_func_types) => {
+                            let (n, idx) = u.choose(import_func_types)?.clone();
+                            (Some(n), idx)
+                        }
+                        None => (None, *u.choose(&m.func_types)?),
+                    };
                     let ty = m.func_type(idx).clone();
-                    Ok(EntityType::Func(idx, ty))
+                    Ok((n, EntityType::Func(idx, ty)))
                 });
             }
             if self.can_add_local_or_import_module() {
                 choices.push(|u, m| {
                     let idx = *u.choose(&m.module_types)?;
                     let ty = m.module_type(idx).clone();
-                    Ok(EntityType::Module(idx, ty.clone()))
+                    Ok((None, EntityType::Module(idx, ty.clone())))
                 });
             }
             if self.can_add_local_or_import_instance() {
                 choices.push(|u, m| {
                     let idx = *u.choose(&m.instance_types)?;
                     let ty = m.instance_type(idx).clone();
-                    Ok(EntityType::Instance(idx, ty))
+                    Ok((None, EntityType::Instance(idx, ty)))
                 });
             }
             if self.can_add_local_or_import_global() {
                 choices.push(|u, m| {
                     let ty = m.arbitrary_global_type(u)?;
-                    Ok(EntityType::Global(ty))
+                    Ok((None, EntityType::Global(ty)))
                 });
             }
             if self.can_add_local_or_import_memory() {
                 choices.push(|u, m| {
                     let ty = m.arbitrary_memtype(u)?;
-                    Ok(EntityType::Memory(ty))
+                    Ok((None, EntityType::Memory(ty)))
                 });
             }
             if self.can_add_local_or_import_table() {
                 choices.push(|u, m| {
                     let ty = m.arbitrary_table_type(u)?;
-                    Ok(EntityType::Table(ty))
+                    Ok((None, EntityType::Table(ty)))
                 });
             }
 
@@ -1308,7 +1355,7 @@ where
             // Generate a type to import, but only actually add the item if the
             // type size budget allows us to.
             let f = u.choose(&choices)?;
-            let ty = f(u, self)?;
+            let (import_name, ty) = f(u, self)?;
             let budget = self.config.max_type_size() - self.type_size;
             if ty.size() + 1 > budget {
                 return Ok(false);
@@ -1320,8 +1367,10 @@ where
             // might be implicitly generating an instance. If that's the case
             // then we need to record the type of this instance.
             let module_linking = self.config.module_linking_enabled() || self.outers.len() > 0;
-            let (module, name) =
-                unique_import_strings(1_000, &mut self.import_names, module_linking, u)?;
+            let (module, name) = match import_name {
+                Some(it) => it,
+                None => unique_import_strings(1_000, &mut self.import_names, module_linking, u)?,
+            };
             if module_linking
                 && name.is_some()
                 && self.import_names[&module].as_ref().unwrap().len() == 1

crates/wasm-smith/tests/tests.rs

@@ -1,7 +1,7 @@
 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, Module, SwarmConfig, ValType};
+use wasmparser::{ImportSectionEntryType, Parser, Validator, WasmFeatures};
 
 fn wasm_features() -> WasmFeatures {
     WasmFeatures {
@@ -68,6 +68,79 @@ fn smoke_test_swarm_config() {
     }
 }
 
+#[test]
+fn smoke_test_imports_config() {
+    #[derive(Clone, Copy)]
+    struct ImportConfig;
+
+    impl Config for ImportConfig {
+        fn available_func_imports(&self) -> Option<Vec<(String, Option<String>, FuncType)>> {
+            Some(vec![
+                (
+                    "env".to_string(),
+                    Some("ping".to_string()),
+                    FuncType {
+                        params: vec![ValType::I32],
+                        results: vec![],
+                    },
+                ),
+                (
+                    "env".to_string(),
+                    Some("pong".to_string()),
+                    FuncType {
+                        params: vec![],
+                        results: vec![ValType::I32],
+                    },
+                ),
+            ])
+        }
+    }
+
+    impl<'a> Arbitrary<'a> for ImportConfig {
+        fn arbitrary(u: &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,