fix(kms-connector): increase gas_limit (#258)

* fix(kms-connector): add gas_limit

* fix(kms-connector): increase gas limit

* fix(kms-connector): infinite gas limit + add tx retry

* fix(kms-connector): add more logs

Author: eudelins-zama

kms-connector/README.md

@@ -156,14 +156,14 @@ The KMS Connector supports flexible configuration through both TOML files and en
    # Set required configuration
    export KMS_CONNECTOR_GATEWAY_URL="ws://localhost:8547"
    export KMS_CONNECTOR_KMS_CORE_ENDPOINT="http://localhost:50052"
-   
+
    # Wallet configuration (one of the following is required)
    export KMS_CONNECTOR_PRIVATE_KEY="0x0000000000000000000000000000000000000000000000000000000000000001"
    # OR for AWS KMS
    export KMS_CONNECTOR_AWS_KMS_CONFIG__KEY_ID="alias/my-kms-key"
    export KMS_CONNECTOR_AWS_KMS_CONFIG__REGION="us-east-1"
    export KMS_CONNECTOR_AWS_KMS_CONFIG__ENDPOINT="http://localhost:4566"
-   
+
    export KMS_CONNECTOR_CHAIN_ID="31337"
    export KMS_CONNECTOR_DECRYPTION_ADDRESS="0x..."
    export KMS_CONNECTOR_GATEWAY_CONFIG_ADDRESS="0x..."
@@ -478,4 +478,4 @@ cargo test
 
 🌟 If you find this project helpful or interesting, please consider giving it a star on GitHub! Your support helps to grow the community and motivates further development.
 
-[![GitHub stars](https://img.shields.io/github/stars/zama-ai/fhevm?style=social)](https://github.com/zama-ai/fhevm/)
+[![GitHub stars](https://img.shields.io/github/stars/zama-ai/fhevm?style=social)](https://github.com/zama-ai/fhevm/)

kms-connector/simple-connector/src/core/connector.rs

@@ -37,6 +37,8 @@ impl<P: Provider + Clone + 'static> KmsCoreConnector<P> {
             config.gateway_config_address,
             event_tx,
         );
+
+        // Possible gas limit
         let decryption = DecryptionAdapter::new(config.decryption_address, provider.clone());
 
         let decryption_handler =

kms-connector/simple-connector/src/gw_adapters/decryption/adapter.rs

@@ -1,11 +1,19 @@
 use crate::error::{Error, Result};
 use alloy::{
+    network::Ethereum,
     primitives::{Address, Bytes, U256},
-    providers::Provider,
+    providers::{PendingTransactionBuilder, Provider},
+    rpc::types::TransactionRequest,
 };
 use fhevm_gateway_rust_bindings::decryption::Decryption;
-use std::sync::Arc;
-use tracing::{debug, info};
+use std::{sync::Arc, time::Duration};
+use tracing::{debug, info, warn};
+
+/// The max value for the `gas_limit` of a transaction.
+const INFINITE_GAS_LIMIT: u64 = u64::MAX;
+
+/// The time to wait between two transactions attempt.
+const TX_INTERVAL: Duration = Duration::from_secs(3);
 
 /// Adapter for decryption operations
 #[derive(Clone)]
@@ -57,18 +65,20 @@ impl<P: Provider + Clone> DecryptionAdapter<P> {
 
         let contract = Decryption::new(self.decryption_address, self.provider.clone());
 
-        // Create and send transaction
-        let call = contract.publicDecryptionResponse(id, result, signature.into());
-        let tx = call
-            .send()
-            .await
-            .map_err(|e| Error::Contract(e.to_string()))?;
+        let call_builder = contract.publicDecryptionResponse(id, result, signature.into());
+        info!(decryption_id = ?id, "public decryption calldata length {}", call_builder.calldata().len());
+
+        let mut call = call_builder.into_transaction_request();
+        self.estimate_gas(id, &mut call).await;
+        let tx = self.send_tx_with_retry(call).await?;
+
         // TODO: optimize for low latency
         let receipt = tx
             .get_receipt()
             .await
             .map_err(|e| Error::Contract(e.to_string()))?;
         info!(decryption_id = ?id, "🎯 Public Decryption response sent with tx receipt: {:?}", receipt);
+        info!(decryption_id = ?id, "⛽ Gas consumed for Public Decryption: {}", receipt.gas_used);
         Ok(())
     }
 
@@ -102,17 +112,58 @@ impl<P: Provider + Clone> DecryptionAdapter<P> {
         let contract = Decryption::new(self.decryption_address, self.provider.clone());
 
         // Create and send transaction
-        let call = contract.userDecryptionResponse(id, result, signature.into());
-        let tx = call
-            .send()
-            .await
-            .map_err(|e| Error::Contract(e.to_string()))?;
+        let call_builder = contract.userDecryptionResponse(id, result, signature.into());
+        info!(decryption_id = ?id, "user decryption calldata length {}", call_builder.calldata().len());
+
+        let mut call = call_builder.into_transaction_request();
+        self.estimate_gas(id, &mut call).await;
+        let tx = self.send_tx_with_retry(call).await?;
+
         // TODO: optimize for low latency
         let receipt = tx
             .get_receipt()
             .await
             .map_err(|e| Error::Contract(e.to_string()))?;
         info!(decryption_id = ?id, "🎯 User Decryption response sent with tx receipt: {:?}", receipt);
+        info!(decryption_id = ?id, "⛽ Gas consumed for User Decryption: {}", receipt.gas_used);
         Ok(())
     }
+
+    /// Estimates the `gas_limit` for the upcoming transaction.
+    async fn estimate_gas(&self, id: U256, call: &mut TransactionRequest) {
+        match self.provider.estimate_gas(call.clone()).await {
+            Ok(gas) => info!(decryption_id = ?id, "Initial gas estimation for the tx: {gas}"),
+            Err(e) => warn!(decryption_id = ?id, "Failed to estimate gas for the tx: {e}"),
+        }
+
+        // TODO: temporary workaround for out-of-gas errors
+        // Our automatic estimation fails during gas pikes.
+        // (see https://zama-ai.slack.com/archives/C0915Q59CKG/p1749843623276629?thread_ts=1749828466.079719&cid=C0915Q59CKG)
+        info!(decryption_id = ?id, "Updating `gas_limit` to max value");
+        call.gas = Some(INFINITE_GAS_LIMIT);
+    }
+
+    /// Sends the requested transactions with one retry.
+    async fn send_tx_with_retry(
+        &self,
+        call: TransactionRequest,
+    ) -> Result<PendingTransactionBuilder<Ethereum>> {
+        match self.provider.send_transaction(call.clone()).await {
+            Ok(tx) => Ok(tx),
+            Err(e) => {
+                warn!(
+                    "Retrying to send transaction in {}s after failure: {}",
+                    TX_INTERVAL.as_secs(),
+                    e
+                );
+
+                tokio::time::sleep(TX_INTERVAL).await;
+
+                self.provider
+                    .send_transaction(call)
+                    .await
+                    .map_err(|e| Error::Contract(e.to_string()))
+            }
+        }
+    }
 }