feat: Add an error code on the metering middleware

Author: jubianchi

examples/metering.rs

@@ -21,7 +21,10 @@ use wasmer::CompilerConfig;
 use wasmer::{imports, wat2wasm, Instance, Module, Store};
 use wasmer_compiler_cranelift::Cranelift;
 use wasmer_engine_jit::JIT;
-use wasmer_middlewares::metering::{get_remaining_points, set_remaining_points, Metering};
+use wasmer_middlewares::{
+    metering::{get_error, get_remaining_points, set_remaining_points},
+    Metering, MeteringError,
+};
 
 fn main() -> anyhow::Result<()> {
     // Let's declare the Wasm module.
@@ -132,19 +135,26 @@ fn main() -> anyhow::Result<()> {
             );
         }
         Err(_) => {
-            println!("Calling `add_one` failed: not enough gas points remaining.");
-        }
-    }
+            println!("Calling `add_one` failed.");
 
-    // Becasue the previous call failed, it did not consume any gas point.
-    // We still have 2 remaining points.
-    let remaining_points_after_third_call = get_remaining_points(&instance);
-    assert_eq!(remaining_points_after_third_call, 2);
+            // Because the last needed more than the remaining gas points, we got an error.
+            let error = get_error(&instance);
 
-    println!(
-        "Remaining points after third call: {:?}",
-        remaining_points_after_third_call
-    );
+            match error {
+                MeteringError::NoError => bail!("No error."),
+                MeteringError::OutOfGas => println!("Not enough gas points remaining."),
+            }
+
+            // There is now 0 remaining points (we can't go below 0).
+            let remaining_points_after_third_call = get_remaining_points(&instance);
+            assert_eq!(remaining_points_after_third_call, 0);
+
+            println!(
+                "Remaining points after third call: {:?}",
+                remaining_points_after_third_call
+            );
+        }
+    }
 
     // Now let's see how we can set a new limit...
     println!("Set new remaining points points to 10");

lib/middlewares/src/lib.rs

@@ -3,3 +3,4 @@ pub mod metering;
 // The most commonly used symbol are exported at top level of the module. Others are available
 // via modules, e.g. `wasmer_middlewares::metering::get_remaining_points`
 pub use metering::Metering;
+pub use metering::{get_error, get_remaining_points, set_remaining_points, MeteringError};

lib/middlewares/src/metering.rs

@@ -1,6 +1,7 @@
 //! `metering` is a middleware for tracking how many operators are executed in total
 //! and putting a limit on the total number of operators executed.
 
+use core::convert::TryFrom;
 use std::convert::TryInto;
 use std::fmt;
 use std::sync::Mutex;
@@ -11,7 +12,7 @@ use wasmer::{
     ExportIndex, FunctionMiddleware, GlobalInit, GlobalType, Instance, LocalFunctionIndex,
     MiddlewareReaderState, ModuleMiddleware, Mutability, Type,
 };
-use wasmer_types::GlobalIndex;
+use wasmer_types::{GlobalIndex, Value};
 use wasmer_vm::ModuleInfo;
 
 /// The module-level metering middleware.
@@ -28,6 +29,9 @@ pub struct Metering<F: Fn(&Operator) -> u64 + Copy + Clone + Send + Sync> {
     /// Function that maps each operator to a cost in "points".
     cost_function: F,
 
+    /// The global index in the current module for the error code.
+    error_code_index: Mutex<Option<GlobalIndex>>,
+
     /// The global index in the current module for remaining points.
     remaining_points_index: Mutex<Option<GlobalIndex>>,
 }
@@ -37,19 +41,52 @@ pub struct FunctionMetering<F: Fn(&Operator) -> u64 + Copy + Clone + Send + Sync
     /// Function that maps each operator to a cost in "points".
     cost_function: F,
 
+    /// The global index in the current module for the error code.
+    error_code_index: GlobalIndex,
+
     /// The global index in the current module for remaining points.
     remaining_points_index: GlobalIndex,
 
     /// Accumulated cost of the current basic block.
     accumulated_cost: u64,
 }
 
