feat(instrument): Add --stub-wasi flag to replace WASI imports with stubs (#104)

* feat(instrument): Add --stub-wasi flag to replace WASI imports with stubs

Problem:
When using `ic-wasm instrument` on a canister WASM compiled with Emscripten,
the resulting instrumented WASM contains WASI imports (fd_close, fd_write,
fd_seek, etc.) that IC rejects at install time:

  Error: Wasm module has an invalid import section.
  Module imports function 'fd_close' from 'wasi_snapshot_preview1'
  that is not exported by the runtime.

Solution:
Add `--stub-wasi` flag that replaces WASI imports with local stub functions
returning 0 (success) or trapping for proc_exit. This allows Emscripten-
generated WASMs to be instrumented and deployed to IC.

Usage:
  ic-wasm canister.wasm -o out.wasm instrument --stub-wasi

* docs(instrument): Reorganize documentation to clarify execution tracing vs WASI stubbing

- Restructure Instrument section intro to list both capabilities
- Add subsection headers: "Execution tracing" and "Stubbing WASI imports"
- Move WASI stubbing docs to its own subsection after execution tracing
- Update CHANGELOG with --stub-wasi flag addition

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

---------

Co-authored-by: Linwei Shang <linwei.shang@dfinity.org>
Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: 3 people

CHANGELOG.md

@@ -6,6 +6,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [unreleased]
 
+* Add `--stub-wasi` flag to `instrument` subcommand to replace WASI imports with stub functions.
+
 ## [0.9.9] - 2025-11-18
 
 * Add support for comments in optional hidden endpoint file for `check-endpoints` command.

README.md

@@ -126,9 +126,13 @@ Usage: `ic-wasm <input.wasm> check-endpoints [--candid <file>] [--hidden <file>]
 
 ### Instrument (experimental)
 
-Instrument canister method to emit execution trace to stable memory.
+Provides instrumentation capabilities for canister WebAssembly modules:
+- **Execution tracing**: Instrument canister methods to emit execution trace to stable memory for performance profiling
+- **WASI compatibility**: Replace WASI imports with stub functions to enable modules compiled with Emscripten or wasi-sdk to run on the Internet Computer
 
-Usage: `ic-wasm <input.wasm> -o <output.wasm> instrument --trace-only func1 --trace-only func2 --start-page 16 --page-limit 30`
+Usage: `ic-wasm <input.wasm> -o <output.wasm> instrument [--trace-only func1] [--start-page 16] [--page-limit 30] [--stub-wasi]`
+
+#### Execution tracing
 
 Instrumented canister has the following additional endpoints:
 
@@ -203,6 +207,30 @@ fn post_upgrade() {
 * We cannot measure query calls.
 * No concurrent calls.
 
+#### Stubbing WASI imports
+
+The `--stub-wasi` flag replaces WASI imports with local stub functions, enabling WASM modules compiled with Emscripten or wasi-sdk to run on the Internet Computer. Without this flag, such modules would fail at install time with errors like:
+
+```
+Error: Wasm module has an invalid import section.
+Module imports function 'fd_close' from 'wasi_snapshot_preview1' that is not exported by the runtime.
+```
+
+The stub functions behave as follows:
+
+| WASI Function | Stub Behavior |
+|---------------|---------------|
+| `fd_close` | Returns 0 (success) |
+| `fd_write` | Writes 0 to nwritten, returns 0 |
+| `fd_read` | Writes 0 to nread, returns 0 |
+| `fd_seek` | Writes 0 to newoffset, returns 0 |
+| `environ_sizes_get` | Writes 0 to both params, returns 0 |
+| `environ_get` | Returns 0 |
+| `proc_exit` | Traps (unreachable) |
+| Others | Returns 0 |
+
+**Note**: This is a workaround for edge cases. The recommended approach is to build without WASI imports (e.g., using `wasm32-unknown-unknown` target with ic-cdk for Rust). Stub functions return success, which may hide real failures if your code depends on WASI functionality.
+
 ## Library
 
 To use `ic-wasm` as a library, add this to your `Cargo.toml`:

src/bin/main.rs

@@ -101,6 +101,9 @@ enum SubCommand {
         /// The number of pages of the preallocated stable memory
         #[clap(short, long, requires("start_page"))]
         page_limit: Option<i32>,
+        /// Replace WASI imports with stub functions that return 0 (success)
+        #[clap(long)]
+        stub_wasi: bool,
     },
     /// Check canister endpoints against provided Candid interface
     #[cfg(feature = "check-endpoints")]
@@ -164,12 +167,14 @@ fn main() -> anyhow::Result<()> {
             trace_only,
             start_page,
             page_limit,
+            stub_wasi,
         } => {
             use ic_wasm::instrumentation::{instrument, Config};
             let config = Config {
                 trace_only_funcs: trace_only.clone().unwrap_or(vec![]),
                 start_address: start_page.map(|page| i64::from(page) * 65536),
                 page_limit: *page_limit,
+                stub_wasi: *stub_wasi,
             };
             instrument(&mut m, config).map_err(|e| anyhow::anyhow!("{e}"))?;
         }

src/instrumentation.rs

@@ -38,6 +38,7 @@ pub struct Config {
     pub trace_only_funcs: Vec<String>,
     pub start_address: Option<i64>,
     pub page_limit: Option<i32>,
+    pub stub_wasi: bool,
 }
 impl Config {
     pub fn is_preallocated(&self) -> bool {
@@ -61,6 +62,9 @@ impl Config {
 /// When trace_only_funcs is not empty, counting and tracing is only enabled for those listed functions per update call.
 /// TODO: doesn't handle recursive entry functions. Need to create a wrapper for the recursive entry function.
 pub fn instrument(m: &mut Module, config: Config) -> Result<(), String> {
+    if config.stub_wasi {
+        stub_wasi_imports(m);
+    }
     let mut trace_only_ids = HashSet::new();
     for name in config.trace_only_funcs.iter() {
         let id = match m.funcs.by_name(name) {
@@ -960,3 +964,223 @@ fn make_toggle_func(m: &mut Module, name: &str, var: GlobalId) {
     let id = builder.finish(vec![], &mut m.funcs);
     m.exports.add(&format!("canister_update {name}"), id);
 }
+
+/// Replace WASI imports with stub functions that return 0 (success) or trap for proc_exit
+fn stub_wasi_imports(m: &mut Module) {
+    use walrus::FunctionBuilder;
+
+    // Find all WASI imports
+    let wasi_imports: Vec<_> = m
+        .imports
+        .iter()
+        .filter(|i| i.module == "wasi_snapshot_preview1")
+        .filter_map(|i| {
+            if let ImportKind::Function(func_id) = i.kind {
+                Some((i.id(), i.name.clone(), func_id))
+            } else {
+                None
+            }
+        })
+        .collect();
+
+    let memory = m.memories.iter().next().map(|mem| mem.id());
+
+    for (import_id, name, old_func_id) in wasi_imports {
+        // Get the function type
+        let func = m.funcs.get(old_func_id);
+        let ty_id = func.ty();
+        let ty = m.types.get(ty_id);
+        let params: Vec<_> = ty.params().to_vec();
+        let results: Vec<_> = ty.results().to_vec();
+
+        // Create stub function
+        let mut builder = FunctionBuilder::new(&mut m.types, &params, &results);
+        builder.name(format!("__wasi_{name}_stub"));
+
+        // Create locals for parameters
+        let param_locals: Vec<_> = params.iter().map(|t| m.locals.add(*t)).collect();
+
+        match name.as_str() {
+            "fd_write" => {
+                // fd_write(fd: i32, iovs: i32, iovs_len: i32, nwritten: i32) -> i32
+                // Write 0 to nwritten and return 0
+                if let Some(mem) = memory {
+                    if param_locals.len() >= 4 {
+                        builder
+                            .func_body()
+                            .local_get(param_locals[3]) // nwritten ptr
+                            .i32_const(0)
+                            .store(
+                                mem,
+                                StoreKind::I32 { atomic: false },
+                                MemArg {
+                                    offset: 0,
+                                    align: 4,
+                                },
+                            )
+                            .i32_const(0);
+                    } else {
+                        builder.func_body().i32_const(0);
+                    }
+                } else {
+                    builder.func_body().i32_const(0);
+                }
+            }
+            "fd_read" => {
+                // fd_read(fd: i32, iovs: i32, iovs_len: i32, nread: i32) -> i32
+                // Write 0 to nread and return 0
+                if let Some(mem) = memory {
+                    if param_locals.len() >= 4 {
+                        builder
+                            .func_body()
+                            .local_get(param_locals[3]) // nread ptr
+                            .i32_const(0)
+                            .store(
+                                mem,
+                                StoreKind::I32 { atomic: false },
+                                MemArg {
+                                    offset: 0,
+                                    align: 4,
+                                },
+                            )
+                            .i32_const(0);
+                    } else {
+                        builder.func_body().i32_const(0);
+                    }
+                } else {
+                    builder.func_body().i32_const(0);
+                }
+            }
+            "fd_seek" => {
+                // fd_seek(fd: i32, offset: i64, whence: i32, newoffset: i32) -> i32
+                // Write 0 to newoffset and return 0
+                if let Some(mem) = memory {
+                    if param_locals.len() >= 4 {
+                        builder
+                            .func_body()
+                            .local_get(param_locals[3]) // newoffset ptr
+                            .i64_const(0)
+                            .store(
+                                mem,
+                                StoreKind::I64 { atomic: false },
+                                MemArg {
+                                    offset: 0,
+                                    align: 8,
+                                },
+                            )
+                            .i32_const(0);
+                    } else {
+                        builder.func_body().i32_const(0);
+                    }
+                } else {
+                    builder.func_body().i32_const(0);
+                }
+            }
+            "fd_close" => {
+                // fd_close(fd: i32) -> i32
+                // Just return 0 (success)
+                builder.func_body().i32_const(0);
+            }
+            "environ_sizes_get" => {
+                // environ_sizes_get(count: i32, buf_size: i32) -> i32
+                // Write 0 to both pointers and return 0
+                if let Some(mem) = memory {
+                    if param_locals.len() >= 2 {
+                        builder
+                            .func_body()
+                            .local_get(param_locals[0]) // count ptr
+                            .i32_const(0)
+                            .store(
+                                mem,
+                                StoreKind::I32 { atomic: false },
+                                MemArg {
+                                    offset: 0,
+                                    align: 4,
+                                },
+                            )
+                            .local_get(param_locals[1]) // buf_size ptr
+                            .i32_const(0)
+                            .store(
+                                mem,
+                                StoreKind::I32 { atomic: false },
+                                MemArg {
+                                    offset: 0,
+                                    align: 4,
+                                },
+                            )
+                            .i32_const(0);
+                    } else {
+                        builder.func_body().i32_const(0);
+                    }
+                } else {
+                    builder.func_body().i32_const(0);
+                }
+            }
+            "environ_get" => {
+                // environ_get(environ: i32, environ_buf: i32) -> i32
+                // Just return 0 (no environment variables)
+                builder.func_body().i32_const(0);
+            }
+            "proc_exit" => {
+                // proc_exit(code: i32) -> !
+                // Trap unconditionally
+                builder.func_body().unreachable();
+            }
+            _ => {
+                // Default: just return 0 for i32 result, or appropriate zero values
+                for result in &results {
+                    match result {
+                        ValType::I32 => {
+                            builder.func_body().i32_const(0);
+                        }
+                        ValType::I64 => {
+                            builder.func_body().i64_const(0);
+                        }
+                        ValType::F32 => {
+                            builder.func_body().f32_const(0.0);
+                        }
+                        ValType::F64 => {
+                            builder.func_body().f64_const(0.0);
+                        }
+                        _ => {}
+                    }
+                }
+            }
+        }
+
+        let stub_func_id = builder.finish(param_locals, &mut m.funcs);
+
+        // Replace all calls to old_func_id with stub_func_id
+        for (_, func) in m.funcs.iter_local_mut() {
+            replace_calls_in_func(func, old_func_id, stub_func_id);
+        }
+
+        // Remove the import
+        m.imports.delete(import_id);
+    }
+}
+
+fn replace_calls_in_func(func: &mut LocalFunction, old_id: FunctionId, new_id: FunctionId) {
+    let mut stack = vec![func.entry_block()];
+    while let Some(seq_id) = stack.pop() {
+        let mut builder = func.builder_mut().instr_seq(seq_id);
+        for (instr, _) in builder.instrs_mut().iter_mut() {
+            match instr {
+                Instr::Call(Call { func }) if *func == old_id => {
+                    *func = new_id;
+                }
+                Instr::Block(Block { seq }) | Instr::Loop(Loop { seq }) => {
+                    stack.push(*seq);
+                }
+                Instr::IfElse(IfElse {
+                    consequent,
+                    alternative,
+                }) => {
+                    stack.push(*consequent);
+                    stack.push(*alternative);
+                }
+                _ => {}
+            }
+        }
+    }
+}

tests/tests.rs

@@ -491,3 +491,47 @@ fn create_tempfile(content: &str) -> NamedTempFile {
     write!(temp_file, "{content}").expect("Failed to write temp file content");
     temp_file
 }
+
+#[test]
+fn stub_wasi() {
+    let path = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests");
+    let out_path = path.join("out-wasi.wasm");
+
+    // First verify input has WASI imports
+    let input_module = walrus::Module::from_file(path.join("wasi-test.wasm")).unwrap();
+    let input_wasi_imports: Vec<_> = input_module
+        .imports
+        .iter()
+        .filter(|i| i.module == "wasi_snapshot_preview1")
+        .collect();
+    assert!(
+        !input_wasi_imports.is_empty(),
+        "Input should have WASI imports"
+    );
+
+    // Test that --stub-wasi removes WASI imports
+    wasm_input("wasi-test.wasm", false)
+        .arg("-o")
+        .arg(&out_path)
+        .arg("instrument")
+        .arg("--stub-wasi")
+        .assert()
+        .success();
+
+    // Verify the output WASM has no WASI imports
+    let module = walrus::Module::from_file(&out_path).unwrap();
+    let wasi_imports: Vec<_> = module
+        .imports
+        .iter()
+        .filter(|i| i.module == "wasi_snapshot_preview1")
+        .collect();
+
+    assert!(
+        wasi_imports.is_empty(),
+        "WASI imports should be removed, but found: {:?}",
+        wasi_imports.iter().map(|i| &i.name).collect::<Vec<_>>()
+    );
+
+    // Clean up
+    fs::remove_file(&out_path).ok();
+}

tests/wasi-test.wasm

@@ -0,0 +1,21 @@
+(module
+  ;; WASI imports
+  (import "wasi_snapshot_preview1" "fd_close" (func $fd_close (param i32) (result i32)))
+  (import "wasi_snapshot_preview1" "fd_write" (func $fd_write (param i32 i32 i32 i32) (result i32)))
+  (import "wasi_snapshot_preview1" "fd_seek" (func $fd_seek (param i32 i64 i32 i32) (result i32)))
+  
+  ;; IC imports
+  (import "ic0" "msg_reply" (func $msg_reply))
+  (import "ic0" "msg_reply_data_append" (func $msg_reply_data_append (param i32 i32)))
+  
+  (memory 1)
+  
+  (func $test (export "canister_query test")
+    ;; Call WASI functions (they will fail at runtime without stubs)
+    i32.const 1
+    call $fd_close
+    drop
+    
+    call $msg_reply
+  )
+)