From 5b1dc876d1b22bb106f714831a8e1e3085152941 Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Tue, 23 Jul 2024 12:42:58 +0200
Subject: [PATCH] runtime-sdk/modules/rofl: Fix gas estimation in client

---
 runtime-sdk/src/modules/rofl/app/client.rs | 40 +++++++++++++---------
 1 file changed, 24 insertions(+), 16 deletions(-)

diff --git a/runtime-sdk/src/modules/rofl/app/client.rs b/runtime-sdk/src/modules/rofl/app/client.rs
index 2997b6e9b6..e6f25fb0c7 100644
--- a/runtime-sdk/src/modules/rofl/app/client.rs
+++ b/runtime-sdk/src/modules/rofl/app/client.rs
@@ -163,6 +163,7 @@ where
         let round = self.latest_round().await?;
 
         // Resolve account nonces.
+        let mut first_signer_address = Default::default();
         for (idx, signer) in signers.iter().enumerate() {
             let sigspec = SignatureAddressSpec::try_from_pk(&signer.public_key())
                 .ok_or(anyhow!("signature scheme not supported"))?;
@@ -171,25 +172,32 @@ where
 
             tx.append_auth_signature(sigspec, nonce);
 
-            // If gas is not set, perform estimation.
-            if idx == 0 && tx.fee_gas() == 0 {
-                let gas = self
-                    .estimate_gas(EstimateGasQuery {
-                        caller: if let PublicKey::Secp256k1(pk) = signer.public_key() {
-                            Some(CallerAddress::EthAddress(
-                                pk.to_eth_address().try_into().unwrap(),
-                            ))
-                        } else {
-                            Some(CallerAddress::Address(address))
-                        },
-                        tx: tx.clone(),
-                        propagate_failures: false,
-                    })
-                    .await?;
-                tx.set_fee_gas(gas);
+            // Store first signer address for gas estimation to avoid rederivation.
+            if idx == 0 {
+                first_signer_address = address;
             }
         }
 
+        // Perform gas estimation after all signer infos have been added as otherwise we may
+        // underestimate the amount of gas needed.
+        if tx.fee_gas() == 0 {
+            let signer = &signers[0]; // Checked to have at least one signer above.
+            let gas = self
+                .estimate_gas(EstimateGasQuery {
+                    caller: if let PublicKey::Secp256k1(pk) = signer.public_key() {
+                        Some(CallerAddress::EthAddress(
+                            pk.to_eth_address().try_into().unwrap(),
+                        ))
+                    } else {
+                        Some(CallerAddress::Address(first_signer_address))
+                    },
+                    tx: tx.clone(),
+                    propagate_failures: false,
+                })
+                .await?;
+            tx.set_fee_gas(gas);
+        }
+
         // Determine gas price. Currently we always use the native denomination.
         let mgp = self.gas_price(round, &token::Denomination::NATIVE).await?;
         let fee = mgp.saturating_mul(tx.fee_gas().into());