Zero the memory of the Wasm VM

src/executor/host.rs

@@ -325,6 +325,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)) => {

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<VirtualMachine, (StartErr, Self)> {
         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);

src/executor/vm/interpreter.rs

@@ -25,6 +25,7 @@ use super::{
 use alloc::{borrow::ToOwned as _, boxed::Box, format, string::ToString as _, sync::Arc, vec::Vec};
 use core::{
     cell::RefCell,
+    cmp,
     convert::{TryFrom, TryInto as _},
     fmt,
 };
@@ -64,6 +65,10 @@ pub struct InterpreterPrototype {
     /// use any memory.
     memory: Option<wasmi::MemoryRef>,
 
+    /// 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
@@ -228,6 +233,7 @@ impl InterpreterPrototype {
         Ok(InterpreterPrototype {
             module,
             memory,
+            memory_zeroed: true,
             indirect_table,
         })
     }
@@ -256,6 +262,7 @@ impl InterpreterPrototype {
         self,
         function_name: &str,
         params: &[WasmValue],
+        zero_memory_above: u32,
     ) -> Result<Interpreter, (StartErr, Self)> {
         let execution = match self.module.export_by_name(function_name) {
             Some(wasmi::ExternVal::Func(f)) => {
@@ -280,6 +287,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,
@@ -540,11 +556,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,
         }
     }

src/executor/vm/jit.rs

@@ -77,6 +77,10 @@ pub struct JitPrototype {
     /// Reference to the memory used by the module, if any.
     memory: Option<wasmtime::Memory>,
 
+    /// 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<wasmtime::Table>,
@@ -245,6 +249,7 @@ impl JitPrototype {
             instance,
             shared,
             memory,
+            memory_zeroed: true,
             indirect_table,
         })
     }
@@ -261,7 +266,12 @@ impl JitPrototype {
     }
 
     /// See [`super::VirtualMachinePrototype::start`].
-    pub fn start(self, function_name: &str, params: &[WasmValue]) -> Result<Jit, (StartErr, Self)> {
+    pub fn start(
+        self,
+        function_name: &str,
+        params: &[WasmValue],
+        zero_memory_above: u32,
+    ) -> Result<Jit, (StartErr, Self)> {
         // Try to start executing `_start`.
         let start_function = match self.instance.get_export(function_name) {
             Some(export) => match export.into_func() {
@@ -292,6 +302,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,
@@ -500,23 +527,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,
         }
     }