forge script Avoid bulk gasEstimation when setting --skip-simulation (#2601)

* Extend test for --skip-simulation script call

* fix: Only estimate gas immediately before tx submission when  is set

Co-authored-by: Matthias Seitz <[email protected]>

Author: ckoopmann

cli/src/cmd/forge/script/broadcast.rs

@@ -9,12 +9,12 @@ use crate::{
     update_progress,
 };
 use ethers::{
-    prelude::{Signer, SignerMiddleware, TxHash},
-    providers::Middleware,
+    prelude::{Provider, Signer, SignerMiddleware, TxHash},
+    providers::{JsonRpcClient, Middleware},
     types::transaction::eip2718::TypedTransaction,
     utils::format_units,
 };
-use eyre::ContextCompat;
+use eyre::{ContextCompat, WrapErr};
 use foundry_common::{get_http_provider, RetryProvider};
 use foundry_config::Chain;
 use futures::StreamExt;
@@ -188,9 +188,9 @@ impl ScriptArgs {
         }
 
         match signer {
-            WalletType::Local(signer) => broadcast(signer, tx).await,
-            WalletType::Ledger(signer) => broadcast(signer, tx).await,
-            WalletType::Trezor(signer) => broadcast(signer, tx).await,
+            WalletType::Local(signer) => self.broadcast(signer, tx).await,
+            WalletType::Ledger(signer) => self.broadcast(signer, tx).await,
+            WalletType::Trezor(signer) => self.broadcast(signer, tx).await,
         }
     }
 
@@ -277,43 +277,106 @@ impl ScriptArgs {
         for mut tx in txes.into_iter() {
             tx.change_type(is_legacy);
 
-            let typed_tx = tx.typed_tx_mut();
+            if !self.skip_simulation {
+                let typed_tx = tx.typed_tx_mut();
 
-            if has_different_gas_calc(chain) || self.skip_simulation {
-                typed_tx.set_gas(
-                    provider.estimate_gas(typed_tx).await? * self.gas_estimate_multiplier / 100,
-                );
-            }
+                if has_different_gas_calc(chain) {
+                    self.estimate_gas(typed_tx, &provider).await?;
+                }
 
-            total_gas += *typed_tx.gas().expect("gas is set");
+                total_gas += *typed_tx.gas().expect("gas is set");
+            }
 
             new_txes.push_back(tx);
         }
 
-        // We don't store it in the transactions, since we want the most updated value. Right before
-        // broadcasting.
-        let per_gas = if let Some(gas_price) = self.with_gas_price {
-            gas_price
-        } else {
-            match new_txes.front().unwrap().typed_tx() {
-                TypedTransaction::Legacy(_) | TypedTransaction::Eip2930(_) => {
-                    provider.get_gas_price().await?
+        if !self.skip_simulation {
+            // We don't store it in the transactions, since we want the most updated value. Right
+            // before broadcasting.
+            let per_gas = if let Some(gas_price) = self.with_gas_price {
+                gas_price
+            } else {
+                match new_txes.front().unwrap().typed_tx() {
+                    TypedTransaction::Legacy(_) | TypedTransaction::Eip2930(_) => {
+                        provider.get_gas_price().await?
+                    }
+                    TypedTransaction::Eip1559(_) => provider.estimate_eip1559_fees(None).await?.0,
                 }
-                TypedTransaction::Eip1559(_) => provider.estimate_eip1559_fees(None).await?.0,
-            }
-        };
+            };
 
-        println!("\n==========================");
-        println!("\nEstimated total gas used for script: {}", total_gas);
-        println!(
-            "\nEstimated amount required: {} ETH",
-            format_units(total_gas.saturating_mul(per_gas), 18)
-                .unwrap_or_else(|_| "[Could not calculate]".to_string())
-                .trim_end_matches('0')
-        );
-        println!("\n==========================");
+            println!("\n==========================");
+            println!("\nEstimated total gas used for script: {}", total_gas);
+            println!(
+                "\nEstimated amount required: {} ETH",
+                format_units(total_gas.saturating_mul(per_gas), 18)
+                    .unwrap_or_else(|_| "[Could not calculate]".to_string())
+                    .trim_end_matches('0')
+            );
+            println!("\n==========================");
+        }
         Ok(new_txes)
     }
+    /// Uses the signer to submit a transaction to the network. If it fails, it tries to retrieve
+    /// the transaction hash that can be used on a later run with `--resume`.
+    async fn broadcast<T, U>(
+        &self,
+        signer: &SignerMiddleware<T, U>,
+        mut legacy_or_1559: TypedTransaction,
+    ) -> Result<TxHash, BroadcastError>
+    where
+        T: Middleware,
+        U: Signer,
+    {
+        tracing::debug!("sending transaction: {:?}", legacy_or_1559);
+
+        // Chains which use `eth_estimateGas` are being sent sequentially and require their gas to
+        // be re-estimated right before broadcasting.
+        if has_different_gas_calc(signer.signer().chain_id()) || self.skip_simulation {
+            // if already set, some RPC endpoints might simply return the gas value that is already
+            // set in the request and omit the estimate altogether, so we remove it here
+            let _ = legacy_or_1559.gas_mut().take();
+
+            self.estimate_gas(&mut legacy_or_1559, signer.provider()).await?;
+        }
+
+        // Signing manually so we skip `fill_transaction` and its `eth_createAccessList` request.
+        let signature = signer
+            .sign_transaction(
+                &legacy_or_1559,
+                *legacy_or_1559.from().expect("Tx should have a `from`."),
+            )
+            .await
+            .map_err(|err| BroadcastError::Simple(err.to_string()))?;
+
+        // Submit the raw transaction
+        let pending = signer
+            .provider()
+            .send_raw_transaction(legacy_or_1559.rlp_signed(&signature))
+            .await
+            .map_err(|err| BroadcastError::Simple(err.to_string()))?;
+
+        Ok(pending.tx_hash())
+    }
+
+    async fn estimate_gas<T>(
+        &self,
+        tx: &mut TypedTransaction,
+        provider: &Provider<T>,
+    ) -> Result<(), BroadcastError>
+    where
+        T: JsonRpcClient,
+    {
+        tx.set_gas(
+            provider
+                .estimate_gas(tx)
+                .await
+                .wrap_err_with(|| format!("Failed to estimate gas for tx: {}", tx.sighash()))
+                .map_err(|err| BroadcastError::Simple(err.to_string()))? *
+                self.gas_estimate_multiplier /
+                100,
+        );
+        Ok(())
+    }
 }
 
 #[derive(thiserror::Error, Debug, Clone)]
