feat(enclave): multi-msgs response + gas simulation support

crates/cli/src/handler/handshake.rs

@@ -1,3 +1,5 @@
+use std::iter;
+
 use async_trait::async_trait;
 use color_eyre::{eyre::eyre, owo_colors::OwoColorize, Report, Result};
 use cw_client::{CliClient, CwClient};
@@ -55,7 +57,7 @@ async fn handshake(args: HandshakeRequest, config: Config) -> Result<String> {
                 &config.chain_id,
                 2000000,
                 &config.tx_sender,
-                json!(res),
+                iter::once(json!(res)),
                 "0untrn"
             )
             .await
@@ -105,7 +107,7 @@ async fn handshake(args: HandshakeRequest, config: Config) -> Result<String> {
                 &config.chain_id,
                 2000000,
                 &config.tx_sender,
-                json!(res),
+                iter::once(json!(res)),
                 "0untrn"
             )
             .await

crates/enclave/core/src/chain_client.rs

@@ -1,7 +1,10 @@
 use std::fmt::Display;
 
+use cosmrs::abci::GasInfo;
 use serde::{de::DeserializeOwned, Serialize};
 
+use crate::chain_client::default::DefaultTxConfig;
+
 pub mod default;
 
 /// Abstraction over a blockchain client.
@@ -19,8 +22,6 @@ pub trait ChainClient: Send + Sync + 'static {
     type Proof: Serialize + Send + Sync + 'static;
     /// The type used to represent query messages.
     type Query: Send + Sync + 'static;
-    /// The configuration type for transactions (e.g. gas fees, parameters).
-    type TxConfig: Send + Sync + 'static;
     /// The output type returned after sending a transaction.
     type TxOutput: Send + Sync + 'static;
 
@@ -63,20 +64,39 @@ pub trait ChainClient: Send + Sync + 'static {
     /// # Parameters
     ///
     /// - `contract`: A reference to the contract identifier.
-    /// - `tx`: The transaction payload, which must be serializable.
+    /// - `msgs`: The transaction messages, which must be serializable.
     /// - `config`: The transaction configuration (e.g., gas, fees).
     ///
     /// # Returns
     ///
     /// A `Result` containing the transaction output of type `Self::TxOutput` on success,
     /// or an error of type `Self::Error` if the transaction fails.
