Implement support for detection of diversified addressess

3 years ago · e41874df40
7 changed files with 118 additions and 30 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@ -3155,7 +3155,7 @@ dependencies = [

 [[package]]
 name = "zecwallet-cli"
-version = "1.4.1"
+version = "1.4.2"
 dependencies = [
 "byteorder",
 "clap",
--- a/cli/Cargo.toml
+++ b/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/" }
-
--- a/cli/src/version.rs
+++ b/cli/src/version.rs
@ -1 +1 @@
-pub const VERSION:&str = "1.4.1";
+pub const VERSION:&str = "1.4.2";
--- a/lib/src/lightclient.rs
+++ b/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();
-
-            // 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();
+            let all_zkeys = wallet.zkeys.read().unwrap();

            // 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()) &&
-                                                    wtx.block <= anchor_height && nd.spent.is_none() && nd.unconfirmed_spent.is_none();
+
+                            let spendable = match zkeys.iter().find(|zk| zk.extfvk == nd.extfvk) {
+                                None => false,
+                                Some(zk) => zk.have_spending_key()
+                            };

                            Some(object!{
                                "created_in_block"   => wtx.block,
--- a/lib/src/lightwallet.rs
+++ b/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 => {}
                }
--- a/lib/src/lightwallet/data.rs
+++ b/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>>,
--- a/lib/src/lightwallet/walletzkey.rs
+++ b/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