Merge pull request #18 from Electric-Coin-Company/server-selection

Server selection logic

src/commands/init.rs

@@ -17,7 +17,7 @@ use zcash_primitives::{
 use crate::{
     data::{get_db_paths, init_wallet_keys, Network},
     error,
-    remote::connect_to_lightwalletd,
+    remote::{connect_to_lightwalletd, Servers},
 };
 
 // Options accepted for the `init` command
@@ -34,6 +34,13 @@ pub(crate) struct Command {
         parse(try_from_str = "Network::parse")
     )]
     network: Network,
+
+    #[options(
+        help = "the server to initialize with (default is \"ecc\")",
+        default = "ecc",
+        parse(try_from_str = "Servers::parse")
+    )]
+    server: Servers,
 }
 
 impl Command {
@@ -42,7 +49,7 @@ impl Command {
         let params = consensus::Network::from(opts.network);
 
         // Get the current chain height (for the wallet's birthday).
-        let mut client = connect_to_lightwalletd(&params).await?;
+        let mut client = connect_to_lightwalletd(opts.server.pick(params)?).await?;
         let birthday = if let Some(birthday) = opts.birthday {
             birthday
         } else {

src/commands/reset.rs

@@ -2,12 +2,19 @@ use gumdrop::Options;
 
 use crate::{
     data::{erase_wallet_state, read_keys},
-    remote::connect_to_lightwalletd,
+    remote::{connect_to_lightwalletd, Servers},
 };
 
 // Options accepted for the `reset` command
 #[derive(Debug, Options)]
-pub(crate) struct Command {}
+pub(crate) struct Command {
+    #[options(
+        help = "the server to re-initialize with (default is \"ecc\")",
+        default = "ecc",
+        parse(try_from_str = "Servers::parse")
+    )]
+    server: Servers,
+}
 
 impl Command {
     pub(crate) async fn run(self, wallet_dir: Option<String>) -> Result<(), anyhow::Error> {
@@ -16,7 +23,7 @@ impl Command {
         let params = keys.network();
 
         // Connect to the client (for re-initializing the wallet).
-        let client = connect_to_lightwalletd(&params).await?;
+        let client = connect_to_lightwalletd(self.server.pick(params)?).await?;
 
         // Erase the wallet state (excluding key material).
         erase_wallet_state(wallet_dir.as_ref()).await;

src/commands/send.rs

@@ -24,7 +24,7 @@ use crate::{
     commands::propose::{parse_fee_rule, FeeRule},
     data::{get_db_paths, read_keys},
     error,
-    remote::connect_to_lightwalletd,
+    remote::{connect_to_lightwalletd, Servers},
     MIN_CONFIRMATIONS,
 };
 
@@ -46,6 +46,13 @@ pub(crate) struct Command {
         parse(try_from_str = "parse_fee_rule")
     )]
     fee_rule: FeeRule,
+
+    #[options(
+        help = "the server to send via (default is \"ecc\")",
+        default = "ecc",
+        parse(try_from_str = "Servers::parse")
+    )]
+    server: Servers,
 }
 
 impl Command {
@@ -71,7 +78,7 @@ impl Command {
             UnifiedSpendingKey::from_seed(&params, keys.seed().expose_secret(), account_index)
                 .map_err(error::Error::from)?;
 
-        let mut client = connect_to_lightwalletd(&params).await?;
+        let mut client = connect_to_lightwalletd(self.server.pick(params)?).await?;
 
         // Create the transaction.
         println!("Creating transaction...");

src/commands/sync.rs

@@ -27,7 +27,7 @@ use zcash_protocol::consensus::{BlockHeight, Parameters};
 use crate::{
     data::{get_block_path, get_db_paths, get_wallet_network},
     error,
-    remote::connect_to_lightwalletd,
+    remote::{connect_to_lightwalletd, Servers},
 };
 
 #[cfg(feature = "tui")]
@@ -41,6 +41,13 @@ const BATCH_SIZE: u32 = 10_000;
 // Options accepted for the `sync` command
 #[derive(Debug, Options)]
 pub(crate) struct Command {
+    #[options(
+        help = "the server to sync with (default is \"ecc\")",
+        default = "ecc",
+        parse(try_from_str = "Servers::parse")
+    )]
+    server: Servers,
+
     #[cfg(feature = "tui")]
     pub(crate) defrag: bool,
 }