+#[derive(Debug, PartialEq)]
+pub enum MeteringError {
+    NoError,
+    OutOfGas,
+}
+
+impl TryFrom<i32> for MeteringError {
+    type Error = ();
+
+    fn try_from(value: i32) -> Result<Self, Self::Error> {
+        match value {
+            value if value == MeteringError::NoError as _ => Ok(MeteringError::NoError),
+            value if value == MeteringError::OutOfGas as _ => Ok(MeteringError::OutOfGas),
+            _ => Err(()),
+        }
+    }
+}
+
+impl<T> TryFrom<Value<T>> for MeteringError {
+    type Error = ();
+
+    fn try_from(v: Value<T>) -> Result<Self, Self::Error> {
+        match v {
+            Value::I32(value) => value.try_into(),
+            _ => Err(()),
+        }
+    }
+}
+
 impl<F: Fn(&Operator) -> u64 + Copy + Clone + Send + Sync> Metering<F> {
     /// Creates a `Metering` middleware.
     pub fn new(initial_limit: u64, cost_function: F) -> Self {
         Self {
             initial_limit,
             cost_function,
+            error_code_index: Mutex::new(None),
             remaining_points_index: Mutex::new(None),
         }
     }
@@ -60,6 +97,7 @@ impl<F: Fn(&Operator) -> u64 + Copy + Clone + Send + Sync> fmt::Debug for Meteri
         f.debug_struct("Metering")
             .field("initial_limit", &self.initial_limit)
             .field("cost_function", &"<function>")
+            .field("error_code_index", &self.error_code_index)
             .field("remaining_points_index", &self.remaining_points_index)
             .finish()
     }
@@ -72,6 +110,11 @@ impl<F: Fn(&Operator) -> u64 + Copy + Clone + Send + Sync + 'static> ModuleMiddl
     fn generate_function_middleware(&self, _: LocalFunctionIndex) -> Box<dyn FunctionMiddleware> {
         Box::new(FunctionMetering {
             cost_function: self.cost_function,
+            error_code_index: self
+                .error_code_index
+                .lock()
+                .unwrap()
+                .expect("Metering::generate_function_middleware: Error code index not set up."),
             remaining_points_index: self.remaining_points_index.lock().unwrap().expect(
                 "Metering::generate_function_middleware: Remaining points index not set up.",
             ),
@@ -81,23 +124,37 @@ impl<F: Fn(&Operator) -> u64 + Copy + Clone + Send + Sync + 'static> ModuleMiddl
 
     /// Transforms a `ModuleInfo` struct in-place. This is called before application on functions begins.
     fn transform_module_info(&self, module_info: &mut ModuleInfo) {
+        let mut error_code_index = self.error_code_index.lock().unwrap();
         let mut remaining_points_index = self.remaining_points_index.lock().unwrap();
-        if remaining_points_index.is_some() {
+
+        if error_code_index.is_some() || remaining_points_index.is_some(){
             panic!("Metering::transform_module_info: Attempting to use a `Metering` middleware from multiple modules.");
         }
 
+        // Append a global for the error code and initialize it.
+        let error_code_global_index = module_info
+            .globals
+            .push(GlobalType::new(Type::I32, Mutability::Var));
+        *error_code_index = Some(error_code_global_index.clone());
+
         // Append a global for remaining points and initialize it.
-        let global_index = module_info
+        let remaining_points_global_index = module_info
             .globals
             .push(GlobalType::new(Type::I64, Mutability::Var));
-        *remaining_points_index = Some(global_index.clone());
+        *remaining_points_index = Some(remaining_points_global_index.clone());
+
         module_info
             .global_initializers
             .push(GlobalInit::I64Const(self.initial_limit as i64));
 
         module_info.exports.insert(
-            "remaining_points".to_string(),
-            ExportIndex::Global(global_index),
+            "wasmer_metering_error_code".to_string(),
+            ExportIndex::Global(error_code_global_index),
+        );
+
+        module_info.exports.insert(
+            "wasmer_metering_remaining_points".to_string(),
+            ExportIndex::Global(remaining_points_global_index),
         );
     }
 }
@@ -106,6 +163,7 @@ impl<F: Fn(&Operator) -> u64 + Copy + Clone + Send + Sync> fmt::Debug for Functi
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("FunctionMetering")
             .field("cost_function", &"<function>")
+            .field("error_code_index", &self.error_code_index)
             .field("remaining_points_index", &self.remaining_points_index)
             .finish()
     }
@@ -143,7 +201,11 @@ impl<F: Fn(&Operator) -> u64 + Copy + Clone + Send + Sync> FunctionMiddleware
                         Operator::I64Const { value: self.accumulated_cost as i64 },
                         Operator::I64LtU,
                         Operator::If { ty: WpTypeOrFuncType::Type(WpType::EmptyBlockType) },
-                        Operator::Unreachable, // FIXME: Signal the error properly.
+                        Operator::I32Const { value: MeteringError::OutOfGas as i32 },
+                        Operator::GlobalSet { global_index: self.error_code_index.as_u32() },
+                        Operator::I64Const { value: 0 },
+                        Operator::GlobalSet { global_index: self.remaining_points_index.as_u32() },
+                        Operator::Unreachable,
                         Operator::End,
 
                         // globals[remaining_points_index] -= self.accumulated_cost;
@@ -164,6 +226,28 @@ impl<F: Fn(&Operator) -> u64 + Copy + Clone + Send + Sync> FunctionMiddleware
     }
 }
 
