diff --git a/src/tests.rs b/src/tests.rs index 4e773cb..ac97da6 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -452,3 +452,53 @@ fn test_smart_contract_audit() { let wasm_ctx = WASMCtx::new(wasm_args); test_wasm_snark_with(wasm_ctx, step_size).unwrap(); } + +#[test] +fn test_memory_truncation_with_host_call() { + pub struct CustomWasmCtx { + wasm_args: crate::wasm_ctx::WASMArgs, + } + + impl ZKWASMCtx for CustomWasmCtx { + type T = (); + + fn create_store(&self, engine: &wasmi::Engine) -> wasmi::Store { + wasmi::Store::new(engine, ()) + } + + fn create_linker( + &self, + engine: &wasmi::Engine, + _module: &wasmi::Module, + ) -> Result, ZKWASMError> { + let mut linker = wasmi::Linker::::new(engine); + + linker + .func_wrap( + "host", + "call", + |_caller: wasmi::Caller<()>, i: u32| -> u32 { i }, + ) + .unwrap(); + + Ok(linker) + } + + fn args(&self) -> &crate::wasm_ctx::WASMArgs { + &self.wasm_args + } + } + + init_logger(); + let step_size = StepSize::new(1000).set_memory_step_size(1000); + + let wasm_args = WASMArgsBuilder::default() + .file_path(PathBuf::from("wasm/host_call_test.wat")) + .unwrap() + .func_args(vec![]) + .invoke("test") + .build(); + + let wasm_ctx = CustomWasmCtx { wasm_args }; + test_wasm_snark_with(wasm_ctx, step_size).unwrap(); +} diff --git a/third-party/wasmi/crates/wasmi/src/tracer/mod.rs b/third-party/wasmi/crates/wasmi/src/tracer/mod.rs index adf0af8..8c026e1 100644 --- a/third-party/wasmi/crates/wasmi/src/tracer/mod.rs +++ b/third-party/wasmi/crates/wasmi/src/tracer/mod.rs @@ -148,7 +148,11 @@ impl Tracer { // Truncate to the largest 8-byte chunk that can accommodate the used highest address. let new_len = highest_address.map_or(0, |addr| (addr / 7) + 1); + self.IS_mem.truncate(new_len as usize); + + self.execution_trace + .retain(|vm| vm.instr != Instruction::HostCallStep || vm.Y < new_len); } /// Push initial heap/linear WASM memory to tracer for MCC diff --git a/wasm/host_call_test.wat b/wasm/host_call_test.wat new file mode 100644 index 0000000..066e2d2 --- /dev/null +++ b/wasm/host_call_test.wat @@ -0,0 +1,10 @@ + (module + (import "host" "call" (func $f (param i32) (result i32))) + (memory 1 2) + (func $test + (i32.const 42) + (call $f) + (drop) + ) + (export "test" (func $test)) + )