wasmi_cli: Support fuel metering (#679)

wasmi_cli: support fuel metering

crates/cli/src/args.rs

@@ -82,6 +82,12 @@ pub struct Args {
     #[clap(long = "invoke", value_name = "FUNCTION")]
     invoke: Option<String>,
 
+    /// Enable execution fiel metering with N units of fuel.
+    ///
+    /// The execution will trap after running out of the N units of fuel.
+    #[clap(long = "fuel", value_name = "N")]
+    fuel: Option<u64>,
+
     /// Arguments given to the Wasm module or the invoked function.
     #[clap(value_name = "ARGS")]
     func_args: Vec<String>,
@@ -103,6 +109,11 @@ impl Args {
         &self.func_args[..]
     }
 
+    /// Returns the amount of fuel given to the CLI app if any.
+    pub fn fuel(&self) -> Option<u64> {
+        self.fuel
+    }
+
     /// Pre-opens all directories given in `--dir` and returns them for use by the [`WasiCtx`].
     ///
     /// # Errors

crates/cli/src/context.rs

@@ -1,7 +1,7 @@
 use crate::utils;
 use anyhow::{anyhow, Error};
 use std::path::Path;
-use wasmi::{ExternType, Func, FuncType, Instance, Module, Store};
+use wasmi::{Config, ExternType, Func, FuncType, Instance, Module, Store};
 use wasmi_wasi::WasiCtx;
 
 /// The [`Context`] for the `wasmi` CLI application.
@@ -23,13 +23,22 @@ impl Context {
     ///
     /// - If parsing, validating, compiling or instantiating the Wasm module failed.
     /// - If adding WASI defintions to the linker failed.
-    pub fn new(wasm_file: &Path, wasi_ctx: WasiCtx) -> Result<Self, Error> {
-        let engine = wasmi::Engine::default();
+    pub fn new(wasm_file: &Path, wasi_ctx: WasiCtx, fuel: Option<u64>) -> Result<Self, Error> {
+        let mut config = Config::default();
+        if fuel.is_some() {
+            config.consume_fuel(true);
+        }
+        let engine = wasmi::Engine::new(&config);
         let wasm_bytes = utils::read_wasm_or_wat(wasm_file)?;
         let module = wasmi::Module::new(&engine, &mut &wasm_bytes[..]).map_err(|error| {
             anyhow!("failed to parse and validate Wasm module {wasm_file:?}: {error}")
         })?;
         let mut store = wasmi::Store::new(&engine, wasi_ctx);
+        if let Some(fuel) = fuel {
+            store.add_fuel(fuel).unwrap_or_else(|error| {
+                panic!("error: fuel metering is enabled but encountered: {error}")
+            });
+        }
         let mut linker = <wasmi::Linker<WasiCtx>>::default();
         wasmi_wasi::define_wasi(&mut linker, &mut store, |ctx| ctx)
             .map_err(|error| anyhow!("failed to add WASI definitions to the linker: {error}"))?;

crates/cli/src/main.rs

@@ -20,7 +20,7 @@ fn main() -> Result<()> {
     let args = Args::parse();
     let wasm_file = args.wasm_file();
     let wasi_ctx = args.wasi_context()?;
-    let mut ctx = Context::new(wasm_file, wasi_ctx)?;
+    let mut ctx = Context::new(wasm_file, wasi_ctx, args.fuel())?;
     let (func_name, func) = get_invoked_func(&args, &ctx)?;
     let ty = func.ty(ctx.store());
     let func_args = utils::decode_func_args(&ty, args.func_args())?;
@@ -39,6 +39,7 @@ fn main() -> Result<()> {
 
     match func.call(ctx.store_mut(), &func_args, &mut func_results) {
         Ok(()) => {
+            print_remaining_fuel(&args, &ctx);
             print_pretty_results(&func_results);
             Ok(())
         }
@@ -48,6 +49,7 @@ fn main() -> Result<()> {
                     // We received an exit code from the WASI program,
                     // therefore we exit with the same exit code after
                     // pretty printing the results.
+                    print_remaining_fuel(&args, &ctx);
                     print_pretty_results(&func_results);
                     process::exit(exit_code)
                 }
@@ -57,6 +59,17 @@ fn main() -> Result<()> {
     }
 }
 
+/// Prints the remaining fuel so far if fuel metering was enabled.
+fn print_remaining_fuel(args: &Args, ctx: &Context) {
+    if let Some(total_fuel) = args.fuel() {
+        let consumed = ctx.store().fuel_consumed().unwrap_or_else(|| {
+            panic!("fuel metering is enabled but could not query consumed fuel")
+        });
+        let remaining = total_fuel - consumed;
+        println!("fuel consumed: {consumed}, fuel remaining: {remaining}");
+    }
+}
+
 /// Performs minor typecheck on the function signature.
 ///
 /// # Note