diff --git a/src/commands/init.rs b/src/commands/init.rs index be01a2a..8d6c63f 100644 --- a/src/commands/init.rs +++ b/src/commands/init.rs @@ -16,9 +16,9 @@ use zcash_primitives::{ }; use crate::{ - data::{get_db_paths, init_wallet_keys}, + data::{get_db_paths, init_wallet_keys, Network}, error, - remote::connect_to_lightwalletd, + remote::{connect_to_lightwalletd, Lightwalletd}, }; // Options accepted for the `init` command @@ -29,18 +29,24 @@ pub(crate) struct Command { #[options(help = "the wallet's birthday (default is current chain height)")] birthday: Option, + + #[options( + help = "the network the wallet will be used with: \"test\" or \"main\" (default is \"test\")", + parse(try_from_str = "Network::parse") + )] + network: Network, } impl Command { pub(crate) async fn run( self, - params: impl Parameters + 'static, + params: impl Parameters + Lightwalletd + 'static, wallet_dir: Option, ) -> Result<(), anyhow::Error> { let opts = self; // Get the current chain height (for the wallet's birthday). - let mut client = connect_to_lightwalletd().await?; + let mut client = connect_to_lightwalletd(¶ms).await?; let birthday = if let Some(birthday) = opts.birthday { birthday } else { @@ -60,7 +66,7 @@ impl Command { }; // Save the wallet keys to disk. - init_wallet_keys(wallet_dir.as_ref(), &mnemonic, birthday)?; + init_wallet_keys(wallet_dir.as_ref(), &mnemonic, birthday, opts.network)?; let seed = { let mut seed = mnemonic.to_seed(""); @@ -69,14 +75,14 @@ impl Command { SecretVec::new(secret) }; - Self::init_dbs(client, params, wallet_dir, seed, birthday).await + Self::init_dbs(client, params, wallet_dir, &seed, birthday).await } pub(crate) async fn init_dbs( mut client: CompactTxStreamerClient, params: impl Parameters + 'static, wallet_dir: Option, - seed: SecretVec, + seed: &SecretVec, birthday: u64, ) -> Result<(), anyhow::Error> { // Initialise the block and wallet DBs. @@ -99,7 +105,7 @@ impl Command { }; // Add one account. - let (account, _) = db_data.create_account(&seed, birthday)?; + let (account, _) = db_data.create_account(seed, birthday)?; assert_eq!(account, AccountId::from(0)); Ok(()) diff --git a/src/commands/reset.rs b/src/commands/reset.rs index c0b38ca..e2f7b1d 100644 --- a/src/commands/reset.rs +++ b/src/commands/reset.rs @@ -3,8 +3,8 @@ use gumdrop::Options; use zcash_primitives::consensus::Parameters; use crate::{ - data::{erase_wallet_state, get_wallet_seed_and_birthday}, - remote::connect_to_lightwalletd, + data::{erase_wallet_state, read_keys}, + remote::{connect_to_lightwalletd, Lightwalletd}, }; // Options accepted for the `reset` command @@ -14,19 +14,26 @@ pub(crate) struct Command {} impl Command { pub(crate) async fn run( self, - params: impl Parameters + 'static, + params: impl Parameters + Lightwalletd + 'static, wallet_dir: Option, ) -> Result<(), anyhow::Error> { // Connect to the client (for re-initializing the wallet). - let client = connect_to_lightwalletd().await?; + let client = connect_to_lightwalletd(¶ms).await?; // Load the wallet seed and birthday from disk. - let (seed, birthday) = get_wallet_seed_and_birthday(wallet_dir.as_ref())?; + let keys = read_keys(wallet_dir.as_ref())?; // Erase the wallet state (excluding key material). erase_wallet_state(wallet_dir.as_ref()).await; // Re-initialize the wallet state. - super::init::Command::init_dbs(client, params, wallet_dir, seed, birthday.into()).await + super::init::Command::init_dbs( + client, + params, + wallet_dir, + keys.seed(), + keys.birthday().into(), + ) + .await } } diff --git a/src/commands/send.rs b/src/commands/send.rs index 4d2190e..6a16bd9 100644 --- a/src/commands/send.rs +++ b/src/commands/send.rs @@ -23,7 +23,7 @@ use crate::{ commands::propose::{parse_fee_rule, FeeRule}, data::{get_db_paths, get_wallet_seed}, error, - remote::connect_to_lightwalletd, + remote::{connect_to_lightwalletd, Lightwalletd}, MIN_CONFIRMATIONS, }; @@ -47,7 +47,7 @@ pub(crate) struct Command { impl Command { pub(crate) async fn run( self, - params: impl Parameters + Copy + 'static, + params: impl Parameters + Lightwalletd + Copy + 'static, wallet_dir: Option, ) -> Result<(), anyhow::Error> { let account = AccountId::from(0); @@ -58,7 +58,7 @@ impl Command { let usk = UnifiedSpendingKey::from_seed(¶ms, seed.expose_secret(), account) .map_err(error::Error::from)?; - let mut client = connect_to_lightwalletd().await?; + let mut client = connect_to_lightwalletd(¶ms).await?; // Create the transaction. println!("Creating transaction..."); diff --git a/src/commands/sync.rs b/src/commands/sync.rs index cb90cb3..c86249b 100644 --- a/src/commands/sync.rs +++ b/src/commands/sync.rs @@ -26,7 +26,7 @@ use zcash_primitives::{ use crate::{ data::{get_block_path, get_db_paths}, error, - remote::connect_to_lightwalletd, + remote::{connect_to_lightwalletd, Lightwalletd}, }; const BATCH_SIZE: u32 = 10_000; @@ -38,14 +38,14 @@ pub(crate) struct Command {} impl Command { pub(crate) async fn run( self, - params: impl Parameters + Copy + Send + 'static, + params: impl Parameters + Lightwalletd + Copy + Send + 'static, wallet_dir: Option, ) -> Result<(), anyhow::Error> { let (fsblockdb_root, db_data) = get_db_paths(wallet_dir.as_ref()); let fsblockdb_root = fsblockdb_root.as_path(); let mut db_cache = FsBlockDb::for_path(fsblockdb_root).map_err(error::Error::from)?; let mut db_data = WalletDb::for_path(db_data, params)?; - let mut client = connect_to_lightwalletd().await?; + let mut client = connect_to_lightwalletd(¶ms).await?; // 1) Download note commitment tree data from lightwalletd // 2) Pass the commitment tree data to the database. diff --git a/src/data.rs b/src/data.rs index fe5b8b5..71b091a 100644 --- a/src/data.rs +++ b/src/data.rs @@ -6,7 +6,10 @@ use secrecy::{SecretVec, Zeroize}; use tracing::error; use zcash_client_sqlite::chain::BlockMeta; -use zcash_primitives::{consensus::BlockHeight, zip339::Mnemonic}; +use zcash_primitives::{ + consensus::{self, BlockHeight}, + zip339::Mnemonic, +}; use crate::error; @@ -15,10 +18,49 @@ const KEYS_FILE: &str = "keys.txt"; const BLOCKS_FOLDER: &str = "blocks"; const DATA_DB: &str = "data.sqlite"; +#[derive(Clone, Copy, Debug)] +pub(crate) enum Network { + Test, + Main, +} + +impl Default for Network { + fn default() -> Self { + Network::Test + } +} + +impl Network { + pub(crate) fn parse(name: &str) -> Result { + match name { + "main" => Ok(Network::Main), + "test" => Ok(Network::Test), + other => Err(format!("Unsupported network: {}", other)), + } + } + + pub(crate) fn name(&self) -> &str { + match self { + Network::Test => "test", + Network::Main => "main", + } + } +} + +impl From for consensus::Network { + fn from(value: Network) -> Self { + match value { + Network::Test => consensus::Network::TestNetwork, + Network::Main => consensus::Network::MainNetwork, + } + } +} + pub(crate) fn init_wallet_keys>( wallet_dir: Option

