From 10212cf0c2db1cc44fc0a19ac7eefd9844ad079c Mon Sep 17 00:00:00 2001 From: Jay Geng Date: Fri, 26 May 2023 17:38:31 -0400 Subject: [PATCH] Add memory fuel metering Add memory fuel metering --- crates/wasmi/src/engine/executor.rs | 14 +++++++- crates/wasmi/src/func/caller.rs | 20 +++++++++++ crates/wasmi/src/store.rs | 35 +++++++++++++++++++ .../tests/e2e/v1/fuel_consumption_mode.rs | 10 +++--- 4 files changed, 74 insertions(+), 5 deletions(-) diff --git a/crates/wasmi/src/engine/executor.rs b/crates/wasmi/src/engine/executor.rs index 64df6d1e5d..d4006345f5 100644 --- a/crates/wasmi/src/engine/executor.rs +++ b/crates/wasmi/src/engine/executor.rs @@ -727,6 +727,17 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> { } } + /// Consume an amount of memory fuel specified by `delta`. + #[inline(always)] + fn consume_mem_fuel(&mut self, bytes: u64) -> Result + { + if self.ctx.engine().config().get_consume_fuel() { + self.ctx.mem_fuel_mut().consume_fuel(bytes) + } else { + Ok(0) + } + } + /// Consume an amount of fuel specified by `delta` if `exec` succeeds. /// /// Prior to executing `exec` it is checked if enough fuel is remaining @@ -1087,9 +1098,10 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> { return self.try_next_instr(); } }; + let delta_in_bytes = delta.to_bytes().unwrap_or(0) as u64; + self.consume_mem_fuel(delta_in_bytes)?; let result = self.consume_fuel_with( |costs| { - let delta_in_bytes = delta.to_bytes().unwrap_or(0) as u64; costs.fuel_for_bytes(delta_in_bytes) }, |this| { diff --git a/crates/wasmi/src/func/caller.rs b/crates/wasmi/src/func/caller.rs index acf7cef0ed..6db1ff89bf 100644 --- a/crates/wasmi/src/func/caller.rs +++ b/crates/wasmi/src/func/caller.rs @@ -100,6 +100,26 @@ impl<'a, T> Caller<'a, T> { pub fn reset_fuel(&mut self) -> Result<(), FuelError> { self.ctx.store.reset_fuel() } + + pub fn add_mem_fuel(&mut self, delta: u64) -> Result<(), FuelError> { + self.ctx.store.add_mem_fuel(delta) + } + + pub fn mem_fuel_consumed(&self) -> Option { + self.ctx.store.mem_fuel_consumed() + } + + pub fn mem_fuel_total(&self) -> Option { + self.ctx.store.mem_fuel_total() + } + + pub fn consume_mem_fuel(&mut self, delta: u64) -> Result { + self.ctx.store.consume_mem_fuel(delta) + } + + pub fn reset_mem_fuel(&mut self) -> Result<(), FuelError> { + self.ctx.store.reset_mem_fuel() + } } impl AsContext for Caller<'_, T> { diff --git a/crates/wasmi/src/store.rs b/crates/wasmi/src/store.rs index 831ccf69cb..70e22706db 100644 --- a/crates/wasmi/src/store.rs +++ b/crates/wasmi/src/store.rs @@ -114,6 +114,7 @@ pub struct StoreInner { engine: Engine, /// The fuel of the [`Store`]. fuel: Fuel, + mem_fuel: Fuel, } #[test] @@ -242,6 +243,7 @@ impl StoreInner { elems: Arena::new(), extern_objects: Arena::new(), fuel: Fuel::default(), + mem_fuel: Fuel::default(), } } @@ -260,6 +262,10 @@ impl StoreInner { &mut self.fuel } + pub fn mem_fuel_mut(&mut self) -> &mut Fuel { + &mut self.mem_fuel + } + /// Wraps an entitiy `Idx` (index type) as a [`Stored`] type. /// /// # Note @@ -796,6 +802,35 @@ impl Store { Ok(self.inner.fuel.reset_fuel()) } + pub fn add_mem_fuel(&mut self, delta: u64) -> Result<(), FuelError> { + self.check_fuel_metering_enabled()?; + self.inner.mem_fuel.add_fuel(delta); + Ok(()) + } + + pub fn mem_fuel_consumed(&self) -> Option { + self.check_fuel_metering_enabled().ok()?; + Some(self.inner.mem_fuel.fuel_consumed()) + } + + pub fn mem_fuel_total(&self) -> Option { + self.check_fuel_metering_enabled().ok()?; + Some(self.inner.mem_fuel.total) + } + + pub fn consume_mem_fuel(&mut self, delta: u64) -> Result { + self.check_fuel_metering_enabled()?; + self.inner + .mem_fuel + .consume_fuel(delta) + .map_err(|_error| FuelError::out_of_fuel()) + } + + pub fn reset_mem_fuel(&mut self) -> Result<(), FuelError> { + self.check_fuel_metering_enabled()?; + Ok(self.inner.mem_fuel.reset_fuel()) + } + /// Allocates a new [`TrampolineEntity`] and returns a [`Trampoline`] reference to it. pub(super) fn alloc_trampoline(&mut self, func: TrampolineEntity) -> Trampoline { let idx = self.trampolines.alloc(func); diff --git a/crates/wasmi/tests/e2e/v1/fuel_consumption_mode.rs b/crates/wasmi/tests/e2e/v1/fuel_consumption_mode.rs index 4736b941da..b6eededdc7 100644 --- a/crates/wasmi/tests/e2e/v1/fuel_consumption_mode.rs +++ b/crates/wasmi/tests/e2e/v1/fuel_consumption_mode.rs @@ -71,23 +71,25 @@ fn test_module() -> &'static str { )"# } -fn check_consumption_mode(mode: FuelConsumptionMode, given_fuel: u64, consumed_fuel: u64) { - assert!(given_fuel >= consumed_fuel); +fn check_consumption_mode(mode: FuelConsumptionMode, given_fuel: u64, consumed_fuel: u64, given_mem_fuel: u64, consumed_mem_fuel: u64) { + assert!(given_fuel >= consumed_fuel && given_mem_fuel >= consumed_mem_fuel); let wasm = wat2wasm(test_module()); let (mut store, func) = default_test_setup(mode, &wasm); let func = func.typed::<(), i32>(&store).unwrap(); // Now add enough fuel, so execution should succeed. store.add_fuel(given_fuel).unwrap(); // this is just enough fuel for a successful `memory.grow` + store.add_mem_fuel(given_mem_fuel).unwrap(); assert_success(func.call(&mut store, ())); assert_eq!(store.fuel_consumed(), Some(consumed_fuel)); + assert_eq!(store.mem_fuel_consumed(), Some(consumed_mem_fuel)); } #[test] fn lazy_consumption_mode() { - check_consumption_mode(FuelConsumptionMode::Lazy, 1030, 4); + check_consumption_mode(FuelConsumptionMode::Lazy, 1030, 4, 65540, 65536); } #[test] fn eager_consumption_mode() { - check_consumption_mode(FuelConsumptionMode::Eager, 1030, 1028); + check_consumption_mode(FuelConsumptionMode::Eager, 1030, 1028, 65540, 65536); }