@@ -332,50 +395,3 @@ impl fmt::Display for BroadcastError {
         }
     }
 }
-
-/// Uses the signer to submit a transaction to the network. If it fails, it tries to retrieve the
-/// transaction hash that can be used on a later run with `--resume`.
-async fn broadcast<T, U>(
-    signer: &SignerMiddleware<T, U>,
-    mut legacy_or_1559: TypedTransaction,
-) -> Result<TxHash, BroadcastError>
-where
-    T: Middleware,
-    U: Signer,
-{
-    tracing::debug!("sending transaction: {:?}", legacy_or_1559);
-
-    // Chains which use `eth_estimateGas` are being sent sequentially and require their gas to be
-    // re-estimated right before broadcasting.
-    if has_different_gas_calc(signer.signer().chain_id()) {
-        // if already set, some RPC endpoints might simply return the gas value that is already set
-        // in the request and omit the estimate altogether, so we remove it here
-        let _ = legacy_or_1559.gas_mut().take();
-
-        legacy_or_1559.set_gas(
-            signer
-                .provider()
-                .estimate_gas(&legacy_or_1559)
-                .await
-                .map_err(|err| BroadcastError::Simple(err.to_string()))?,
-        );
-    }
-
-    // Signing manually so we skip `fill_transaction` and its `eth_createAccessList` request.
-    let signature = signer
-        .sign_transaction(
-            &legacy_or_1559,
-            *legacy_or_1559.from().expect("Tx should have a `from`."),
-        )
-        .await
-        .map_err(|err| BroadcastError::Simple(err.to_string()))?;
-
-    // Submit the raw transaction
-    let pending = signer
-        .provider()
-        .send_raw_transaction(legacy_or_1559.rlp_signed(&signature))
-        .await
-        .map_err(|err| BroadcastError::Simple(err.to_string()))?;
-
-    Ok(pending.tx_hash())
-}

cli/tests/it/script.rs

@@ -1,5 +1,6 @@
 //! Contains various tests related to forge script
 use anvil::{spawn, NodeConfig};
