Skip to main content

foundry_wallets/wallet_multi/
mod.rs

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