+/// Get the error code in an `Instance`.
+///
+/// When instance execution traps, this error code will help know if it was caused by
+/// remaining points being exhausted.
+///
+/// This can be used in a headless engine after an ahead-of-time compilation
+/// as all required state lives in the instance.
+///
+/// # Panic
+///
+/// The instance Module must have been processed with the [`Metering`] middleware
+/// at compile time, otherwise this will panic.
+pub fn get_error(instance: &Instance) -> MeteringError {
+    instance
+        .exports
+        .get_global("wasmer_metering_error_code")
+        .expect("Can't get `wasmer_metering_error_code` from Instance")
+        .get()
+        .try_into()
+        .expect("`wasmer_metering_error_code` from Instance has wrong type")
+}
+
 /// Get the remaining points in an `Instance`.
 ///
 /// This can be used in a headless engine after an ahead-of-time compilation
@@ -176,11 +260,11 @@ impl<F: Fn(&Operator) -> u64 + Copy + Clone + Send + Sync> FunctionMiddleware
 pub fn get_remaining_points(instance: &Instance) -> u64 {
     instance
         .exports
-        .get_global("remaining_points")
-        .expect("Can't get `remaining_points` from Instance")
+        .get_global("wasmer_metering_remaining_points")
+        .expect("Can't get `wasmer_metering_remaining_points` from Instance")
         .get()
         .try_into()
-        .expect("`remaining_points` from Instance has wrong type")
+        .expect("`wasmer_metering_remaining_points` from Instance has wrong type")
 }
 
 /// Set the provided remaining points in an `Instance`.
@@ -195,10 +279,10 @@ pub fn get_remaining_points(instance: &Instance) -> u64 {
 pub fn set_remaining_points(instance: &Instance, points: u64) {
     instance
         .exports
-        .get_global("remaining_points")
-        .expect("Can't get `remaining_points` from Instance")
+        .get_global("wasmer_metering_remaining_points")
+        .expect("Can't get `wasmer_metering_remaining_points` from Instance")
         .set(points.into())
-        .expect("Can't set `remaining_points` in Instance");
+        .expect("Can't set `wasmer_metering_remaining_points` in Instance");
 }
 
 #[cfg(test)]
@@ -258,16 +342,17 @@ mod tests {
             .unwrap();
         add_one.call(1).unwrap();
         assert_eq!(get_remaining_points(&instance), 6);
+        assert_eq!(get_error(&instance), MeteringError::NoError);
 
         // Second call
         add_one.call(1).unwrap();
         assert_eq!(get_remaining_points(&instance), 2);
+        assert_eq!(get_error(&instance), MeteringError::NoError);
 
         // Third call fails due to limit
         assert!(add_one.call(1).is_err());
-        // TODO: what do we expect now? 0 or 2? See https://github.com/wasmerio/wasmer/issues/1931
-        // assert_eq!(metering.get_remaining_points(&instance), 2);
-        // assert_eq!(metering.get_remaining_points(&instance), 0);
+        assert_eq!(get_remaining_points(&instance), 0);
+        assert_eq!(get_error(&instance), MeteringError::OutOfGas);
     }
 
     #[test]
@@ -294,9 +379,14 @@ mod tests {
         // Ensure we can use the new points now
         add_one.call(1).unwrap();
         assert_eq!(get_remaining_points(&instance), 8);
+        assert_eq!(get_error(&instance), MeteringError::NoError);
+
         add_one.call(1).unwrap();
         assert_eq!(get_remaining_points(&instance), 4);
+        assert_eq!(get_error(&instance), MeteringError::NoError);
+
         add_one.call(1).unwrap();
         assert_eq!(get_remaining_points(&instance), 0);
+        assert_eq!(get_error(&instance), MeteringError::NoError);
     }
 }