, mnemonic: &Mnemonic, birthday: u64, + network: Network, ) -> Result<(), anyhow::Error> { // Create the wallet directory. let wallet_dir = wallet_dir @@ -39,6 +81,7 @@ pub(crate) fn init_wallet_keys>( mnemonic.phrase() )?; writeln!(&mut keys_file, "{} # wallet birthday", birthday)?; + writeln!(&mut keys_file, "{} # network", network.name())?; Ok(()) } @@ -55,16 +98,37 @@ pub(crate) fn get_keys_file>( Ok(BufReader::new(File::open(p)?)) } +pub(crate) struct WalletKeys { + network: consensus::Network, + seed: SecretVec, + birthday: BlockHeight, +} + +impl WalletKeys { + pub(crate) fn seed(&self) -> &SecretVec { + &self.seed + } + + pub(crate) fn birthday(&self) -> BlockHeight { + self.birthday + } +} + +pub(crate) fn get_wallet_network>( + wallet_dir: Option

, +) -> Result { + Ok(read_keys(wallet_dir)?.network) +} + pub(crate) fn get_wallet_seed>( wallet_dir: Option

, ) -> Result, anyhow::Error> { - let (seed, _) = get_wallet_seed_and_birthday(wallet_dir)?; - Ok(seed) + Ok(read_keys(wallet_dir)?.seed) } -pub(crate) fn get_wallet_seed_and_birthday>( +pub(crate) fn read_keys>( wallet_dir: Option

, -) -> Result<(SecretVec, BlockHeight), anyhow::Error> { +) -> Result { let keys_file = get_keys_file(wallet_dir)?; let mut keys_file_lines = keys_file.lines(); @@ -77,9 +141,9 @@ pub(crate) fn get_wallet_seed_and_birthday>( .ok_or(error::Error::InvalidKeysFile)? .trim(), )?; - let mut seed = mnemonic.to_seed(""); - let secret = SecretVec::new(seed.to_vec()); - seed.zeroize(); + let mut seed_bytes = mnemonic.to_seed(""); + let seed = SecretVec::new(seed_bytes.to_vec()); + seed_bytes.zeroize(); let birthday = keys_file_lines .next() @@ -92,7 +156,24 @@ pub(crate) fn get_wallet_seed_and_birthday>( .map(BlockHeight::from) .map_err(|_| error::Error::InvalidKeysFile)?; - Ok((secret, birthday)) + let network = keys_file_lines.next().transpose()?.map_or_else( + || Ok(consensus::Network::TestNetwork), + |line| { + let network_name = line + .split('#') + .next() + .ok_or(error::Error::InvalidKeysFile)?; + Network::parse(network_name.trim()) + .map(consensus::Network::from) + .map_err(|_| error::Error::InvalidKeysFile) + }, + )?; + + Ok(WalletKeys { + network, + seed, + birthday, + }) } pub(crate) fn get_db_paths>(wallet_dir: Option

) -> (PathBuf, PathBuf) { diff --git a/src/main.rs b/src/main.rs index 1ada5ec..e299c93 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,10 +7,9 @@ use std::env; use std::num::NonZeroU32; use std::sync::atomic::{AtomicUsize, Ordering}; +use data::get_wallet_network; use gumdrop::Options; -use zcash_primitives::consensus::TEST_NETWORK; - mod commands; mod data; mod error; @@ -81,7 +80,7 @@ fn main() -> Result<(), anyhow::Error> { }) .build()?; - let params = TEST_NETWORK; + let params = get_wallet_network(opts.wallet_dir.as_ref())?; runtime.block_on(async { match opts.command { diff --git a/src/remote.rs b/src/remote.rs index 5da6ec8..5f1824c 100644 --- a/src/remote.rs +++ b/src/remote.rs @@ -1,23 +1,37 @@ use tonic::transport::{Channel, ClientTlsConfig}; use zcash_client_backend::proto::service::compact_tx_streamer_client::CompactTxStreamerClient; +use zcash_primitives::consensus; -const LIGHTWALLETD_HOST: &str = "lightwalletd.testnet.electriccoin.co"; -const LIGHTWALLETD_PORT: u16 = 9067; +pub(crate) trait Lightwalletd { + fn host(&self) -> &str; + fn port(&self) -> u16; +} + +impl Lightwalletd for consensus::Network { + fn host(&self) -> &str { + match self { + consensus::Network::MainNetwork => "mainnet.lightwalletd.com", + consensus::Network::TestNetwork => "lightwalletd.testnet.electriccoin.co", + } + } + + fn port(&self) -> u16 { + 9067 + } +} pub(crate) async fn connect_to_lightwalletd( + network: &impl Lightwalletd, ) -> Result, anyhow::Error> { - println!("Connecting to {}:{}", LIGHTWALLETD_HOST, LIGHTWALLETD_PORT); + println!("Connecting to {}:{}", network.host(), network.port()); - let tls = ClientTlsConfig::new().domain_name(LIGHTWALLETD_HOST); + let tls = ClientTlsConfig::new().domain_name(network.host()); - let channel = Channel::from_shared(format!( - "https://{}:{}", - LIGHTWALLETD_HOST, LIGHTWALLETD_PORT - ))? - .tls_config(tls)? - .connect() - .await?; + let channel = Channel::from_shared(format!("https://{}:{}", network.host(), network.port()))? + .tls_config(tls)? + .connect() + .await?; Ok(CompactTxStreamerClient::new(channel)) }