feat(zk-toolbox): add balance check

zk_toolbox/crates/common/src/ethereum.rs

@@ -1,25 +1,33 @@
 use std::{ops::Add, time::Duration};
 
 use ethers::{
+    core::k256::ecdsa::SigningKey,
     middleware::MiddlewareBuilder,
-    prelude::{Http, LocalWallet, Provider, Signer},
+    prelude::{Http, LocalWallet, Provider},
+    prelude::{SignerMiddleware, H256},
     providers::Middleware,
     types::{Address, TransactionRequest},
 };
 
 use crate::wallets::Wallet;
 
+pub fn create_ethers_client(
+    private_key: H256,
+    l1_rpc: String,
+) -> anyhow::Result<SignerMiddleware<Provider<Http>, ethers::prelude::Wallet<SigningKey>>> {
+    let wallet = LocalWallet::from_bytes(private_key.as_bytes())?;
+    let client = Provider::<Http>::try_from(l1_rpc)?.with_signer(wallet);
+    Ok(client)
+}
+
 pub async fn distribute_eth(
     main_wallet: Wallet,
     addresses: Vec<Address>,
     l1_rpc: String,
     chain_id: u32,
     amount: u128,
 ) -> anyhow::Result<()> {
-    let wallet = LocalWallet::from_bytes(main_wallet.private_key.unwrap().as_bytes())?
-        .with_chain_id(chain_id);
-    let client = Provider::<Http>::try_from(l1_rpc)?.with_signer(wallet);
-
+    let client = create_ethers_client(main_wallet.private_key.unwrap(), l1_rpc)?;
     let mut pending_txs = vec![];
     let mut nonce = client.get_transaction_count(client.address(), None).await?;
     for address in addresses {

zk_toolbox/crates/common/src/forge.rs

@@ -1,12 +1,17 @@
 use std::path::{Path, PathBuf};
+use std::str::FromStr;
 
 use clap::Parser;
+use ethers::abi::Address;
+use ethers::middleware::Middleware;
+use ethers::prelude::{LocalWallet, Signer, U256};
 use ethers::{abi::AbiEncode, types::H256};
 use serde::{Deserialize, Serialize};
 use strum_macros::Display;
 use xshell::{cmd, Shell};
 
 use crate::cmd::Cmd;
+use crate::ethereum::create_ethers_client;
 
 /// Forge is a wrapper around the forge binary.
 pub struct Forge {
@@ -93,6 +98,43 @@ impl ForgeScript {
         });
         self
     }
+    // Do not start the script if balance is not enough
+    pub fn private_key(&self) -> Option<H256> {
+        self.args.args.iter().find_map(|a| {
+            if let ForgeScriptArg::PrivateKey { private_key } = a {
+                Some(H256::from_str(private_key).unwrap())
+            } else {
+                None
+            }
+        })
+    }
+
+    pub fn rpc_url(&self) -> Option<String> {
+        self.args.args.iter().find_map(|a| {
+            if let ForgeScriptArg::RpcUrl { url } = a {
+                Some(url.clone())
+            } else {
+                None
+            }
+        })
+    }
+
+    pub fn address(&self) -> Option<Address> {
+        self.private_key()
+            .flat_map(|a| LocalWallet::from_bytes(a.as_bytes()).map(|a| a.address()))
+    }
+
+    pub async fn check_the_balance(&self, minimum_value: U256) -> anyhow::Result<bool> {
+        let Some(rpc_url) = self.rpc_url() else {
+            return Ok(true)
+        };
+        let Some(private_key) = self.private_key() else {
+            return Ok(true)
+        };
+        let client = create_ethers_client(private_key, rpc_url)?;
+        let balance = client.get_balance(client.address(), None).await?;
+        Ok(balance > minimum_value))
+    }
 }
 
 const PROHIBITED_ARGS: [&str; 10] = [

zk_toolbox/crates/common/src/prompt/confirm.rs

@@ -1,11 +1,12 @@
 use cliclack::Confirm;
+use std::fmt::Display;
 
 pub struct PromptConfirm {
     inner: Confirm,
 }
 
 impl PromptConfirm {
-    pub fn new(question: &str) -> Self {
+    pub fn new(question: impl Display) -> Self {
         Self {
             inner: Confirm::new(question),
         }

zk_toolbox/crates/zk_inception/src/accept_ownership.rs

@@ -5,6 +5,7 @@ use common::{
 use ethers::{abi::Address, types::H256};
 use xshell::Shell;
 
+use crate::forge_utils::check_the_balance;
 use crate::{
     configs::{
         forge_interface::accept_ownership::AcceptOwnershipInput, EcosystemConfig, SaveConfig,
@@ -13,7 +14,7 @@ use crate::{
     forge_utils::fill_forge_private_key,
 };
 
-pub fn accept_admin(
+pub async fn accept_admin(
     shell: &Shell,
     ecosystem_config: &EcosystemConfig,
     governor_contract: Address,
@@ -36,9 +37,10 @@ pub fn accept_admin(
         target_address,
         forge,
     )
+    .await
 }
 
-pub fn accept_owner(
+pub async fn accept_owner(
     shell: &Shell,
     ecosystem_config: &EcosystemConfig,
     governor_contract: Address,
@@ -61,9 +63,10 @@ pub fn accept_owner(
         target_address,
         forge,
     )
+    .await
 }
 
-fn accept_ownership(
+async fn accept_ownership(
     shell: &Shell,
     ecosystem_config: &EcosystemConfig,
     governor_contract: Address,
@@ -82,6 +85,7 @@ fn accept_ownership(
 
     forge = fill_forge_private_key(forge, governor)?;
 
+    check_the_balance(&forge).await?;
     let spinner = Spinner::new("Accepting governance");
     forge.run(shell)?;
     spinner.finish();

zk_toolbox/crates/zk_inception/src/commands/chain/deploy_paymaster.rs

@@ -6,6 +6,7 @@ use common::{
 };
 use xshell::Shell;
 
+use crate::forge_utils::check_the_balance;
 use crate::{
     configs::{
         forge_interface::paymaster::{DeployPaymasterInput, DeployPaymasterOutput},
@@ -15,16 +16,16 @@ use crate::{
     forge_utils::fill_forge_private_key,
 };
 
-pub fn run(args: ForgeScriptArgs, shell: &Shell) -> anyhow::Result<()> {
+pub async fn run(args: ForgeScriptArgs, shell: &Shell) -> anyhow::Result<()> {
     let chain_name = global_config().chain_name.clone();
     let ecosystem_config = EcosystemConfig::from_file(shell)?;
     let chain_config = ecosystem_config
         .load_chain(chain_name)
         .context("Chain not initialized. Please create a chain first")?;
-    deploy_paymaster(shell, &chain_config, &ecosystem_config, args)
+    deploy_paymaster(shell, &chain_config, &ecosystem_config, args).await
 }
 
-pub fn deploy_paymaster(
+pub async fn deploy_paymaster(
     shell: &Shell,
     chain_config: &ChainConfig,
     ecosystem_config: &EcosystemConfig,
@@ -46,6 +47,7 @@ pub fn deploy_paymaster(
     )?;
 
     let spinner = Spinner::new("Deploying paymaster");
+    check_the_balance(&forge).await?;
     forge.run(shell)?;
     spinner.finish();
 

zk_toolbox/crates/zk_inception/src/commands/chain/init.rs

@@ -8,6 +8,7 @@ use common::{
 use xshell::Shell;
 
 use super::args::init::InitArgsFinal;
+use crate::forge_utils::check_the_balance;
 use crate::{
     accept_ownership::accept_admin,
     commands::chain::{
@@ -72,23 +73,26 @@ pub async fn init(
         chain_config.get_wallets_config()?.governor_private_key(),
         contracts_config.l1.diamond_proxy_addr,
         &init_args.forge_args.clone(),
-    )?;
+    )
+    .await?;
     spinner.finish();
 
     initialize_bridges::initialize_bridges(
         shell,
         chain_config,
         ecosystem_config,
         init_args.forge_args.clone(),
-    )?;
+    )
+    .await?;
 
     if init_args.deploy_paymaster {
         deploy_paymaster::deploy_paymaster(
             shell,
             chain_config,
             ecosystem_config,
             init_args.forge_args.clone(),
-        )?;
+        )
+        .await?;
     }
 
     genesis(
@@ -124,7 +128,7 @@ async fn register_chain(
         .with_broadcast();
 
     forge = fill_forge_private_key(forge, config.get_wallets()?.governor_private_key())?;
-
+    check_the_balance(&forge).await?;
     forge.run(shell)?;
 
     let register_chain_output =

zk_toolbox/crates/zk_inception/src/commands/chain/initialize_bridges.rs

@@ -9,6 +9,7 @@ use common::{
 };
 use xshell::{cmd, Shell};
 
+use crate::forge_utils::check_the_balance;
 use crate::{
     configs::{
         forge_interface::initialize_bridges::{
@@ -20,21 +21,21 @@ use crate::{
     forge_utils::fill_forge_private_key,
 };
 
-pub fn run(args: ForgeScriptArgs, shell: &Shell) -> anyhow::Result<()> {
+pub async fn run(args: ForgeScriptArgs, shell: &Shell) -> anyhow::Result<()> {
     let chain_name = global_config().chain_name.clone();
     let ecosystem_config = EcosystemConfig::from_file(shell)?;
     let chain_config = ecosystem_config
         .load_chain(chain_name)
         .context("Chain not initialized. Please create a chain first")?;
 
     let spinner = Spinner::new("Initializing bridges");
-    initialize_bridges(shell, &chain_config, &ecosystem_config, args)?;
+    initialize_bridges(shell, &chain_config, &ecosystem_config, args).await?;
     spinner.finish();
 
     Ok(())
 }
 
-pub fn initialize_bridges(
+pub async fn initialize_bridges(
     shell: &Shell,
     chain_config: &ChainConfig,
     ecosystem_config: &EcosystemConfig,
@@ -56,6 +57,7 @@ pub fn initialize_bridges(
         ecosystem_config.get_wallets()?.governor_private_key(),
     )?;
 
+    check_the_balance(&forge).await?;
     forge.run(shell)?;
 
     let output =

zk_toolbox/crates/zk_inception/src/commands/chain/mod.rs

@@ -32,7 +32,7 @@ pub(crate) async fn run(shell: &Shell, args: ChainCommands) -> anyhow::Result<()
         ChainCommands::Create(args) => create::run(args, shell),
         ChainCommands::Init(args) => init::run(args, shell).await,
         ChainCommands::Genesis(args) => genesis::run(args, shell).await,
-        ChainCommands::InitializeBridges(args) => initialize_bridges::run(args, shell),
-        ChainCommands::DeployPaymaster(args) => deploy_paymaster::run(args, shell),
+        ChainCommands::InitializeBridges(args) => initialize_bridges::run(args, shell).await,
+        ChainCommands::DeployPaymaster(args) => deploy_paymaster::run(args, shell).await,
     }
 }