diff --git a/src/executor/host.rs b/src/executor/host.rs index c9334b140e..f9adb9037e 100644 --- a/src/executor/host.rs +++ b/src/executor/host.rs @@ -324,6 +324,7 @@ impl HostVmPrototype { vm::WasmValue::I32(i32::from_ne_bytes(self.heap_base.to_ne_bytes())), vm::WasmValue::I32(i32::from_ne_bytes(data_len_u32.to_ne_bytes())), ], + self.heap_base.saturating_add(data_len_u32), ) { Ok(vm) => vm, Err((error, vm_proto)) => { diff --git a/src/executor/vm.rs b/src/executor/vm.rs index c973b7404e..54b3a0de1b 100644 --- a/src/executor/vm.rs +++ b/src/executor/vm.rs @@ -183,16 +183,20 @@ impl VirtualMachinePrototype { /// Turns this prototype into an actual virtual machine. This requires choosing which function /// to execute. + /// + /// The function will zero the memory starting at offset `zero_memory_above` and above. Pass + /// `0` in order to zero the entire memory, or `u32::max_value()` to zero nothing. pub fn start( mut self, function_name: &str, params: &[WasmValue], + zero_memory_above: u32, ) -> Result { Ok(VirtualMachine { inner: match self.inner { #[cfg(all(target_arch = "x86_64", feature = "std"))] VirtualMachinePrototypeInner::Jit(inner) => { - match inner.start(function_name, params) { + match inner.start(function_name, params, zero_memory_above) { Ok(vm) => VirtualMachineInner::Jit(vm), Err((err, proto)) => { self.inner = VirtualMachinePrototypeInner::Jit(proto); @@ -201,7 +205,7 @@ impl VirtualMachinePrototype { } } VirtualMachinePrototypeInner::Interpreter(inner) => { - match inner.start(function_name, params) { + match inner.start(function_name, params, zero_memory_above) { Ok(vm) => VirtualMachineInner::Interpreter(vm), Err((err, proto)) => { self.inner = VirtualMachinePrototypeInner::Interpreter(proto); diff --git a/src/executor/vm/interpreter.rs b/src/executor/vm/interpreter.rs index a080133d04..cca8b6a545 100644 --- a/src/executor/vm/interpreter.rs +++ b/src/executor/vm/interpreter.rs @@ -23,7 +23,12 @@ use super::{ }; use alloc::{borrow::ToOwned as _, boxed::Box, format, string::ToString as _, sync::Arc, vec::Vec}; -use core::{cell::RefCell, fmt}; +use core::{ + cell::RefCell, + cmp, + convert::{TryFrom, TryInto as _}, + fmt, +}; use wasmi::memory_units::ByteSize as _; /// See [`super::Module`]. @@ -60,6 +65,10 @@ pub struct InterpreterPrototype { /// use any memory. memory: Option, + /// If `true`, the memory in [`InterpreterPrototype::memory`] is zero-ed. Any additional + /// zero-ing can be skipped as an optimization. + memory_zeroed: bool, + /// Table of the indirect function calls. /// /// In Wasm, function pointers are in reality indices in a table called @@ -224,6 +233,7 @@ impl InterpreterPrototype { Ok(InterpreterPrototype { module, memory, + memory_zeroed: true, indirect_table, }) } @@ -252,6 +262,7 @@ impl InterpreterPrototype { self, function_name: &str, params: &[WasmValue], + zero_memory_above: u32, ) -> Result { let execution = match self.module.export_by_name(function_name) { Some(wasmi::ExternVal::Func(f)) => { @@ -274,6 +285,15 @@ impl InterpreterPrototype { _ => return Err((StartErr::NotAFunction, self)), }; + // Zero the memory if necessary. + if !self.memory_zeroed { + if let Some(memory) = self.memory.as_ref() { + let size = memory.current_size().0 * wasmi::memory_units::Pages::byte_size().0; + let offset = cmp::min(size, usize::try_from(zero_memory_above).unwrap()); + memory.zero(offset, size - offset).unwrap(); // Can error only if out of bounds. + } + } + Ok(Interpreter { _module: self.module, memory: self.memory, @@ -527,11 +547,10 @@ impl Interpreter { /// See [`super::VirtualMachine::into_prototype`]. pub fn into_prototype(self) -> InterpreterPrototype { - // TODO: zero the memory - InterpreterPrototype { module: self._module, memory: self.memory, + memory_zeroed: false, indirect_table: self.indirect_table, } } diff --git a/src/executor/vm/jit.rs b/src/executor/vm/jit.rs index cbbed7616b..f6f7e41027 100644 --- a/src/executor/vm/jit.rs +++ b/src/executor/vm/jit.rs @@ -74,6 +74,10 @@ pub struct JitPrototype { /// Reference to the memory used by the module, if any. memory: Option, + /// If `true`, the memory in [`JitPrototype::memory`] is zero-ed. Any additional zero-ing can + /// be skipped as an optimization. + memory_zeroed: bool, + /// Reference to the table of indirect functions, in case we need to access it. /// `None` if the module doesn't export such table. indirect_table: Option, @@ -254,6 +258,7 @@ impl JitPrototype { instance, shared, memory, + memory_zeroed: true, indirect_table, }) } @@ -270,7 +275,12 @@ impl JitPrototype { } /// See [`super::VirtualMachinePrototype::start`]. - pub fn start(self, function_name: &str, params: &[WasmValue]) -> Result { + pub fn start( + self, + function_name: &str, + params: &[WasmValue], + zero_memory_above: u32, + ) -> Result { // Try to start executing `_start`. let start_function = match self.instance.get_export(function_name) { Some(export) => match export.into_func() { @@ -301,6 +311,23 @@ impl JitPrototype { Ok(result.get(0).map(|v| TryFrom::try_from(v).unwrap())) }); + // Zero the memory if necessary. + if !self.memory_zeroed { + if let Some(memory) = self.memory.as_ref() { + let offset = cmp::min( + memory.data_size(), + usize::try_from(zero_memory_above).unwrap(), + ); + + // Soundness: the documentation of wasmtime precisely explains what is safe or not. + // Basically, we are safe as long as we are sure that we don't potentially grow the + // buffer (which would invalidate the buffer pointer). + unsafe { + memory.data_unchecked_mut()[offset..].fill(0); + } + } + } + Ok(Jit { function_call: Some(function_call), instance: self.instance, @@ -509,23 +536,11 @@ impl Jit { pub fn into_prototype(self) -> JitPrototype { // TODO: how do we handle if the coroutine was within a host function? - // TODO: necessary? - /*// Zero-ing the memory. - if let Some(memory) = &self.memory { - // Soundness: the documentation of wasmtime precisely explains what is safe or not. - // Basically, we are safe as long as we are sure that we don't potentially grow the - // buffer (which would invalidate the buffer pointer). - unsafe { - for byte in memory.data_unchecked_mut() { - *byte = 0; - } - } - }*/ - JitPrototype { instance: self.instance, shared: self.shared, memory: self.memory, + memory_zeroed: false, indirect_table: self.indirect_table, } }