-    async fn send_tx<T: Serialize + Send + Sync>(
+    async fn send_tx<M: Serialize>(
         &self,
         contract: &Self::Contract,
-        tx: T,
-        config: Self::TxConfig,
+        msgs: impl Iterator<Item = M> + Send + Sync,
+        config: DefaultTxConfig,
     ) -> Result<Self::TxOutput, Self::Error>;
 
+    /// Simulates a transaction returning the gas_info.
+    ///
+    /// # Parameters
+    ///
+    /// - `contract`: A reference to the contract identifier.
+    /// - `msgs`: The transaction messages, which must be serializable.
+    /// - `config`: The transaction configuration (e.g., gas, fees).
+    ///
+    /// # Returns
+    ///
+    /// A `Result` containing the `GasInfo` on success,
+    /// or an error of type `Self::Error` if the transaction fails.
+    async fn simulate_tx<M: Serialize>(
+        &self,
+        contract: &Self::Contract,
+        msgs: impl Iterator<Item = M> + Send + Sync,
+        config: DefaultTxConfig,
+    ) -> Result<GasInfo, Self::Error>;
+
     /// Waits for a specified number of blocks to be produced on the blockchain.
     ///
     /// # Parameters

crates/enclave/core/src/chain_client/default.rs

@@ -1,5 +1,5 @@
 use anyhow::anyhow;
-use cosmrs::{crypto::secp256k1::SigningKey, AccountId};
+use cosmrs::{abci::GasInfo, crypto::secp256k1::SigningKey, AccountId};
 use cw_client::{CwClient, GrpcClient};
 use futures_util::StreamExt;
 use log::{debug, error, info, trace};
@@ -21,12 +21,12 @@ use crate::chain_client::ChainClient;
 ///     - websocket for waiting for blocks
 ///     - tendermint HTTP RPC for generating light client proofs
 pub struct DefaultChainClient {
-    chain_id: TmChainId,
-    grpc_client: GrpcClient,
-    node_url: Url,
-    ws_url: Url,
-    trusted_height: Height,
-    trusted_hash: Hash,
+    pub chain_id: TmChainId,
+    pub grpc_client: GrpcClient,
+    pub node_url: Url,
+    pub ws_url: Url,
+    pub trusted_height: Height,
+    pub trusted_hash: Hash,
 }
 
 impl DefaultChainClient {
@@ -74,7 +74,6 @@ impl ChainClient for DefaultChainClient {
     type Error = anyhow::Error;
     type Proof = ProofOutput;
     type Query = Query;
-    type TxConfig = DefaultTxConfig;
     type TxOutput = String;
 
     async fn query_contract<R: DeserializeOwned + Default + Send>(
@@ -129,11 +128,11 @@ impl ChainClient for DefaultChainClient {
         Ok(proof_output)
     }
 
-    async fn send_tx<T: Serialize + Send + Sync>(
+    async fn send_tx<M: Serialize>(
         &self,
         contract: &Self::Contract,
-        tx: T,
-        config: Self::TxConfig,
+        msgs: impl Iterator<Item = M> + Send + Sync,
+        config: DefaultTxConfig,
     ) -> Result<Self::TxOutput, Self::Error> {
         debug!(
             "Sending transaction to contract {contract} with gas {}",
@@ -145,7 +144,29 @@ impl ChainClient for DefaultChainClient {
                 &self.chain_id,
                 config.gas,
                 "",
-                json!(tx),
+                msgs.map(|m| json!(m)),
+                &config.amount,
+            )
+            .await
+    }
+
+    async fn simulate_tx<M: Serialize>(
+        &self,
+        contract: &Self::Contract,
+        msgs: impl Iterator<Item = M> + Send + Sync,
+        config: DefaultTxConfig,
+    ) -> Result<GasInfo, Self::Error> {
+        debug!(
+            "Simulating a transaction to contract {contract} with gas {}",
+            config.gas
+        );
+        self.grpc_client
+            .tx_simulate(
+                contract,
+                &self.chain_id,
+                config.gas,
+                "",
+                msgs.map(|m| json!(m)),
                 &config.amount,
             )
             .await
@@ -185,3 +206,26 @@ pub struct DefaultTxConfig {
     pub gas: u64,
     pub amount: String,
 }
+
+impl DefaultTxConfig {
+    pub fn new(gas_used: u64, multiplier: f64, base_gas_price: f64, denom: &str) -> Self {
+        let gas = scale_gas(gas_used, multiplier);
+        let amount_num = scale_gas(gas, base_gas_price);
+        Self {
+            gas,
+            amount: format!("{amount_num}{denom}"),
+        }
+    }
+}
+
+pub fn scale_gas(gas: u64, multiplier: f64) -> u64 {
+    if !multiplier.is_finite() || multiplier < 0.0 {
+        return gas;
+    }
+    let prod = (gas as f64) * multiplier;
+    if prod >= (u64::MAX as f64) {
+        u64::MAX
+    } else {
+        prod.ceil() as u64
+    }
+}

crates/enclave/core/src/host.rs

@@ -23,7 +23,10 @@ use tonic_health::server::health_reporter;
 
 use crate::{
     backup_restore::Backup,
-    chain_client::{default::DefaultChainClient, ChainClient},
+    chain_client::{
+        default::{DefaultChainClient, DefaultTxConfig},
+        ChainClient,
+    },
     event::QuartzEvent,
     handler::Handler,
     store::Store,
@@ -156,10 +159,11 @@ where
     C: ChainClient<Contract = AccountId, Error = anyhow::Error>,
     <C as ChainClient>::TxOutput: Display,
     R: Handler<E, Error = Status> + Debug,
-    <R as Handler<E>>::Response: Serialize + Send + Sync + 'static,
+    <R as Handler<E>>::Response: Iterator + Send + Sync,
+    <<R as Handler<E>>::Response as Iterator>::Item: Serialize + Send + Sync + 'static,
     EV: Handler<C, Response = R, Error = anyhow::Error>,
     EV: TryFrom<TmEvent, Error = anyhow::Error>,
-    GF: Fn(&<R as Handler<E>>::Response) -> <C as ChainClient>::TxConfig + Send + Sync + 'static,
+    GF: GasProvider<<R as Handler<E>>::Response, C> + Send + Sync + 'static,
 {
     type ChainClient = C;
     type Enclave = E;
@@ -280,10 +284,13 @@ where
             };
 
             // submit response to the chain
-            let tx_config = (self.gas_fn)(&response);
+            let gas_info = self
+                .gas_fn
+                .gas_for_tx(&response, &self.chain_client, &contract)
+                .await?;
             let output = self
                 .chain_client
-                .send_tx(&contract, response, tx_config)
+                .send_tx(&contract, response, gas_info)
                 .await;
             match output {
                 Ok(o) => info!("tx output: {o}"),
@@ -315,3 +322,13 @@ fn busy_wait_iters(mut iters: u64) {
         iters -= 1;
     }
 }
+
+#[async_trait::async_trait]
+pub trait GasProvider<Tx, CC> {
+    async fn gas_for_tx(
+        &self,
+        tx: &Tx,
+        chain_client: &CC,
+        contract: &AccountId,
+    ) -> Result<DefaultTxConfig, anyhow::Error>;
+}

crates/utils/cw-client/src/cli.rs

@@ -1,7 +1,7 @@
 use std::process::Command;
 
 use color_eyre::{eyre::eyre, Help, Report, Result};
-use cosmrs::{tendermint::chain::Id, AccountId};
+use cosmrs::{abci::GasInfo, tendermint::chain::Id, AccountId};
 use reqwest::Url;
 use serde::de::DeserializeOwned;
 
@@ -145,20 +145,24 @@ impl CwClient for CliClient {
         Ok(query_result)
     }
 
-    async fn tx_execute<M: ToString + Send>(
+    async fn tx_execute<M: ToString>(
         &self,
         contract: &Self::Address,
         chain_id: &Id,
         gas: u64,
         sender: &str,
-        msg: M,
+        msgs: impl Iterator<Item = M> + Send + Sync,
         pay_amount: &str,
     ) -> Result<String, Self::Error> {
         let gas_amount = match gas {
             0 => "auto",
             _ => &gas.to_string(),
         };
 
+        // only support one message for now
+        let msgs = msgs.collect::<Vec<_>>();
+        let msg = msgs.first().ok_or(eyre!("No messages provided"))?;
+
         let mut command = self.new_command()?;
         let command = command
             .args(["--node", self.url.as_str()])
@@ -183,6 +187,49 @@ impl CwClient for CliClient {
         Ok((String::from_utf8(output.stdout)?).to_string())
     }
 
+    async fn tx_simulate<M: ToString + Send + Sync>(
+        &self,
+        contract: &Self::Address,
+        chain_id: &Id,
+        gas: u64,
+        sender: &str,
+        msgs: impl Iterator<Item = M> + Send + Sync,
+        pay_amount: &str,
+    ) -> std::result::Result<GasInfo, Self::Error> {
+        let gas_amount = match gas {
+            0 => "auto",
+            _ => &gas.to_string(),
+        };
+
+        // only support one message for now
+        let msgs = msgs.collect::<Vec<_>>();
+        let msg = msgs.first().ok_or(eyre!("No messages provided"))?;
+
+        let mut command = self.new_command()?;
+        let command = command
+            .args(["--node", self.url.as_str()])
+            .args(["--chain-id", chain_id.as_ref()])
+            .args(["tx", "wasm"])
+            .args(["execute", contract.as_ref(), &msg.to_string()])
+            .args(["--amount", pay_amount])
+            .args(["--gas", gas_amount])
+            .args(["--gas-adjustment", "1.3"])
+            .args(["--gas-prices", "0.025untrn"])
+            .args(["--from", sender])
+            .args(["--output", "json"])
+            .arg("--dry-run")
+            .arg("-y");
+
+        let output = command.output()?;
+
+        if !output.status.success() {
+            return Err(eyre!("{:?}", output));
+        }
+
+        let gas_info: GasInfo = serde_json::from_slice(&output.stdout)?;
+        Ok(gas_info)
+    }
+
     fn deploy<M: ToString>(
         &self,
         chain_id: &Id,