Add memory fuel metering

Add memory fuel metering

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<u64, TrapCode>
+    {
+        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| {

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<u64> {
+        self.ctx.store.mem_fuel_consumed()
+    }
+
+    pub fn mem_fuel_total(&self) -> Option<u64> {
+        self.ctx.store.mem_fuel_total()
+    }
+
+    pub fn consume_mem_fuel(&mut self, delta: u64) -> Result<u64, FuelError> {
+        self.ctx.store.consume_mem_fuel(delta)
+    }
+
+    pub fn reset_mem_fuel(&mut self) -> Result<(), FuelError> {
+        self.ctx.store.reset_mem_fuel()
+    }
 }
 
 impl<T> AsContext for Caller<'_, T> {

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<Idx>`] type.
     ///
     /// # Note
@@ -796,6 +802,35 @@ impl<T> Store<T> {
         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<u64> {
+        self.check_fuel_metering_enabled().ok()?;
+        Some(self.inner.mem_fuel.fuel_consumed())
+    }
+
+    pub fn mem_fuel_total(&self) -> Option<u64> {
+        self.check_fuel_metering_enabled().ok()?;
+        Some(self.inner.mem_fuel.total)
+    }
+
+    pub fn consume_mem_fuel(&mut self, delta: u64) -> Result<u64, FuelError> {
+        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<T>) -> Trampoline {
         let idx = self.trampolines.alloc(func);

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);
 }