From c51e19ffa701b10c51a0e50b51a1ec07e9c7172a Mon Sep 17 00:00:00 2001 From: wasm-forge <122647775+wasm-forge@users.noreply.github.com> Date: Sat, 7 Oct 2023 23:28:33 +0200 Subject: [PATCH] new tests --- src/main.rs | 36 ++++++++++-------- src/tests.rs | 101 +++++++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 118 insertions(+), 19 deletions(-) diff --git a/src/main.rs b/src/main.rs index abc6fc6..592b992 100644 --- a/src/main.rs +++ b/src/main.rs @@ -253,12 +253,28 @@ fn remove_start_export(module: &mut walrus::Module) { } -fn main() -> anyhow::Result<()> { +fn do_module_replacements(module: &mut walrus::Module) { + // find corresponding IDs for replacements + let fn_replacement_ids = gather_replacement_ids(&module); - let exe_name = std::env::current_exe().unwrap().file_name().unwrap().to_str().unwrap().to_owned(); + // do recursive call replacement + replace_calls(module, &fn_replacement_ids); + + // add start entry (this is needed to do initialization) + add_start_entry(module); + // remove the _start export to clean up the module exports + remove_start_export(module); + + // clean-up unused imports + walrus::passes::gc::run(module); +} + +fn main() -> anyhow::Result<()> { env_logger::init(); + let exe_name = std::env::current_exe().unwrap().file_name().unwrap().to_str().unwrap().to_owned(); + let input_wasm = std::env::args() .nth(1) .ok_or_else(|| anyhow::anyhow!("The launch parameters are incorrect, try: {} [output.wasm]", {exe_name}))?; @@ -266,20 +282,8 @@ fn main() -> anyhow::Result<()> { // load Wasm module from file let mut module = walrus::Module::from_file(&input_wasm)?; - // find corresponding IDs for replacements - let fn_replacement_ids = gather_replacement_ids(&module); - - // do recursive call replacement - replace_calls(&mut module, &fn_replacement_ids); - - // add start entry (this is needed to do initialization) - add_start_entry(&mut module); - - // remove the _start export to clean up the module exports - remove_start_export(&mut module); - - // clean-up unused imports - walrus::passes::gc::run(&mut module); + // do the substitution workflow + do_module_replacements(&mut module); // try store as binary let wasm = module.emit_wasm(); diff --git a/src/tests.rs b/src/tests.rs index b466bd8..1093428 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -1,7 +1,7 @@ use crate::*; -#[warn(dead_code)] +#[allow(dead_code)] fn print_wasm(module: &mut walrus::Module) { let wasm = module.emit_wasm(); let text = wasmprinter::print_bytes(wasm).unwrap(); @@ -113,12 +113,11 @@ fn test_remove_start_export() { } - #[test] fn test_gather_replacement_ids() { let wat = r#" - (module + (module (type (;0;) (func)) (type (;1;) (func (param i32))) (type (;2;) (func (param i32 i32))) @@ -170,3 +169,99 @@ fn test_gather_replacement_ids() { assert!(id_reps[&3] == 7); } + +#[test] +fn test_do_module_replacements() { + + let wat = r#" + (module + (type (;0;) (func)) + (type (;1;) (func (param i32))) + (type (;2;) (func (param i32 i32))) + (type (;3;) (func (param i32 i32) (result i32))) + (type (;4;) (func (param i32 i32 i32))) + (type (;5;) (func (param i32 i32 i32 i32) (result i32))) + (import "ic0" "debug_print" (func $_dprint (;0;) (type 2))) + (import "ic0" "msg_reply" (func $_msg_reply (;1;) (type 0))) + (import "wasi_snapshot_preview1" "fd_write" (func $_wasi_snapshot_preview_fd_write (;2;) (type 5))) + (import "wasi_snapshot_preview1" "random_get" (func $_wasi_snapshot_preview_random_get (;3;) (type 3))) + (import "wasi_snapshot_preview1" "environ_get" (func $__imported_wasi_snapshot_preview1_environ_get (;4;) (type 3))) + (import "wasi_snapshot_preview1" "proc_exit" (func $__imported_wasi_snapshot_preview1_proc_exit (;5;) (type 1))) + + (func $_start (;6;) (type 0) + i32.const 1 + i32.const 2 + call $_wasi_snapshot_preview_random_get + + (block $test_block + i32.const 0 + (if + (then + (block $test_block3 + i32.const 1 + i32.const 2 + + call $_wasi_snapshot_preview_random_get + + drop + ) + ) + (else + (loop $test_loop + i32.const 1 + i32.const 2 + + call $_wasi_snapshot_preview_random_get + + br_if $test_loop + ) + ) + ) + ) + + i32.const 3 + i32.const 4 + i32.const 5 + call $_wasi_snapshot_preview_fd_write + + call $__imported_wasi_snapshot_preview1_proc_exit + ) + + (func $__ic_custom_random_get (;7;) (type 3) (param i32 i32) (result i32) + call $_msg_reply + i32.const 421 + ) + + (func $__ic_custom_fd_write (;8;) (type 5) (param i32 i32 i32 i32) (result i32) + i32.const 0 + i32.const 0 + call $_dprint + i32.const 42 + ) + + (export "_start" (func $_start)) + ) + "#; + + let binary = wat::parse_str(wat).unwrap(); + let mut module = walrus::Module::from_buffer(&binary).unwrap(); + + do_module_replacements(&mut module); + + // we expect random_get and fd_write to be replaced, environ_get to be removed and the calls to the proc_exit to remain + let imports = module.imports; + + let result = imports.find("ic0", "debug_print"); + assert!(result.is_some()); + let result = imports.find("ic0", "msg_reply"); + assert!(result.is_some()); + let result = imports.find("wasi_snapshot_preview1", "proc_exit"); + assert!(result.is_some()); + let result = imports.find("wasi_snapshot_preview1", "fd_write"); + assert!(None == result); + let result = imports.find("wasi_snapshot_preview1", "random_get"); + assert!(None == result); + let result = imports.find("wasi_snapshot_preview1", "environ_get"); + assert!(None == result); + +}