Implement support for detection of diversified addressess

Cargo.lock

@@ -3155,7 +3155,7 @@ dependencies = [
 
 [[package]]
 name = "zecwallet-cli"
-version = "1.4.1"
+version = "1.4.2"
 dependencies = [
  "byteorder",
  "clap",

cli/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "zecwallet-cli"
-version = "1.4.1"
+version = "1.4.2"
 edition = "2018"
 
 [dependencies]
@@ -15,4 +15,3 @@ byteorder = "1"
 tiny-bip39 = "0.6.2"
 
 zecwalletlitelib = { path = "../lib/" }
-

cli/src/version.rs

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

lib/src/lightclient.rs

@@ -1,4 +1,5 @@
 use crate::lightwallet::LightWallet;
+use crate::lightwallet::walletzkey::WalletDiversifiers;
 
 use rand::{rngs::OsRng, seq::SliceRandom};
 
@@ -6,7 +7,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::{HashSet, HashMap};
+use std::collections::{HashMap};
 use std::cmp::{max, min};
 use std::io;
 use std::io::prelude::*;
@@ -584,6 +585,33 @@ impl LightClient {
             sync_status     : Arc::new(RwLock::new(WalletStatus::new())),
         };
 
+        //Load Diversified Addresses from SaplingNotes
+        {
+            let note_wallet = lc.wallet.write().unwrap();
+            let txs = note_wallet.txs.read().unwrap();
+            for (_t, tx) in txs.iter() {
+                for n in tx.notes.iter() {
+                    match LightWallet::note_address(lc.config.hrp_sapling_address(), &n) {
+                        Some(a) => {
+                            //Add diversified addresses to address list
+                            let mut zaddrs = note_wallet.zaddresses.write().unwrap();
+                            let mut found = false;
+                            for z in zaddrs.iter() {
+                                if z.zaddress == a {
+                                    found = true;
+                                }
+                            }
+
+                            if !found {
+                                zaddrs.push(WalletDiversifiers{extfvk: n.extfvk.clone(), diversifier: n.diversifier.clone(), zaddress: a});
+                            }
+                        },
+                        None => {}
+                    }
+                }
+            }
+        }
+        
         #[cfg(feature = "embed_params")]
         lc.read_sapling_params();
 
@@ -611,6 +639,33 @@ impl LightClient {
             sync_status     : Arc::new(RwLock::new(WalletStatus::new())),
         };
 
+        //Load Diversified Addresses from SaplingNotes
+        {
+            let note_wallet = lc.wallet.write().unwrap();
+            let txs = note_wallet.txs.read().unwrap();
+            for (_t, tx) in txs.iter() {
+                for n in tx.notes.iter() {
+                    match LightWallet::note_address(lc.config.hrp_sapling_address(), &n) {
+                        Some(a) => {
+                            //Add diversified addresses to address list
+                            let mut zaddrs = note_wallet.zaddresses.write().unwrap();
+                            let mut found = false;
+                            for z in zaddrs.iter() {
+                                if z.zaddress == a {
+                                    found = true;
+                                }
+                            }
+
+                            if !found {
+                                zaddrs.push(WalletDiversifiers{extfvk: n.extfvk.clone(), diversifier: n.diversifier.clone(), zaddress: a});
+                            }
+                        },
+                        None => {}
+                    }
+                }
+            }
+        }
+
         #[cfg(feature = "embed_params")]
         lc.read_sapling_params();
 
@@ -905,29 +960,24 @@ 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;
-
         {
             let wallet = self.wallet.read().unwrap();
-
+            let all_zkeys = wallet.zkeys.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();
+                    let zkeys = all_zkeys.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()) &&
+                            let spendable = match zkeys.iter().find(|zk| zk.extfvk == nd.extfvk) {
-                                                    wtx.block <= anchor_height && nd.spent.is_none() && nd.unconfirmed_spent.is_none();
+                                None => false,
+                                Some(zk) => zk.have_spending_key()
+                            };
 
                             Some(object!{
                                 "created_in_block"   => wtx.block,

lib/src/lightwallet.rs

@@ -55,11 +55,11 @@ mod extended_key;
 mod utils;
 mod address;
 mod prover;
-mod walletzkey;
+pub mod walletzkey;
 
 use data::{BlockData, WalletTx, Utxo, SaplingNoteData, SpendableNote, OutgoingTxMetadata};
 use extended_key::{KeyIndex, ExtendedPrivKey};
-use walletzkey::{WalletZKey, WalletZKeyType};
+use walletzkey::{WalletZKey, WalletZKeyType, WalletDiversifiers};
 
 pub const MAX_REORG: usize = 100;
 
@@ -116,7 +116,8 @@ pub struct LightWallet {
 
     // List of keys, actually in this wallet. This is a combination of HD keys derived from the seed,
     // viewing keys and imported spending keys.
-    zkeys: Arc<RwLock<Vec<WalletZKey>>>,
+    pub zkeys: Arc<RwLock<Vec<WalletZKey>>>,
+    pub zaddresses: Arc<RwLock<Vec<WalletDiversifiers>>>,
 
     // Transparent keys. If the wallet is locked, then the secret keys will be encrypted,
     // but the addresses will be present.
@@ -229,6 +230,7 @@ impl LightWallet {
             nonce:       vec![],
             seed:        seed_bytes,
             zkeys:       Arc::new(RwLock::new(vec![WalletZKey::new_hdkey(hdkey_num, extsk)])),
+            zaddresses:  Arc::new(RwLock::new(vec![])),
             tkeys:       Arc::new(RwLock::new(vec![])),
             taddresses:  Arc::new(RwLock::new(vec![])),
             blocks:      Arc::new(RwLock::new(vec![])),
@@ -345,18 +347,18 @@ impl LightWallet {
             Vector::read(&mut reader, |r| WalletZKey::read(r))?
         };
 
-        let tkeys = Vector::read(&mut reader, |r| {
+        let _tkeys = Vector::read(&mut reader, |r| {
             let mut tpk_bytes = [0u8; 32];
             r.read_exact(&mut tpk_bytes)?;
             secp256k1::SecretKey::from_slice(&tpk_bytes).map_err(|e| io::Error::new(ErrorKind::InvalidData, e))
         })?;
 
-        let taddresses = if version >= 4 {
+        let _taddresses = if version >= 4 {
             // Read the addresses
             Vector::read(&mut reader, |r| utils::read_string(r))?
         } else {
             // Calculate the addresses
-            tkeys.iter().map(|sk| LightWallet::address_from_prefix_sk(&config.base58_pubkey_address(), sk)).collect()
+            _tkeys.iter().map(|sk| LightWallet::address_from_prefix_sk(&config.base58_pubkey_address(), sk)).collect()
         };
 
         let blocks = Vector::read(&mut reader, |r| BlockData::read(r))?;
@@ -385,6 +387,7 @@ impl LightWallet {
             nonce:       nonce,
             seed:        seed_bytes,
             zkeys:       Arc::new(RwLock::new(zkeys)),
+            zaddresses:  Arc::new(RwLock::new(vec![])),
             tkeys:       Arc::new(RwLock::new(vec![])),
             taddresses:  Arc::new(RwLock::new(vec![])),
             blocks:      Arc::new(RwLock::new(blocks)),
@@ -753,9 +756,23 @@ impl LightWallet {
     }
 
     pub fn get_all_zaddresses(&self) -> Vec<String> {
-        self.zkeys.read().unwrap().iter().map( |zk| {
+        let mut zaddrs: Vec<String> = self.zkeys.read().unwrap().iter().map( |zk| {
             encode_payment_address(self.config.hrp_sapling_address(), &zk.zaddress)
-        }).collect()
+        }).collect();
+
+        let dzaddrs = self.zaddresses.read().unwrap();
+        for z in dzaddrs.iter() {
+            let mut found = false;
+            for ud in zaddrs.iter() {
+                if ud == &z.zaddress.clone() {
+                    found = true;
+                }
+            }
+            if !found {
+                zaddrs.push(z.zaddress.clone());
+            }
+        }
+        zaddrs
     }
 
     pub fn address_from_prefix_sk(prefix: &[u8; 2], sk: &secp256k1::SecretKey) -> String {
@@ -869,7 +886,7 @@ impl LightWallet {
         let bip39_seed = bip39::Seed::new(&Mnemonic::from_entropy(&seed, Language::English).unwrap(), "");
 
         // Transparent keys
-        let mut tkeys = vec![];
+        let tkeys = vec![];
         // for pos in 0..self.taddresses.read().unwrap().len() {
         //     let sk = LightWallet::get_taddr_from_bip39seed(&self.config, &bip39_seed.as_bytes(), pos as u32);
         //     let address = self.address_from_sk(&sk);
@@ -1965,7 +1982,21 @@ impl LightWallet {
                 match LightWallet::note_address(self.config.hrp_sapling_address(), &new_note) {
                     Some(a) => {
                         info!("Received sapling output to {}", a);
-                        self.ensure_hd_zaddresses(&a);
+                        // self.ensure_hd_zaddresses(&a);
+
+                        //Add diversified addresses to address list
+                        let mut zaddrs = self.zaddresses.write().unwrap();
+                        let mut found = false;
+                        for z in zaddrs.iter() {
+                            if z.zaddress == a {
+                                found = true;
+                            }
+                        }
+
+                        if !found {
+                            zaddrs.push(WalletDiversifiers{extfvk: new_note.extfvk.clone(), diversifier: new_note.diversifier.clone(), zaddress: a});
+                        }
+
                     },
                     None => {}
                 }

lib/src/lightwallet/data.rs

@@ -62,7 +62,7 @@ impl BlockData {
 
 pub struct SaplingNoteData {
     pub(super) account: usize,
-    pub(super) extfvk: ExtendedFullViewingKey, // Technically, this should be recoverable from the account number, but we're going to refactor this in the future, so I'll write it again here.
+    pub extfvk: ExtendedFullViewingKey, // Technically, this should be recoverable from the account number, but we're going to refactor this in the future, so I'll write it again here.
     pub diversifier: Diversifier,
     pub note: Note<Bls12>,
     pub(super) witnesses: Vec<IncrementalWitness<Node>>,

lib/src/lightwallet/walletzkey.rs

@@ -9,7 +9,7 @@ use sodiumoxide::crypto::secretbox;
 use zcash_primitives::{
     serialize::{Vector, Optional},
     zip32::{ExtendedFullViewingKey, ExtendedSpendingKey},
-    primitives::{PaymentAddress},
+    primitives::{Diversifier, PaymentAddress},
 };
 
 use crate::lightclient::{LightClientConfig};
@@ -22,13 +22,21 @@ pub enum WalletZKeyType {
     ImportedViewKey = 2
 }
 
+// A struct that holds diversified addresses
+#[derive(Clone, Debug, PartialEq)]
+pub struct WalletDiversifiers {
+  pub extfvk: ExtendedFullViewingKey,
+  pub diversifier: Diversifier,
+  pub zaddress: String
+}
+
 // 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 extfvk: ExtendedFullViewingKey,
   pub(super) zaddress: PaymentAddress<Bls12>,
 
   // If this is a HD key, what is the key number