Add 5 unused addresses all the time

lib/src/lightwallet.rs

@@ -52,6 +52,7 @@ use data::{BlockData, WalletTx, Utxo, SaplingNoteData, SpendableNote, OutgoingTx
 use extended_key::{KeyIndex, ExtendedPrivKey};
 
 pub const MAX_REORG: usize = 100;
+pub const GAP_RULE_UNUSED_ADDRESSES: usize = 5;
 
 fn now() -> f64 {
     SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs() as f64
@@ -879,29 +880,45 @@ impl LightWallet {
         }
     }
 
-    // If the last taddress was used, ensure we add the next HD taddress to the wallet. 
+    // If one of the last 'n' taddress was used, ensure we add the next HD taddress to the wallet. 
-    pub fn ensure_hd_taddresses(&self, address: &String) {
+    pub fn ensure_hd_taddresses(&self, address: &String) {        
-        let last_address = {
+        let last_addresses = {
-            self.taddresses.read().unwrap().last().unwrap().clone()
+            self.taddresses.read().unwrap().iter().rev().take(GAP_RULE_UNUSED_ADDRESSES).map(|s| s.clone()).collect::<Vec<String>>()
         };
         
-        if *last_address == *address {
+        match last_addresses.iter().position(|s| *s == *address) {
-            // If the wallet is locked, this is a no-op. That is fine, since we really
+            None => return,
-            // need to only add new addresses when restoring a new wallet, when it will not be locked.
+            Some(pos) => {
-            // Also, if it is locked, the user can't create new addresses anyway. 
+                // If it in the last unused, addresses, create that many more
-            self.add_taddr();
+                for _ in 0..(GAP_RULE_UNUSED_ADDRESSES - pos) {
+                    // If the wallet is locked, this is a no-op. That is fine, since we really
+                    // need to only add new addresses when restoring a new wallet, when it will not be locked.
+                    // Also, if it is locked, the user can't create new addresses anyway. 
+                    self.add_taddr();
+                }
+            }
         }
     }
 
-    // If the last zaddress was used, ensure we add the next HD zaddress to the wallet
+    // 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_address = encode_payment_address(self.config.hrp_sapling_address(), self.zaddress.read().unwrap().last().unwrap());
+        let last_addresses = {
+            self.zaddress.read().unwrap().iter().rev().take(GAP_RULE_UNUSED_ADDRESSES)
+                .map(|s| encode_payment_address(self.config.hrp_sapling_address(), s))
+                .collect::<Vec<String>>()
+        };
         
-        if last_address == *address {
+        match last_addresses.iter().position(|s| *s == *address) {
-            // If the wallet is locked, this is a no-op. That is fine, since we really
+            None => return,
-            // need to only add new addresses when restoring a new wallet, when it will not be locked.
+            Some(pos) => {
-            // Also, if it is locked, the user can't create new addresses anyway. 
+                // If it in the last unused, addresses, create that many more
-            self.add_zaddr();
+                for _ in 0..(GAP_RULE_UNUSED_ADDRESSES - pos) {
+                    // If the wallet is locked, this is a no-op. That is fine, since we really
+                    // need to only add new addresses when restoring a new wallet, when it will not be locked.
+                    // Also, if it is locked, the user can't create new addresses anyway. 
+                    self.add_zaddr();
+                }
+            }
         }
     }
 

lib/src/lightwallet/tests.rs

@@ -796,7 +796,7 @@ fn test_multi_z() {
     const AMOUNT1: u64 = 50000;
     let (wallet, txid1, block_hash) = get_test_wallet(AMOUNT1);
 
-    let zaddr2 = wallet.add_zaddr();
+    let zaddr2 = wallet.add_zaddr();   // This is acually address #6, since there are 5 initial addresses in the wallet
 
     const AMOUNT_SENT: u64 = 20;
 
@@ -845,7 +845,7 @@ fn test_multi_z() {
         assert_eq!(LightWallet::memo_str(&txs[&sent_txid].notes[change_note_number].memo), None);
 
         assert_eq!(txs[&sent_txid].notes[ext_note_number].note.value, AMOUNT_SENT);
-        assert_eq!(txs[&sent_txid].notes[ext_note_number].account, 2);
+        assert_eq!(txs[&sent_txid].notes[ext_note_number].account, 6);  // The new addr is added after the change addresses
         assert_eq!(txs[&sent_txid].notes[ext_note_number].is_change, false);
         assert_eq!(txs[&sent_txid].notes[ext_note_number].spent, None);
         assert_eq!(txs[&sent_txid].notes[ext_note_number].unconfirmed_spent, None);
@@ -1140,9 +1140,9 @@ fn test_add_new_zt_hd_after_incoming() {
     let branch_id = u32::from_str_radix("2bb40e60", 16).unwrap();
     let (ss, so) = get_sapling_params().unwrap();
 
-    assert_eq!(wallet.zaddress.read().unwrap().len(), 2);   // Starts with 2 addresses
+    assert_eq!(wallet.zaddress.read().unwrap().len(), 6);   // Starts with 1+5 addresses
 
-    // Create a tx and send to address
+    // Create a tx and send to the last address
     let raw_tx = wallet.send_to_address(branch_id, &ss, &so,
                             vec![(&my_address, AMOUNT1 - fee, None)]).unwrap();
     let sent_tx = Transaction::read(&raw_tx[..]).unwrap();
@@ -1152,21 +1152,24 @@ fn test_add_new_zt_hd_after_incoming() {
     cb3.add_tx(&sent_tx);
     wallet.scan_block(&cb3.as_bytes()).unwrap();
 
-    assert_eq!(wallet.zaddress.read().unwrap().len(), 3);   // Now has a new address
+    // NOw, 5 new addresses should be created
-    
+    assert_eq!(wallet.zaddress.read().unwrap().len(), 6+5);     
 
     let mut rng = OsRng;
     let secp = Secp256k1::new();
     // Send a fake transaction to the last taddr
     let pk = PublicKey::from_secret_key(&secp, &wallet.tkeys.read().unwrap().last().unwrap());
 
-    assert_eq!(wallet.taddresses.read().unwrap().len(), 1); // Start with 1 taddr
+    // Start with 1 taddr
+    assert_eq!(wallet.taddresses.read().unwrap().len(), 1); 
 
+    // Send a Tx to the last address
     let mut tx = FakeTransaction::new(&mut rng);
     tx.add_t_output(&pk, AMOUNT1);
-
     wallet.scan_full_tx(&tx.get_tx(), 3, 0);  
-    assert_eq!(wallet.taddresses.read().unwrap().len(), 2); // Now there should be 2 addrs
+
+    // Now, 5 new addresses should be created. 
+    assert_eq!(wallet.taddresses.read().unwrap().len(), 1+5); 
 }
 
 #[test]
@@ -1359,11 +1362,11 @@ fn test_multi_spends() {
     const AMOUNT1: u64 = 50000;
     let (wallet, txid1, block_hash) = get_test_wallet(AMOUNT1);
 
-    let zaddr2 = wallet.add_zaddr();
+    let zaddr2 = wallet.add_zaddr();    // Address number 6
     const ZAMOUNT2:u64 = 30;
     let outgoing_memo2 = "Outgoing Memo2".to_string();
 
-    let zaddr3 = wallet.add_zaddr();
+    let zaddr3 = wallet.add_zaddr();    // Address number 7
     const ZAMOUNT3:u64 = 40;
     let outgoing_memo3 = "Outgoing Memo3".to_string();
 
@@ -1412,7 +1415,7 @@ fn test_multi_spends() {
 
         // Find zaddr2
         let zaddr2_note = txs[&sent_txid].notes.iter().find(|n| n.note.value == ZAMOUNT2).unwrap();
-        assert_eq!(zaddr2_note.account, 2);
+        assert_eq!(zaddr2_note.account, 6);
         assert_eq!(zaddr2_note.is_change, false);
         assert_eq!(zaddr2_note.spent, None);
         assert_eq!(zaddr2_note.unconfirmed_spent, None);
@@ -1420,7 +1423,7 @@ fn test_multi_spends() {
 
         // Find zaddr3
         let zaddr3_note = txs[&sent_txid].notes.iter().find(|n| n.note.value == ZAMOUNT3).unwrap();
-        assert_eq!(zaddr3_note.account, 3);
+        assert_eq!(zaddr3_note.account, 7);
         assert_eq!(zaddr3_note.is_change, false);
         assert_eq!(zaddr3_note.spent, None);
         assert_eq!(zaddr3_note.unconfirmed_spent, None);
@@ -1781,8 +1784,8 @@ fn test_lock_unlock() {
     // Add some addresses
     let zaddr0 = encode_payment_address(config.hrp_sapling_address(), 
                                         &wallet.extfvks.read().unwrap()[0].default_address().unwrap().1);
-    let zaddr1 = wallet.add_zaddr(); // This is actually address at index 2
+    let zaddr1 = wallet.add_zaddr(); // This is actually address at index 6
-    let zaddr2 = wallet.add_zaddr(); // This is actually address at index 3
+    let zaddr2 = wallet.add_zaddr(); // This is actually address at index 7
 
     let taddr0 = wallet.address_from_sk(&wallet.tkeys.read().unwrap()[0]);
     let taddr1 = wallet.add_taddr();
@@ -1813,15 +1816,15 @@ fn test_lock_unlock() {
     {
         let extsks = wallet.extsks.read().unwrap();
         let tkeys = wallet.tkeys.read().unwrap();
-        assert_eq!(extsks.len(), 4); // 3 zaddrs + 1 added originally in get_test_wallet()
+        assert_eq!(extsks.len(), 8); // 3 zaddrs + 1 original + 4 extra HD addreses
         assert_eq!(tkeys.len(), 3);
 
         assert_eq!(zaddr0, encode_payment_address(config.hrp_sapling_address(), 
                                 &ExtendedFullViewingKey::from(&extsks[0]).default_address().unwrap().1));
         assert_eq!(zaddr1, encode_payment_address(config.hrp_sapling_address(), 
-                                &ExtendedFullViewingKey::from(&extsks[2]).default_address().unwrap().1));
+                                &ExtendedFullViewingKey::from(&extsks[6]).default_address().unwrap().1));
         assert_eq!(zaddr2, encode_payment_address(config.hrp_sapling_address(), 
-                                &ExtendedFullViewingKey::from(&extsks[3]).default_address().unwrap().1));
+                                &ExtendedFullViewingKey::from(&extsks[7]).default_address().unwrap().1));
 
         assert_eq!(taddr0, wallet.address_from_sk(&tkeys[0]));
         assert_eq!(taddr1, wallet.address_from_sk(&tkeys[1]));
@@ -1849,15 +1852,15 @@ fn test_lock_unlock() {
     {
         let extsks = wallet2.extsks.read().unwrap();
         let tkeys = wallet2.tkeys.read().unwrap();
-        assert_eq!(extsks.len(), 4);
+        assert_eq!(extsks.len(), 8);
         assert_eq!(tkeys.len(), 3);
 
         assert_eq!(zaddr0, encode_payment_address(wallet2.config.hrp_sapling_address(), 
                                 &ExtendedFullViewingKey::from(&extsks[0]).default_address().unwrap().1));
         assert_eq!(zaddr1, encode_payment_address(wallet2.config.hrp_sapling_address(), 
-                                &ExtendedFullViewingKey::from(&extsks[2]).default_address().unwrap().1));
+                                &ExtendedFullViewingKey::from(&extsks[6]).default_address().unwrap().1));
         assert_eq!(zaddr2, encode_payment_address(wallet2.config.hrp_sapling_address(), 
-                                &ExtendedFullViewingKey::from(&extsks[3]).default_address().unwrap().1));
+                                &ExtendedFullViewingKey::from(&extsks[7]).default_address().unwrap().1));
 
         assert_eq!(taddr0, wallet2.address_from_sk(&tkeys[0]));
         assert_eq!(taddr1, wallet2.address_from_sk(&tkeys[1]));
@@ -2003,7 +2006,7 @@ fn test_encrypted_zreceive() {
     wallet.unlock(password.clone()).unwrap();
 
     // Second z address
-    let zaddr2 = wallet.add_zaddr();
+    let zaddr2 = wallet.add_zaddr();    // This is address number 6
     const ZAMOUNT2:u64 = 30;
     let outgoing_memo2 = "Outgoing Memo2".to_string();
 
@@ -2040,7 +2043,7 @@ fn test_encrypted_zreceive() {
 
         // Find zaddr2
         let zaddr2_note = txs[&txid2].notes.iter().find(|n| n.note.value == ZAMOUNT2).unwrap();
-        assert_eq!(zaddr2_note.account, 2);
+        assert_eq!(zaddr2_note.account, 6);
         assert_eq!(zaddr2_note.is_change, false);
         assert_eq!(zaddr2_note.spent, None);
         assert_eq!(zaddr2_note.unconfirmed_spent, None);