Repository: bytecodealliance/wasm-tools
Pull request #1412 (2024-02-10 17:57:56)

Remove usage of `AtomicU64` in `wasmparser`
Instead substitute this with `AtomicUsize` which panics on overflow to catch any possible issues on 32-bit platforms. This is motivated by https://github.com/rust-lang/rust/pull/120588 which is pulling `wasmparser` into the Rust compiler and `AtomicU64` is not available on all the platforms the compiler itself is built for.

I plan on backporting this to the `0.118.x` release track as well which is what's used by the `object` crate currently to avoid the need for a whole bunch of cascading updates.


Patch:
diff --git a/crates/wasm-smith/Cargo.toml b/crates/wasm-smith/Cargo.toml
index 037833dd..5fbcfa3e 100644
--- a/crates/wasm-smith/Cargo.toml
+++ b/crates/wasm-smith/Cargo.toml
@@ -19,7 +19,7 @@ harness = false
 workspace = true
 
 [dependencies]
-anyhow = { workspace = true, optional = true }
+anyhow = { workspace = true }
 arbitrary = { workspace = true, features = ["derive"] }
 clap = { workspace = true, optional = true }
 flagset = "0.4"
@@ -42,5 +42,5 @@ wat = { path = "../wat" }
 libfuzzer-sys = { workspace = true }
 
 [features]
-_internal_cli = ["anyhow", "clap", "flagset/serde", "serde", "serde_derive", "wasmparser", "wat"]
+_internal_cli = ["clap", "flagset/serde", "serde", "serde_derive", "wasmparser", "wat"]
 wasmparser = ['dep:wasmparser', 'wasm-encoder/wasmparser']
diff --git a/crates/wasm-smith/src/config.rs b/crates/wasm-smith/src/config.rs
index b183530e..6efe88d9 100644
--- a/crates/wasm-smith/src/config.rs
+++ b/crates/wasm-smith/src/config.rs
@@ -481,6 +481,18 @@ define_config! {
         ///
         /// Defaults to `false`.
         pub threads_enabled: bool = false,
+
+        /// Indicates whether wasm-smith is allowed to generate invalid function
+        /// bodies.
+        ///
+        /// When enabled this option will enable taking raw bytes from the input
+        /// byte stream and using them as a wasm function body. This means that
+        /// the output module is not guaranteed to be valid but can help tickle
+        /// various parts of validation/compilation in some circumstances as
+        /// well.
+        ///
+        /// Defaults to `false`.
+        pub allow_invalid_funcs: bool = false,
     }
 }
 
@@ -617,6 +629,7 @@ impl<'a> Arbitrary<'a> for Config {
             tail_call_enabled: false,
             gc_enabled: false,
             generate_custom_sections: false,
+            allow_invalid_funcs: false,
         })
     }
 }
diff --git a/crates/wasm-smith/src/core.rs b/crates/wasm-smith/src/core.rs
index 6a2836ae..ee1b8d2e 100644
--- a/crates/wasm-smith/src/core.rs
+++ b/crates/wasm-smith/src/core.rs
@@ -172,7 +172,7 @@ impl Module {
         duplicate_imports_behavior: DuplicateImportsBehavior,
     ) -> Result<Self> {
         let mut module = Module::empty(config, duplicate_imports_behavior);
-        module.build(u, false)?;
+        module.build(u)?;
         Ok(module)
     }
 
@@ -213,30 +213,6 @@ impl Module {
     }
 }
 
