piecrust: allow specifying HostQuery price

Host queries are modified to be able to specify their price. The calling
of a such a query proceeds in two stages - deserializing and pricing,
and then execution. In between the two stages, `piecrust` is able to
make the choice of whether to continue with execution or not, based on
the gas remaining for the execution.

To achieve this, the `HostQuery` trait is modified by adding two
functions - `deserialize_and_price` and `execute` - and by removing the
`Fn(&mut [u8], u32) -> u32` bound. This structure allows the implementer
to price queries fairly, as well as prevent double deserialization by
leveraging `Box<dyn Any>` to pass data between the two functions.

`HostQuery` remains implemented for all `Fn(&mut [u8], u32) -> u32` that
are also `Send` and `Sync`, so this is not a breaking change.

See-also: #359

piecrust/CHANGELOG.md

@@ -7,8 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## Added
+
+- Apply charging mechanism for host queries [#359]
+- Add `HostQuery::execute` and `HostQuery::deserialize_and_price` [#359]
+
 ## Changed
 
+- Drop `Fn(&mut [u8], u32) -> u32` bound for `HostQuery` [#359]
 - Make storage instructions cost 4 gas per byte [#359]
 - Upgrade `dusk-wasmtime` to version `21.0.0-alpha` [#359]
 

piecrust/src/imports.rs

@@ -7,6 +7,7 @@
 mod wasm32;
 mod wasm64;
 
+use std::any::Any;
 use std::sync::Arc;
 
 use dusk_wasmtime::{
@@ -161,9 +162,28 @@ pub(crate) fn hq(
             .map(ToOwned::to_owned)
     })?;
 
-    Ok(instance
-        .with_arg_buf_mut(|buf| env.host_query(&name, buf, arg_len))
-        .ok_or(Error::MissingHostQuery(name))?)
+    // Get the host query if it exists.
+    let host_query =
+        env.host_query(&name).ok_or(Error::MissingHostQuery(name))?;
+    let mut arg: Box<dyn Any> = Box::new(());
+
+    // Price the query, allowing for an early exit if the gas is insufficient.
+    let query_cost = instance.with_arg_buf(|arg_buf| {
+        let arg_len = arg_len as usize;
+        let arg_buf = &arg_buf[..arg_len];
+        host_query.deserialize_and_price(arg_buf, &mut arg)
+    });
+
+    // If the gas is insufficient, return an error.
+    let gas_remaining = instance.get_remaining_gas();
+    if gas_remaining < query_cost {
+        instance.set_remaining_gas(0);
+        Err(Error::OutOfGas)?;
+    }
+    instance.set_remaining_gas(gas_remaining - query_cost);
+
+    // Execute the query and return the result.
+    Ok(instance.with_arg_buf_mut(|arg_buf| host_query.execute(&arg, arg_buf)))
 }
 
 pub(crate) fn hd(

piecrust/src/session.rs

@@ -30,7 +30,7 @@ use crate::error::Error::{self, InitalizationError, PersistenceError};
 use crate::instance::WrappedInstance;
 use crate::store::{ContractSession, PageOpening, PAGE_SIZE};
 use crate::types::StandardBufSerializer;
-use crate::vm::HostQueries;
+use crate::vm::{HostQueries, HostQuery};
 
 const MAX_META_SIZE: usize = ARGBUF_LEN;
 pub const INIT_METHOD: &str = "init";
@@ -522,7 +522,7 @@ impl Session {
             .inner
             .contract_session
             .contract(contract_id)
-            .map_err(|err| Error::PersistenceError(Arc::new(err)))?
+            .map_err(|err| PersistenceError(Arc::new(err)))?
             .map(|data| data.memory.current_len))
     }
 
@@ -612,13 +612,8 @@ impl Session {
         Ok(instance)
     }
 
-    pub(crate) fn host_query(
-        &self,
-        name: &str,
-        buf: &mut [u8],
-        arg_len: u32,
-    ) -> Option<u32> {
-        self.inner.host_queries.call(name, buf, arg_len)
+    pub(crate) fn host_query(&self, name: &str) -> Option<&dyn HostQuery> {
+        self.inner.host_queries.get(name)
     }
 
     pub(crate) fn nth_from_top(&self, n: usize) -> Option<CallTreeElem> {

piecrust/src/vm.rs

@@ -4,6 +4,7 @@
 //
 // Copyright (c) DUSK NETWORK. All rights reserved.
 
+use std::any::Any;
 use std::borrow::Cow;
 use std::collections::BTreeMap;
 use std::fmt::{self, Debug, Formatter};
@@ -269,8 +270,8 @@ impl HostQueries {
         self.map.insert(name.into(), Arc::new(query));
     }
 
-    pub fn call(&self, name: &str, buf: &mut [u8], len: u32) -> Option<u32> {
-        self.map.get(name).map(|host_query| host_query(buf, len))
+    pub fn get(&self, name: &str) -> Option<&dyn HostQuery> {
+        self.map.get(name).map(|q| q.as_ref())
     }
 }
 
@@ -281,5 +282,54 @@ impl HostQueries {
 /// function, and should be processed first. Once this is done, the implementor
 /// should emplace the return of the query in the same buffer, and return the
 /// length written.
-pub trait HostQuery: Send + Sync + Fn(&mut [u8], u32) -> u32 {}
-impl<F> HostQuery for F where F: Send + Sync + Fn(&mut [u8], u32) -> u32 {}
+///
+/// Implementers of `Fn(&mut [u8], u32) -> u32` can be used as a `HostQuery`,
+/// but the cost will be 0.
+pub trait HostQuery: Send + Sync {
+    /// Deserialize the argument buffer and return the price of the query.
+    ///
+    /// The buffer passed will be of the length of the argument the contract
+    /// used to call the query.
+    ///
+    /// Any information needed to perform the query after deserializing the
+    /// argument should be stored in `arg`, and will be passed to [`execute`],
+    /// if there's enough gas to execute the query.
+    ///
+    /// [`execute`]: HostQuery::execute
+    fn deserialize_and_price(
+        &self,
+        arg_buf: &[u8],
+        arg: &mut Box<dyn Any>,
+    ) -> u64;
+
+    /// Perform the query and return the length of the result written to the
+    /// argument buffer.
+    ///
+    /// The whole argument buffer is passed, together with any information
+    /// stored in `arg` previously, during [`deserialize_and_price`].
+    ///
+    /// [`deserialize_and_price`]: HostQuery::deserialize_and_price
+    fn execute(&self, arg: &Box<dyn Any>, arg_buf: &mut [u8]) -> u32;
+}
+
+/// An implementer of `Fn(&mut [u8], u32) -> u32` can be used as a `HostQuery`,
+/// and the cost will be 0.
+impl<F> HostQuery for F
+where
+    F: Send + Sync + Fn(&mut [u8], u32) -> u32,
+{
+    fn deserialize_and_price(
+        &self,
+        arg_buf: &[u8],
+        arg: &mut Box<dyn Any>,
+    ) -> u64 {
+        let len = Box::new(arg_buf.len() as u32);
+        *arg = len;
+        0
+    }
+
+    fn execute(&self, arg: &Box<dyn Any>, arg_buf: &mut [u8]) -> u32 {
+        let arg_len = *arg.downcast_ref::<u32>().unwrap();
+        self(arg_buf, arg_len)
+    }
+}

piecrust/tests/host.rs

@@ -6,9 +6,12 @@
 
 use dusk_plonk::prelude::*;
 use once_cell::sync::Lazy;
-use piecrust::{contract_bytecode, ContractData, Error, SessionData, VM};
+use piecrust::{
+    contract_bytecode, ContractData, Error, HostQuery, SessionData, VM,
+};
 use rand::rngs::OsRng;
 use rkyv::Deserialize;
+use std::any::Any;
 
 const OWNER: [u8; 32] = [0u8; 32];
 const LIMIT: u64 = 1_000_000;
@@ -58,10 +61,27 @@ fn verify_proof(buf: &mut [u8], len: u32) -> u32 {
     valid_bytes.len() as u32
 }
 
+struct VeryExpensiveQuery;
+
+impl HostQuery for VeryExpensiveQuery {
+    fn deserialize_and_price(
+        &self,
+        _arg_buf: &[u8],
+        _arg: &mut Box<dyn Any>,
+    ) -> u64 {
+        u64::MAX
+    }
+
+    fn execute(&self, _arg: &Box<dyn Any>, _arg_buf: &mut [u8]) -> u32 {
+        unreachable!("Query will never be executed since its price is too high")
+    }
+}
+
 fn new_ephemeral_vm() -> Result<VM, Error> {
     let mut vm = VM::ephemeral()?;
     vm.register_host_query("hash", hash);
     vm.register_host_query("verify_proof", verify_proof);
+    vm.register_host_query("very_expensive", VeryExpensiveQuery);
     Ok(vm)
 }
 
@@ -87,6 +107,27 @@ pub fn host_hash() -> Result<(), Error> {
     Ok(())
 }
 
+#[test]
+pub fn host_very_expensive_oog() -> Result<(), Error> {
+    let vm = new_ephemeral_vm()?;
+
+    let mut session = vm.session(SessionData::builder())?;
+
+    let id = session.deploy(
+        contract_bytecode!("host"),
+        ContractData::builder().owner(OWNER),
+        LIMIT,
+    )?;
+
+    let err = session
+        .call::<_, String>(id, "host_very_expensive", &(), LIMIT)
+        .expect_err("query should fail since it's too expensive");
+
+    assert!(matches!(err, Error::OutOfGas));
+
+    Ok(())
+}
+
 /// Proves that we know a number `c` such that `a + b = c`.
 #[derive(Default)]
 struct TestCircuit {