foundry_wallets/
multi_wallet.rs

Help
1use crate::{
2    utils,
3    wallet_signer::{PendingSigner, WalletSigner},
4};
5use alloy_primitives::{map::AddressHashMap, Address};
6use alloy_signer::Signer;
7use clap::Parser;
8use derive_builder::Builder;
9use eyre::Result;
10use foundry_config::Config;
11use serde::Serialize;
12use std::path::PathBuf;
13
14/// Container for multiple wallets.
15#[derive(Debug, Default)]
16pub struct MultiWallet {
17    /// Vector of wallets that require an action to be unlocked.
18    /// Those are lazily unlocked on the first access of the signers.
19    pending_signers: Vec<PendingSigner>,
20    /// Contains unlocked signers.
21    signers: AddressHashMap<WalletSigner>,
22}
23
24impl MultiWallet {
25    pub fn new(pending_signers: Vec<PendingSigner>, signers: Vec<WalletSigner>) -> Self {
26        let signers = signers.into_iter().map(|signer| (signer.address(), signer)).collect();
27        Self { pending_signers, signers }
28    }
29
30    fn maybe_unlock_pending(&mut self) -> Result<()> {
31        for pending in self.pending_signers.drain(..) {
32            let signer = pending.unlock()?;
33            self.signers.insert(signer.address(), signer);
34        }
35        Ok(())
36    }
37
38    pub fn signers(&mut self) -> Result<&AddressHashMap<WalletSigner>> {
39        self.maybe_unlock_pending()?;
40        Ok(&self.signers)
41    }
42
43    pub fn into_signers(mut self) -> Result<AddressHashMap<WalletSigner>> {
44        self.maybe_unlock_pending()?;
45        Ok(self.signers)
46    }
47
48    pub fn add_signer(&mut self, signer: WalletSigner) {
49        self.signers.insert(signer.address(), signer);
50    }
51}
52
53/// A macro that initializes multiple wallets
54///
55/// Should be used with a [`MultiWallet`] instance
56macro_rules! create_hw_wallets {
57    ($self:ident, $create_signer:expr, $signers:ident) => {
58        let mut $signers = vec![];
59
60        if let Some(hd_paths) = &$self.hd_paths {
61            for path in hd_paths {
62                let hw = $create_signer(Some(path), 0).await?;
63                $signers.push(hw);
64            }
65        }
66
67        if let Some(mnemonic_indexes) = &$self.mnemonic_indexes {
68            for index in mnemonic_indexes {
69                let hw = $create_signer(None, *index).await?;
70                $signers.push(hw);
71            }
72        }
73
74        if $signers.is_empty() {
75            let hw = $create_signer(None, 0).await?;
76            $signers.push(hw);
77        }
78    };
79}
80
81/// The wallet options can either be:
82/// 1. Ledger
83/// 2. Trezor
84/// 3. Mnemonics (via file path)
85/// 4. Keystores (via file path)
86/// 5. Private Keys (cleartext in CLI)
87/// 6. Private Keys (interactively via secure prompt)
88/// 7. AWS KMS
89#[derive(Builder, Clone, Debug, Default, Serialize, Parser)]
90#[command(next_help_heading = "Wallet options", about = None, long_about = None)]
91pub struct MultiWalletOpts {
92    /// The sender accounts.
93    #[arg(
94        long,
95        short = 'a',
96        help_heading = "Wallet options - raw",
97        value_name = "ADDRESSES",
98        env = "ETH_FROM",
99        num_args(0..),
100    )]
101    #[builder(default = "None")]
102    pub froms: Option<Vec<Address>>,
103
104    /// Open an interactive prompt to enter your private key.
105    ///
106    /// Takes a value for the number of keys to enter.
107    #[arg(
108        long,
109        short,
110        help_heading = "Wallet options - raw",
111        default_value = "0",
112        value_name = "NUM"
113    )]
114    pub interactives: u32,
115
116    /// Use the provided private keys.
117    #[arg(long, help_heading = "Wallet options - raw", value_name = "RAW_PRIVATE_KEYS")]
118    #[builder(default = "None")]
119    pub private_keys: Option<Vec<String>>,
120
121    /// Use the provided private key.
122    #[arg(
123        long,
124        help_heading = "Wallet options - raw",
125        conflicts_with = "private_keys",
126        value_name = "RAW_PRIVATE_KEY"
127    )]
128    #[builder(default = "None")]
129    pub private_key: Option<String>,
130
131    /// Use the mnemonic phrases of mnemonic files at the specified paths.
132    #[arg(long, alias = "mnemonic-paths", help_heading = "Wallet options - raw")]
133    #[builder(default = "None")]
134    pub mnemonics: Option<Vec<String>>,
135
136    /// Use a BIP39 passphrases for the mnemonic.
137    #[arg(long, help_heading = "Wallet options - raw", value_name = "PASSPHRASE")]
138    #[builder(default = "None")]
139    pub mnemonic_passphrases: Option<Vec<String>>,
140
141    /// The wallet derivation path.
142    ///
143    /// Works with both --mnemonic-path and hardware wallets.
144    #[arg(
145        long = "mnemonic-derivation-paths",
146        alias = "hd-paths",
147        help_heading = "Wallet options - raw",
148        value_name = "PATH"
149    )]
150    #[builder(default = "None")]
151    pub hd_paths: Option<Vec<String>>,
152
153    /// Use the private key from the given mnemonic index.
154    ///
155    /// Can be used with --mnemonics, --ledger, --aws and --trezor.
156    #[arg(
157        long,
158        conflicts_with = "hd_paths",
159        help_heading = "Wallet options - raw",
160        default_value = "0",
161        value_name = "INDEXES"
162    )]
163    pub mnemonic_indexes: Option<Vec<u32>>,
164
165    /// Use the keystore by its filename in the given folder.
166    #[arg(
167        long = "keystore",
168        visible_alias = "keystores",
169        help_heading = "Wallet options - keystore",
170        value_name = "PATHS",
171        env = "ETH_KEYSTORE"
172    )]
173    #[builder(default = "None")]
174    pub keystore_paths: Option<Vec<String>>,
175
176    /// Use a keystore from the default keystores folder (~/.foundry/keystores) by its filename.
177    #[arg(
178        long = "account",
179        visible_alias = "accounts",
180        help_heading = "Wallet options - keystore",
181        value_name = "ACCOUNT_NAMES",
182        env = "ETH_KEYSTORE_ACCOUNT",
183        conflicts_with = "keystore_paths"
184    )]
185    #[builder(default = "None")]
186    pub keystore_account_names: Option<Vec<String>>,
187
188    /// The keystore password.
189    ///
190    /// Used with --keystore.
191    #[arg(
192        long = "password",
193        help_heading = "Wallet options - keystore",
194        requires = "keystore_paths",
195        value_name = "PASSWORDS"
196    )]
197    #[builder(default = "None")]
198    pub keystore_passwords: Option<Vec<String>>,
199
200    /// The keystore password file path.
201    ///
202    /// Used with --keystore.
203    #[arg(
204        long = "password-file",
205        help_heading = "Wallet options - keystore",
206        requires = "keystore_paths",
207        value_name = "PATHS",
208        env = "ETH_PASSWORD"
209    )]
210    #[builder(default = "None")]
211    pub keystore_password_files: Option<Vec<String>>,
212
213    /// Use a Ledger hardware wallet.
214    #[arg(long, short, help_heading = "Wallet options - hardware wallet")]
215    pub ledger: bool,
216
217    /// Use a Trezor hardware wallet.
218    #[arg(long, short, help_heading = "Wallet options - hardware wallet")]
219    pub trezor: bool,
220
221    /// Use AWS Key Management Service.
222    #[arg(long, help_heading = "Wallet options - remote", hide = !cfg!(feature = "aws-kms"))]
223    pub aws: bool,
224
225    /// Use Google Cloud Key Management Service.
226    #[arg(long, help_heading = "Wallet options - remote", hide = !cfg!(feature = "gcp-kms"))]
227    pub gcp: bool,
228}
229
230impl MultiWalletOpts {
231    /// Returns [MultiWallet] container configured with provided options.
232    pub async fn get_multi_wallet(&self) -> Result<MultiWallet> {
233        let mut pending = Vec::new();
234        let mut signers: Vec<WalletSigner> = Vec::new();
235
236        if let Some(ledgers) = self.ledgers().await? {
237            signers.extend(ledgers);
238        }
239        if let Some(trezors) = self.trezors().await? {
240            signers.extend(trezors);
241        }
242        if let Some(aws_signers) = self.aws_signers().await? {
243            signers.extend(aws_signers);
244        }
245        if let Some(gcp_signer) = self.gcp_signers().await? {
246            signers.extend(gcp_signer);
247        }
248        if let Some((pending_keystores, unlocked)) = self.keystores()? {
249            pending.extend(pending_keystores);
250            signers.extend(unlocked);
251        }
252        if let Some(pks) = self.private_keys()? {
253            signers.extend(pks);
254        }
255        if let Some(mnemonics) = self.mnemonics()? {
256            signers.extend(mnemonics);
257        }
258        if self.interactives > 0 {
259            pending.extend(std::iter::repeat_n(
260                PendingSigner::Interactive,
261                self.interactives as usize,
262            ));
263        }
264
265        Ok(MultiWallet::new(pending, signers))
266    }
267
268    pub fn private_keys(&self) -> Result<Option<Vec<WalletSigner>>> {
269        let mut pks = vec![];
270        if let Some(private_key) = &self.private_key {
271            pks.push(private_key);
272        }
273        if let Some(private_keys) = &self.private_keys {
274            for pk in private_keys {
275                pks.push(pk);
276            }
277        }
278        if !pks.is_empty() {
279            let wallets = pks
280                .into_iter()
281                .map(|pk| utils::create_private_key_signer(pk))
282                .collect::<Result<Vec<_>>>()?;
283            Ok(Some(wallets))
284        } else {
285            Ok(None)
286        }
287    }
288
289    fn keystore_paths(&self) -> Result<Option<Vec<PathBuf>>> {
290        if let Some(keystore_paths) = &self.keystore_paths {
291            return Ok(Some(keystore_paths.iter().map(PathBuf::from).collect()));
292        }
293        if let Some(keystore_account_names) = &self.keystore_account_names {
294            let default_keystore_dir = Config::foundry_keystores_dir()
295                .ok_or_else(|| eyre::eyre!("Could not find the default keystore directory."))?;
296            return Ok(Some(
297                keystore_account_names
298                    .iter()
299                    .map(|keystore_name| default_keystore_dir.join(keystore_name))
300                    .collect(),
301            ));
302        }
303        Ok(None)
304    }
305
306    /// Returns all wallets read from the provided keystores arguments
307    ///
308    /// Returns `Ok(None)` if no keystore provided.
309    pub fn keystores(&self) -> Result<Option<(Vec<PendingSigner>, Vec<WalletSigner>)>> {
310        if let Some(keystore_paths) = self.keystore_paths()? {
311            let mut pending = Vec::new();
312            let mut signers = Vec::new();
313
314            let mut passwords_iter =
315                self.keystore_passwords.clone().unwrap_or_default().into_iter();
316
317            let mut password_files_iter =
318                self.keystore_password_files.clone().unwrap_or_default().into_iter();
319
320            for path in &keystore_paths {
321                let (maybe_signer, maybe_pending) = utils::create_keystore_signer(
322                    path,
323                    passwords_iter.next().as_deref(),
324                    password_files_iter.next().as_deref(),
325                )?;
326                if let Some(pending_signer) = maybe_pending {
327                    pending.push(pending_signer);
328                } else if let Some(signer) = maybe_signer {
329                    signers.push(signer);
330                }
331            }
332            return Ok(Some((pending, signers)));
333        }
334        Ok(None)
335    }
336
337    pub fn mnemonics(&self) -> Result<Option<Vec<WalletSigner>>> {
338        if let Some(ref mnemonics) = self.mnemonics {
339            let mut wallets = vec![];
340
341            let mut hd_paths_iter = self.hd_paths.clone().unwrap_or_default().into_iter();
342
343            let mut passphrases_iter =
344                self.mnemonic_passphrases.clone().unwrap_or_default().into_iter();
345
346            let mut indexes_iter = self.mnemonic_indexes.clone().unwrap_or_default().into_iter();
347
348            for mnemonic in mnemonics {
349                let wallet = utils::create_mnemonic_signer(
350                    mnemonic,
351                    passphrases_iter.next().as_deref(),
352                    hd_paths_iter.next().as_deref(),
353                    indexes_iter.next().unwrap_or(0),
354                )?;
355                wallets.push(wallet);
356            }
357            return Ok(Some(wallets));
358        }
359        Ok(None)
360    }
361
362    pub async fn ledgers(&self) -> Result<Option<Vec<WalletSigner>>> {
363        if self.ledger {
364            let mut args = self.clone();
365
366            if let Some(paths) = &args.hd_paths {
367                if paths.len() > 1 {
368                    eyre::bail!("Ledger only supports one signer.");
369                }
370                args.mnemonic_indexes = None;
371            }
372
373            create_hw_wallets!(args, utils::create_ledger_signer, wallets);
374            return Ok(Some(wallets));
375        }
376        Ok(None)
377    }
378
379    pub async fn trezors(&self) -> Result<Option<Vec<WalletSigner>>> {
380        if self.trezor {
381            create_hw_wallets!(self, utils::create_trezor_signer, wallets);
382            return Ok(Some(wallets));
383        }
384        Ok(None)
385    }
386
387    pub async fn aws_signers(&self) -> Result<Option<Vec<WalletSigner>>> {
388        #[cfg(feature = "aws-kms")]
389        if self.aws {
390            let mut wallets = vec![];
391            let aws_keys = std::env::var("AWS_KMS_KEY_IDS")
392                .or(std::env::var("AWS_KMS_KEY_ID"))?
393                .split(',')
394                .map(|k| k.to_string())
395                .collect::<Vec<_>>();
396
397            for key in aws_keys {
398                let aws_signer = WalletSigner::from_aws(key).await?;
399                wallets.push(aws_signer)
400            }
401
402            return Ok(Some(wallets));
403        }
404
405        Ok(None)
406    }
407
408    /// Returns a list of GCP signers if the GCP flag is set.
409    ///
410    /// The GCP signers are created from the following environment variables:
411    /// - GCP_PROJECT_ID: The GCP project ID. e.g. `my-project-123456`.
412    /// - GCP_LOCATION: The GCP location. e.g. `us-central1`.
413    /// - GCP_KEY_RING: The GCP key ring name. e.g. `my-key-ring`.
414    /// - GCP_KEY_NAME: The GCP key name. e.g. `my-key`.
415    /// - GCP_KEY_VERSION: The GCP key version. e.g. `1`.
416    ///
417    /// For more information on GCP KMS, see the [official documentation](https://cloud.google.com/kms/docs).
418    pub async fn gcp_signers(&self) -> Result<Option<Vec<WalletSigner>>> {
419        #[cfg(feature = "gcp-kms")]
420        if self.gcp {
421            let mut wallets = vec![];
422
423            let project_id = std::env::var("GCP_PROJECT_ID")?;
424            let location = std::env::var("GCP_LOCATION")?;
425            let key_ring = std::env::var("GCP_KEY_RING")?;
426            let key_names = std::env::var("GCP_KEY_NAME")?;
427            let key_version = std::env::var("GCP_KEY_VERSION")?;
428
429            let gcp_signer = WalletSigner::from_gcp(
430                project_id,
431                location,
432                key_ring,
433                key_names,
434                key_version.parse()?,
435            )
436            .await?;
437            wallets.push(gcp_signer);
438
439            return Ok(Some(wallets));
440        }
441
442        Ok(None)
443    }
444}
445
446#[cfg(test)]
447mod tests {
448    use super::*;
449    use alloy_primitives::address;
450    use std::path::Path;
451
452    #[test]
453    fn parse_keystore_args() {
454        let args: MultiWalletOpts =
455            MultiWalletOpts::parse_from(["foundry-cli", "--keystores", "my/keystore/path"]);
456        assert_eq!(args.keystore_paths, Some(vec!["my/keystore/path".to_string()]));
457
458        std::env::set_var("ETH_KEYSTORE", "MY_KEYSTORE");
459        let args: MultiWalletOpts = MultiWalletOpts::parse_from(["foundry-cli"]);
460        assert_eq!(args.keystore_paths, Some(vec!["MY_KEYSTORE".to_string()]));
461
462        std::env::remove_var("ETH_KEYSTORE");
463    }
464
465    #[test]
466    fn parse_keystore_password_file() {
467        let keystore =
468            Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/../cast/tests/fixtures/keystore"));
469        let keystore_file = keystore
470            .join("UTC--2022-12-20T10-30-43.591916000Z--ec554aeafe75601aaab43bd4621a22284db566c2");
471
472        let keystore_password_file = keystore.join("password-ec554").into_os_string();
473
474        let args: MultiWalletOpts = MultiWalletOpts::parse_from([
475            "foundry-cli",
476            "--keystores",
477            keystore_file.to_str().unwrap(),
478            "--password-file",
479            keystore_password_file.to_str().unwrap(),
480        ]);
481        assert_eq!(
482            args.keystore_password_files,
483            Some(vec![keystore_password_file.to_str().unwrap().to_string()])
484        );
485
486        let (_, unlocked) = args.keystores().unwrap().unwrap();
487        assert_eq!(unlocked.len(), 1);
488        assert_eq!(unlocked[0].address(), address!("0xec554aeafe75601aaab43bd4621a22284db566c2"));
489    }
490
491    // https://github.com/foundry-rs/foundry/issues/5179
492    #[test]
493    fn should_not_require_the_mnemonics_flag_with_mnemonic_indexes() {
494        let wallet_options = vec![
495            ("ledger", "--mnemonic-indexes", 1),
496            ("trezor", "--mnemonic-indexes", 2),
497            ("aws", "--mnemonic-indexes", 10),
498        ];
499
500        for test_case in wallet_options {
501            let args: MultiWalletOpts = MultiWalletOpts::parse_from([
502                "foundry-cli",
503                &format!("--{}", test_case.0),
504                test_case.1,
505                &test_case.2.to_string(),
506            ]);
507
508            match test_case.0 {
509                "ledger" => assert!(args.ledger),
510                "trezor" => assert!(args.trezor),
511                "aws" => assert!(args.aws),
512                _ => panic!("Should have matched one of the previous wallet options"),
513            }
514
515            assert_eq!(
516                args.mnemonic_indexes.expect("--mnemonic-indexes should have been set")[0],
517                test_case.2
518            )
519        }
520    }
521}
foundry_wallets/multi_wallet.rs

foundry_wallets/
multi_wallet.rs
