feat: add ability to limit the number of instructions executed by a plugin (#754)

*Note*: this will be limited for the time being as not all runtimes
support this yet

runtime/extism.h

@@ -195,6 +195,17 @@ ExtismPlugin *extism_plugin_new(const uint8_t *wasm,
                                 bool with_wasi,
                                 char **errmsg);
 
+/**
+ * Create a new plugin and set the number of instructions a plugin is allowed to execute
+ */
+ExtismPlugin *extism_plugin_new_with_fuel_limit(const uint8_t *wasm,
+                                                ExtismSize wasm_size,
+                                                const ExtismFunction **functions,
+                                                ExtismSize n_functions,
+                                                bool with_wasi,
+                                                uint64_t fuel_limit,
+                                                char **errmsg);
+
 /**
  * Free the error returned by `extism_plugin_new`, errors returned from `extism_plugin_error` don't need to be freed
  */

runtime/src/plugin.rs

@@ -83,6 +83,8 @@ pub struct Plugin {
     pub(crate) debug_options: DebugOptions,
 
     pub(crate) error_msg: Option<Vec<u8>>,
+
+    pub(crate) fuel: Option<u64>,
 }
 
 unsafe impl Send for Plugin {}
@@ -292,7 +294,14 @@ impl Plugin {
         imports: impl IntoIterator<Item = Function>,
         with_wasi: bool,
     ) -> Result<Plugin, Error> {
-        Self::build_new(wasm.into(), imports, with_wasi, Default::default(), None)
+        Self::build_new(
+            wasm.into(),
+            imports,
+            with_wasi,
+            Default::default(),
+            None,
+            None,
+        )
     }
 
     pub(crate) fn build_new(
@@ -301,6 +310,7 @@ impl Plugin {
         with_wasi: bool,
         debug_options: DebugOptions,
         cache_dir: Option<Option<PathBuf>>,
+        fuel: Option<u64>,
     ) -> Result<Plugin, Error> {
         // Setup wasmtime types
         let mut config = Config::new();
@@ -313,6 +323,10 @@ impl Plugin {
             .wasm_function_references(true)
             .wasm_gc(true);
 
+        if fuel.is_some() {
+            config.consume_fuel(true);
+        }
+
         match cache_dir {
             Some(None) => (),
             Some(Some(path)) => {
@@ -346,6 +360,9 @@ impl Plugin {
             CurrentPlugin::new(manifest, with_wasi, available_pages, id)?,
         );
         store.set_epoch_deadline(1);
+        if let Some(fuel) = fuel {
+            store.set_fuel(fuel)?;
+        }
 
         let imports: Vec<Function> = imports.into_iter().collect();
         let (instance_pre, linker) = relink(&engine, &mut store, &imports, &modules, with_wasi)?;
@@ -367,6 +384,7 @@ impl Plugin {
             debug_options,
             _functions: imports,
             error_msg: None,
+            fuel,
         };
 
         plugin.current_plugin_mut().store = &mut plugin.store;
@@ -400,6 +418,10 @@ impl Plugin {
             );
             self.store.set_epoch_deadline(1);
 
+            if let Some(fuel) = self.fuel {
+                self.store.set_fuel(fuel)?;
+            }
+
             let (instance_pre, linker) = relink(
                 &engine,
                 &mut self.store,
@@ -752,6 +774,9 @@ impl Plugin {
             .expect("Timer should start");
         self.store.epoch_deadline_trap();
         self.store.set_epoch_deadline(1);
+        if let Some(fuel) = self.fuel {
+            self.store.set_fuel(fuel).map_err(|x| (x, -1))?;
+        }
         self.current_plugin_mut().start_time = std::time::Instant::now();
 
         // Call the function

runtime/src/plugin_builder.rs

@@ -39,6 +39,7 @@ pub struct PluginBuilder<'a> {
     functions: Vec<Function>,
     debug_options: DebugOptions,
     cache_config: Option<Option<PathBuf>>,
+    fuel: Option<u64>,
 }
 
 impl<'a> PluginBuilder<'a> {
@@ -50,6 +51,7 @@ impl<'a> PluginBuilder<'a> {
             functions: vec![],
             debug_options: DebugOptions::default(),
             cache_config: None,
+            fuel: None,
         }
     }
 
@@ -148,6 +150,12 @@ impl<'a> PluginBuilder<'a> {
         self
     }
 
+    // Limit the number of instructions that can be executed
+    pub fn with_fuel_limit(mut self, fuel: u64) -> Self {
+        self.fuel = Some(fuel);
+        self
+    }
+
     /// Generate a new plugin with the configured settings
     pub fn build(self) -> Result<Plugin, Error> {
         Plugin::build_new(
@@ -156,6 +164,7 @@ impl<'a> PluginBuilder<'a> {
             self.wasi,
             self.debug_options,
             self.cache_config,
+            self.fuel,
         )
     }
 }

runtime/src/sdk.rs

@@ -334,6 +334,66 @@ pub unsafe extern "C" fn extism_plugin_new(
     }
 }
 
+/// Create a new plugin and set the number of instructions a plugin is allowed to execute
+#[no_mangle]
+pub unsafe extern "C" fn extism_plugin_new_with_fuel_limit(
+    wasm: *const u8,
+    wasm_size: Size,
+    functions: *mut *const ExtismFunction,
+    n_functions: Size,
+    with_wasi: bool,
+    fuel_limit: u64,
+    errmsg: *mut *mut std::ffi::c_char,
+) -> *mut Plugin {
+    trace!(
+        "Call to extism_plugin_new_with_fuel_limit with wasm pointer {:?}",
+        wasm
+    );
+    let data = std::slice::from_raw_parts(wasm, wasm_size as usize);
+    let mut funcs = vec![];
+
+    if !functions.is_null() {
+        for i in 0..n_functions {
+            unsafe {
+                let f = *functions.add(i as usize);
+                if f.is_null() {
+                    continue;
+                }
+                if let Some(f) = (*f).0.take() {
+                    funcs.push(f);
+                } else {
+                    let e = std::ffi::CString::new(
+                        "Function cannot be registered with multiple different Plugins",
+                    )
+                    .unwrap();
+                    *errmsg = e.into_raw();
+                }
+            }
+        }
+    }
+
+    let plugin = Plugin::build_new(
+        data.into(),
+        funcs,
+        with_wasi,
+        Default::default(),
+        None,
+        Some(fuel_limit),
+    );
+
+    match plugin {
+        Err(e) => {
+            if !errmsg.is_null() {
+                let e = std::ffi::CString::new(format!("Unable to create Extism plugin: {}", e))
+                    .unwrap();
+                *errmsg = e.into_raw();
+            }
+            std::ptr::null_mut()
+        }
+        Ok(p) => Box::into_raw(Box::new(p)),
+    }
+}
+
 /// Free the error returned by `extism_plugin_new`, errors returned from `extism_plugin_error` don't need to be freed
 #[no_mangle]
 pub unsafe extern "C" fn extism_plugin_new_error_free(err: *mut std::ffi::c_char) {

runtime/src/tests/runtime.rs

@@ -241,6 +241,22 @@ fn test_timeout() {
     assert!(err == "timeout");
 }
 
+#[test]
+fn test_fuel() {
+    let manifest = Manifest::new([extism_manifest::Wasm::data(WASM_LOOP)]);
+    let mut plugin = PluginBuilder::new(manifest)
+        .with_wasi(true)
+        .with_fuel_limit(1)
+        .build()
+        .unwrap();
+    for _ in 0..10001 {
+        let output: Result<&[u8], Error> = plugin.call("loop_forever", "abc123");
+        let err = output.unwrap_err().root_cause().to_string();
+        println!("Fuel limited plugin exited with error: {:?}", &err);
+        assert!(err.contains("fuel"));
+    }
+}
+
 #[test]
 #[cfg(feature = "http")]
 fn test_http_timeout() {