Merge commit '6c0a672d6643b82b336160da5d3c05cbd2167967' into merge_upstream

Cargo.lock

@@ -2504,7 +2504,7 @@ dependencies = [
 
 [[package]]
 name = "zecwallet-cli"
-version = "1.3.3"
+version = "1.4.1"
 dependencies = [
  "byteorder 1.3.2 (registry+https://github.com/rust-lang/crates.io-index)",
  "clap 2.33.0 (registry+https://github.com/rust-lang/crates.io-index)",

cli/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "zecwallet-cli"
-version = "1.3.3"
+version = "1.4.1"
 edition = "2018"
 
 [dependencies]

cli/src/version.rs

@@ -1 +1 @@
-pub const VERSION:&str = "1.3.3";
+pub const VERSION:&str = "1.4.1";

lib/Cargo.toml

@@ -29,7 +29,7 @@ ring = "0.16.9"
 libflate = "0.1"
 subtle = "2"
 threadpool = "1.8.0"
-num_cpus = "1.13.0"
+num_cpus = "1.12.0"
 
 tonic = { version = "0.2.1", features = ["tls", "tls-roots"] }
 bytes = "0.4"

lib/src/commands.rs

@@ -307,16 +307,6 @@ impl Command for EncryptCommand {
             return self.help();
         }
 
-        // Refuse to encrypt if the bip39 bug has not been fixed
-        use crate::lightwallet::bugs::BugBip39Derivation;
-        if BugBip39Derivation::has_bug(lightclient) {
-            let mut h = vec![];
-            h.push("It looks like your wallet has the bip39bug. Please run 'fixbip39bug' to fix it");
-            h.push("before encrypting your wallet.");
-            h.push("ERROR: Cannot encrypt while wallet has the bip39bug.");
-            return h.join("\n");
-        }
-
         let passwd = args[0].to_string();
 
         match lightclient.wallet.write().unwrap().encrypt(passwd) {
@@ -624,8 +614,9 @@ impl Command for TransactionsCommand {
         let mut h = vec![];
         h.push("List all incoming and outgoing transactions from this wallet");
         h.push("Usage:");
-        h.push("list");
+        h.push("list [allmemos]");
         h.push("");
+        h.push("If you include the 'allmemos' argument, all memos are returned in their raw hex format");
 
         h.join("\n")
     }
@@ -634,8 +625,107 @@ impl Command for TransactionsCommand {
         "List all transactions in the wallet".to_string()
     }
 
-    fn exec(&self, _args: &[&str], lightclient: &LightClient) -> String {
+    fn exec(&self, args: &[&str], lightclient: &LightClient) -> String {
-        format!("{}", lightclient.do_list_transactions().pretty(2))
+        if args.len() > 1 {
+            return format!("Didn't understand arguments\n{}", self.help());
+        }
+
+        let include_memo_hex = if args.len() == 1 {
+            if args[0] == "allmemos" || args[0] == "true" || args[0] == "yes" {
+                true
+            } else {
+                return format!("Couldn't understand first argument '{}'\n{}", args[0], self.help());
+            }
+        } else {
+            false
+        };
+
+        format!("{}", lightclient.do_list_transactions(include_memo_hex).pretty(2))
+    }
+}
+
+struct ImportCommand {}
+impl Command for ImportCommand {
+    fn help(&self) -> String {
+        let mut h = vec![];
+        h.push("Import an external spending or viewing key into the wallet");
+        h.push("Usage:");
+        h.push("import <spending_key | viewing_key> <birthday> [norescan]");
+        h.push("OR");
+        h.push("import '{'key': <spending_key or viewing_key>, 'birthday': <birthday>, 'norescan': <true>}'");
+        h.push("");
+        h.push("Birthday is the earliest block number that has transactions belonging to the imported key. Rescanning will start from this block. If not sure, you can specify '0', which will start rescanning from the first sapling block.");
+        h.push("Note that you can import only the full spending (private) key or the full viewing key.");
+
+        h.join("\n")
+    }
+
+
+    fn short_help(&self) -> String {
+        "Import spending or viewing keys into the wallet".to_string()
+    }
+
+    fn exec(&self, args: &[&str], lightclient: &LightClient) -> String {
+        if args.len() == 0 || args.len() > 3 {
+            return format!("Insufficient arguments\n\n{}", self.help());
+        }
+
+        let (key, birthday, rescan) = if args.len() == 1 {
+            // If only one arg, parse it as JSON
+            let json_args = match json::parse(&args[0]) {
+                Ok(j) => j,
+                Err(e) => {
+                    let es = format!("Couldn't understand JSON: {}", e);
+                    return format!("{}\n{}", es, self.help());
+                }
+            };
+
+            if !json_args.is_object() {
+                return format!("Couldn't parse argument as a JSON object\n{}", self.help());
+            }
+
+            if !json_args.has_key("key") {
+                return format!("'key' field is required in the JSON, containing the spending or viewing key to import\n{}", self.help());
+            }
+
+            if !json_args.has_key("birthday") {
+                return format!("'birthday' field is required in the JSON, containing the birthday of the spending or viewing key\n{}", self.help());
+            }
+
+            (json_args["key"].as_str().unwrap().to_string(), json_args["birthday"].as_u64().unwrap(), !json_args.has_key("norescan"))
+        } else {
+            let key = args[0];
+            let birthday = match args[1].parse::<u64>() {
+                Ok(b) => b,
+                Err(_) => return format!("Couldn't parse {} as birthday. Please specify an integer. Ok to use '0'", args[1]),
+            };
+    
+            let rescan = if args.len() == 3 {
+                if args[2] == "norescan" || args[2] == "false" || args[2] == "no" { 
+                    false 
+                } else {
+                    return format!("Couldn't undestand the argument '{}'. Please pass 'norescan' to prevent rescanning the wallet", args[2]);
+                }
+            } else {
+                true
+            };
+
+            (key.to_string(), birthday, rescan)
+        };
+
+        let r = match lightclient.do_import_key(key, birthday) {
+            Ok(r) => r.pretty(2),
+            Err(e) => return format!("Error: {}", e),
+        };
+
+        if rescan {
+            match lightclient.do_rescan() {
+                Ok(_) => {},
+                Err(e) => return format!("Error: Rescan failed: {}", e),
+            };
+        }
+
+        return r;
     }
 }
 
@@ -645,7 +735,7 @@ impl Command for HeightCommand {
         let mut h = vec![];
         h.push("Get the latest block height that the wallet is at.");
         h.push("Usage:");
-        h.push("height [do_sync = true | false]");
+        h.push("height");
         h.push("");
         h.push("Pass 'true' (default) to sync to the server to get the latest block height. Pass 'false' to get the latest height in the wallet without checking with the server.");
 
@@ -656,11 +746,7 @@ impl Command for HeightCommand {
         "Get the latest block height that the wallet is at".to_string()
     }
 
-    fn exec(&self, args: &[&str], lightclient: &LightClient) -> String {
+    fn exec(&self, _args: &[&str], lightclient: &LightClient) -> String {
-        if args.len() > 1 {
-            return format!("Didn't understand arguments\n{}", self.help());
-        }
-
         format!("{}", object! { "height" => lightclient.last_scanned_height()}.pretty(2))
     }
 }
@@ -733,29 +819,6 @@ impl Command for NotesCommand {
     }
 }
 
-struct FixBip39BugCommand {}
-impl Command for FixBip39BugCommand {
-    fn help(&self)  -> String {
-        let mut h = vec![];
-        h.push("Detect if the wallet has the Bip39 derivation bug, and fix it automatically");
-        h.push("Usage:");
-        h.push("fixbip39bug");
-        h.push("");
-
-        h.join("\n")
-    }
-
-    fn short_help(&self) -> String {
-        "Detect if the wallet has the Bip39 derivation bug, and fix it automatically".to_string()
-    }
-
-    fn exec(&self, _args: &[&str], lightclient: &LightClient) -> String {
-        use crate::lightwallet::bugs::BugBip39Derivation;
-
-        BugBip39Derivation::fix_bug(lightclient)
-    }
-}
-
 struct QuitCommand {}
 impl Command for QuitCommand {
     fn help(&self)  -> String {
@@ -792,6 +855,7 @@ pub fn get_commands() -> Box<HashMap<String, Box<dyn Command>>> {
     map.insert("balance".to_string(),           Box::new(BalanceCommand{}));
     map.insert("addresses".to_string(),         Box::new(AddressCommand{}));
     map.insert("height".to_string(),            Box::new(HeightCommand{}));
+    map.insert("import".to_string(),            Box::new(ImportCommand{}));
     map.insert("export".to_string(),            Box::new(ExportCommand{}));
     map.insert("info".to_string(),              Box::new(InfoCommand{}));
     map.insert("send".to_string(),              Box::new(SendCommand{}));
@@ -805,7 +869,6 @@ pub fn get_commands() -> Box<HashMap<String, Box<dyn Command>>> {
     map.insert("decrypt".to_string(),           Box::new(DecryptCommand{}));
     map.insert("unlock".to_string(),            Box::new(UnlockCommand{}));
     map.insert("lock".to_string(),              Box::new(LockCommand{}));
-    map.insert("fixbip39bug".to_string(),       Box::new(FixBip39BugCommand{}));
 
     Box::new(map)
 }

lib/src/lightclient.rs

@@ -6,7 +6,7 @@ use std::sync::{Arc, RwLock, Mutex, mpsc::channel};
 use std::sync::atomic::{AtomicI32, AtomicUsize, Ordering};
 use std::path::{Path, PathBuf};
 use std::fs::File;
-use std::collections::HashMap;
+use std::collections::{HashSet, HashMap};
 use std::cmp::{max, min};
 use std::io;
 use std::io::prelude::*;
@@ -18,9 +18,7 @@ use threadpool::ThreadPool;
 
 use json::{object, array, JsonValue};
 use zcash_primitives::transaction::{TxId, Transaction};
-use zcash_client_backend::{
+use zcash_client_backend::{constants::testnet, constants::mainnet, constants::regtest,};
-    constants::testnet, constants::mainnet, constants::regtest, encoding::encode_payment_address,
-};
 
 use log::{info, warn, error, LevelFilter};
 use log4rs::append::rolling_file::RollingFileAppender;
@@ -148,6 +146,9 @@ impl LightClientConfig {
                 zcash_data_location = dirs::data_dir().expect("Couldn't determine app data directory!");
                 zcash_data_location.push("Zero");
             } else {
+                if dirs::home_dir().is_none() {
+                    info!("Couldn't determine home dir!");
+                }
                 zcash_data_location = dirs::home_dir().expect("Couldn't determine home directory!");
                 zcash_data_location.push(".zero");
             };
@@ -175,6 +176,28 @@ impl LightClientConfig {
         zcash_data_location.into_boxed_path()
     }
 
+    pub fn get_zcash_params_path(&self) -> io::Result<Box<Path>> {
+        if dirs::home_dir().is_none() {
+            return Err(io::Error::new(io::ErrorKind::InvalidData, "Couldn't determine Home Dir"));
+        }
+
+        let mut zcash_params = self.get_zcash_data_path().into_path_buf();
+        zcash_params.push("..");
+        if cfg!(target_os="macos") || cfg!(target_os="windows") {
+            zcash_params.push("ZcashParams");
+        } else {
+            zcash_params.push(".zcash-params");
+        }
+
+        match std::fs::create_dir_all(zcash_params.clone()) {
+            Ok(_) => Ok(zcash_params.into_boxed_path()),
+            Err(e) => {
+                eprintln!("Couldn't create zcash params directory\n{}", e);
+                Err(e)
+            }
+        }
+    }
+
     pub fn get_wallet_path(&self) -> Box<Path> {
         let mut wallet_location = self.get_zcash_data_path().into_path_buf();
         wallet_location.push(WALLET_NAME);
@@ -253,6 +276,15 @@ impl LightClientConfig {
         }
     }
 
+    pub fn hrp_sapling_viewing_key(&self) -> &str {
+        match &self.chain_name[..] {
+            "main"    => mainnet::HRP_SAPLING_EXTENDED_FULL_VIEWING_KEY,
+            "test"    => testnet::HRP_SAPLING_EXTENDED_FULL_VIEWING_KEY,
+            "regtest" => regtest::HRP_SAPLING_EXTENDED_FULL_VIEWING_KEY,
+            c         => panic!("Unknown chain {}", c)
+        }
+    }
+
     pub fn base58_pubkey_address(&self) -> [u8; 2] {
         match &self.chain_name[..] {
             "main"    => mainnet::B58_PUBKEY_ADDRESS_PREFIX,
@@ -308,6 +340,17 @@ impl LightClient {
         };
     }
 
+    fn write_file_if_not_exists(dir: &Box<Path>, name: &str, bytes: &[u8]) -> io::Result<()> {
+        let mut file_path = dir.to_path_buf();
+        file_path.push(name);
+        if !file_path.exists() {
+            let mut file = File::create(&file_path)?;
+            file.write_all(bytes)?;
+        }
+
+        Ok(())
+    }
+
     #[cfg(feature = "embed_params")]
     fn read_sapling_params(&mut self) {
         // Read Sapling Params
@@ -339,6 +382,27 @@ impl LightClient {
             self.sapling_spend.extend_from_slice(sapling_spend);
         }
 
+        // Ensure that the sapling params are stored on disk properly as well. Only on desktop
+        if cfg!(all(not(target_os="ios"), not(target_os="android"))) {
+            match self.config.get_zcash_params_path() {
+                Ok(zcash_params_dir) => {
+                    // Create the sapling output and spend params files
+                    match LightClient::write_file_if_not_exists(&zcash_params_dir, "sapling-output.params", &self.sapling_output) {
+                        Ok(_) => {},
+                        Err(e) => eprintln!("Warning: Couldn't write the output params!\n{}", e)
+                    };
+
+                    match LightClient::write_file_if_not_exists(&zcash_params_dir, "sapling-spend.params", &self.sapling_spend) {
+                        Ok(_) => {},
+                        Err(e) => eprintln!("Warning: Couldn't write the output params!\n{}", e)
+                    }
+                },
+                Err(e) => {
+                    eprintln!("{}", e);
+                }
+            };
+        }
+
         Ok(())
     }
 
@@ -477,12 +541,6 @@ impl LightClient {
         info!("Read wallet with birthday {}", lc.wallet.read().unwrap().get_first_tx_block());
         info!("Created LightClient to {}", &config.server);
 
-        if crate::lightwallet::bugs::BugBip39Derivation::has_bug(&lc) {
-            let m = format!("WARNING!!!\nYour wallet has a bip39derivation bug that's showing incorrect addresses.\nPlease run 'fixbip39bug' to automatically fix the address derivation in your wallet!\nPlease see: https://github.com/adityapk00/zecwallet-light-cli/blob/master/bip39bug.md");
-             info!("{}", m);
-             println!("{}", m);
-        }
-
         Ok(lc)
     }
 
@@ -577,11 +635,12 @@ impl LightClient {
         let wallet = self.wallet.read().unwrap();
         // Go over all z addresses
         let z_keys = wallet.get_z_private_keys().iter()
-            .filter( move |(addr, _)| address.is_none() || address.as_ref() == Some(addr))
+            .filter( move |(addr, _, _)| address.is_none() || address.as_ref() == Some(addr))
-            .map( |(addr, pk)|
+            .map( |(addr, pk, vk)|
                 object!{
                     "address"     => addr.clone(),
-                    "private_key" => pk.clone()
+                    "private_key" => pk.clone(),
+                    "viewing_key" => vk.clone(),
                 }
             ).collect::<Vec<JsonValue>>();
 
@@ -609,9 +668,7 @@ impl LightClient {
         let wallet = self.wallet.read().unwrap();
 
         // Collect z addresses
-        let z_addresses = wallet.zaddress.read().unwrap().iter().map( |ad| {
+        let z_addresses = wallet.get_all_zaddresses();
-            encode_payment_address(self.config.hrp_sapling_address(), &ad)
-        }).collect::<Vec<String>>();
 
         // Collect t addresses
         let t_addresses = wallet.taddresses.read().unwrap().iter().map( |a| a.clone() )
@@ -627,12 +684,13 @@ impl LightClient {
         let wallet = self.wallet.read().unwrap();
 
         // Collect z addresses
-        let z_addresses = wallet.zaddress.read().unwrap().iter().map( |ad| {
+        let z_addresses = wallet.get_all_zaddresses().iter().map(|zaddress| {
-            let address = encode_payment_address(self.config.hrp_sapling_address(), &ad);
             object!{
-                "address" => address.clone(),
+                "address" => zaddress.clone(),
-                "zbalance" => wallet.zbalance(Some(address.clone())),
+                "zbalance" => wallet.zbalance(Some(zaddress.clone())),
-                "verified_zbalance" => wallet.verified_zbalance(Some(address)),
+                "verified_zbalance"  => wallet.verified_zbalance(Some(zaddress.clone())),
+                "spendable_zbalance" => wallet.spendable_zbalance(Some(zaddress.clone())),
+                "unverified_zbalance"   => wallet.unverified_zbalance(Some(zaddress.clone()))
             }
         }).collect::<Vec<JsonValue>>();
 
@@ -650,6 +708,8 @@ impl LightClient {
         object!{
             "zbalance"           => wallet.zbalance(None),
             "verified_zbalance"  => wallet.verified_zbalance(None),
+            "spendable_zbalance" => wallet.spendable_zbalance(None),
+            "unverified_zbalance"   => wallet.unverified_zbalance(None),
             "tbalance"           => wallet.tbalance(None),
             "z_addresses"        => z_addresses,
             "t_addresses"        => t_addresses,
@@ -769,23 +829,40 @@ impl LightClient {
         let mut spent_notes  : Vec<JsonValue> = vec![];
         let mut pending_notes: Vec<JsonValue> = vec![];
 
+        let anchor_height: i32 = self.wallet.read().unwrap().get_anchor_height() as i32;
+
         {
-            // Collect Sapling notes
             let wallet = self.wallet.read().unwrap();
+
+            // First, collect all extfvk's that are spendable (i.e., we have the private key)
+            let spendable_address: HashSet<String> = wallet.get_all_zaddresses().iter()
+                .filter(|address| wallet.have_spending_key_for_zaddress(address))
+                .map(|address| address.clone())
+                .collect();
+
+            // Collect Sapling notes
             wallet.txs.read().unwrap().iter()
                 .flat_map( |(txid, wtx)| {
+                    let spendable_address = spendable_address.clone();
                     wtx.notes.iter().filter_map(move |nd|
                         if !all_notes && nd.spent.is_some() {
                             None
                         } else {
+                            let address = LightWallet::note_address(self.config.hrp_sapling_address(), nd);
+                            let spendable = address.is_some() &&
+                                                    spendable_address.contains(&address.clone().unwrap()) &&
+                                                    wtx.block <= anchor_height && nd.spent.is_none() && nd.unconfirmed_spent.is_none();
+
                             Some(object!{
                                 "created_in_block"   => wtx.block,
                                 "datetime"           => wtx.datetime,
                                 "created_in_txid"    => format!("{}", txid),
                                 "value"              => nd.note.value,
                                 "is_change"          => nd.is_change,
-                                "address"            => LightWallet::note_address(self.config.hrp_sapling_address(), nd),
+                                "address"            => address,
+                                "spendable"          => spendable,
                                 "spent"              => nd.spent.map(|spent_txid| format!("{}", spent_txid)),
+                                "spent_at_height"    => nd.spent_at_height.map(|h| format!("{}", h)),
                                 "unconfirmed_spent"  => nd.unconfirmed_spent.map(|spent_txid| format!("{}", spent_txid)),
                             })
                         }
@@ -862,7 +939,7 @@ impl LightClient {
         }
     }
 
-    pub fn do_list_transactions(&self) -> JsonValue {
+    pub fn do_list_transactions(&self, include_memo_hex: bool) -> JsonValue {
         let wallet = self.wallet.read().unwrap();
 
         // Create a list of TransactionItems from wallet txns
@@ -883,22 +960,39 @@ impl LightClient {
                 let mut incoming_json = v.notes.iter()
                     .filter( |nd| !nd.is_change )
                     .enumerate()
-                    .map ( |(_i, nd)|
+                    .map ( |(_i, nd)| {
-                        object! {
+                        let mut o = object! {
                             "address"      => LightWallet::note_address(self.config.hrp_sapling_address(), nd),
                             "value"       => nd.note.value as i64,
                             "memo"         => LightWallet::memo_str(&nd.memo),
+                        };
+
+                        if include_memo_hex {
+                            o.insert("memohex", match &nd.memo {
+                                Some(m) => hex::encode(m.as_bytes()),
+                                _ => "".to_string(),
+                            }).unwrap();
+                        }
+
+                        return o;
                     })
                     .collect::<Vec<JsonValue>>();
 
                 let incoming_t_json = v.utxos.iter()
                     .filter(|u| !change_addresses.contains(&u.address))
-                    .map( |uo|
+                    .map( |uo| {
-                        object! {
+                        let mut o = object! {
                             "address"       => uo.address.clone(),
                             "value"         => uo.value.clone() as i64,
                             "memo"          => None::<String>,
-                        })
+                        };
+
+                        if include_memo_hex {
+                            o.insert("memohex", None::<String>).unwrap();
+                        }
+
+                        return o;
+                    })
                     .collect::<Vec<JsonValue>>();
 
                 for json in incoming_t_json {
@@ -909,22 +1003,39 @@ impl LightClient {
                 let mut incoming_change_json = v.notes.iter()
                     .filter( |nd| nd.is_change )
                     .enumerate()
-                    .map ( |(_i, nd)|
+                    .map ( |(_i, nd)| {
-                        object! {
+                        let mut o = object! {
                             "address"      => LightWallet::note_address(self.config.hrp_sapling_address(), nd),
                             "value"       => nd.note.value as i64,
                             "memo"         => LightWallet::memo_str(&nd.memo),
+                        };
+
+                        if include_memo_hex {
+                            o.insert("memohex", match &nd.memo {
+                                Some(m) => hex::encode(m.as_bytes()),
+                                _ => "".to_string(),
+                            }).unwrap();
+                        }
+
+                        return o;
                     })
                     .collect::<Vec<JsonValue>>();
 
                 let incoming_t_change_json = v.utxos.iter()
                     .filter(|u| change_addresses.contains(&u.address))
-                    .map( |uo|
+                    .map( |uo| {
-                        object! {
+                        let mut o = object! {
                             "address"       => uo.address.clone(),
                             "value"         => uo.value.clone() as i64,
                             "memo"          => None::<String>,
-                        })
+                        };
+
+                        if include_memo_hex {
+                            o.insert("memohex", None::<String>).unwrap();
+                        }
+
+                        return o;
+                    })
                     .collect::<Vec<JsonValue>>();
 
                 for json in incoming_t_change_json {
@@ -933,21 +1044,35 @@ impl LightClient {
 
                 // Collect outgoing metadata
                 let outgoing_json = v.outgoing_metadata.iter()
-                    .map(|om|
+                    .map(|om| {
-                        object!{
+                        let mut o = object!{
                             "address" => om.address.clone(),
                             "value"   => om.value,
                             "memo"    => LightWallet::memo_str(&Some(om.memo.clone())),
+                        };
+
+                        if include_memo_hex {
+                            o.insert("memohex", hex::encode(om.memo.as_bytes())).unwrap();
+                        }
+
+                        return o;
                     })
                     .collect::<Vec<JsonValue>>();
 
                 // Collect outgoing metadata change
                 let outgoing_change_json = v.outgoing_metadata_change.iter()
-                    .map(|om|
+                    .map(|om| {
-                        object!{
+                        let mut o = object!{
                             "address" => om.address.clone(),
                             "value"   => om.value,
                             "memo"    => LightWallet::memo_str(&Some(om.memo.clone())),
+                        };
+
+                        if include_memo_hex {
+                            o.insert("memohex", hex::encode(om.memo.as_bytes())).unwrap();
+                        }
+
+                        return o;
                     })
                     .collect::<Vec<JsonValue>>();
 
@@ -979,12 +1104,20 @@ impl LightClient {
 
             // Collect outgoing metadata
             let outgoing_json = wtx.outgoing_metadata.iter()
-                .map(|om|
+                .map(|om| {
-                    object!{
+                    let mut o = object!{
                         "address" => om.address.clone(),
                         "value"   => om.value,
                         "memo"    => LightWallet::memo_str(&Some(om.memo.clone())),
-                }).collect::<Vec<JsonValue>>();
+                    };
+
+                    if include_memo_hex {
+                        o.insert("memohex", hex::encode(om.memo.as_bytes())).unwrap();
+                    }
+
+                    return o;
+                })
+                .collect::<Vec<JsonValue>>();
 
             object! {
                 "block_height" => wtx.block,
@@ -1017,7 +1150,7 @@ impl LightClient {
         let new_address = {
             let wallet = self.wallet.write().unwrap();
 
-            match addr_type {
+            let addr = match addr_type {
                 "z" => wallet.add_zaddr(),
                 "t" => wallet.add_taddr(),
                 _   => {
@@ -1025,7 +1158,15 @@ impl LightClient {
                     error!("{}", e);
                     return Err(e);
                 }
+            };
+
+            if addr.starts_with("Error") {
+                let e = format!("Error creating new address: {}", addr);
+                    error!("{}", e);
+                    return Err(e);
             }
+
+            addr
         };
 
         self.do_save()?;
@@ -1033,6 +1174,69 @@ impl LightClient {
         Ok(array![new_address])
     }
 
+    /// Convinence function to determine what type of key this is and import it
+    pub fn do_import_key(&self, key: String, birthday: u64) -> Result<JsonValue, String> {
+        if key.starts_with(self.config.hrp_sapling_private_key()) {
+            self.do_import_sk(key, birthday)
+        } else if key.starts_with(self.config.hrp_sapling_viewing_key()) {
+            self.do_import_vk(key, birthday)
+        } else {
+            Err(format!("'{}' was not recognized as either a spending key or a viewing key because it didn't start with either '{}' or '{}'",
+                key, self.config.hrp_sapling_private_key(), self.config.hrp_sapling_viewing_key()))
+        }
+    }
+
+    /// Import a new private key
+    pub fn do_import_sk(&self, sk: String, birthday: u64) -> Result<JsonValue, String> {
+        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
+            error!("Wallet is locked");
+            return Err("Wallet is locked".to_string());
+        }
+
+        let new_address = {
+            let mut wallet = self.wallet.write().unwrap();
+
+            let addr = wallet.add_imported_sk(sk, birthday);
+            if addr.starts_with("Error") {
+                let e = format!("Error creating new address{}", addr);
+                    error!("{}", e);
+                    return Err(e);
+            }
+
+            addr
+        };
+
+        self.do_save()?;
+
+        Ok(array![new_address])
+    }
+
+    /// Import a new viewing key
+    pub fn do_import_vk(&self, vk: String, birthday: u64) -> Result<JsonValue, String> {
+        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
+            error!("Wallet is locked");
+            return Err("Wallet is locked".to_string());
+        }
+
+        let new_address = {
+            let mut wallet = self.wallet.write().unwrap();
+
+            let addr = wallet.add_imported_vk(vk, birthday);
+            if addr.starts_with("Error") {
+                let e = format!("Error creating new address{}", addr);
+                    error!("{}", e);
+                    return Err(e);
+            }
+
+            addr
+        };
+
+        self.do_save()?;
+
+        Ok(array![new_address])
+    }
+
+
     pub fn clear_state(&self) {
         // First, clear the state from the wallet
         self.wallet.read().unwrap().clear_blocks();
@@ -1377,20 +1581,20 @@ impl LightClient {
 
         info!("Creating transaction");
 
-        let rawtx = self.wallet.write().unwrap().send_to_address(
+        let result = {
-            u32::from_str_radix(&self.config.consensus_branch_id, 16).unwrap(),
+            let _lock = self.sync_lock.lock().unwrap();
-            &self.sapling_spend, &self.sapling_output,
+
-            from, addrs, fee
+            self.wallet.write().unwrap().send_to_address(
-        );
+                u32::from_str_radix(&self.config.consensus_branch_id, 16).unwrap(),
+                &self.sapling_spend, &self.sapling_output,
+                from, addrs, fee,
+                |txbytes| broadcast_raw_tx(&self.get_server_uri(), txbytes)
+            )
+        };
 
         info!("Transaction Complete");
 
-
+        result.map(|(txid, _)| txid)
-
-        match rawtx {
-            Ok(txbytes)   => broadcast_raw_tx(&self.get_server_uri(), txbytes),
-            Err(e)        => Err(format!("Error: No Tx to broadcast. Error was: {}", e))
-        }
     }
 }
 
@@ -1433,6 +1637,14 @@ pub mod tests {
         assert!(!lc.do_new_address("z").is_err());
     }
 
+    #[test]
+    pub fn test_bad_import() {
+        let lc = super::LightClient::unconnected(TEST_SEED.to_string(), None).unwrap();
+
+        assert!(lc.do_import_sk("bad_priv_key".to_string(), 0).is_err());
+        assert!(lc.do_import_vk("bad_view_key".to_string(), 0).is_err());
+    }
+
     #[test]
     pub fn test_addresses() {
         let lc = super::LightClient::unconnected(TEST_SEED.to_string(), None).unwrap();

lib/src/lightwallet.rs

@@ -23,8 +23,10 @@ use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use pairing::bls12_381::{Bls12};
 use sha2::{Sha256, Digest};
 
+use sodiumoxide::crypto::secretbox;
+
 use zcash_client_backend::{
-    encoding::{encode_payment_address, encode_extended_spending_key},
+    encoding::{encode_payment_address, encode_extended_spending_key, encode_extended_full_viewing_key, decode_extended_spending_key, decode_extended_full_viewing_key},
     proto::compact_formats::{CompactBlock, CompactOutput},
     wallet::{WalletShieldedOutput, WalletShieldedSpend}
 };
@@ -55,13 +57,15 @@ mod extended_key;
 mod utils;
 mod address;
 mod prover;
-pub mod bugs;
+mod walletzkey;
 
 use data::{BlockData, WalletTx, Utxo, SaplingNoteData, SpendableNote, OutgoingTxMetadata};
 use extended_key::{KeyIndex, ExtendedPrivKey};
+use walletzkey::{WalletZKey, WalletZKeyType};
 
 pub const MAX_REORG: usize = 100;
-pub const GAP_RULE_UNUSED_ADDRESSES: usize = 5;
+
+pub const GAP_RULE_UNUSED_ADDRESSES: usize = 0;
 
 fn now() -> f64 {
     SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs() as f64
@@ -112,12 +116,9 @@ pub struct LightWallet {
 
     seed: [u8; 32],    // Seed phrase for this wallet. If wallet is locked, this is 0
 
-    // List of keys, actually in this wallet. If the wallet is locked, the `extsks` will be
+    // List of keys, actually in this wallet. This is a combination of HD keys derived from the seed,
-    // encrypted (but the fvks are not encrpyted)
+    // viewing keys and imported spending keys.
-    extsks:  Arc<RwLock<Vec<ExtendedSpendingKey>>>,
+    zkeys: Arc<RwLock<Vec<WalletZKey>>>,
-    extfvks: Arc<RwLock<Vec<ExtendedFullViewingKey>>>,
-
-    pub zaddress: Arc<RwLock<Vec<PaymentAddress<Bls12>>>>,
 
     // Transparent keys. If the wallet is locked, then the secret keys will be encrypted,
     // but the addresses will be present.
@@ -143,7 +144,7 @@ pub struct LightWallet {
 
 impl LightWallet {
     pub fn serialized_version() -> u64 {
-        return 6;
+        return 8;
     }
 
     fn get_taddr_from_bip39seed(config: &LightClientConfig, bip39_seed: &[u8], pos: u32) -> SecretKey {
@@ -219,8 +220,9 @@ impl LightWallet {
 
         // TODO: We need to monitor addresses, and always keep 1 "free" address, so
         // users can import a seed phrase and automatically get all used addresses
-        let (extsk, extfvk, address)
+        let hdkey_num = 0;
-            = LightWallet::get_zaddr_from_bip39seed(&config, &bip39_seed.as_bytes(), 0);
+        let (extsk, _, _)
+            = LightWallet::get_zaddr_from_bip39seed(&config, &bip39_seed.as_bytes(), hdkey_num);
 
         let lw = LightWallet {
             encrypted:   false,
@@ -228,9 +230,7 @@ impl LightWallet {
             enc_seed:    [0u8; 48],
             nonce:       vec![],
             seed:        seed_bytes,
-            extsks:      Arc::new(RwLock::new(vec![extsk])),
+            zkeys:       Arc::new(RwLock::new(vec![WalletZKey::new_hdkey(hdkey_num, extsk)])),
-            extfvks:     Arc::new(RwLock::new(vec![extfvk])),
-            zaddress:    Arc::new(RwLock::new(vec![address])),
             tkeys:       Arc::new(RwLock::new(vec![tpk])),
             taddresses:  Arc::new(RwLock::new(vec![taddr])),
             blocks:      Arc::new(RwLock::new(vec![])),
@@ -294,22 +294,59 @@ impl LightWallet {
         let mut seed_bytes = [0u8; 32];
         reader.read_exact(&mut seed_bytes)?;
 
-        // Read the spending keys
+        let zkeys = if version <= 6 {
-        let extsks = Vector::read(&mut reader, |r| ExtendedSpendingKey::read(r))?;
+            // Up until version 6, the wallet keys were written out individually
+            // Read the spending keys
+            let extsks = Vector::read(&mut reader, |r| ExtendedSpendingKey::read(r))?;
 
-        let extfvks = if version >= 4 {
+            let extfvks = if version >= 4 {
-            // Read the viewing keys
+                // Read the viewing keys
-            Vector::read(&mut reader, |r| ExtendedFullViewingKey::read(r))?
+                Vector::read(&mut reader, |r| ExtendedFullViewingKey::read(r))?
-        } else {
+            } else {
-            // Calculate the viewing keys
+                // Calculate the viewing keys
-            extsks.iter().map(|sk| ExtendedFullViewingKey::from(sk))
+                extsks.iter().map(|sk| ExtendedFullViewingKey::from(sk))
-                .collect::<Vec<ExtendedFullViewingKey>>()
+                    .collect::<Vec<ExtendedFullViewingKey>>()
+            };
+
+            // Calculate the addresses
+            let addresses = extfvks.iter().map( |fvk| fvk.default_address().unwrap().1 )
+                .collect::<Vec<PaymentAddress<Bls12>>>();
+
+            // If extsks is of len 0, then this wallet is locked
+            let zkeys_result = if extsks.len() == 0 {
+                // Wallet is locked, so read only the viewing keys.
+                extfvks.iter().zip(addresses.iter()).enumerate().map(|(i, (extfvk, payment_address))| {
+                let zk = WalletZKey::new_locked_hdkey(i as u32, extfvk.clone());
+                if zk.zaddress != *payment_address {
+                    Err(io::Error::new(ErrorKind::InvalidData, "Payment address didn't match"))
+                } else {
+                    Ok(zk)
+                }
+                }).collect::<Vec<io::Result<WalletZKey>>>()
+            } else {
+                // Wallet is unlocked, read the spending keys as well
+                extsks.into_iter().zip(extfvks.into_iter().zip(addresses.iter())).enumerate()
+                .map(|(i, (extsk, (extfvk, payment_address)))| {
+                    let zk = WalletZKey::new_hdkey(i as u32, extsk);
+                    if zk.zaddress != *payment_address {
+                    return Err(io::Error::new(ErrorKind::InvalidData, "Payment address didn't match"));
+                    }
+
+                    if zk.extfvk != extfvk {
+                    return Err(io::Error::new(ErrorKind::InvalidData, "Full View key didn't match"));
+                    }
+
+                    Ok(zk)
+                }).collect::<Vec<io::Result<WalletZKey>>>()
+            };
+
+            // Convert vector of results into result of vector, returning an error if any one of the keys failed the checks above
+            zkeys_result.into_iter().collect::<io::Result<_>>()?
+        }  else {
+            // After version 6, we read the WalletZKey structs directly
+            Vector::read(&mut reader, |r| WalletZKey::read(r))?
         };
 
-        // Calculate the addresses
-        let addresses = extfvks.iter().map( |fvk| fvk.default_address().unwrap().1 )
-            .collect::<Vec<PaymentAddress<Bls12>>>();
-
         let tkeys = Vector::read(&mut reader, |r| {
             let mut tpk_bytes = [0u8; 32];
             r.read_exact(&mut tpk_bytes)?;
@@ -343,15 +380,13 @@ impl LightWallet {
 
         let birthday = reader.read_u64::<LittleEndian>()?;
 
-        Ok(LightWallet{
+        let lw = LightWallet{
             encrypted:   encrypted,
             unlocked:    !encrypted, // When reading from disk, if wallet is encrypted, it starts off locked.
             enc_seed:    enc_seed,
             nonce:       nonce,
             seed:        seed_bytes,
-            extsks:      Arc::new(RwLock::new(extsks)),
+            zkeys:       Arc::new(RwLock::new(zkeys)),
-            extfvks:     Arc::new(RwLock::new(extfvks)),
-            zaddress:    Arc::new(RwLock::new(addresses)),
             tkeys:       Arc::new(RwLock::new(tkeys)),
             taddresses:  Arc::new(RwLock::new(taddresses)),
             blocks:      Arc::new(RwLock::new(blocks)),
@@ -360,7 +395,14 @@ impl LightWallet {
             config:      config.clone(),
             birthday,
             total_scan_duration: Arc::new(RwLock::new(vec![Duration::new(0, 0)])),
-        })
+        };
+
+        // Do a one-time fix of the spent_at_height for older wallets
+        if version <= 7 {
+            lw.fix_spent_at_height();
+        }
+
+        Ok(lw)
     }
 
     pub fn write<W: Write>(&self, mut writer: W) -> io::Result<()> {
@@ -387,14 +429,9 @@ impl LightWallet {
         // Flush after writing the seed, so in case of a disaster, we can still recover the seed.
         writer.flush()?;
 
-        // Write all the spending keys
+        // Write all the wallet's keys
-        Vector::write(&mut writer, &self.extsks.read().unwrap(),
+        Vector::write(&mut writer, &self.zkeys.read().unwrap(),
-             |w, sk| sk.write(w)
+             |w, zk| zk.write(w)
-        )?;
-
-        // Write the FVKs
-        Vector::write(&mut writer, &self.extfvks.read().unwrap(),
-             |w, fvk| fvk.write(w)
         )?;
 
         // Write the transparent private keys
@@ -458,14 +495,20 @@ impl LightWallet {
             .unwrap_or(&cmp::max(self.birthday, self.config.sapling_activation_height))
     }
 
-    // Get all z-address private keys. Returns a Vector of (address, privatekey)
+    // Get all z-address private keys. Returns a Vector of (address, privatekey, viewkey)
-    pub fn get_z_private_keys(&self) -> Vec<(String, String)> {
+    pub fn get_z_private_keys(&self) -> Vec<(String, String, String)> {
-        self.extsks.read().unwrap().iter().map(|sk| {
+        let keys = self.zkeys.read().unwrap().iter().map(|k| {
-            (encode_payment_address(self.config.hrp_sapling_address(),
+            let pkey = match k.extsk.clone().map(|extsk| encode_extended_spending_key(self.config.hrp_sapling_private_key(), &extsk)) {
-                                    &ExtendedFullViewingKey::from(sk).default_address().unwrap().1),
+                Some(pk) => pk,
-             encode_extended_spending_key(self.config.hrp_sapling_private_key(), &sk)
+                None => "".to_string()
-            )
+            };
-        }).collect::<Vec<(String, String)>>()
+
+            let vkey = encode_extended_full_viewing_key(self.config.hrp_sapling_viewing_key(), &k.extfvk);
+
+            (encode_payment_address(self.config.hrp_sapling_address(),&k.zaddress), pkey, vkey)
+        }).collect::<Vec<(String, String, String)>>();
+
+        keys
     }
 
     /// Get all t-address private keys. Returns a Vector of (address, secretkey)
@@ -481,29 +524,34 @@ impl LightWallet {
     /// NOTE: This does NOT rescan
     pub fn add_zaddr(&self) -> String {
         if !self.unlocked {
-            return "".to_string();
+            return "Error: Can't add key while wallet is locked".to_string();
         }
 
-        let pos = self.extsks.read().unwrap().len() as u32;
+        // Find the highest pos we have
+        let pos = self.zkeys.read().unwrap().iter()
+            .filter(|zk| zk.hdkey_num.is_some())
+            .max_by(|zk1, zk2| zk1.hdkey_num.unwrap().cmp(&zk2.hdkey_num.unwrap()))
+            .map_or(0, |zk| zk.hdkey_num.unwrap() + 1);
+
+
         let bip39_seed = bip39::Seed::new(&Mnemonic::from_entropy(&self.seed, Language::English).unwrap(), "");
 
-        let (extsk, extfvk, address) =
+        let (extsk, _, _) =
             LightWallet::get_zaddr_from_bip39seed(&self.config, &bip39_seed.as_bytes(), pos);
 
-        let zaddr = encode_payment_address(self.config.hrp_sapling_address(), &address);
+        // let zaddr = encode_payment_address(self.config.hrp_sapling_address(), &address);
-        self.extsks.write().unwrap().push(extsk);
+        let newkey = WalletZKey::new_hdkey(pos, extsk);
-        self.extfvks.write().unwrap().push(extfvk);
+        self.zkeys.write().unwrap().push(newkey.clone());
-        self.zaddress.write().unwrap().push(address);
 
-        zaddr
+        encode_payment_address(self.config.hrp_sapling_address(), &newkey.zaddress)
     }
 
     /// Add a new t address to the wallet. This will derive a new address from the seed
     /// at the next position.
-    /// NOTE: This is not rescan the wallet
+    /// NOTE: This will not rescan the wallet
     pub fn add_taddr(&self) -> String {
         if !self.unlocked {
-            return "".to_string();
+            return "Error: Can't add key while wallet is locked".to_string();
         }
 
         let pos = self.tkeys.read().unwrap().len() as u32;
@@ -518,6 +566,81 @@ impl LightWallet {
         address
     }
 
+    // Add a new imported spending key to the wallet
+    /// NOTE: This will not rescan the wallet
+    pub fn add_imported_sk(&mut self, sk: String, birthday: u64) -> String {
+        if self.encrypted {
+            return "Error: Can't import spending key while wallet is encrypted".to_string();
+        }
+
+        // First, try to interpret the key
+        let extsk = match decode_extended_spending_key(self.config.hrp_sapling_private_key(), &sk) {
+            Ok(Some(k)) => k,
+            Ok(None) => return format!("Error: Couldn't decode spending key"),
+            Err(e) => return format!("Error importing spending key: {}", e)
+        };
+
+        // Make sure the key doesn't already exist
+        if self.zkeys.read().unwrap().iter().find(|&wk| wk.extsk.is_some() && wk.extsk.as_ref().unwrap() == &extsk.clone()).is_some() {
+            return "Error: Key already exists".to_string();
+        }
+
+        let extfvk = ExtendedFullViewingKey::from(&extsk);
+        let zaddress = {
+            let mut zkeys = self.zkeys.write().unwrap();
+            let maybe_existing_zkey = zkeys.iter_mut().find(|wk| wk.extfvk == extfvk);
+
+            // If the viewing key exists, and is now being upgraded to the spending key, replace it in-place
+            if maybe_existing_zkey.is_some() {
+                let mut existing_zkey = maybe_existing_zkey.unwrap();
+                existing_zkey.extsk = Some(extsk);
+                existing_zkey.keytype = WalletZKeyType::ImportedSpendingKey;
+                existing_zkey.zaddress.clone()
+            } else {
+                let newkey = WalletZKey::new_imported_sk(extsk);
+                zkeys.push(newkey.clone());
+                newkey.zaddress
+            }
+        };
+
+        // Adjust wallet birthday
+        if birthday < self.birthday {
+            self.birthday = if birthday < self.config.sapling_activation_height {self.config.sapling_activation_height} else {birthday};
+        }
+
+        encode_payment_address(self.config.hrp_sapling_address(), &zaddress)
+    }
+
+    // Add a new imported viewing key to the wallet
+    /// NOTE: This will not rescan the wallet
+    pub fn add_imported_vk(&mut self, vk: String, birthday: u64) -> String {
+        if !self.unlocked {
+            return "Error: Can't add key while wallet is locked".to_string();
+        }
+
+        // First, try to interpret the key
+        let extfvk = match decode_extended_full_viewing_key(self.config.hrp_sapling_viewing_key(), &vk) {
+            Ok(Some(k)) => k,
+            Ok(None) => return format!("Error: Couldn't decode viewing key"),
+            Err(e) => return format!("Error importing viewing key: {}", e)
+        };
+
+        // Make sure the key doesn't already exist
+        if self.zkeys.read().unwrap().iter().find(|wk| wk.extfvk == extfvk.clone()).is_some() {
+            return "Error: Key already exists".to_string();
+        }
+
+        let newkey = WalletZKey::new_imported_viewkey(extfvk);
+        self.zkeys.write().unwrap().push(newkey.clone());
+
+        // Adjust wallet birthday
+        if birthday < self.birthday {
+            self.birthday = if birthday < self.config.sapling_activation_height {self.config.sapling_activation_height} else {birthday};
+        }
+
+        encode_payment_address(self.config.hrp_sapling_address(), &newkey.zaddress)
+    }
+
     /// Clears all the downloaded blocks and resets the state back to the initial block.
     /// After this, the wallet's initial state will need to be set
     /// and the wallet will need to be rescanned
@@ -611,6 +734,14 @@ impl LightWallet {
         }
     }
 
+    /// Get the height of the anchor block
+    pub fn get_anchor_height(&self) -> u32 {
+        match self.get_target_height_and_anchor_offset() {
+            Some((height, anchor_offset)) => height - anchor_offset as u32 - 1,
+            None => return 0,
+        }
+    }
+
     pub fn memo_str(memo: &Option<Memo>) -> Option<String> {
         match memo {
             Some(memo) => {
@@ -623,6 +754,12 @@ impl LightWallet {
         }
     }
 
+    pub fn get_all_zaddresses(&self) -> Vec<String> {
+        self.zkeys.read().unwrap().iter().map( |zk| {
+            encode_payment_address(self.config.hrp_sapling_address(), &zk.zaddress)
+        }).collect()
+    }
+
     pub fn address_from_prefix_sk(prefix: &[u8; 2], sk: &secp256k1::SecretKey) -> String {
         let secp = secp256k1::Secp256k1::new();
         let pk = secp256k1::PublicKey::from_secret_key(&secp, &sk);
@@ -661,8 +798,6 @@ impl LightWallet {
     }
 
     pub fn encrypt(&mut self, passwd: String) -> io::Result<()> {
-        use sodiumoxide::crypto::secretbox;
-
         if self.encrypted {
             return Err(io::Error::new(ErrorKind::AlreadyExists, "Wallet is already encrypted"));
         }
@@ -674,8 +809,12 @@ impl LightWallet {
         let cipher = secretbox::seal(&self.seed, &nonce, &key);
 
         self.enc_seed.copy_from_slice(&cipher);
-        self.nonce = vec![];
+        self.nonce = nonce.as_ref().to_vec();
-        self.nonce.extend_from_slice(nonce.as_ref());
+
+        // Encrypt the individual keys
+        self.zkeys.write().unwrap().iter_mut()
+            .map(|k| k.encrypt(&key))
+            .collect::<io::Result<Vec<()>>>()?;
 
         self.encrypted = true;
         self.lock()?;
@@ -695,7 +834,11 @@ impl LightWallet {
         // Empty the seed and the secret keys
         self.seed.copy_from_slice(&[0u8; 32]);
         self.tkeys = Arc::new(RwLock::new(vec![]));
-        self.extsks = Arc::new(RwLock::new(vec![]));
+
+        // Remove all the private key from the zkeys
+        self.zkeys.write().unwrap().iter_mut().map(|zk| {
+            zk.lock()
+        }).collect::<io::Result<Vec<_>>>()?;
 
         self.unlocked = false;
 
@@ -703,8 +846,6 @@ impl LightWallet {
     }
 
     pub fn unlock(&mut self, passwd: String) -> io::Result<()> {
-        use sodiumoxide::crypto::secretbox;
-
         if !self.encrypted {
             return Err(Error::new(ErrorKind::AlreadyExists, "Wallet is not encrypted"));
         }
@@ -729,26 +870,6 @@ impl LightWallet {
         // we need to get the 64 byte bip39 entropy
         let bip39_seed = bip39::Seed::new(&Mnemonic::from_entropy(&seed, Language::English).unwrap(), "");
 
-        // Sapling keys
-        let mut extsks = vec![];
-        for pos in 0..self.zaddress.read().unwrap().len() {
-            let (extsk, extfvk, address) =
-                LightWallet::get_zaddr_from_bip39seed(&self.config, &bip39_seed.as_bytes(), pos as u32);
-
-            if address != self.zaddress.read().unwrap()[pos] {
-                return Err(io::Error::new(ErrorKind::InvalidData,
-                        format!("zaddress mismatch at {}. {:?} vs {:?}", pos, address, self.zaddress.read().unwrap()[pos])));
-            }
-
-            if extfvk != self.extfvks.read().unwrap()[pos] {
-                return Err(io::Error::new(ErrorKind::InvalidData,
-                            format!("fvk mismatch at {}. {:?} vs {:?}", pos, extfvk, self.extfvks.read().unwrap()[pos])));
-            }
-
-            // Don't add it to self yet, we'll do that at the end when everything is verified
-            extsks.push(extsk);
-        }
-
         // Transparent keys
         let mut tkeys = vec![];
         for pos in 0..self.taddresses.read().unwrap().len() {
@@ -763,8 +884,12 @@ impl LightWallet {
             tkeys.push(sk);
         }
 
+        // Go over the zkeys, and add the spending keys again
+        self.zkeys.write().unwrap().iter_mut().map(|zk| {
+            zk.unlock(&self.config, bip39_seed.as_bytes(), &key)
+        }).collect::<io::Result<Vec<()>>>()?;
+
         // Everything checks out, so we'll update our wallet with the decrypted values
-        self.extsks = Arc::new(RwLock::new(extsks));
         self.tkeys = Arc::new(RwLock::new(tkeys));
         self.seed.copy_from_slice(&seed);
 
@@ -786,6 +911,11 @@ impl LightWallet {
             self.unlock(passwd)?;
         }
 
+        // Remove encryption from individual zkeys
+        self.zkeys.write().unwrap().iter_mut().map(|zk| {
+            zk.remove_encryption()
+        }).collect::<io::Result<Vec<()>>>()?;
+
         // Permanantly remove the encryption
         self.encrypted = false;
         self.nonce = vec![];
@@ -847,6 +977,48 @@ impl LightWallet {
             .sum::<u64>()
     }
 
+    pub fn unverified_zbalance(&self, addr: Option<String>) -> u64 {
+        let anchor_height = match self.get_target_height_and_anchor_offset() {
+            Some((height, anchor_offset)) => height - anchor_offset as u32 - 1,
+            None => return 0,
+        };
+
+        self.txs
+            .read()
+            .unwrap()
+            .values()
+            .map(|tx| {
+                tx.notes
+                    .iter()
+                    .filter(|nd| nd.spent.is_none() && nd.unconfirmed_spent.is_none())
+                    .filter(|nd| {
+                        // Check to see if we have this note's spending key.
+                        self.have_spendingkey_for_extfvk(&nd.extfvk)
+                    })
+                    .filter(|nd| {  // TODO, this whole section is shared with verified_balance. Refactor it.
+                        match addr.clone() {
+                            Some(a) => a == encode_payment_address(
+                                                self.config.hrp_sapling_address(),
+                                                &nd.extfvk.fvk.vk
+                                                    .to_payment_address(nd.diversifier, &JUBJUB).unwrap()
+                                            ),
+                            None    => true
+                        }
+                    })
+                    .map(|nd| {
+                        if tx.block as u32 <= anchor_height {
+                            // If confirmed, then unconfirmed is 0
+                            0
+                        } else {
+                            // If confirmed but dont have anchor yet, it is unconfirmed
+                            nd.note.value
+                        }
+                    })
+                    .sum::<u64>()
+            })
+            .sum::<u64>()
+    }
+
     pub fn verified_zbalance(&self, addr: Option<String>) -> u64 {
         let anchor_height = match self.get_target_height_and_anchor_offset() {
             Some((height, anchor_offset)) => height - anchor_offset as u32 - 1,
@@ -861,6 +1033,7 @@ impl LightWallet {
                 if tx.block as u32 <= anchor_height {
                     tx.notes
                         .iter()
+                        .filter(|nd| nd.spent.is_none() && nd.unconfirmed_spent.is_none())
                         .filter(|nd| {  // TODO, this whole section is shared with verified_balance. Refactor it.
                             match addr.clone() {
                                 Some(a) => a == encode_payment_address(
@@ -871,7 +1044,7 @@ impl LightWallet {
                                 None    => true
                             }
                         })
-                        .map(|nd| if nd.spent.is_none() && nd.unconfirmed_spent.is_none() { nd.note.value } else { 0 })
+                        .map(|nd| nd.note.value)
                         .sum::<u64>()
                 } else {
                     0
@@ -880,6 +1053,57 @@ impl LightWallet {
             .sum::<u64>()
     }
 
+    pub fn spendable_zbalance(&self, addr: Option<String>) -> u64 {
+        let anchor_height = self.get_anchor_height();
+
+        self.txs
+            .read()
+            .unwrap()
+            .values()
+            .map(|tx| {
+                if tx.block as u32 <= anchor_height {
+                    tx.notes
+                        .iter()
+                        .filter(|nd| nd.spent.is_none() && nd.unconfirmed_spent.is_none())
+                        .filter(|nd| {
+                            // Check to see if we have this note's spending key.
+                            self.have_spendingkey_for_extfvk(&nd.extfvk)
+                        })
+                        .filter(|nd| {  // TODO, this whole section is shared with verified_balance. Refactor it.
+                            match addr.clone() {
+                                Some(a) => a == encode_payment_address(
+                                                    self.config.hrp_sapling_address(),
+                                                    &nd.extfvk.fvk.vk
+                                                        .to_payment_address(nd.diversifier, &JUBJUB).unwrap()
+                                                ),
+                                None    => true
+                            }
+                        })
+                        .map(|nd| nd.note.value)
+                        .sum::<u64>()
+                } else {
+                    0
+                }
+            })
+            .sum::<u64>()
+    }
+
+    pub fn have_spendingkey_for_extfvk(&self, extfvk: &ExtendedFullViewingKey) -> bool {
+        match self.zkeys.read().unwrap().iter().find(|zk| zk.extfvk == *extfvk) {
+            None => false,
+            Some(zk) => zk.have_spending_key()
+        }
+    }
+
+    pub fn have_spending_key_for_zaddress(&self, address: &String) -> bool {
+        match self.zkeys.read().unwrap().iter()
+            .find(|zk| encode_payment_address(self.config.hrp_sapling_address(), &zk.zaddress) == *address)
+            {
+                None => false,
+                Some(zk) => zk.have_spending_key()
+            }
+    }
+
     fn add_toutput_to_wtx(&self, height: i32, timestamp: u64, txid: &TxId, vout: &TxOut, n: u64) {
         let mut txs = self.txs.write().unwrap();
 
@@ -945,8 +1169,11 @@ impl LightWallet {
     // If one of the last 'n' zaddress was used, ensure we add the next HD zaddress to the wallet
     pub fn ensure_hd_zaddresses(&self, address: &String) {
         let last_addresses = {
-            self.zaddress.read().unwrap().iter().rev().take(GAP_RULE_UNUSED_ADDRESSES)
+            self.zkeys.read().unwrap().iter()
-                .map(|s| encode_payment_address(self.config.hrp_sapling_address(), s))
+                .filter(|zk| zk.keytype == WalletZKeyType::HdKey)
+                .rev()
+                .take(GAP_RULE_UNUSED_ADDRESSES)
+                .map(|s| encode_payment_address(self.config.hrp_sapling_address(), &s.zaddress))
                 .collect::<Vec<String>>()
         };
 
@@ -1071,18 +1298,18 @@ impl LightWallet {
 
         // Scan shielded sapling outputs to see if anyone of them is us, and if it is, extract the memo
         for output in tx.shielded_outputs.iter() {
-            let ivks: Vec<_> = self.extfvks.read().unwrap().iter().map(
+            let ivks: Vec<_> = self.zkeys.read().unwrap().iter()
-                |extfvk| extfvk.fvk.vk.ivk().clone()
+                .map(|zk| zk.extfvk.fvk.vk.ivk()
-            ).collect();
+                ).collect();
 
             let cmu = output.cmu;
             let ct  = output.enc_ciphertext;
 
             // Search all of our keys
-            for (_account, ivk) in ivks.iter().enumerate() {
+            for ivk in ivks {
                 let epk_prime = output.ephemeral_key.as_prime_order(&JUBJUB).unwrap();
 
-                let (note, _to, memo) = match try_sapling_note_decryption(ivk, &epk_prime, &cmu, &ct) {
+                let (note, _to, memo) = match try_sapling_note_decryption(&ivk, &epk_prime, &cmu, &ct) {
                     Some(ret) => ret,
                     None => continue,
                 };
@@ -1098,7 +1325,10 @@ impl LightWallet {
                         .and_then(|t| {
                             t.notes.iter_mut().find(|nd| nd.note == note)
                         }) {
-                            None => (),
+                            None => {
+                                info!("No txid matched for incoming sapling funds while updating memo");
+                                ()
+                            },
                             Some(nd) => {
                                 nd.memo = Some(memo)
                             }
@@ -1156,12 +1386,13 @@ impl LightWallet {
 
 
             // Search all ovks that we have
-            let ovks: Vec<_> = self.extfvks.read().unwrap().iter().map(
+            let ovks: Vec<_> = self.zkeys.read().unwrap().iter()
-                |extfvk| extfvk.fvk.ovk.clone()
+                .map(|zk| zk.extfvk.fvk.ovk.clone())
-            ).collect();
+                .collect();
 
-            for (_account, ovk) in ovks.iter().enumerate() {
+            for ovk in ovks {
-                match try_sapling_output_recovery(ovk,
+                match try_sapling_output_recovery(
+                    &ovk,
                     &output.cv,
                     &output.cmu,
                     &output.ephemeral_key.as_prime_order(&JUBJUB).unwrap(),
@@ -1579,6 +1810,16 @@ impl LightWallet {
             // Create a write lock
             let mut txs = self.txs.write().unwrap();
 
+            // Trim the older witnesses
+            txs.values_mut().for_each(|wtx| {
+                wtx.notes
+                    .iter_mut()
+                    .filter(|nd| nd.spent.is_some() && nd.spent_at_height.is_some() && nd.spent_at_height.unwrap() < height - (MAX_REORG as i32) - 1)
+                    .for_each(|nd| {
+                        nd.witnesses.clear()
+                    })
+            });
+
             // Create a Vec containing all unspent nullifiers.
             // Include only the confirmed spent nullifiers, since unconfirmed ones still need to be included
             // during scan_block below.
@@ -1614,17 +1855,30 @@ impl LightWallet {
 
             new_txs = {
                 let nf_refs = nfs.iter().map(|(nf, account, _)| (nf.to_vec(), *account)).collect::<Vec<_>>();
+                let extfvks: Vec<ExtendedFullViewingKey> = self.zkeys.read().unwrap().iter().map(|zk| zk.extfvk.clone()).collect();
 
-                // Create a single mutable slice of all the newly-added witnesses.
+                // Create a single mutable slice of all the wallet's note's witnesses.
                 let mut witness_refs: Vec<_> = txs
                     .values_mut()
-                    .map(|tx| tx.notes.iter_mut().filter_map(|nd| nd.witnesses.last_mut()))
+                    .map(|tx|
+                        tx.notes.iter_mut()
+                            .filter_map(|nd|
+                                // Note was not spent
+                                if nd.spent.is_none() && nd.unconfirmed_spent.is_none() {
+                                    nd.witnesses.last_mut()
+                                } else if nd.spent.is_some() && nd.spent_at_height.is_some() && nd.spent_at_height.unwrap() < height - (MAX_REORG as i32) - 1 {
+                                   // Note was spent in the last 100 blocks
+                                    nd.witnesses.last_mut()
+                                } else {
+                                    // If note was old (spent NOT in the last 100 blocks)
+                                    None
+                                }))
                     .flatten()
                     .collect();
 
                 self.scan_block_internal(
                     block.clone(),
-                    &self.extfvks.read().unwrap(),
+                    &extfvks,
                     nf_refs,
                     &mut block_data.tree,
                     &mut witness_refs[..],
@@ -1671,6 +1925,7 @@ impl LightWallet {
                 // Mark the note as spent, and remove the unconfirmed part of it
                 info!("Marked a note as spent");
                 spent_note.spent = Some(tx.txid);
+                spent_note.spent_at_height = Some(height);
                 spent_note.unconfirmed_spent = None::<TxId>;
 
                 total_shielded_value_spent += spent_note.note.value;
@@ -1687,7 +1942,7 @@ impl LightWallet {
             // Save notes.
             for output in tx.shielded_outputs
             {
-                let new_note = SaplingNoteData::new(&self.extfvks.read().unwrap()[output.account], output);
+                let new_note = SaplingNoteData::new(&self.zkeys.read().unwrap()[output.account].extfvk, output);
                 match LightWallet::note_address(self.config.hrp_sapling_address(), &new_note) {
                     Some(a) => {
                         info!("Received sapling output to {}", a);
@@ -1733,15 +1988,34 @@ impl LightWallet {
         Ok(all_txs)
     }
 
-    pub fn send_to_address(
+    // Add the spent_at_height for each sapling note that has been spent. This field was added in wallet version 8,
+    // so for older wallets, it will need to be added
+    pub fn fix_spent_at_height(&self) {
+        // First, build an index of all the txids and the heights at which they were spent.
+        let spent_txid_map: HashMap<_, _> = self.txs.read().unwrap().iter().map(|(txid, wtx)| (txid.clone(), wtx.block)).collect();
+
+        // Go over all the sapling notes that might need updating
+        self.txs.write().unwrap().values_mut().for_each(|wtx| {
+            wtx.notes.iter_mut()
+                .filter(|nd| nd.spent.is_some() && nd.spent_at_height.is_none())
+                .for_each(|nd| {
+                    nd.spent_at_height = spent_txid_map.get(&nd.spent.unwrap()).map(|b| *b);
+                })
+        });
+    }
+
+    pub fn send_to_address<F> (
         &self,
         consensus_branch_id: u32,
         spend_params: &[u8],
         output_params: &[u8],
         from: &str,
         tos: Vec<(&str, u64, Option<String>)>,
-        fee: &u64
+        fee: &u64,
-    ) -> Result<Box<[u8]>, String> {
+        broadcast_fn: F
+    ) -> Result<(String, Vec<u8>), String>
+        where F: Fn(Box<[u8]>) -> Result<String, String>
+    {
         if !self.unlocked {
             return Err("Cannot spend while wallet is locked".to_string());
         }
@@ -1788,14 +2062,30 @@ impl LightWallet {
 
         // Select notes to cover the target value
         println!("{}: Selecting notes", now() - start_time);
-        let target_value = Amount::from_u64(total_value).unwrap() + Amount::from_u64(*fee).unwrap() ;
+        let target_value = Amount::from_u64(total_value).unwrap() + Amount::from_u64(*fee).unwrap();
-        let notes: Vec<_> = self.txs.read().unwrap().iter()
+
+        // Select the candidate notes that are eligible to be spent
+        let mut candidate_notes: Vec<_> = self.txs.read().unwrap().iter()
             .map(|(txid, tx)| tx.notes.iter().map(move |note| (*txid, note)))
             .flatten()
-            .filter(|(_txid, note)|LightWallet::note_address(self.config.hrp_sapling_address(), note).unwrap() == from)
+            .filter_map(|(txid, note)| {
-            .filter_map(|(txid, note)|
+                // Filter out notes that are already spent
-                SpendableNote::from(txid, note, anchor_offset, &self.extsks.read().unwrap()[note.account])
+                if note.spent.is_some() || note.unconfirmed_spent.is_some() {
-            )
+                    None
+                } else {
+                    // Get the spending key for the selected fvk, if we have it
+                    let extsk = self.zkeys.read().unwrap().iter()
+                        .find(|zk| zk.extfvk == note.extfvk)
+                        .and_then(|zk| zk.extsk.clone());
+                    SpendableNote::from(txid, note, anchor_offset, &extsk)
+                }
+            }).collect();
+
+        // Sort by highest value-notes first.
+        candidate_notes.sort_by(|a, b| b.note.value.cmp(&a.note.value));
+
+        // Select the minimum number of notes required to satisfy the target value
+        let notes: Vec<_> = candidate_notes.iter()
             .scan(0, |running_total, spendable| {
                 let value = spendable.note.value;
                 let ret = if *running_total < u64::from(target_value) {
@@ -1890,7 +2180,7 @@ impl LightWallet {
         // Use the ovk belonging to the address being sent from, if not using any notes
         // use the first address in the wallet for the ovk.
         let ovk = if notes.len() == 0 {
-            self.extfvks.read().unwrap()[0].fvk.ovk
+            self.zkeys.read().unwrap()[0].extfvk.fvk.ovk
         } else {
             ExtendedFullViewingKey::from(&notes[0].extsk).fvk.ovk
         };
@@ -1928,13 +2218,16 @@ impl LightWallet {
             // Compute memo if it exists
             let encoded_memo = match memo {
                 None => None,
-                Some(s) => match Memo::from_bytes(s.as_bytes()) {
+                Some(s) => {
-                    None => {
+                    // If the string starts with an "0x", and contains only hex chars ([a-f0-9]+) then
-                        let e = format!("Error creating output. Memo {:?} is too long", s);
+                    // interpret it as a hex
-                        error!("{}", e);
+                    match utils::interpret_memo_string(&s) {
-                        return Err(e);
+                        Ok(m) => Some(m),
-                    },
+                        Err(e) => {
-                    Some(m) => Some(m)
+                            error!("{}", e);
+                            return Err(e);
+                        }
+                    }
                 }
             };
 
@@ -1970,7 +2263,11 @@ impl LightWallet {
         println!("{}: Transaction created", now() - start_time);
         println!("Transaction ID: {}", tx.txid());
 
+        // Create the TX bytes
+        let mut raw_tx = vec![];
+        tx.write(&mut raw_tx).unwrap();
 
+        let txid = broadcast_fn(raw_tx.clone().into_boxed_slice())?;
 
         // Mark notes as spent.
         {
@@ -2008,10 +2305,16 @@ impl LightWallet {
                                 None    => Memo::default(),
                                 Some(s) => {
                                     // If the address is not a z-address, then drop the memo
-                                    if LightWallet::is_shielded_address(&addr.to_string(), &self.config) {
+                                    if !LightWallet::is_shielded_address(&addr.to_string(), &self.config) {
-                                            Memo::from_bytes(s.as_bytes()).unwrap()
-                                    } else {
                                         Memo::default()
+                                    } else {
+                                        match utils::interpret_memo_string(s) {
+                                            Ok(m) => m,
+                                            Err(e) => {
+                                                error!("{}", e);
+                                                Memo::default()
+                                            }
+                                        }
                                     }
                                 }
                             },
@@ -2032,10 +2335,7 @@ impl LightWallet {
             }
         }
 
-        // Return the encoded transaction, so the caller can send it.
+        Ok((txid, raw_tx))
-        let mut raw_tx = vec![];
-        tx.write(&mut raw_tx).unwrap();
-        Ok(raw_tx.into_boxed_slice())
     }
 
     // After some blocks have been mined, we need to remove the Txns from the mempool_tx structure

lib/src/lightwallet/data.rs

@@ -68,6 +68,7 @@ pub struct SaplingNoteData {
     pub(super) witnesses: Vec<IncrementalWitness<Node>>,
     pub(super) nullifier: [u8; 32],
     pub spent: Option<TxId>,             // If this note was confirmed spent
+    pub spent_at_height: Option<i32>,        // The height at which this note was spent
     pub unconfirmed_spent: Option<TxId>, // If this note was spent in a send, but has not yet been confirmed.
     pub memo:  Option<Memo>,
     pub is_change: bool,
@@ -106,7 +107,7 @@ pub fn read_note<R: Read>(mut reader: R) -> io::Result<(u64, Fs)> {
 
 impl SaplingNoteData {
     fn serialized_version() -> u64 {
-        1
+        2
     }
 
     pub fn new(
@@ -132,6 +133,7 @@ impl SaplingNoteData {
             witnesses: vec![witness],
             nullifier: nf,
             spent: None,
+            spent_at_height: None,
             unconfirmed_spent: None,
             memo: None,
             is_change: output.is_change,
@@ -141,10 +143,9 @@ impl SaplingNoteData {
     // Reading a note also needs the corresponding address to read from.
     pub fn read<R: Read>(mut reader: R) -> io::Result<Self> {
         let version = reader.read_u64::<LittleEndian>()?;
-        assert_eq!(version, SaplingNoteData::serialized_version());
 
         let account = reader.read_u64::<LittleEndian>()? as usize;
-
+        
         let extfvk = ExtendedFullViewingKey::read(&mut reader)?;
 
         let mut diversifier_bytes = [0u8; 11];
@@ -176,6 +177,12 @@ impl SaplingNoteData {
             Ok(TxId{0: txid_bytes})
         })?;
 
+        let spent_at_height = if version >=2 {
+            Optional::read(&mut reader, |r| r.read_i32::<LittleEndian>())?
+        } else {
+            None
+        };
+
         let memo = Optional::read(&mut reader, |r| {
             let mut memo_bytes = [0u8; 512];
             r.read_exact(&mut memo_bytes)?;
@@ -195,6 +202,7 @@ impl SaplingNoteData {
             witnesses,
             nullifier,
             spent,
+            spent_at_height,
             unconfirmed_spent: None,
             memo,
             is_change,
@@ -222,6 +230,8 @@ impl SaplingNoteData {
         writer.write_all(&self.nullifier)?;
         Optional::write(&mut writer, &self.spent, |w, t| w.write_all(&t.0))?;
 
+        Optional::write(&mut writer, &self.spent_at_height, |w, h| w.write_i32::<LittleEndian>(*h))?;
+
         Optional::write(&mut writer, &self.memo, |w, m| w.write_all(m.as_bytes()))?;
 
         writer.write_u8(if self.is_change {1} else {0})?;
@@ -512,9 +522,9 @@ pub struct SpendableNote {
 }
 
 impl SpendableNote {
-    pub fn from(txid: TxId, nd: &SaplingNoteData, anchor_offset: usize, extsk: &ExtendedSpendingKey) -> Option<Self> {
+    pub fn from(txid: TxId, nd: &SaplingNoteData, anchor_offset: usize, extsk: &Option<ExtendedSpendingKey>) -> Option<Self> {
         // Include only notes that haven't been spent, or haven't been included in an unconfirmed spend yet.
-        if nd.spent.is_none() && nd.unconfirmed_spent.is_none() &&
+        if nd.spent.is_none() && nd.unconfirmed_spent.is_none() && extsk.is_some() &&
                 nd.witnesses.len() >= (anchor_offset + 1) {
             let witness = nd.witnesses.get(nd.witnesses.len() - anchor_offset - 1);
 
@@ -524,7 +534,7 @@ impl SpendableNote {
                 diversifier: nd.diversifier,
                 note: nd.note.clone(),
                 witness: w.clone(),
-                extsk: extsk.clone(),
+                extsk: extsk.clone().unwrap(),
             })
         } else {
             None

lib/src/lightwallet/utils.rs

@@ -1,5 +1,6 @@
 use std::io::{self, Read, Write};
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+use zcash_primitives::note_encryption::Memo;
 
 pub fn read_string<R: Read>(mut reader: R) -> io::Result<String> {
     // Strings are written as <littleendian> len + bytes
@@ -18,4 +19,26 @@ pub fn write_string<W: Write>(mut writer: W, s: &String) -> io::Result<()> {
     // Strings are written as len + utf8
     writer.write_u64::<LittleEndian>(s.as_bytes().len() as u64)?;
     writer.write_all(s.as_bytes())
+}
+
+// Interpret a string or hex-encoded memo, and return a Memo object
+pub fn interpret_memo_string(memo_str: &String) -> Result<Memo, String> {
+    // If the string starts with an "0x", and contains only hex chars ([a-f0-9]+) then
+    // interpret it as a hex
+    let s_bytes = if memo_str.to_lowercase().starts_with("0x") {
+        match hex::decode(&memo_str[2..memo_str.len()]) {
+            Ok(data) => data,
+            Err(_) => Vec::from(memo_str.as_bytes())
+        }
+    } else {
+        Vec::from(memo_str.as_bytes())
+    };
+
+    match Memo::from_bytes(&s_bytes) {
+        None => {
+            let e = format!("Error creating output. Memo {:?} is too long", memo_str);
+            return Err(e);
+        },
+        Some(m) => Ok(m)
+    }
 }

lib/src/lightwallet/walletzkey.rs

@@ -0,0 +1,425 @@
+use std::io::{self, Read, Write};
+use std::io::{Error, ErrorKind};
+
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+use pairing::bls12_381::{Bls12};
+
+use sodiumoxide::crypto::secretbox;
+
+use zcash_primitives::{
+    serialize::{Vector, Optional},
+    zip32::{ExtendedFullViewingKey, ExtendedSpendingKey},
+    primitives::{PaymentAddress},
+};
+
+use crate::lightclient::{LightClientConfig};
+use crate::lightwallet::LightWallet;
+
+#[derive(PartialEq, Debug, Clone)]
+pub enum WalletZKeyType {
+    HdKey = 0,
+    ImportedSpendingKey = 1,
+    ImportedViewKey = 2
+}
+
+// A struct that holds z-address private keys or view keys
+#[derive(Clone, Debug, PartialEq)]
+pub struct WalletZKey {
+  pub(super) keytype: WalletZKeyType,
+  locked: bool,
+  pub(super) extsk: Option<ExtendedSpendingKey>,
+  pub(super) extfvk: ExtendedFullViewingKey,
+  pub(super) zaddress: PaymentAddress<Bls12>,
+
+  // If this is a HD key, what is the key number
+  pub(super) hdkey_num: Option<u32>,
+
+  // If locked, the encrypted private key is stored here
+  enc_key: Option<Vec<u8>>,
+  nonce: Option<Vec<u8>>,
+}
+
+impl WalletZKey {
+  pub fn new_hdkey(hdkey_num: u32, extsk: ExtendedSpendingKey) -> Self {
+    let extfvk = ExtendedFullViewingKey::from(&extsk);
+    let zaddress = extfvk.default_address().unwrap().1;
+
+    WalletZKey {
+        keytype: WalletZKeyType::HdKey,
+        locked: false,
+        extsk: Some(extsk),
+        extfvk,
+        zaddress,
+        hdkey_num: Some(hdkey_num),
+        enc_key: None,
+        nonce: None,
+    }
+  }
+
+  pub fn new_locked_hdkey(hdkey_num: u32, extfvk: ExtendedFullViewingKey) -> Self {
+    let zaddress = extfvk.default_address().unwrap().1;
+
+    WalletZKey {
+      keytype: WalletZKeyType::HdKey,
+      locked: true,
+      extsk: None,
+      extfvk,
+      zaddress,
+      hdkey_num: Some(hdkey_num),
+      enc_key: None,
+      nonce: None
+    }
+  }
+
+  pub fn new_imported_sk(extsk: ExtendedSpendingKey) -> Self {
+      let extfvk = ExtendedFullViewingKey::from(&extsk);
+      let zaddress = extfvk.default_address().unwrap().1;
+
+      WalletZKey {
+          keytype: WalletZKeyType::ImportedSpendingKey,
+          locked: false,
+          extsk: Some(extsk),
+          extfvk,
+          zaddress,
+          hdkey_num: None,
+          enc_key: None,
+          nonce: None,
+      }
+  }
+
+  pub fn new_imported_viewkey(extfvk: ExtendedFullViewingKey) -> Self {
+    let zaddress = extfvk.default_address().unwrap().1;
+
+    WalletZKey {
+      keytype: WalletZKeyType::ImportedViewKey,
+      locked: false,
+      extsk: None,
+      extfvk,
+      zaddress,
+      hdkey_num: None,
+      enc_key: None,
+      nonce: None,
+    }
+  }
+
+  pub fn have_spending_key(&self) -> bool {
+    self.extsk.is_some() || self.enc_key.is_some() || self.hdkey_num.is_some()
+  }
+
+  fn serialized_version() -> u8 {
+      return 1;
+  }
+
+  pub fn read<R: Read>(mut inp: R) -> io::Result<Self> {
+    let version = inp.read_u8()?;
+    assert!(version <= Self::serialized_version());
+
+    let keytype: WalletZKeyType = match inp.read_u32::<LittleEndian>()? {
+      0 => Ok(WalletZKeyType::HdKey),
+      1 => Ok(WalletZKeyType::ImportedSpendingKey),
+      2 => Ok(WalletZKeyType::ImportedViewKey),
+      n => Err(io::Error::new(ErrorKind::InvalidInput, format!("Unknown zkey type {}", n)))
+    }?;
+
+    let locked = inp.read_u8()? > 0;
+
+    let extsk = Optional::read(&mut inp, |r| ExtendedSpendingKey::read(r))?;
+    let extfvk = ExtendedFullViewingKey::read(&mut inp)?;
+    let zaddress = extfvk.default_address().unwrap().1;
+
+    let hdkey_num = Optional::read(&mut inp, |r| r.read_u32::<LittleEndian>())?;
+
+    let enc_key = Optional::read(&mut inp, |r| 
+        Vector::read(r, |r| r.read_u8()))?;
+    let nonce = Optional::read(&mut inp, |r| 
+        Vector::read(r, |r| r.read_u8()))?;
+
+    Ok(WalletZKey {
+      keytype,
+      locked,
+      extsk,
+      extfvk,
+      zaddress,
+      hdkey_num,
+      enc_key,
+      nonce,
+    })
+  }
+
+  pub fn write<W: Write>(&self, mut out: W) -> io::Result<()> {
+    out.write_u8(Self::serialized_version())?;
+
+    out.write_u32::<LittleEndian>(self.keytype.clone() as u32)?;
+
+    out.write_u8(self.locked as u8)?;
+
+    Optional::write(&mut out, &self.extsk, |w, sk| ExtendedSpendingKey::write(sk, w))?;
+
+    ExtendedFullViewingKey::write(&self.extfvk, &mut out)?;
+
+    Optional::write(&mut out, &self.hdkey_num, |o, n| o.write_u32::<LittleEndian>(*n))?;
+    
+    // Write enc_key
+    Optional::write(&mut out, &self.enc_key, |o, v| 
+        Vector::write(o, v, |o,n| o.write_u8(*n)))?;
+
+    // Write nonce
+    Optional::write(&mut out, &self.nonce, |o, v| 
+        Vector::write(o, v, |o,n| o.write_u8(*n)))
+  }
+
+  pub fn lock(&mut self) -> io::Result<()> {
+    match self.keytype {
+        WalletZKeyType::HdKey => {
+            // For HD keys, just empty out the keys, since they will be reconstructed from the hdkey_num
+            self.extsk = None;
+            self.locked = true;
+        },
+        WalletZKeyType::ImportedSpendingKey => {
+            // For imported keys, encrypt the key into enckey
+            // assert that we have the encrypted key. 
+            if self.enc_key.is_none() {
+              return Err(Error::new(ErrorKind::InvalidInput, "Can't lock when imported key is not encrypted"));
+            }
+            self.extsk = None;
+            self.locked = true;
+        },
+        WalletZKeyType::ImportedViewKey => {
+            // For viewing keys, there is nothing to lock, so just return true
+            self.locked = true;
+        }
+    }
+
+    Ok(())
+  }
+
+  pub fn unlock(&mut self, config: &LightClientConfig, bip39_seed: &[u8], key: &secretbox::Key) -> io::Result<()> {
+    match self.keytype {
+      WalletZKeyType::HdKey => {
+        let (extsk, extfvk, address) =
+            LightWallet::get_zaddr_from_bip39seed(&config, &bip39_seed, self.hdkey_num.unwrap());
+
+        if address != self.zaddress {
+            return Err(io::Error::new(ErrorKind::InvalidData, 
+                    format!("zaddress mismatch at {}. {:?} vs {:?}", self.hdkey_num.unwrap(), address, self.zaddress)));
+        }
+
+        if extfvk != self.extfvk {
+            return Err(io::Error::new(ErrorKind::InvalidData, 
+                        format!("fvk mismatch at {}. {:?} vs {:?}", self.hdkey_num.unwrap(), extfvk, self.extfvk)));
+        }
+
+        self.extsk = Some(extsk);
+      },
+      WalletZKeyType::ImportedSpendingKey => {
+        // For imported keys, we need to decrypt from the encrypted key
+        let nonce = secretbox::Nonce::from_slice(&self.nonce.as_ref().unwrap()).unwrap();
+        let extsk_bytes = match secretbox::open(&self.enc_key.as_ref().unwrap(), &nonce, &key) {
+            Ok(s) => s,
+            Err(_) => {return Err(io::Error::new(ErrorKind::InvalidData, "Decryption failed. Is your password correct?"));}
+        };
+
+        self.extsk = Some(ExtendedSpendingKey::read(&extsk_bytes[..])?);
+      },
+      WalletZKeyType::ImportedViewKey => {
+        // Viewing key unlocking is basically a no op
+      }
+    };
+
+    self.locked = false;
+    Ok(())
+  }
+
+  pub fn encrypt(&mut self, key: &secretbox::Key) -> io::Result<()> {
+    match self.keytype {
+        WalletZKeyType::HdKey => {
+            // For HD keys, we don't need to do anything, since the hdnum has all the info to recreate this key
+        },
+        WalletZKeyType::ImportedSpendingKey => {
+            // For imported keys, encrypt the key into enckey
+            let nonce = secretbox::gen_nonce();
+
+            let mut sk_bytes = vec![];
+            self.extsk.as_ref().unwrap().write(&mut sk_bytes)?;
+
+            self.enc_key = Some(secretbox::seal(&sk_bytes, &nonce, &key));
+            self.nonce = Some(nonce.as_ref().to_vec());
+        },
+        WalletZKeyType::ImportedViewKey => {
+            // Encrypting a viewing key is a no-op
+        }
+    }
+
+    // Also lock after encrypt
+    self.lock()
+  }
+
+  pub fn remove_encryption(&mut self) -> io::Result<()> {
+    if self.locked {
+      return Err(Error::new(ErrorKind::InvalidInput, "Can't remove encryption while locked"));
+    }
+
+    match self.keytype {
+      WalletZKeyType::HdKey => {
+          // For HD keys, we don't need to do anything, since the hdnum has all the info to recreate this key
+          Ok(())
+      },
+      WalletZKeyType::ImportedSpendingKey => {
+          self.enc_key = None;
+          self.nonce = None;
+          Ok(())
+      },
+      WalletZKeyType::ImportedViewKey => {
+          // Removing encryption is a no-op for viewing keys
+          Ok(())
+      }
+    }
+  }
+}
+
+#[cfg(test)]
+pub mod tests {
+  use zcash_client_backend::{
+    encoding::{encode_payment_address, decode_extended_spending_key, decode_extended_full_viewing_key}
+  };
+  use sodiumoxide::crypto::secretbox;
+
+  use crate::lightclient::LightClientConfig;
+  use super::WalletZKey;
+
+  fn get_config() -> LightClientConfig {
+    LightClientConfig {
+      server: "0.0.0.0:0".parse().unwrap(),
+      chain_name: "main".to_string(),
+      sapling_activation_height: 0,
+      consensus_branch_id: "000000".to_string(),
+      anchor_offset: 0,
+      data_dir: None,
+    }
+  }
+
+  #[test]
+  fn test_serialize() {
+    let config = get_config();
+
+    // Priv Key's address is "zs1fxgluwznkzm52ux7jkf4st5znwzqay8zyz4cydnyegt2rh9uhr9458z0nk62fdsssx0cqhy6lyv"
+    let privkey = "secret-extended-key-main1q0p44m9zqqqqpqyxfvy5w2vq6ahvxyrwsk2w4h2zleun4cft4llmnsjlv77lhuuknv6x9jgu5g2clf3xq0wz9axxxq8klvv462r5pa32gjuj5uhxnvps6wsrdg6xll05unwks8qpgp4psmvy5e428uxaggn4l29duk82k3sv3njktaaj453fdmfmj2fup8rls4egqxqtj2p5a3yt4070khn99vzxj5ag5qjngc4v2kq0ctl9q2rpc2phu4p3e26egu9w88mchjf83sqgh3cev";
+
+    let esk = decode_extended_spending_key(config.hrp_sapling_private_key(), privkey).unwrap().unwrap();
+    let wzk = WalletZKey::new_imported_sk(esk);
+    assert_eq!(encode_payment_address(config.hrp_sapling_address(), &wzk.zaddress), "zs1fxgluwznkzm52ux7jkf4st5znwzqay8zyz4cydnyegt2rh9uhr9458z0nk62fdsssx0cqhy6lyv".to_string());
+
+    let mut v: Vec<u8> = vec![];
+    // Serialize
+    wzk.write(&mut v).unwrap();
+    // Read it right back
+    let wzk2 = WalletZKey::read(&v[..]).unwrap();
+
+    {
+      assert_eq!(wzk, wzk2);
+      assert_eq!(wzk.extsk, wzk2.extsk);
+      assert_eq!(wzk.extfvk, wzk2.extfvk);
+      assert_eq!(wzk.zaddress, wzk2.zaddress);
+    }
+  }
+
+  #[test]
+  fn test_encrypt_decrypt_sk() {
+    let config = get_config();
+
+    // Priv Key's address is "zs1fxgluwznkzm52ux7jkf4st5znwzqay8zyz4cydnyegt2rh9uhr9458z0nk62fdsssx0cqhy6lyv"
+    let privkey = "secret-extended-key-main1q0p44m9zqqqqpqyxfvy5w2vq6ahvxyrwsk2w4h2zleun4cft4llmnsjlv77lhuuknv6x9jgu5g2clf3xq0wz9axxxq8klvv462r5pa32gjuj5uhxnvps6wsrdg6xll05unwks8qpgp4psmvy5e428uxaggn4l29duk82k3sv3njktaaj453fdmfmj2fup8rls4egqxqtj2p5a3yt4070khn99vzxj5ag5qjngc4v2kq0ctl9q2rpc2phu4p3e26egu9w88mchjf83sqgh3cev";
+
+    let esk = decode_extended_spending_key(config.hrp_sapling_private_key(), privkey).unwrap().unwrap();
+    let mut wzk = WalletZKey::new_imported_sk(esk);
+    assert_eq!(encode_payment_address(config.hrp_sapling_address(), &wzk.zaddress), "zs1fxgluwznkzm52ux7jkf4st5znwzqay8zyz4cydnyegt2rh9uhr9458z0nk62fdsssx0cqhy6lyv".to_string());
+
+    // Can't lock without encryption
+    assert!(wzk.lock().is_err());
+
+    // Encryption key
+    let key = secretbox::Key::from_slice(&[0; 32]).unwrap();
+
+    // Encrypt, but save the extsk first
+    let orig_extsk = wzk.extsk.clone().unwrap();
+    wzk.encrypt(&key).unwrap();
+    {
+      assert!(wzk.enc_key.is_some());
+      assert!(wzk.nonce.is_some());
+    }
+
+    // Now lock
+    assert!(wzk.lock().is_ok());
+    {
+      assert!(wzk.extsk.is_none());
+      assert_eq!(wzk.locked, true);
+      assert_eq!(wzk.zaddress, wzk.extfvk.default_address().unwrap().1);
+    }
+
+    // Can't remove encryption without unlocking
+    assert!(wzk.remove_encryption().is_err());
+
+    // Unlock
+    assert!(wzk.unlock(&config, &[], &key).is_ok());
+    {
+      assert_eq!(wzk.extsk, Some(orig_extsk));
+    }
+
+    // Remove encryption
+    assert!(wzk.remove_encryption().is_ok());
+    {
+      assert_eq!(wzk.enc_key, None);
+      assert_eq!(wzk.nonce, None);
+    }
+  }
+
+
+  #[test]
+  fn test_encrypt_decrypt_vk() {
+    let config = get_config();
+
+    // Priv Key's address is "zs1va5902apnzlhdu0pw9r9q7ca8s4vnsrp2alr6xndt69jnepn2v2qrj9vg3wfcnjyks5pg65g9dc" 
+    let viewkey = "zxviews1qvvx7cqdqyqqpqqte7292el2875kw2fgvnkmlmrufyszlcy8xgstwarnumqye3tr3d9rr3ydjm9zl9464majh4pa3ejkfy779dm38sfnkar67et7ykxkk0z9rfsmf9jclfj2k85xt2exkg4pu5xqyzyxzlqa6x3p9wrd7pwdq2uvyg0sal6zenqgfepsdp8shestvkzxuhm846r2h3m4jvsrpmxl8pfczxq87886k0wdasppffjnd2eh47nlmkdvrk6rgyyl0ekh3ycqtvvje";
+
+    let extfvk = decode_extended_full_viewing_key(config.hrp_sapling_viewing_key(), viewkey).unwrap().unwrap();
+    let mut wzk = WalletZKey::new_imported_viewkey(extfvk);
+    
+    assert_eq!(encode_payment_address(config.hrp_sapling_address(), &wzk.zaddress), "zs1va5902apnzlhdu0pw9r9q7ca8s4vnsrp2alr6xndt69jnepn2v2qrj9vg3wfcnjyks5pg65g9dc".to_string());
+
+    // Encryption key
+    let key = secretbox::Key::from_slice(&[0; 32]).unwrap();
+
+    // Encrypt
+    wzk.encrypt(&key).unwrap();
+    {
+      assert!(wzk.enc_key.is_none());
+      assert!(wzk.nonce.is_none());
+    }
+
+    // Now lock
+    assert!(wzk.lock().is_ok());
+    {
+      assert!(wzk.extsk.is_none());
+      assert_eq!(wzk.locked, true);
+      assert_eq!(wzk.zaddress, wzk.extfvk.default_address().unwrap().1);
+    }
+
+    // Can't remove encryption without unlocking
+    assert!(wzk.remove_encryption().is_err());
+
+    // Unlock
+    assert!(wzk.unlock(&config, &[], &key).is_ok());
+    {
+      assert_eq!(wzk.extsk, None);
+    }
+
+    // Remove encryption
+    assert!(wzk.remove_encryption().is_ok());
+    {
+      assert_eq!(wzk.enc_key, None);
+      assert_eq!(wzk.nonce, None);
+    }
+  }
+
+
+}