From d62d7621c67a1cd910c3a727f7c36b886a72098d Mon Sep 17 00:00:00 2001 From: Pierre Krieger Date: Wed, 16 Jun 2021 10:13:00 +0200 Subject: [PATCH 1/2] Zero the memory of the Wasm VM --- src/executor/host.rs | 1 + src/executor/vm.rs | 8 ++++-- src/executor/vm/interpreter.rs | 19 ++++++++++++-- src/executor/vm/jit.rs | 45 +++++++++++++++++++++++----------- 4 files changed, 55 insertions(+), 18 deletions(-) diff --git a/src/executor/host.rs b/src/executor/host.rs index e567b5ad93..b4129bb599 100644 --- a/src/executor/host.rs +++ b/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)) => { diff --git a/src/executor/vm.rs b/src/executor/vm.rs index 78318e0696..dbfbf15ef0 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 45c29342c8..0daadf553c 100644 --- a/src/executor/vm/interpreter.rs +++ b/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, + /// 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 { 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, } } diff --git a/src/executor/vm/jit.rs b/src/executor/vm/jit.rs index 86bb6f8f9e..4b8ec35a5b 100644 --- a/src/executor/vm/jit.rs +++ b/src/executor/vm/jit.rs @@ -72,6 +72,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, @@ -240,6 +244,7 @@ impl JitPrototype { instance, shared, memory, + memory_zeroed: true, indirect_table, }) } @@ -256,7 +261,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() { @@ -287,6 +297,25 @@ 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 { + for byte in &mut memory.data_unchecked_mut()[offset..] { + *byte = 0; + } + } + } + } + Ok(Jit { function_call: Some(function_call), instance: self.instance, @@ -495,23 +524,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, } } From 4f0e918f1efb228b146c5c6a1f519f70a3761157 Mon Sep 17 00:00:00 2001 From: Pierre Krieger Date: Wed, 16 Jun 2021 10:18:52 +0200 Subject: [PATCH 2/2] Use `fill` instead of manually writing zeroes --- src/executor/vm/jit.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/executor/vm/jit.rs b/src/executor/vm/jit.rs index 4b8ec35a5b..022c6e07b4 100644 --- a/src/executor/vm/jit.rs +++ b/src/executor/vm/jit.rs @@ -309,9 +309,7 @@ impl JitPrototype { // 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 &mut memory.data_unchecked_mut()[offset..] { - *byte = 0; - } + memory.data_unchecked_mut()[offset..].fill(0); } } }