Skip to main content

foundry_wallets/wallet_multi/
mod.rs

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