@@ -57,7 +64,7 @@ impl Command {
         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(&params).await?;
+        let mut client = connect_to_lightwalletd(self.server.pick(params)?).await?;
 
         #[cfg(feature = "tui")]
         let tui_handle = if self.defrag {
@@ -278,8 +285,6 @@ async fn update_subtree_roots<P: Parameters>(
 ) -> Result<(), anyhow::Error> {
     let mut request = service::GetSubtreeRootsArg::default();
     request.set_shielded_protocol(service::ShieldedProtocol::Sapling);
-    // Hack to work around a bug in the initial lightwalletd implementation.
-    request.max_entries = 65536;
     let sapling_roots: Vec<CommitmentTreeRoot<sapling::Node>> = client
         .get_subtree_roots(request)
         .await?
@@ -299,8 +304,6 @@ async fn update_subtree_roots<P: Parameters>(
 
     let mut request = service::GetSubtreeRootsArg::default();
     request.set_shielded_protocol(service::ShieldedProtocol::Orchard);
-    // Hack to work around a bug in the initial lightwalletd implementation.
-    request.max_entries = 65536;
     let orchard_roots: Vec<CommitmentTreeRoot<MerkleHashOrchard>> = client
         .get_subtree_roots(request)
         .await?

src/main.rs

@@ -70,7 +70,7 @@ fn main() -> Result<(), anyhow::Error> {
     let log_configured = false;
     #[cfg(feature = "tui")]
     let log_configured =
-        if let Some(Command::Sync(commands::sync::Command { defrag: true })) = opts.command {
+        if let Some(Command::Sync(commands::sync::Command { defrag: true, .. })) = opts.command {
             use tracing_subscriber::layer::SubscriberExt;
 
             tracing::subscriber::set_global_default(

src/remote.rs

@@ -1,35 +1,136 @@
+use std::{borrow::Cow, fmt};
+
+use anyhow::anyhow;
 use tonic::transport::{Channel, ClientTlsConfig};
 
 use tracing::info;
 use zcash_client_backend::proto::service::compact_tx_streamer_client::CompactTxStreamerClient;
-use zcash_primitives::consensus;
+use zcash_protocol::consensus::Network;
+
+const ECC_TESTNET: &[Server<'_>] = &[Server::fixed("lightwalletd.testnet.electriccoin.co", 9067)];
+
+const YWALLET_MAINNET: &[Server<'_>] = &[
+    Server::fixed("lwd1.zcash-infra.com", 9067),
+    Server::fixed("lwd2.zcash-infra.com", 9067),
+    Server::fixed("lwd3.zcash-infra.com", 9067),
+    Server::fixed("lwd4.zcash-infra.com", 9067),
+    Server::fixed("lwd5.zcash-infra.com", 9067),
+    Server::fixed("lwd6.zcash-infra.com", 9067),
+    Server::fixed("lwd7.zcash-infra.com", 9067),
+    Server::fixed("lwd8.zcash-infra.com", 9067),
+];
+
+const ZEC_ROCKS_MAINNET: &[Server<'_>] = &[
+    Server::fixed("zec.rocks", 443),
+    Server::fixed("ap.zec.rocks", 443),
+    Server::fixed("eu.zec.rocks", 443),
+    Server::fixed("na.zec.rocks", 443),
+    Server::fixed("sa.zec.rocks", 443),
+];
+const ZEC_ROCKS_TESTNET: &[Server<'_>] = &[Server::fixed("testnet.zec.rocks", 443)];
 
-pub(crate) trait Lightwalletd {
-    fn host(&self) -> &str;
-    fn port(&self) -> u16;
+#[derive(Debug)]
+pub(crate) enum ServerOperator {
+    Ecc,
+    YWallet,
+    ZecRocks,
 }
 
-impl Lightwalletd for consensus::Network {
-    fn host(&self) -> &str {
+impl ServerOperator {
+    fn servers(&self, network: Network) -> &[Server<'_>] {
+        match (self, network) {
+            (ServerOperator::Ecc, Network::MainNetwork) => &[],
+            (ServerOperator::Ecc, Network::TestNetwork) => ECC_TESTNET,
+            (ServerOperator::YWallet, Network::MainNetwork) => YWALLET_MAINNET,
+            (ServerOperator::YWallet, Network::TestNetwork) => &[],
+            (ServerOperator::ZecRocks, Network::MainNetwork) => ZEC_ROCKS_MAINNET,
+            (ServerOperator::ZecRocks, Network::TestNetwork) => ZEC_ROCKS_TESTNET,
+        }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) enum Servers {
+    Hosted(ServerOperator),
+    Custom(Vec<Server<'static>>),
+}
+
+impl Servers {
+    pub(crate) fn parse(s: &str) -> anyhow::Result<Self> {
+        match s {
+            "ecc" => Ok(Self::Hosted(ServerOperator::Ecc)),
+            "ywallet" => Ok(Self::Hosted(ServerOperator::YWallet)),
+            "zecrocks" => Ok(Self::Hosted(ServerOperator::ZecRocks)),
+            _ => s
+                .split(',')
+                .map(|sub| {
+                    sub.split_once(':').and_then(|(host, port_str)| {
+                        port_str
+                            .parse()
+                            .ok()
+                            .map(|port| Server::custom(host.into(), port))
+                    })
+                })
+                .collect::<Option<_>>()
+                .map(Self::Custom)
+                .ok_or(anyhow!("'{}' must be one of ['ecc', 'ywallet', 'zecrocks'], or a comma-separated list of host:port", s)),
+        }
+    }
+
+    pub(crate) fn pick(&self, network: Network) -> anyhow::Result<&Server<'_>> {
+        // For now just use the first server in the list.
         match self {
-            consensus::Network::MainNetwork => "mainnet.lightwalletd.com",
-            consensus::Network::TestNetwork => "lightwalletd.testnet.electriccoin.co",
+            Servers::Hosted(server_operator) => server_operator
+                .servers(network)
+                .first()
+                .ok_or(anyhow!("{:?} doesn't serve {:?}", server_operator, network)),
+            Servers::Custom(servers) => Ok(servers.first().expect("not empty")),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct Server<'a> {
+    host: Cow<'a, str>,
+    port: u16,
+}
+
+impl<'a> fmt::Display for Server<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.host, self.port)
+    }
+}
+
+impl Server<'static> {
+    const fn fixed(host: &'static str, port: u16) -> Self {
+        Self {
+            host: Cow::Borrowed(host),
+            port,
+        }
+    }
+}
+
+impl<'a> Server<'a> {
+    fn custom(host: String, port: u16) -> Self {
+        Self {
+            host: Cow::Owned(host),
+            port,
         }
     }
 
-    fn port(&self) -> u16 {
-        9067
+    fn endpoint(&self) -> String {
+        format!("https://{}:{}", self.host, self.port)
     }
 }
 
 pub(crate) async fn connect_to_lightwalletd(
-    network: &impl Lightwalletd,
+    server: &Server<'_>,
 ) -> Result<CompactTxStreamerClient<Channel>, anyhow::Error> {
-    info!("Connecting to {}:{}", network.host(), network.port());
+    info!("Connecting to {}", server);
 
-    let tls = ClientTlsConfig::new().domain_name(network.host());
+    let tls = ClientTlsConfig::new().domain_name(server.host.to_string());
 
-    let channel = Channel::from_shared(format!("https://{}:{}", network.host(), network.port()))?
+    let channel = Channel::from_shared(server.endpoint())?
         .tls_config(tls)?
         .connect()
         .await?;