Allow the network for a wallet to be specified in the keys.txt file.

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<u64>,
+
+    #[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<String>,
     ) -> 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(&params).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<Channel>,
         params: impl Parameters + 'static,
         wallet_dir: Option<String>,
-        seed: SecretVec<u8>,
+        seed: &SecretVec<u8>,
         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(())

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<String>,
     ) -> Result<(), anyhow::Error> {
         // Connect to the client (for re-initializing the wallet).
-        let client = connect_to_lightwalletd().await?;
+        let client = connect_to_lightwalletd(&params).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
     }
 }

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<String>,
     ) -> Result<(), anyhow::Error> {
         let account = AccountId::from(0);
@@ -58,7 +58,7 @@ impl Command {
         let usk = UnifiedSpendingKey::from_seed(&params, seed.expose_secret(), account)
             .map_err(error::Error::from)?;
 
-        let mut client = connect_to_lightwalletd().await?;
+        let mut client = connect_to_lightwalletd(&params).await?;
 
         // Create the transaction.
         println!("Creating transaction...");

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<String>,
     ) -> 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(&params).await?;
 
         // 1) Download note commitment tree data from lightwalletd
         // 2) Pass the commitment tree data to the database.

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<Network, String> {
+        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<Network> 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<P: AsRef<Path>>(
     wallet_dir: Option<P>,
     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<P: AsRef<Path>>(
         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<P: AsRef<Path>>(
     Ok(BufReader::new(File::open(p)?))
 }
 
+pub(crate) struct WalletKeys {
+    network: consensus::Network,
+    seed: SecretVec<u8>,
+    birthday: BlockHeight,
+}
+
+impl WalletKeys {
+    pub(crate) fn seed(&self) -> &SecretVec<u8> {
+        &self.seed
+    }
+
+    pub(crate) fn birthday(&self) -> BlockHeight {
+        self.birthday
+    }
+}
+
+pub(crate) fn get_wallet_network<P: AsRef<Path>>(
+    wallet_dir: Option<P>,
+) -> Result<consensus::Network, anyhow::Error> {
+    Ok(read_keys(wallet_dir)?.network)
+}
+
 pub(crate) fn get_wallet_seed<P: AsRef<Path>>(
     wallet_dir: Option<P>,
 ) -> Result<SecretVec<u8>, 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<P: AsRef<Path>>(
+pub(crate) fn read_keys<P: AsRef<Path>>(
     wallet_dir: Option<P>,
-) -> Result<(SecretVec<u8>, BlockHeight), anyhow::Error> {
+) -> Result<WalletKeys, anyhow::Error> {
     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<P: AsRef<Path>>(
             .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<P: AsRef<Path>>(
         .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<P: AsRef<Path>>(wallet_dir: Option<P>) -> (PathBuf, PathBuf) {

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 {

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<CompactTxStreamerClient<Channel>, 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))
 }