foundry_wallets/
multi_wallet.rs

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

foundry_wallets/
multi_wallet.rs
