Repository: bytecodealliance/wasm-tools
Pull request #1616 (2024-06-13 17:05:04)

`wasm-smith`: fix generation of memories with custom page sizes


Patch:
diff --git a/crates/wasm-smith/Cargo.toml b/crates/wasm-smith/Cargo.toml
index 16e4810d..49d93793 100644
--- a/crates/wasm-smith/Cargo.toml
+++ b/crates/wasm-smith/Cargo.toml
@@ -35,7 +35,7 @@ wat = { workspace = true, optional = true }
 [dev-dependencies]
 criterion = { workspace = true }
 rand = { workspace = true }
-wasmparser = { workspace = true }
+wasmparser = { workspace = true, features = ["validate"] }
 wasmprinter = { path = "../wasmprinter" }
 wat = { path = "../wat" }
 
diff --git a/crates/wasm-smith/src/config.rs b/crates/wasm-smith/src/config.rs
index 59ec8166..f628b7c7 100644
--- a/crates/wasm-smith/src/config.rs
+++ b/crates/wasm-smith/src/config.rs
@@ -379,17 +379,21 @@ define_config! {
         /// wasm proposal.
         pub max_memories: usize = 1,
 
-        /// The maximum, in 64k Wasm pages, of any 32-bit memory's initial or
-        /// maximum size.
+        /// The maximum, in bytes, of any 32-bit memory's initial or maximum
+        /// size.
         ///
-        /// Defaults to 2^16.
-        pub max_memory32_pages: u64 = 1 << 16,
+        /// May not be larger than `2**32`.
+        ///
+        /// Defaults to `2**32`.
+        pub max_memory32_bytes: u64 = u32::MAX as u64 + 1,
 
-        /// The maximum, in 64k Wasm pages, of any 64-bit memory's initial or
-        /// maximum size.
+        /// The maximum, in bytes, of any 64-bit memory's initial or maximum
+        /// size.
+        ///
+        /// May not be larger than `2**64`.
         ///
-        /// Defaults to 2^48.
-        pub max_memory64_pages: u64 = 1 << 48,
+        /// Defaults to `2**64`.
+        pub max_memory64_bytes: u128 = u64::MAX as u128 + 1,
 
         /// The maximum number of modules to use. Defaults to 10.
         ///
@@ -664,8 +668,8 @@ impl<'a> Arbitrary<'a> for Config {
             max_instructions: u.int_in_range(0..=MAX_MAXIMUM)?,
             max_memories: u.int_in_range(0..=100)?,
             max_tables,
-            max_memory32_pages: u.int_in_range(0..=1 << 16)?,
-            max_memory64_pages: u.int_in_range(0..=1 << 48)?,
+            max_memory32_bytes: u.int_in_range(0..=u32::MAX as u64 + 1)?,
+            max_memory64_bytes: u.int_in_range(0..=u64::MAX as u128 + 1)?,
             min_uleb_size: u.int_in_range(0..=5)?,
             bulk_memory_enabled: u.arbitrary()?,
             reference_types_enabled,
diff --git a/crates/wasm-smith/src/core.rs b/crates/wasm-smith/src/core.rs
index ede06314..e35059ac 100644
--- a/crates/wasm-smith/src/core.rs
+++ b/crates/wasm-smith/src/core.rs
@@ -2558,17 +2558,22 @@ pub(crate) fn arbitrary_memtype(u: &mut Unstructured, config: &Config) -> Result
     // We want to favor memories <= 1gb in size, allocate at most 16k pages,
     // depending on the maximum number of memories.
     let memory64 = config.memory64_enabled && u.arbitrary()?;
-    let page_size = if config.custom_page_sizes_enabled && u.arbitrary()? {
-        Some(1 << u.int_in_range(0..=16)?)
+    let page_size_log2 = if config.custom_page_sizes_enabled && u.arbitrary()? {
+        Some(if u.arbitrary()? { 0 } else { 16 })
     } else {
         None
     };
     let max_inbounds = 16 * 1024 / u64::try_from(config.max_memories).unwrap();
     let min_pages = if config.disallow_traps { Some(1) } else { None };
     let max_pages = min_pages.unwrap_or(0).max(if memory64 {
-        config.max_memory64_pages
+        u64::try_from(config.max_memory64_bytes >> page_size_log2.unwrap_or(16))
+            // Can only fail when we have a custom page size of 1 byte and a
+            // memory size of `2**64 == u64::MAX + 1`. In this case, just
+            // saturate to `u64::MAX`.
+            .unwrap_or(u64::MAX as u64)
     } else {
-        config.max_memory32_pages
+        // Unlike above, this can never fail.
+        u64::try_from(config.max_memory32_bytes >> page_size_log2.unwrap_or(16)).unwrap()
     });
     let (minimum, maximum) = arbitrary_limits64(
         u,
@@ -2582,7 +2587,7 @@ pub(crate) fn arbitrary_memtype(u: &mut Unstructured, config: &Config) -> Result
         maximum,
         memory64,
         shared,
-        page_size_log2: page_size,
+        page_size_log2,
     })
 }
 
diff --git a/fuzz/src/lib.rs b/fuzz/src/lib.rs
index 2cf513f1..32c2f8f6 100644
--- a/fuzz/src/lib.rs
+++ b/fuzz/src/lib.rs
@@ -28,6 +28,7 @@ pub fn generate_valid_module(
     config.exceptions_enabled = u.arbitrary()?;
     config.canonicalize_nans = u.arbitrary()?;
     config.tail_call_enabled = u.arbitrary()?;
+    config.custom_page_sizes_enabled = u.arbitrary()?;
 
     config.gc_enabled = u.arbitrary()?;
     config.reference_types_enabled = config.reference_types_enabled || config.gc_enabled;
@@ -91,6 +92,10 @@ pub fn validator_for_config(config: &Config) -> wasmparser::Validator {
     features.set(WasmFeatures::MEMORY64, config.memory64_enabled);
     features.set(WasmFeatures::THREADS, config.threads_enabled);
     features.set(WasmFeatures::EXCEPTIONS, config.exceptions_enabled);
+    features.set(
+        WasmFeatures::CUSTOM_PAGE_SIZES,
+        config.custom_page_sizes_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
diff --git a/fuzz/src/no_traps.rs b/fuzz/src/no_traps.rs
index 90e13431..782ff0ca 100644
--- a/fuzz/src/no_traps.rs
+++ b/fuzz/src/no_traps.rs
@@ -16,8 +16,8 @@ pub fn run(u: &mut Unstructured<'_>) -> Result<()> {
         config.threads_enabled = false;
         config.exceptions_enabled = false;
         config.gc_enabled = false;
-        config.max_memory32_pages = config.max_memory32_pages.min(100);
-        config.max_memory64_pages = config.max_memory64_pages.min(100);
+        config.max_memory32_bytes = config.max_memory32_bytes.min(1 << 18);
+        config.max_memory64_bytes = config.max_memory64_bytes.min(1 << 18);
 
         // NB: should re-enable once wasmtime implements the table64 extension
         // to the memory64 proposal.


Test patch:
diff --git a/crates/wasm-smith/tests/core.rs b/crates/wasm-smith/tests/core.rs
index 9b18fcde..d7972f29 100644
--- a/crates/wasm-smith/tests/core.rs
+++ b/crates/wasm-smith/tests/core.rs
@@ -142,6 +142,25 @@ fn smoke_test_wasm_gc() {
     }
 }
 
+#[test]
+fn smoke_test_wasm_custom_page_sizes() {
+    let mut rng = SmallRng::seed_from_u64(0);
+    let mut buf = vec![0; 2048];
+    for _ in 0..1024 {
+        rng.fill_bytes(&mut buf);
+        let mut u = Unstructured::new(&buf);
+        let config = Config {
+            custom_page_sizes_enabled: true,
+            ..Config::default()
+        };
+        if let Ok(module) = Module::new(config, &mut u) {
+            let wasm_bytes = module.to_bytes();
+            let mut validator = Validator::new_with_features(wasm_features());
+            validate(&mut validator, &wasm_bytes);
+        }
+    }
+}
+
 fn wasm_features() -> WasmFeatures {
     WasmFeatures::all()
 }