refactor: wallet management (#7141)

* refactor foundry-wallets

* Use MultiWallet in cheats

* Add comments

* clippy + fmt

* maybe_load_private_key

* fmt

* fix ci

* clippy

* refactor

* Wallet -> WalletOpts

* rm blank lines

* Review fixes

* fix comment

* comments

* review fixes

* clippy

* fixes

* fmt

crates/cast/bin/cmd/send.rs

@@ -1,7 +1,7 @@
 use cast::{Cast, TxBuilder};
 use clap::Parser;
 use ethers_core::types::NameOrAddress;
-use ethers_middleware::MiddlewareBuilder;
+use ethers_middleware::SignerMiddleware;
 use ethers_providers::Middleware;
 use ethers_signers::Signer;
 use eyre::Result;
@@ -170,7 +170,7 @@ impl SendTxArgs {
         // enough information to sign and we must bail.
         } else {
             // Retrieve the signer, and bail if it can't be constructed.
-            let signer = eth.wallet.signer(chain.id()).await?;
+            let signer = eth.wallet.signer().await?;
             let from = signer.address();
 
             // prevent misconfigured hwlib from sending a transaction that defies
@@ -191,7 +191,7 @@ corresponds to the sender, or let foundry automatically detect it by not specify
                 tx.nonce = Some(provider.get_transaction_count(from, None).await?.to_alloy());
             }
 
-            let provider = provider.with_signer(signer);
+            let provider = SignerMiddleware::new_with_provider_chain(provider, signer).await?;
 
             cast_send(
                 provider,

crates/cast/bin/cmd/wallet/mod.rs

@@ -9,7 +9,7 @@ use ethers_signers::Signer;
 use eyre::{Context, Result};
 use foundry_common::{fs, types::ToAlloy};
 use foundry_config::Config;
-use foundry_wallets::{RawWallet, Wallet};
+use foundry_wallets::{RawWalletOpts, WalletOpts, WalletSigner};
 use rand::thread_rng;
 use serde_json::json;
 use std::{path::Path, str::FromStr};
@@ -72,7 +72,7 @@ pub enum WalletSubcommands {
         private_key_override: Option<String>,
 
         #[clap(flatten)]
-        wallet: Wallet,
+        wallet: WalletOpts,
     },
 
     /// Sign a message or typed data.
@@ -106,7 +106,7 @@ pub enum WalletSubcommands {
         no_hash: bool,
 
         #[clap(flatten)]
-        wallet: Wallet,
+        wallet: WalletOpts,
     },
 
     /// Verify the signature of a message.
@@ -133,7 +133,7 @@ pub enum WalletSubcommands {
         #[clap(long, short)]
         keystore_dir: Option<String>,
         #[clap(flatten)]
-        raw_wallet_options: RawWallet,
+        raw_wallet_options: RawWalletOpts,
     },
     /// List all the accounts in the keystore default directory
     #[clap(visible_alias = "ls")]
@@ -241,18 +241,18 @@ impl WalletSubcommands {
             }
             WalletSubcommands::Address { wallet, private_key_override } => {
                 let wallet = private_key_override
-                    .map(|pk| Wallet {
-                        raw: RawWallet { private_key: Some(pk), ..Default::default() },
+                    .map(|pk| WalletOpts {
+                        raw: RawWalletOpts { private_key: Some(pk), ..Default::default() },
                         ..Default::default()
                     })
                     .unwrap_or(wallet)
-                    .signer(0)
+                    .signer()
                     .await?;
                 let addr = wallet.address();
                 println!("{}", addr.to_alloy().to_checksum(None));
             }
             WalletSubcommands::Sign { message, data, from_file, no_hash, wallet } => {
-                let wallet = wallet.signer(0).await?;
+                let wallet = wallet.signer().await?;
                 let sig = if data {
                     let typed_data: TypedData = if from_file {
                         // data is a file name, read json from file
@@ -297,16 +297,21 @@ impl WalletSubcommands {
                 }
 
                 // get wallet
-                let wallet: Wallet = raw_wallet_options.into();
-                let wallet = wallet.try_resolve_local_wallet()?.ok_or_else(|| {
-                    eyre::eyre!(
-                        "\
+                let wallet = raw_wallet_options
+                    .signer()?
+                    .and_then(|s| match s {
+                        WalletSigner::Local(s) => Some(s),
+                        _ => None,
+                    })
+                    .ok_or_else(|| {
+                        eyre::eyre!(
+                            "\
 Did you set a private key or mnemonic?
 Run `cast wallet import --help` and use the corresponding CLI
 flag to set your key via:
 --private-key, --mnemonic-path or --interactive."
-                    )
-                })?;
+                        )
+                    })?;
 
                 let private_key = wallet.signer().to_bytes();
                 let password = rpassword::prompt_password("Enter password: ")?;

crates/cheatcodes/Cargo.toml

@@ -17,6 +17,7 @@ foundry-common.workspace = true
 foundry-compilers.workspace = true
 foundry-config.workspace = true
 foundry-evm-core.workspace = true
+foundry-wallets.workspace = true
 
 alloy-dyn-abi.workspace = true
 alloy-json-abi.workspace = true
@@ -26,6 +27,7 @@ alloy-sol-types.workspace = true
 alloy-providers.workspace = true
 alloy-rpc-types.workspace = true
 alloy-signer = { workspace = true, features = ["mnemonic", "keystore"] }
+parking_lot = "0.12"
 
 
 eyre.workspace = true

crates/cheatcodes/src/config.rs

@@ -1,5 +1,5 @@
 use super::Result;
-use crate::Vm::Rpc;
+use crate::{script::ScriptWallets, Vm::Rpc};
 use alloy_primitives::Address;
 use foundry_common::fs::normalize_path;
 use foundry_compilers::{utils::canonicalize, ProjectPathsConfig};
@@ -36,11 +36,13 @@ pub struct CheatsConfig {
     pub evm_opts: EvmOpts,
     /// Address labels from config
     pub labels: HashMap<Address, String>,
+    /// Script wallets
+    pub script_wallets: Option<ScriptWallets>,
 }
 
 impl CheatsConfig {
     /// Extracts the necessary settings from the Config
-    pub fn new(config: &Config, evm_opts: EvmOpts) -> Self {
+    pub fn new(config: &Config, evm_opts: EvmOpts, script_wallets: Option<ScriptWallets>) -> Self {
         let mut allowed_paths = vec![config.__root.0.clone()];
         allowed_paths.extend(config.libs.clone());
         allowed_paths.extend(config.allow_paths.clone());
@@ -58,6 +60,7 @@ impl CheatsConfig {
             allowed_paths,
             evm_opts,
             labels: config.labels.clone(),
+            script_wallets,
         }
     }
 
@@ -172,6 +175,7 @@ impl Default for CheatsConfig {
             allowed_paths: vec![],
             evm_opts: Default::default(),
             labels: Default::default(),
+            script_wallets: None,
         }
     }
 }
@@ -185,6 +189,7 @@ mod tests {
         CheatsConfig::new(
             &Config { __root: PathBuf::from(root).into(), fs_permissions, ..Default::default() },
             Default::default(),
+            None,
         )
     }
 

crates/cheatcodes/src/error.rs

@@ -5,6 +5,7 @@ use alloy_sol_types::SolError;
 use foundry_common::errors::FsPathError;
 use foundry_config::UnresolvedEnvVarError;
 use foundry_evm_core::backend::DatabaseError;
+use foundry_wallets::error::WalletSignerError;
 use k256::ecdsa::signature::Error as SignatureError;
 use std::{borrow::Cow, fmt};
 
@@ -298,6 +299,7 @@ impl_from!(
     UnresolvedEnvVarError,
     WalletError,
     SignerError,
+    WalletSignerError,
 );
 
 #[cfg(test)]

crates/cheatcodes/src/inspector.rs

@@ -7,7 +7,7 @@ use crate::{
         prank::Prank,
         DealRecord, RecordAccess,
     },
-    script::Broadcast,
+    script::{Broadcast, ScriptWallets},
     test::expect::{
         self, ExpectedCallData, ExpectedCallTracker, ExpectedCallType, ExpectedEmit,
         ExpectedRevert, ExpectedRevertKind,
@@ -16,7 +16,6 @@ use crate::{
 };
 use alloy_primitives::{Address, Bytes, B256, U256, U64};
 use alloy_rpc_types::request::{TransactionInput, TransactionRequest};
-use alloy_signer::LocalWallet;
 use alloy_sol_types::{SolInterface, SolValue};
 use foundry_common::{evm::Breakpoints, provider::alloy::RpcUrl};
 use foundry_evm_core::{
@@ -127,7 +126,7 @@ pub struct Cheatcodes {
     pub labels: HashMap<Address, String>,
 
     /// Remembered private keys
-    pub script_wallets: Vec<LocalWallet>,
+    pub script_wallets: Option<ScriptWallets>,
 
     /// Prank information
     pub prank: Option<Prank>,
@@ -218,7 +217,8 @@ impl Cheatcodes {
     #[inline]
     pub fn new(config: Arc<CheatsConfig>) -> Self {
         let labels = config.labels.clone();
-        Self { config, fs_commit: true, labels, ..Default::default() }
+        let script_wallets = config.script_wallets.clone();
+        Self { config, fs_commit: true, labels, script_wallets, ..Default::default() }
     }
 
     fn apply_cheatcode<DB: DatabaseExt>(

crates/cheatcodes/src/lib.rs

@@ -36,6 +36,7 @@ mod string;
 mod test;
 mod utils;
 
+pub use script::ScriptWallets;
 pub use test::expect::ExpectedCallTracker;
 
 /// Cheatcode implementation.

crates/cheatcodes/src/script.rs

@@ -2,8 +2,11 @@
 
 use crate::{Cheatcode, CheatsCtxt, DatabaseExt, Result, Vm::*};
 use alloy_primitives::{Address, U256};
-use alloy_signer::Signer;
+use alloy_signer::{LocalWallet, Signer};
 use foundry_config::Config;
+use foundry_wallets::{multi_wallet::MultiWallet, WalletSigner};
+use parking_lot::Mutex;
+use std::sync::Arc;
 
 impl Cheatcode for broadcast_0Call {
     fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
@@ -72,6 +75,44 @@ pub struct Broadcast {
     pub single_call: bool,
 }
 
+/// Contains context for wallet management.
+#[derive(Debug)]
+pub struct ScriptWalletsInner {
+    /// All signers in scope of the script.
+    pub multi_wallet: MultiWallet,
+    /// Optional signer provided as `--sender` flag.
+    pub provided_sender: Option<Address>,
+}
+
+/// Clonable wrapper around [ScriptWalletsInner].
+#[derive(Debug, Clone)]
+pub struct ScriptWallets {
+    /// Inner data.
+    pub inner: Arc<Mutex<ScriptWalletsInner>>,
+}
+
+impl ScriptWallets {
+    #[allow(missing_docs)]
+    pub fn new(multi_wallet: MultiWallet, provided_sender: Option<Address>) -> Self {
+        Self { inner: Arc::new(Mutex::new(ScriptWalletsInner { multi_wallet, provided_sender })) }
+    }
+
+    /// Consumes [ScriptWallets] and returns [MultiWallet].
+    ///
+    /// Panics if [ScriptWallets] is still in use.
+    pub fn into_multi_wallet(self) -> MultiWallet {
+        Arc::into_inner(self.inner)
+            .map(|m| m.into_inner().multi_wallet)
+            .unwrap_or_else(|| panic!("not all instances were dropped"))
+    }
+
+    /// Locks inner Mutex and adds a signer to the [MultiWallet].
+    pub fn add_signer(&self, private_key: impl AsRef<[u8]>) -> Result {
+        self.inner.lock().multi_wallet.add_signer(WalletSigner::from_private_key(private_key)?);
+        Ok(Default::default())
+    }
+}
+
 /// Sets up broadcasting from a script using `new_origin` as the sender.
 fn broadcast<DB: DatabaseExt>(
     ccx: &mut CheatsCtxt<DB>,
@@ -86,8 +127,25 @@ fn broadcast<DB: DatabaseExt>(
 
     correct_sender_nonce(ccx)?;
 
+    let mut new_origin = new_origin.cloned();
+
+    if new_origin.is_none() {
+        if let Some(script_wallets) = &ccx.state.script_wallets {
+            let mut script_wallets = script_wallets.inner.lock();
+            if let Some(provided_sender) = script_wallets.provided_sender {
+                new_origin = Some(provided_sender);
+            } else {
+                let signers = script_wallets.multi_wallet.signers()?;
+                if signers.len() == 1 {
+                    let address = signers.keys().next().unwrap();
+                    new_origin = Some(*address);
+                }
+            }
+        }
+    }
+
     let broadcast = Broadcast {
-        new_origin: *new_origin.unwrap_or(&ccx.data.env.tx.caller),
+        new_origin: new_origin.unwrap_or(ccx.data.env.tx.caller),
         original_caller: ccx.caller,
         original_origin: ccx.data.env.tx.caller,
         depth: ccx.data.journaled_state.depth(),
@@ -106,13 +164,15 @@ fn broadcast_key<DB: DatabaseExt>(
     private_key: &U256,
     single_call: bool,
 ) -> Result {
-    let mut wallet = super::utils::parse_wallet(private_key)?;
-    wallet.set_chain_id(Some(ccx.data.env.cfg.chain_id));
-    let new_origin = &wallet.address();
+    let key = super::utils::parse_private_key(private_key)?;
+    let new_origin = LocalWallet::from(key.clone()).address();
+
+    let result = broadcast(ccx, Some(&new_origin), single_call);
 
-    let result = broadcast(ccx, Some(new_origin), single_call);
     if result.is_ok() {
-        ccx.state.script_wallets.push(wallet);
+        if let Some(script_wallets) = &ccx.state.script_wallets {
+            script_wallets.add_signer(key.to_bytes())?;
+        }
     }
     result
 }

crates/cheatcodes/src/utils.rs

@@ -87,9 +87,11 @@ impl Cheatcode for deriveKey_3Call {
 impl Cheatcode for rememberKeyCall {
     fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { privateKey } = self;
-        let wallet = parse_wallet(privateKey)?.with_chain_id(Some(ccx.data.env.cfg.chain_id));
-        let address = wallet.address();
-        ccx.state.script_wallets.push(wallet);
+        let key = parse_private_key(privateKey)?;
+        let address = LocalWallet::from(key.clone()).address();
+        if let Some(script_wallets) = &ccx.state.script_wallets {
+            script_wallets.add_signer(key.to_bytes())?;
+        }
         Ok(address.abi_encode())
     }
 }