diff --git a/crates/c-api/include/wasmtime/store.h b/crates/c-api/include/wasmtime/store.h index e4e76173ee..f8298c10fa 100644 --- a/crates/c-api/include/wasmtime/store.h +++ b/crates/c-api/include/wasmtime/store.h @@ -139,6 +139,20 @@ WASM_API_EXTERN wasmtime_error_t *wasmtime_context_add_fuel(wasmtime_context_t * */ WASM_API_EXTERN bool wasmtime_context_fuel_consumed(const wasmtime_context_t *context, uint64_t *fuel); +/** + * \brief Attempt to manually consume fuel from the store. + * + * If fuel consumption is not enabled via #wasmtime_config_consume_fuel_set then + * this function will return an error. Otherwise this will attempt to consume + * the specified amount of `fuel` from the store. If successful the remaining + * amount of fuel is stored into `remaining`. If `fuel` couldn't be consumed + * then an error is returned. + * + * Also note that fuel, if enabled, must be originally configured via + * #wasmtime_context_add_fuel. + */ +WASM_API_EXTERN wasmtime_error_t *wasmtime_context_consume_fuel(wasmtime_context_t *context, uint64_t fuel, uint64_t *remaining); + /** * \brief Configres WASI state within the specified store. * diff --git a/crates/c-api/src/store.rs b/crates/c-api/src/store.rs index 86af26b318..ea8994ddc9 100644 --- a/crates/c-api/src/store.rs +++ b/crates/c-api/src/store.rs @@ -145,6 +145,17 @@ pub extern "C" fn wasmtime_context_fuel_consumed(store: CStoreContext<'_>, fuel: } } +#[no_mangle] +pub extern "C" fn wasmtime_context_consume_fuel( + mut store: CStoreContextMut<'_>, + fuel: u64, + remaining_fuel: &mut u64, +) -> Option> { + crate::handle_result(store.consume_fuel(fuel), |remaining| { + *remaining_fuel = remaining; + }) +} + #[repr(C)] pub struct wasmtime_interrupt_handle_t { handle: InterruptHandle, diff --git a/crates/wasmtime/src/func.rs b/crates/wasmtime/src/func.rs index f11e201179..586857afc9 100644 --- a/crates/wasmtime/src/func.rs +++ b/crates/wasmtime/src/func.rs @@ -1621,6 +1621,13 @@ impl Caller<'_, T> { self.store.add_fuel(fuel) } + /// Synthetically consumes fuel from the store. + /// + /// For more information see [`Store::consume_fuel`](crate::Store::consume_fuel) + pub fn consume_fuel(&mut self, fuel: u64) -> Result { + self.store.consume_fuel(fuel) + } + /// Configures this `Store` to trap whenever fuel runs out. /// /// For more information see diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index 116d1d5cfd..e42020b3ba 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -614,6 +614,29 @@ impl Store { self.inner.add_fuel(fuel) } + /// Synthetically consumes fuel from this [`Store`]. + /// + /// For this method to work fuel consumption must be enabled via + /// [`Config::consume_fuel`](crate::Config::consume_fuel). + /// + /// WebAssembly execution will automatically consume fuel but if so desired + /// the embedder can also consume fuel manually to account for relative + /// costs of host functions, for example. + /// + /// This function will attempt to consume `fuel` units of fuel from within + /// this store. If the remaining amount of fuel allows this then `Ok(N)` is + /// returned where `N` is the amount of remaining fuel. Otherwise an error + /// is returned and no fuel is consumed. + /// + /// # Errors + /// + /// This function will return an either either if fuel consumption via + /// [`Config`](crate::Config) is disabled or if `fuel` exceeds the amount + /// of remaining fuel within this store. + pub fn consume_fuel(&mut self, fuel: u64) -> Result { + self.inner.consume_fuel(fuel) + } + /// Configures a [`Store`] to generate a [`Trap`] whenever it runs out of /// fuel. /// @@ -748,6 +771,13 @@ impl<'a, T> StoreContextMut<'a, T> { self.0.add_fuel(fuel) } + /// Synthetically consume fuel from this store. + /// + /// For more information see [`Store::consume_fuel`] + pub fn consume_fuel(&mut self, fuel: u64) -> Result { + self.0.consume_fuel(fuel) + } + /// Configures this `Store` to trap whenever fuel runs out. /// /// For more information see [`Store::out_of_fuel_trap`] @@ -1028,6 +1058,20 @@ impl StoreOpaque { Ok(()) } + fn consume_fuel(&mut self, fuel: u64) -> Result { + let consumed_ptr = unsafe { &mut *self.interrupts.fuel_consumed.get() }; + match i64::try_from(fuel) + .ok() + .and_then(|fuel| consumed_ptr.checked_add(fuel)) + { + Some(consumed) if consumed < 0 => { + *consumed_ptr = consumed; + Ok(u64::try_from(-consumed).unwrap()) + } + _ => bail!("not enough fuel remaining in store"), + } + } + #[inline] pub fn signal_handler(&self) -> Option<*const SignalHandler<'static>> { let handler = self.signal_handler.as_ref()?; diff --git a/tests/all/fuel.rs b/tests/all/fuel.rs index 8478e542ef..5b2a9ab2e1 100644 --- a/tests/all/fuel.rs +++ b/tests/all/fuel.rs @@ -122,3 +122,70 @@ fn iloop() { ); } } + +#[test] +fn manual_fuel() { + let mut config = Config::new(); + config.consume_fuel(true); + let engine = Engine::new(&config).unwrap(); + let mut store = Store::new(&engine, ()); + store.add_fuel(10_000).unwrap(); + assert_eq!(store.fuel_consumed(), Some(0)); + assert_eq!(store.consume_fuel(1).unwrap(), 9_999); + assert_eq!(store.fuel_consumed(), Some(1)); + assert!(store.consume_fuel(10_000).is_err()); + assert_eq!(store.consume_fuel(999).unwrap(), 9_000); + assert!(store.consume_fuel(10_000).is_err()); + assert_eq!(store.consume_fuel(8998).unwrap(), 2); + assert!(store.consume_fuel(2).is_err()); + assert_eq!(store.consume_fuel(1).unwrap(), 1); + assert!(store.consume_fuel(1).is_err()); + assert_eq!(store.consume_fuel(0).unwrap(), 1); +} + +#[test] +fn host_function_consumes_all() { + const FUEL: u64 = 10_000; + let mut config = Config::new(); + config.consume_fuel(true); + let engine = Engine::new(&config).unwrap(); + let module = Module::new( + &engine, + r#" + (module + (import "" "" (func)) + (func (export "") + call 0 + call $other) + (func $other)) + "#, + ) + .unwrap(); + let mut store = Store::new(&engine, ()); + store.add_fuel(FUEL).unwrap(); + let func = Func::wrap(&mut store, |mut caller: Caller<'_, ()>| { + let consumed = caller.fuel_consumed().unwrap(); + assert_eq!(caller.consume_fuel((FUEL - consumed) - 1).unwrap(), 1); + }); + + let instance = Instance::new(&mut store, &module, &[func.into()]).unwrap(); + let export = instance + .get_typed_func::<(), (), _>(&mut store, "") + .unwrap(); + let trap = export.call(&mut store, ()).err().unwrap().to_string(); + assert!(trap.contains("all fuel consumed"), "bad error: {}", trap); +} + +#[test] +fn manual_edge_cases() { + let mut config = Config::new(); + config.consume_fuel(true); + let engine = Engine::new(&config).unwrap(); + let mut store = Store::new(&engine, ()); + store.add_fuel(u64::MAX).unwrap(); + assert_eq!(store.fuel_consumed(), Some(0)); + assert!(store.consume_fuel(u64::MAX).is_err()); + assert!(store.consume_fuel(i64::MAX as u64).is_err()); + assert!(store.consume_fuel(i64::MAX as u64 + 1).is_err()); + assert_eq!(store.consume_fuel(i64::MAX as u64 - 1).unwrap(), 1); +}