+use cast::SimpleCast;
 use ethers::abi::Address;
 use foundry_cli_test_utils::{
     forgetest, forgetest_async, forgetest_init,
@@ -8,6 +9,7 @@ use foundry_cli_test_utils::{
 };
 use foundry_utils::rpc;
 use regex::Regex;
+use serde_json::Value;
 use std::{env, path::PathBuf, str::FromStr};
 
 // Tests that fork cheat codes can be used in script
@@ -180,10 +182,10 @@ forgetest_async!(
     |prj: TestProject, mut cmd: TestCommand| async move {
         foundry_cli_test_utils::util::initialize(prj.root());
         // This example script would fail in on-chain simulation
-        let script = prj
+        let deploy_script = prj
             .inner()
             .add_source(
-                "Foo",
+                "DeployScript",
                 r#"
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity 0.8.10;
@@ -196,36 +198,90 @@ contract HashChecker {
         require(newHash != lastHash, "Hash didn't change");
         lastHash = newHash;
     }
+
+    function checkLastHash() public {
+        require(lastHash != bytes32(0),  "Hash shouldn't be zero");
+    }
 }
-contract Demo is Script {
+contract DeployScript is Script {
     function run() external returns (uint256 result, uint8) {
         vm.startBroadcast();
         HashChecker hashChecker = new HashChecker();
-        uint numUpdates = 8;
-        vm.roll(block.number - numUpdates);
-        for(uint i = 0; i < numUpdates; i++) {
-            vm.roll(block.number + 1);
-            hashChecker.update();
-        }
     }
 }"#,
             )
             .unwrap();
 
+        let deploy_contract = deploy_script.display().to_string() + ":DeployScript";
+
         let node_config = NodeConfig::test()
             .with_eth_rpc_url(Some(rpc::next_http_archive_rpc_endpoint()))
             .silent();
-
         let (_api, handle) = spawn(node_config).await;
-        let target_contract = script.display().to_string() + ":Demo";
-
-        let wallet = handle.dev_wallets().next().unwrap();
-        let private_key = hex::encode(wallet.signer().to_bytes());
+        let private_key =
+            "ac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80".to_string();
         cmd.set_current_dir(prj.root());
 
         cmd.args([
             "script",
-            &target_contract,
+            &deploy_contract,
+            "--root",
+            prj.root().to_str().unwrap(),
+            "--fork-url",
+            &handle.http_endpoint(),
+            "-vvvvv",
+            "--broadcast",
+            "--slow",
+            "--skip-simulation",
+            "--private-key",
+            &private_key,
+        ]);
+
+        let output = cmd.stdout_lossy();
+
+        assert!(output.contains("SKIPPING ON CHAIN SIMULATION"));
+        assert!(output.contains("ONCHAIN EXECUTION COMPLETE & SUCCESSFUL"));
+
+        let run_log =
+            std::fs::read_to_string("broadcast/DeployScript.sol/1/run-latest.json").unwrap();
+        let run_object: Value = serde_json::from_str(&run_log).unwrap();
+        let contract_address = SimpleCast::checksum_address(
+            &ethers::prelude::H160::from_str(
+                run_object["receipts"][0]["contractAddress"].as_str().unwrap(),
+            )
+            .unwrap(),
+        )
+        .unwrap();
+
+        let run_code = r#"
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity 0.8.10;
+import "forge-std/Script.sol";
+import { HashChecker } from "./DeployScript.sol";
+
+contract RunScript is Script {
+    function run() external returns (uint256 result, uint8) {
+        vm.startBroadcast();
+        HashChecker hashChecker = HashChecker(CONTRACT_ADDRESS);
+        uint numUpdates = 8;
+        vm.roll(block.number - numUpdates);
+        for(uint i = 0; i < numUpdates; i++) {
+            vm.roll(block.number + 1);
+            hashChecker.update();
+            hashChecker.checkLastHash();
+        }
+    }
+}"#
+        .replace("CONTRACT_ADDRESS", &contract_address);
+
+        let run_script = prj.inner().add_source("RunScript", run_code).unwrap();
+        let run_contract = run_script.display().to_string() + ":RunScript";
+
+        cmd.forge_fuse();
+        cmd.set_current_dir(prj.root());
+        cmd.args([
+            "script",
+            &run_contract,
             "--root",
             prj.root().to_str().unwrap(),
             "--fork-url",
@@ -241,7 +297,6 @@ contract Demo is Script {
         ]);
 
         let output = cmd.stdout_lossy();
-
         assert!(output.contains("SKIPPING ON CHAIN SIMULATION"));
         assert!(output.contains("ONCHAIN EXECUTION COMPLETE & SUCCESSFUL"));
     }