fix: better error handling when plugin runs out of fuel (#762)

- Checks error results to determine if a plugin ran out of fuel, if the
fuel is 0 after we get the error, then we return an out of fuel error
message instead.
- Updates `extism_plugin_error` to check `Plugin::error_msg` regardless
of the kernel's error state

runtime/src/current_plugin.rs

@@ -1,3 +1,5 @@
+use anyhow::Context;
+
 use crate::*;
 
 /// CurrentPlugin stores data that is available to the caller in PDK functions, this should
@@ -220,9 +222,13 @@ impl CurrentPlugin {
         let (linker, mut store) = self.linker_and_store();
         let output = &mut [Val::I64(0)];
         if let Some(f) = linker.get(&mut store, EXTISM_ENV_MODULE, "alloc") {
-            f.into_func()
-                .unwrap()
-                .call(&mut *store, &[Val::I64(n as i64)], output)?;
+            catch_out_of_fuel!(
+                &store,
+                f.into_func()
+                    .unwrap()
+                    .call(&mut *store, &[Val::I64(n as i64)], output)
+                    .context("failed to allocate extism memory")
+            )?;
         } else {
             anyhow::bail!("{} unable to allocate memory", self.id);
         }
@@ -246,9 +252,13 @@ impl CurrentPlugin {
     pub fn memory_free(&mut self, handle: MemoryHandle) -> Result<(), Error> {
         let (linker, store) = self.linker_and_store();
         if let Some(f) = linker.get(&mut *store, EXTISM_ENV_MODULE, "free") {
-            f.into_func()
-                .unwrap()
-                .call(&mut *store, &[Val::I64(handle.offset as i64)], &mut [])?;
+            catch_out_of_fuel!(
+                &store,
+                f.into_func()
+                    .unwrap()
+                    .call(&mut *store, &[Val::I64(handle.offset as i64)], &mut [])
+                    .context("failed to free extism memory")
+            )?;
         } else {
             anyhow::bail!("unable to locate an extism kernel function: free",)
         }
@@ -259,9 +269,13 @@ impl CurrentPlugin {
         let (linker, store) = self.linker_and_store();
         let output = &mut [Val::I64(0)];
         if let Some(f) = linker.get(&mut *store, EXTISM_ENV_MODULE, "length") {
-            f.into_func()
-                .unwrap()
-                .call(&mut *store, &[Val::I64(offs as i64)], output)?;
+            catch_out_of_fuel!(
+                &store,
+                f.into_func()
+                    .unwrap()
+                    .call(&mut *store, &[Val::I64(offs as i64)], output)
+                    .context("failed to get length of extism memory handle")
+            )?;
         } else {
             anyhow::bail!("unable to locate an extism kernel function: length",)
         }
@@ -279,9 +293,13 @@ impl CurrentPlugin {
         let (linker, store) = self.linker_and_store();
         let output = &mut [Val::I64(0)];
         if let Some(f) = linker.get(&mut *store, EXTISM_ENV_MODULE, "length_unsafe") {
-            f.into_func()
-                .unwrap()
-                .call(&mut *store, &[Val::I64(offs as i64)], output)?;
+            catch_out_of_fuel!(
+                &store,
+                f.into_func()
+                    .unwrap()
+                    .call(&mut *store, &[Val::I64(offs as i64)], output)
+                    .context("failed to get length of extism memory using length_unsafe")
+            )?;
         } else {
             anyhow::bail!("unable to locate an extism kernel function: length_unsafe",)
         }
@@ -357,7 +375,7 @@ impl CurrentPlugin {
 
         let memory_limiter = if let Some(pgs) = available_pages {
             let n = pgs as usize * 65536;
-            Some(crate::current_plugin::MemoryLimiter {
+            Some(MemoryLimiter {
                 max_bytes: n,
                 bytes_left: n,
             })
@@ -456,13 +474,15 @@ impl CurrentPlugin {
     pub fn set_error(&mut self, s: impl AsRef<str>) -> Result<(u64, u64), Error> {
         let s = s.as_ref();
         debug!(plugin = self.id.to_string(), "set error: {:?}", s);
-        let handle = self.current_plugin_mut().memory_new(s)?;
+        let handle = self.memory_new(s)?;
         let (linker, store) = self.linker_and_store();
         if let Some(f) = linker.get(&mut *store, EXTISM_ENV_MODULE, "error_set") {
-            f.into_func().unwrap().call(
-                &mut *store,
-                &[Val::I64(handle.offset() as i64)],
-                &mut [],
+            catch_out_of_fuel!(
+                &store,
+                f.into_func()
+                    .unwrap()
+                    .call(&mut *store, &[Val::I64(handle.offset() as i64)], &mut [])
+                    .context("failed to set extism error")
             )?;
             Ok((handle.offset(), s.len() as u64))
         } else {
@@ -474,7 +494,10 @@ impl CurrentPlugin {
         let (linker, store) = self.linker_and_store();
         let output = &mut [Val::I64(0)];
         if let Some(f) = linker.get(&mut *store, EXTISM_ENV_MODULE, "error_get") {
-            if let Err(e) = f.into_func().unwrap().call(&mut *store, &[], output) {
+            if let Err(e) = catch_out_of_fuel!(
+                &store,
+                f.into_func().unwrap().call(&mut *store, &[], output)
+            ) {
                 error!(
                     plugin = self.id.to_string(),
                     "unable to call extism:host/env::error_get: {:?}", e

runtime/src/lib.rs

@@ -1,6 +1,17 @@
 // Makes proc-macros able to resolve `::extism` correctly
 extern crate self as extism;
 
+macro_rules! catch_out_of_fuel {
+    ($store: expr, $x:expr) => {{
+        let y = $x;
+        if y.is_err() && $store.get_fuel().is_ok_and(|x| x == 0) {
+            Err(Error::msg("plugin ran out of fuel"))
+        } else {
+            y
+        }
+    }};
+}
+
 pub(crate) use extism_convert::*;
 pub(crate) use std::collections::BTreeMap;
 use std::str::FromStr;

runtime/src/plugin.rs

@@ -4,6 +4,8 @@ use std::{
     path::PathBuf,
 };
 
+use anyhow::Context;
+
 use crate::*;
 
 pub const EXTISM_ENV_MODULE: &str = "extism:host/env";
@@ -554,10 +556,16 @@ impl Plugin {
             .linker
             .get(&mut self.store, EXTISM_ENV_MODULE, "input_set")
         {
-            f.into_func().unwrap().call(
-                &mut self.store,
-                &[Val::I64(handle.offset() as i64), Val::I64(len as i64)],
-                &mut [],
+            catch_out_of_fuel!(
+                &self.store,
+                f.into_func()
+                    .unwrap()
+                    .call(
+                        &mut self.store,
+                        &[Val::I64(handle.offset() as i64), Val::I64(len as i64)],
+                        &mut [],
+                    )
+                    .context("unable to set extism input")
             )?;
         }
 
@@ -565,7 +573,8 @@ impl Plugin {
             self.linker
                 .get(&mut self.store, EXTISM_ENV_MODULE, "extism_context")
         {
-            ctxt.set(&mut self.store, Val::ExternRef(host_context))?;
+            ctxt.set(&mut self.store, Val::ExternRef(host_context))
+                .context("unable to set extism host context")?;
         }
 
         Ok(())
@@ -576,7 +585,13 @@ impl Plugin {
         let id = self.id.to_string();
 
         if let Some(f) = self.linker.get(&mut self.store, EXTISM_ENV_MODULE, "reset") {
-            f.into_func().unwrap().call(&mut self.store, &[], &mut [])?;
+            catch_out_of_fuel!(
+                &self.store,
+                f.into_func()
+                    .unwrap()
+                    .call(&mut self.store, &[], &mut [])
+                    .context("extism reset failed")
+            )?;
         } else {
             error!(plugin = &id, "call to extism:host/env::reset failed");
         }
@@ -658,13 +673,22 @@ impl Plugin {
             match runtime {
                 GuestRuntime::Haskell { init, reactor_init } => {
                     if let Some(reactor_init) = reactor_init {
-                        reactor_init.call(&mut *store, &[], &mut [])?;
+                        catch_out_of_fuel!(
+                            &store,
+                            reactor_init
+                                .call(&mut *store, &[], &mut [])
+                                .context("failed to initialize Haskell reactor runtime")
+                        )?;
                     }
                     let mut results = vec![Val::I32(0); init.ty(&*store).results().len()];
-                    init.call(
-                        &mut *store,
-                        &[Val::I32(0), Val::I32(0)],
-                        results.as_mut_slice(),
+                    catch_out_of_fuel!(
+                        &store,
+                        init.call(
+                            &mut *store,
+                            &[Val::I32(0), Val::I32(0)],
+                            results.as_mut_slice(),
+                        )
+                        .context("failed to initialize Haskell using hs_init")
                     )?;
                     debug!(
                         plugin = self.id.to_string(),
@@ -672,7 +696,11 @@ impl Plugin {
                     );
                 }
                 GuestRuntime::Wasi { init } => {
-                    init.call(&mut *store, &[], &mut [])?;
+                    catch_out_of_fuel!(
+                        &store,
+                        init.call(&mut *store, &[], &mut [])
+                            .context("failed to initialize wasi runtime")
+                    )?;
                     debug!(plugin = self.id.to_string(), "initialied WASI runtime");
                 }
             }
@@ -690,7 +718,13 @@ impl Plugin {
             .linker
             .get(&mut *store, EXTISM_ENV_MODULE, "output_offset")
         {
-            f.into_func().unwrap().call(&mut *store, &[], out)?;
+            catch_out_of_fuel!(
+                &store,
+                f.into_func()
+                    .unwrap()
+                    .call(&mut *store, &[], out)
+                    .context("call to set extism output offset failed")
+            )?;
         } else {
             anyhow::bail!("unable to set output")
         }
@@ -698,7 +732,13 @@ impl Plugin {
             .linker
             .get(&mut *store, EXTISM_ENV_MODULE, "output_length")
         {
-            f.into_func().unwrap().call(&mut *store, &[], out_len)?;
+            catch_out_of_fuel!(
+                &store,
+                f.into_func()
+                    .unwrap()
+                    .call(&mut *store, &[], out_len)
+                    .context("call to set extism output length failed")
+            )?;
         } else {
             anyhow::bail!("unable to set output length")
         }
@@ -712,10 +752,10 @@ impl Plugin {
     fn output<'a, T: FromBytes<'a>>(&'a mut self) -> Result<T, Error> {
         let offs = self.output.offset;
         let len = self.output.length;
-        T::from_bytes(
-            self.current_plugin_mut()
-                .memory_bytes(unsafe { MemoryHandle::new(offs, len) })?,
-        )
+        let x = self
+            .current_plugin_mut()
+            .memory_bytes(unsafe { MemoryHandle::new(offs, len) })?;
+        T::from_bytes(x)
     }
 
     // Cache output memory and error information after call is complete
@@ -750,13 +790,12 @@ impl Plugin {
         let name = name.as_ref();
         let input = input.as_ref();
 
-        if let Err(e) = self.reset_store(lock) {
-            error!(
-                plugin = self.id.to_string(),
-                "call to Plugin::reset_store failed: {e:?}"
-            );
+        if let Some(fuel) = self.fuel {
+            self.store.set_fuel(fuel).map_err(|x| (x, -1))?;
         }
 
+        catch_out_of_fuel!(&self.store, self.reset_store(lock)).map_err(|x| (x, -1))?;
+
         self.instantiate(lock).map_err(|e| (e, -1))?;
 
         // Set host context
@@ -805,9 +844,6 @@ 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
@@ -828,34 +864,43 @@ impl Plugin {
         let _ = self.timer_tx.send(TimerAction::Stop { id: self.id });
         self.store_needs_reset = name == "_start";
 
-        // Get extism error
-        self.get_output_after_call().map_err(|x| (x, -1))?;
         let mut rc = 0;
-        if !results.is_empty() {
-            rc = results[0].i32().unwrap_or(-1);
-            debug!(plugin = self.id.to_string(), "got return code: {}", rc);
-        }
+        if self.store.get_fuel().is_ok_and(|x| x == 0) {
+            res = Err(Error::msg("plugin ran out of fuel"));
+            rc = -1;
+        } else {
+            // Get extism error
+            self.get_output_after_call().map_err(|x| (x, -1))?;
+            if !results.is_empty() {
+                rc = results[0].i32().unwrap_or(-1);
+                debug!(plugin = self.id.to_string(), "got return code: {}", rc);
+            }
 
-        if self.output.error_offset != 0 && self.output.error_length != 0 {
-            let handle = MemoryHandle {
-                offset: self.output.error_offset,
-                length: self.output.error_length,
-            };
-            if let Ok(e) = self.current_plugin_mut().memory_str(handle) {
-                let x = e.to_string();
-                error!(
-                    plugin = self.id.to_string(),
-                    "call to {name} returned with error message: {}", x
-                );
-                if let Err(e) = res {
-                    res = Err(Error::msg(x).context(e));
-                } else {
-                    res = Err(Error::msg(x))
+            if self.output.error_offset != 0 && self.output.error_length != 0 {
+                let handle = MemoryHandle {
+                    offset: self.output.error_offset,
+                    length: self.output.error_length,
+                };
+                match self.current_plugin_mut().memory_str(handle) {
+                    Ok(e) => {
+                        let x = e.to_string();
+                        error!(
+                            plugin = self.id.to_string(),
+                            "call to {name} returned with error message: {}", x
+                        );
+                        if let Err(e) = res {
+                            res = Err(Error::msg(x).context(e));
+                        } else {
+                            res = Err(Error::msg(x))
+                        }
+                    }
+                    Err(msg) => {
+                        res = Err(Error::msg(format!(
+                            "Call to Extism plugin function {name} encountered an error: {}",
+                            msg,
+                        )));
+                    }
                 }
-            } else {
-                res = Err(Error::msg(format!(
-                    "Call to Extism plugin function {name} encountered an error"
-                )));
             }
         }
 
@@ -1022,42 +1067,21 @@ impl Plugin {
 
     pub(crate) fn clear_error(&mut self) -> Result<(), Error> {
         trace!(plugin = self.id.to_string(), "clearing error");
-        let (linker, store) = self.linker_and_store();
+        self.error_msg = None;
+        let (linker, mut store) = self.linker_and_store();
+        #[allow(clippy::needless_borrows_for_generic_args)]
         if let Some(f) = linker.get(&mut *store, EXTISM_ENV_MODULE, "error_set") {
-            f.into_func()
+            let x = f
+                .into_func()
                 .unwrap()
-                .call(store, &[Val::I64(0)], &mut [])?;
+                .call(&mut store, &[Val::I64(0)], &mut [])
+                .context("unable to clear error message");
+            catch_out_of_fuel!(&store, x)?;
             Ok(())
         } else {
             anyhow::bail!("Plugin::clear_error failed, extism:host/env::error_set not found")
         }
     }
-
-    // A convenience method to set the plugin error and return a value
-    pub(crate) fn return_error<E>(
-        &mut self,
-        instance_lock: &mut std::sync::MutexGuard<Option<Instance>>,
-        e: impl std::fmt::Display,
-        x: E,
-    ) -> E {
-        if instance_lock.is_none() {
-            error!(
-                plugin = self.id.to_string(),
-                "no instance, unable to set error: {}", e
-            );
-            return x;
-        }
-        match self.current_plugin_mut().set_error(e.to_string()) {
-            Ok((a, b)) => {
-                self.output.error_offset = a;
-                self.output.error_length = b;
-            }
-            Err(e) => {
-                error!(plugin = self.id.to_string(), "unable to set error: {e:?}")
-            }
-        }
-        x
-    }
 }
 
 // Enumerates the PDK languages that need some additional initialization

runtime/src/sdk.rs

@@ -11,6 +11,12 @@ pub struct ExtismFunction(std::cell::Cell<Option<Function>>);
 /// The return code used to specify a successful plugin call
 pub static EXTISM_SUCCESS: i32 = 0;
 
+fn make_error_msg(s: String) -> Vec<u8> {
+    let mut s = s.into_bytes();
+    s.push(0);
+    s
+}
+
 /// A union type for host function argument/return values
 #[repr(C)]
 pub union ValUnion {
@@ -470,8 +476,6 @@ pub unsafe extern "C" fn extism_plugin_config(
         return false;
     }
     let plugin = &mut *plugin;
-    let _lock = plugin.instance.clone();
-    let mut lock = _lock.lock().unwrap();
 
     trace!(
         plugin = plugin.id.to_string(),
@@ -482,8 +486,8 @@ pub unsafe extern "C" fn extism_plugin_config(
     let json: std::collections::BTreeMap<String, Option<String>> =
         match serde_json::from_slice(data) {
             Ok(x) => x,
-            Err(e) => {
-                return plugin.return_error(&mut lock, e, false);
+            Err(_) => {
+                return false;
             }
         };
 
@@ -516,9 +520,6 @@ pub unsafe extern "C" fn extism_plugin_function_exists(
         return false;
     }
     let plugin = &mut *plugin;
-    let _lock = plugin.instance.clone();
-    let mut lock = _lock.lock().unwrap();
-
     let name = std::ffi::CStr::from_ptr(func_name);
     trace!(
         plugin = plugin.id.to_string(),
@@ -528,8 +529,8 @@ pub unsafe extern "C" fn extism_plugin_function_exists(
 
     let name = match name.to_str() {
         Ok(x) => x,
-        Err(e) => {
-            return plugin.return_error(&mut lock, e, false);
+        Err(_) => {
+            return false;
         }
     };
 
@@ -581,20 +582,14 @@ pub unsafe extern "C" fn extism_plugin_call_with_host_context(
     let lock = plugin.instance.clone();
     let mut lock = lock.lock().unwrap();
 
-    if let Err(e) = plugin.reset_store(&mut lock) {
-        error!(
-            plugin = plugin.id.to_string(),
-            "call to Plugin::reset_store failed: {e:?}"
-        );
-    }
-
-    plugin.error_msg = None;
-
     // Get function name
     let name = std::ffi::CStr::from_ptr(func_name);
     let name = match name.to_str() {
         Ok(name) => name,
-        Err(e) => return plugin.return_error(&mut lock, e, -1),
+        Err(e) => {
+            plugin.error_msg = Some(make_error_msg(e.to_string()));
+            return -1;
+        }
     };
 
     trace!(
@@ -610,7 +605,10 @@ pub unsafe extern "C" fn extism_plugin_call_with_host_context(
     };
     let res = plugin.raw_call(&mut lock, name, input, r);
     match res {
-        Err((e, rc)) => plugin.return_error(&mut lock, e, rc),
+        Err((e, rc)) => {
+            plugin.error_msg = Some(make_error_msg(e.to_string()));
+            rc
+        }
         Ok(x) => x,
     }
 }
@@ -633,6 +631,9 @@ pub unsafe extern "C" fn extism_plugin_error(plugin: *mut Plugin) -> *const c_ch
     let _lock = _lock.lock().unwrap();
 
     if plugin.output.error_offset == 0 {
+        if let Some(err) = &plugin.error_msg {
+            return err.as_ptr() as *const _;
+        }
         trace!(plugin = plugin.id.to_string(), "error is NULL");
         return std::ptr::null();
     }