-/// Same as [`Module`], but may be invalid.
-///
-/// This module generates function bodies differnetly than `Module` to try to
-/// better explore wasm decoders and such.
-#[derive(Debug)]
-pub struct MaybeInvalidModule {
-    module: Module,
-}
-
-impl MaybeInvalidModule {
-    /// Encode this Wasm module into bytes.
-    pub fn to_bytes(&self) -> Vec<u8> {
-        self.module.to_bytes()
-    }
-}
-
-impl<'a> Arbitrary<'a> for MaybeInvalidModule {
-    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
-        let mut module = Module::empty(Config::default(), DuplicateImportsBehavior::Allowed);
-        module.build(u, true)?;
-        Ok(MaybeInvalidModule { module })
-    }
-}
-
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub(crate) struct RecGroup {
     pub(crate) types: Vec<SubType>,
@@ -398,7 +374,7 @@ pub(crate) enum GlobalInitExpr {
 }
 
 impl Module {
-    fn build(&mut self, u: &mut Unstructured, allow_invalid: bool) -> Result<()> {
+    fn build(&mut self, u: &mut Unstructured) -> Result<()> {
         self.valtypes = configured_valtypes(&self.config);
 
         // We attempt to figure out our available imports *before* creating the types section here,
@@ -425,7 +401,7 @@ impl Module {
         self.arbitrary_start(u)?;
         self.arbitrary_elems(u)?;
         self.arbitrary_data(u)?;
-        self.arbitrary_code(u, allow_invalid)?;
+        self.arbitrary_code(u)?;
         Ok(())
     }
 
@@ -1795,11 +1771,11 @@ impl Module {
         )
     }
 
-    fn arbitrary_code(&mut self, u: &mut Unstructured, allow_invalid: bool) -> Result<()> {
+    fn arbitrary_code(&mut self, u: &mut Unstructured) -> Result<()> {
         self.code.reserve(self.num_defined_funcs);
         let mut allocs = CodeBuilderAllocations::new(self);
         for (_, ty) in self.funcs[self.funcs.len() - self.num_defined_funcs..].iter() {
-            let body = self.arbitrary_func_body(u, ty, &mut allocs, allow_invalid)?;
+            let body = self.arbitrary_func_body(u, ty, &mut allocs)?;
             self.code.push(body);
         }
         allocs.finish(u, self)?;
@@ -1811,11 +1787,10 @@ impl Module {
         u: &mut Unstructured,
         ty: &FuncType,
         allocs: &mut CodeBuilderAllocations,
-        allow_invalid: bool,
     ) -> Result<Code> {
         let mut locals = self.arbitrary_locals(u)?;
         let builder = allocs.builder(ty, &mut locals);
-        let instructions = if allow_invalid && u.arbitrary().unwrap_or(false) {
+        let instructions = if self.config.allow_invalid_funcs && u.arbitrary().unwrap_or(false) {
             Instructions::Arbitrary(arbitrary_vec_u8(u)?)
         } else {
             Instructions::Generated(builder.arbitrary(u, self)?)
diff --git a/crates/wasm-smith/src/core/terminate.rs b/crates/wasm-smith/src/core/terminate.rs
index adcfeed5..813135c1 100644
--- a/crates/wasm-smith/src/core/terminate.rs
+++ b/crates/wasm-smith/src/core/terminate.rs
@@ -1,4 +1,5 @@
 use super::*;
+use anyhow::{bail, Result};
 use std::mem;
 use wasm_encoder::BlockType;
 
@@ -12,7 +13,14 @@ impl Module {
     ///
     /// The index of the fuel global is returned, so that you may control how
     /// much fuel the module is given.
-    pub fn ensure_termination(&mut self, default_fuel: u32) -> u32 {
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if any function body was generated with
+    /// possibly-invalid bytes rather than being generated by wasm-smith. In
+    /// such a situation this pass does not parse the input bytes and inject
+    /// instructions, instead it returns an error.
+    pub fn ensure_termination(&mut self, default_fuel: u32) -> Result<u32> {
         let fuel_global = self.globals.len() as u32;
         self.globals.push(GlobalType {
             val_type: ValType::I32,
@@ -41,10 +49,12 @@ impl Module {
 
             let instrs = match &mut code.instructions {
                 Instructions::Generated(list) => list,
-                // only present on modules contained within
-                // `MaybeInvalidModule`, which doesn't expose its internal
-                // `Module`.
-                Instructions::Arbitrary(_) => unreachable!(),
+                Instructions::Arbitrary(_) => {
+                    bail!(
+                        "failed to ensure that a function generated due to it \
+                         containing arbitrary instructions"
+                    )
+                }
             };
             let mut new_insts = Vec::with_capacity(instrs.len() * 2);
 
@@ -65,6 +75,6 @@ impl Module {
             *instrs = new_insts;
         }
 
-        fuel_global
+        Ok(fuel_global)
     }
 }
diff --git a/crates/wasm-smith/src/lib.rs b/crates/wasm-smith/src/lib.rs
index 132ccaf2..b985791a 100644
--- a/crates/wasm-smith/src/lib.rs
+++ b/crates/wasm-smith/src/lib.rs
@@ -58,7 +58,7 @@ mod component;
 mod config;
 mod core;
 
-pub use crate::core::{InstructionKind, InstructionKinds, MaybeInvalidModule, Module};
+pub use crate::core::{InstructionKind, InstructionKinds, Module};
 use arbitrary::{Result, Unstructured};
 pub use component::Component;
 pub use config::{Config, MemoryOffsetChoices};
diff --git a/crates/wasmparser/src/validator/types.rs b/crates/wasmparser/src/validator/types.rs
index acab322d..043469f5 100644
--- a/crates/wasmparser/src/validator/types.rs
+++ b/crates/wasmparser/src/validator/types.rs
@@ -14,7 +14,7 @@ use indexmap::{IndexMap, IndexSet};
 use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet};
 use std::ops::{Index, Range};
-use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::{
     borrow::Borrow,
     hash::{Hash, Hasher},
@@ -1451,7 +1451,7 @@ pub struct ResourceId {
     // per resource id is probably too expensive. To amortize that cost each
     // top-level wasm component gets a single globally unique identifier, and
     // then within a component contextually unique identifiers are handed out.
-    globally_unique_id: u64,
+    globally_unique_id: usize,
 
     // A contextually unique id within the globally unique id above. This is
     // allocated within a `TypeAlloc` with its own counter, and allocations of
@@ -2925,7 +2925,7 @@ pub(crate) struct TypeAlloc {
 
     // This is assigned at creation of a `TypeAlloc` and then never changed.
     // It's used in one entry for all `ResourceId`s contained within.
-    globally_unique_id: u64,
+    globally_unique_id: usize,
 
     // This is a counter that's incremeneted each time `alloc_resource_id` is
     // called.
@@ -2934,10 +2934,17 @@ pub(crate) struct TypeAlloc {
 
 impl Default for TypeAlloc {
     fn default() -> TypeAlloc {
-        static NEXT_GLOBAL_ID: AtomicU64 = AtomicU64::new(0);
+        static NEXT_GLOBAL_ID: AtomicUsize = AtomicUsize::new(0);
         let mut ret = TypeAlloc {
             list: TypeList::default(),
-            globally_unique_id: NEXT_GLOBAL_ID.fetch_add(1, Ordering::Relaxed),
+            globally_unique_id: {
+                let id = NEXT_GLOBAL_ID.fetch_add(1, Ordering::Relaxed);
+                if id > usize::MAX - 10_000 {
+                    NEXT_GLOBAL_ID.store(usize::MAX - 10_000, Ordering::Relaxed);
+                    panic!("overflow on the global id counter");
+                }
+                id
+            },
             next_resource_id: 0,
         };
         ret.list.canonical_rec_groups = Some(Default::default());
diff --git a/fuzz/src/lib.rs b/fuzz/src/lib.rs
index 885a55dc..bf4806c8 100644
--- a/fuzz/src/lib.rs
+++ b/fuzz/src/lib.rs
@@ -42,7 +42,7 @@ pub fn generate_valid_module(
     // still produce a valid module.
     if u.ratio(1, 10)? {
         log::debug!("ensuring termination with 100 fuel");
-        module.ensure_termination(100);
+        let _ = module.ensure_termination(100);
     }
 
     log_wasm(&bytes, &config);
@@ -76,6 +76,27 @@ pub fn generate_valid_component(
     Ok((bytes, config))
 }
 
+pub fn validator_for_config(config: &Config) -> wasmparser::Validator {
+    wasmparser::Validator::new_with_features(wasmparser::WasmFeatures {
+        multi_value: config.multi_value_enabled,
+        multi_memory: config.max_memories > 1,
+        bulk_memory: config.bulk_memory_enabled,
+        reference_types: config.reference_types_enabled,
+        simd: config.simd_enabled,
+        relaxed_simd: config.relaxed_simd_enabled,
+        memory64: config.memory64_enabled,
+        threads: config.threads_enabled,
+        exceptions: config.exceptions_enabled,
+        // TODO: determine our larger story for function-references in
+        // wasm-tools and whether we should just have a Wasm GC flag since
+        // function-references is effectively part of the Wasm GC proposal at
+        // this point.
+        function_references: config.gc_enabled,
+        gc: config.gc_enabled,
+        ..wasmparser::WasmFeatures::default()
+    })
+}
+
 // Optionally log the module and its configuration if we've gotten this
 // far. Note that we don't do this unconditionally to avoid slowing down
 // fuzzing, but this is expected to be enabled when debugging a failing
diff --git a/fuzz/src/validate.rs b/fuzz/src/validate.rs
index 13859dc6..c9657a7b 100644
--- a/fuzz/src/validate.rs
+++ b/fuzz/src/validate.rs
@@ -1,8 +1,33 @@
-use arbitrary::{Arbitrary, Result, Unstructured};
-use wasm_smith::MaybeInvalidModule;
+use arbitrary::{Result, Unstructured};
 use wasmparser::{Validator, WasmFeatures};
 
 pub fn run(u: &mut Unstructured<'_>) -> Result<()> {
+    // Either use `wasm-smith` to generate a module with possibly invalid
+    // functions or try validating raw bytes from the input itself.
+    if u.arbitrary()? {
+        validate_maybe_invalid_module(u)?;
+    } else {
+        validate_raw_bytes(u)?;
+    }
+    Ok(())
+}
+
+pub fn validate_maybe_invalid_module(u: &mut Unstructured<'_>) -> Result<()> {
+    // Generate a "valid" module but specifically allow invalid functions which
+    // means that some functions may be defined from the input bytes raw. This
+    // means that most of the module is valid but only some functions may be
+    // invalid which can help stress various bits and pieces of validation.
+    let (wasm, config) = crate::generate_valid_module(u, |config, _| {
+        config.allow_invalid_funcs = true;
+        Ok(())
+    })?;
+    let mut validator = crate::validator_for_config(&config);
+    drop(validator.validate_all(&wasm));
+    Ok(())
+}
+
+pub fn validate_raw_bytes(u: &mut Unstructured<'_>) -> Result<()> {
+    // Enable arbitrary combinations of features to validate the input bytes.
     let mut validator = Validator::new_with_features(WasmFeatures {
         reference_types: u.arbitrary()?,
         multi_value: u.arbitrary()?,
@@ -26,18 +51,8 @@ pub fn run(u: &mut Unstructured<'_>) -> Result<()> {
         component_model_values: u.arbitrary()?,
         component_model_nested_names: u.arbitrary()?,
     });
-    let use_maybe_invalid = u.arbitrary()?;
-
-    if use_maybe_invalid {
-        if let Ok(module) = MaybeInvalidModule::arbitrary(u) {
-            let wasm = module.to_bytes();
-            crate::log_wasm(&wasm, "");
-            drop(validator.validate_all(&wasm));
-        }
-    } else {
-        let wasm = u.bytes(u.len())?;
-        crate::log_wasm(wasm, "");
-        drop(validator.validate_all(wasm));
-    }
+    let wasm = u.bytes(u.len())?;
+    crate::log_wasm(wasm, "");
+    drop(validator.validate_all(wasm));
     Ok(())
 }
diff --git a/fuzz/src/validate_valid_module.rs b/fuzz/src/validate_valid_module.rs
index 85592543..95510225 100644
--- a/fuzz/src/validate_valid_module.rs
+++ b/fuzz/src/validate_valid_module.rs
@@ -25,25 +25,7 @@ pub fn run(u: &mut Unstructured<'_>) -> Result<()> {
     };
 
     // Validate the module or component and assert that it passes validation.
-    let mut validator = wasmparser::Validator::new_with_features(wasmparser::WasmFeatures {
-        component_model: generate_component,
-        multi_value: config.multi_value_enabled,
-        multi_memory: config.max_memories > 1,
-        bulk_memory: config.bulk_memory_enabled,
-        reference_types: config.reference_types_enabled,
-        simd: config.simd_enabled,
-        relaxed_simd: config.relaxed_simd_enabled,
-        memory64: config.memory64_enabled,
-        threads: config.threads_enabled,
-        exceptions: config.exceptions_enabled,
-        // TODO: determine our larger story for function-references in
-        // wasm-tools and whether we should just have a Wasm GC flag since
-        // function-references is effectively part of the Wasm GC proposal at
-        // this point.
-        function_references: config.gc_enabled,
-        gc: config.gc_enabled,
-        ..wasmparser::WasmFeatures::default()
-    });
+    let mut validator = crate::validator_for_config(&config);
     if let Err(e) = validator.validate_all(&wasm_bytes) {
         let component_or_module = if generate_component {
             "component"
diff --git a/src/bin/wasm-tools/smith.rs b/src/bin/wasm-tools/smith.rs
index dc1c32f1..2a6a1aaf 100644
--- a/src/bin/wasm-tools/smith.rs
+++ b/src/bin/wasm-tools/smith.rs
@@ -1,10 +1,9 @@
 use anyhow::{Context, Result};
-use arbitrary::Arbitrary;
 use clap::Parser;
 use std::io::{stdin, Read};
 use std::path::PathBuf;
 use std::process;
-use wasm_smith::{MaybeInvalidModule, Module};
+use wasm_smith::Module;
 
 /// A WebAssembly test case generator.
 ///
@@ -52,23 +51,22 @@ pub struct Opts {
     /// and in function prologues. When the fuel reaches 0, a trap is raised to
     /// terminate execution. Control the default amount of fuel with the
     /// `--fuel` flag.
-    #[clap(long = "ensure-termination")]
+    ///
+    /// Note that when combined with `--allow-invalid-funcs` this subcommand may
+    /// return an error because arbitrary function bodies taken from input bytes
+    /// cannot be validated to ensure they always terminate.
+    #[clap(long)]
     ensure_termination: bool,
 
-    /// Indicates that the generated module may contain invalid wasm functions,
-    /// taken raw from the input DNA.
-    #[clap(long = "maybe-invalid")]
-    maybe_invalid: bool,
-
     /// The default amount of fuel used with `--ensure-termination`.
     ///
     /// This is roughly the number of loop iterations and function calls that
     /// will be executed before a trap is raised to prevent infinite loops.
-    #[clap(short = 'f', long = "fuel")]
+    #[clap(short, long)]
     fuel: Option<u32>,
 
     /// JSON configuration file with settings to control the wasm output.
-    #[clap(short = 'c', long = "config")]
+    #[clap(short, long)]
     config: Option<PathBuf>,
 
     #[clap(flatten)]
@@ -98,36 +96,31 @@ impl Opts {
         };
 
         let mut u = arbitrary::Unstructured::new(&seed);
-        let wasm_bytes = if self.maybe_invalid {
-            MaybeInvalidModule::arbitrary(&mut u)
-                .unwrap_or_else(|e| {
-                    eprintln!("error: failed to generate module: {}", e);
-                    process::exit(2);
-                })
-                .to_bytes()
-        } else {
-            let json = match &self.config {
-                Some(path) => {
-                    let json = std::fs::read_to_string(&path).with_context(|| {
-                        format!("failed to read json config: {}", path.display())
-                    })?;
-                    serde_json::from_str(&json).with_context(|| {
-                        format!("failed to decode json config: {}", path.display())
-                    })?
-                }
-                None => wasm_smith::InternalOptionalConfig::default(),
-            };
-            let config = self.module_config.clone().or(json);
-            let config = wasm_smith::Config::try_from(config)?;
-            let mut module = Module::new(config, &mut u).unwrap_or_else(|e| {
-                eprintln!("error: failed to generate module: {}", e);
-                process::exit(2);
-            });
-            if self.ensure_termination {
-                module.ensure_termination(self.fuel.unwrap_or(100));
+        let json = match &self.config {
+            Some(path) => {
+                let json = std::fs::read_to_string(&path)
+                    .with_context(|| format!("failed to read json config: {}", path.display()))?;
+                serde_json::from_str(&json)
+                    .with_context(|| format!("failed to decode json config: {}", path.display()))?
             }
-            module.to_bytes()
+            None => wasm_smith::InternalOptionalConfig::default(),
         };
+        let config = self.module_config.clone().or(json);
+        let config = wasm_smith::Config::try_from(config)?;
+        let mut module = Module::new(config, &mut u).unwrap_or_else(|e| {
+            eprintln!("error: failed to generate module: {}", e);
+            process::exit(2);
+        });
+        if self.ensure_termination {
+            module
+                .ensure_termination(self.fuel.unwrap_or(100))
+                .context(
+                    "failed to ensure module always terminates, \
+                    `--ensure-termination` is incompatible with \
+                    `--allow-invalid-funcs`",
+                )?;
+        }
+        let wasm_bytes = module.to_bytes();
 
         self.output.output(wasm_tools::Output::Wasm {
             bytes: &wasm_bytes,


Test patch:
diff --git a/crates/wasm-smith/tests/core.rs b/crates/wasm-smith/tests/core.rs
index 5816f8e3..36db8adb 100644
--- a/crates/wasm-smith/tests/core.rs
+++ b/crates/wasm-smith/tests/core.rs
@@ -27,7 +27,7 @@ fn smoke_test_ensure_termination() {
         rng.fill_bytes(&mut buf);
         let u = Unstructured::new(&buf);
         if let Ok(mut module) = Module::arbitrary_take_rest(u) {
-            module.ensure_termination(10);
+            module.ensure_termination(10).unwrap();
             let wasm_bytes = module.to_bytes();
 
             let mut validator = Validator::new_with_features(wasm_features());