Merge pull request #539 from nagisa/nagisa/available-imports

wasm-smith: Allow specifying the set of valid imports

crates/wasm-smith/Cargo.toml

@@ -18,17 +18,18 @@ harness = false
 [dependencies]
 arbitrary = { version = "1.0.0", features = ["derive"] }
 flagset = "0.4"
-leb128 = "0.2.4"
-wasm-encoder = { version = "0.10.0", path = "../wasm-encoder" }
 indexmap = "1.6"
+leb128 = "0.2.4"
 serde = { version = "1", features = ['derive'], optional = true }
+wasm-encoder = { version = "0.10.0", path = "../wasm-encoder" }
+wasmparser = { version = "0.83", path = "../wasmparser" }
 
 [dev-dependencies]
 criterion = "0.3.3"
 libfuzzer-sys = "0.4.0"
 rand = { version = "0.8.0", features = ["small_rng"] }
-wasmparser = { path = "../wasmparser" }
 wasmprinter = { path = "../wasmprinter" }
+wat = { path = "../wat" }
 
 [features]
 _internal_cli = ["serde", "flagset/serde"]

crates/wasm-smith/src/component.rs

@@ -312,7 +312,7 @@ impl Component {
                     let module = crate::limited_string(100, u)?;
                     let existing_module_imports = imports.entry(module.clone()).or_default();
                     let field = crate::unique_string(100, existing_module_imports, u)?;
-                    let entity_ty = match self.arbitrary_core_entity_type(
+                    let entity_type = match self.arbitrary_core_entity_type(
                         u,
                         &types,
                         &mut entity_choices,
@@ -321,9 +321,11 @@ impl Component {
                         None => return Ok(false),
                         Some(x) => x,
                     };
-                    defs.push(ModuleTypeDef::Import(crate::core::Import(
-                        module, field, entity_ty,
-                    )));
+                    defs.push(ModuleTypeDef::Import(crate::core::Import {
+                        module: module.into(),
+                        field: field.into(),
+                        entity_type,
+                    }));
                 }
 
                 // Export.

crates/wasm-smith/src/component/encode.rs

@@ -76,9 +76,9 @@ impl Type {
                         }
                         ModuleTypeDef::Import(imp) => {
                             enc_mod_ty.import(
-                                &imp.0,
-                                &imp.1,
-                                crate::core::encode::translate_entity_type(&imp.2),
+                                &imp.module,
+                                &imp.field,
+                                crate::core::encode::translate_entity_type(&imp.entity_type),
                             );
                         }
                         ModuleTypeDef::Export(name, ty) => {

crates/wasm-smith/src/config.rs

@@ -2,6 +2,7 @@
 
 use crate::InstructionKinds;
 use arbitrary::{Arbitrary, Result, Unstructured};
+use std::borrow::Cow;
 
 /// Configuration for a generated module.
 ///
@@ -59,6 +60,41 @@ pub trait Config: 'static + std::fmt::Debug {
         100
     }
 
+    /// The imports that may be used when generating the module.
+    ///
+    /// Defaults to `None` which means that any arbitrary import can be generated.
+    ///
+    /// To only allow specific imports, override this method to return a WebAssembly module which
+    /// describes the imports allowed.
+    ///
+    /// Note that [`Self::min_imports`] is ignored when `available_imports` are enabled.
+    ///
+    /// # Panics
+    ///
+    /// The returned value must be a valid binary encoding of a WebAssembly module. `wasm-smith`
+    /// will panic if the module cannot be parsed.
+    ///
+    /// # Example
+    ///
+    /// An implementation of this method could use the `wat` crate to provide a human-readable and
+    /// maintainable description:
+    ///
+    /// ```rust
+    /// Some(wat::parse_str(r#"
+    ///     (module
+    ///         (import "env" "ping" (func (param i32)))
+    ///         (import "env" "pong" (func (result i32)))
+    ///         (import "env" "memory" (memory 1))
+    ///         (import "env" "table" (table 1))
+    ///         (import "env" "tag" (tag (param i32)))
+    ///     )
+    /// "#))
+    /// # ;
+    /// ```
+    fn available_imports(&self) -> Option<Cow<'_, [u8]>> {
+        None
+    }
+
     /// The minimum number of functions to generate. Defaults to 0.  This
     /// includes imported functions.
     fn min_funcs(&self) -> usize {
@@ -390,6 +426,7 @@ impl Config for DefaultConfig {}
 #[allow(missing_docs)]
 pub struct SwarmConfig {
     pub allow_start_export: bool,
+    pub available_imports: Option<Vec<u8>>,
     pub bulk_memory_enabled: bool,
     pub canonicalize_nans: bool,
     pub exceptions_enabled: bool,
@@ -489,6 +526,7 @@ impl<'a> Arbitrary<'a> for SwarmConfig {
             memory64_enabled: false,
             max_type_size: 1000,
             canonicalize_nans: false,
+            available_imports: None,
         })
     }
 }
@@ -510,6 +548,12 @@ impl Config for SwarmConfig {
         self.max_imports
     }
 
+    fn available_imports(&self) -> Option<Cow<'_, [u8]>> {
+        self.available_imports
+            .as_ref()
+            .map(|is| Cow::Borrowed(&is[..]))
+    }
+
     fn min_funcs(&self) -> usize {
         self.min_funcs
     }
@@ -650,14 +694,14 @@ impl Config for SwarmConfig {
         self.max_nesting_depth
     }
 
-    fn memory64_enabled(&self) -> bool {
-        self.memory64_enabled
-    }
-
     fn max_type_size(&self) -> u32 {
         self.max_type_size
     }
 
+    fn memory64_enabled(&self) -> bool {
+        self.memory64_enabled
+    }
+
     fn canonicalize_nans(&self) -> bool {
         self.canonicalize_nans
     }