diff --git a/lib/src/lightclient.rs b/lib/src/lightclient.rs index f0b36d1..3089437 100644 --- a/lib/src/lightclient.rs +++ b/lib/src/lightclient.rs @@ -852,6 +852,10 @@ impl LightClient { } pub fn do_rescan(&self) -> Result { + if !self.wallet.read().unwrap().is_unlocked_for_spending() { + warn!("Wallet is locked, new HD addresses won't be added!"); + } + info!("Rescan starting"); self.clear_state(); diff --git a/lib/src/lightwallet.rs b/lib/src/lightwallet.rs index 809ac21..9cb1dbd 100644 --- a/lib/src/lightwallet.rs +++ b/lib/src/lightwallet.rs @@ -879,6 +879,32 @@ impl LightWallet { } } + // If the last taddress was used, ensure we add the next HD taddress to the wallet. + pub fn ensure_hd_taddresses(&self, address: &String) { + let last_address = { + self.taddresses.read().unwrap().last().unwrap().clone() + }; + + if *last_address == *address { + // 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 + 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()); + + if last_address == *address { + // 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(); + } + } + // Scan the full Tx and update memos for incoming shielded transactions. pub fn scan_full_tx(&self, tx: &Transaction, height: i32, datetime: u64) { let mut total_transparent_spend: u64 = 0; @@ -934,6 +960,9 @@ impl LightWallet { if address == hash.to_base58check(&self.config.base58_pubkey_address(), &[]) { // This is our address. Add this as an output to the txid self.add_toutput_to_wtx(height, datetime, &tx.txid(), &vout, n as u64); + + // Ensure that we add any new HD addresses + self.ensure_hd_taddresses(&address); } }, _ => {} @@ -1303,9 +1332,15 @@ impl LightWallet { // Save notes. for output in tx.shielded_outputs { - info!("Received sapling output"); - let new_note = SaplingNoteData::new(&self.extfvks.read().unwrap()[output.account], output); + match LightWallet::note_address(self.config.hrp_sapling_address(), &new_note) { + Some(a) => { + info!("Received sapling output to {}", a); + self.ensure_hd_zaddresses(&a); + }, + None => {} + } + match tx_entry.notes.iter().find(|nd| nd.nullifier == new_note.nullifier) { None => tx_entry.notes.push(new_note), Some(_) => warn!("Tried to insert duplicate note for Tx {}", tx.txid) diff --git a/lib/src/lightwallet/tests.rs b/lib/src/lightwallet/tests.rs index 8a1e9ba..65a22c6 100644 --- a/lib/src/lightwallet/tests.rs +++ b/lib/src/lightwallet/tests.rs @@ -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, 1); + assert_eq!(txs[&sent_txid].notes[ext_note_number].account, 2); 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); @@ -1085,7 +1085,7 @@ fn test_t_spend_to_z() { } } - #[test] +#[test] fn test_z_incoming_memo() { const AMOUNT1: u64 = 50000; let (wallet, _txid1, block_hash) = get_test_wallet(AMOUNT1); @@ -1125,7 +1125,51 @@ fn test_z_incoming_memo() { } } - #[test] +#[test] +fn test_add_new_zt_hd_after_incoming() { + // When an address recieves funds, a new, unused address should automatically get added + const AMOUNT1: u64 = 50000; + let (wallet, _txid1, block_hash) = get_test_wallet(AMOUNT1); + + // Get the last address + let my_address = encode_payment_address(wallet.config.hrp_sapling_address(), + &wallet.extfvks.read().unwrap().last().unwrap().default_address().unwrap().1); + + let fee: u64 = DEFAULT_FEE.try_into().unwrap(); + + 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 + + // Create a tx and send to 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(); + + // Add it to a block + let mut cb3 = FakeCompactBlock::new(2, block_hash); + 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 + + + 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 + + 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 +} + +#[test] fn test_z_to_t_withinwallet() { const AMOUNT: u64 = 500000; const AMOUNT_SENT: u64 = 20000; @@ -1182,7 +1226,7 @@ fn test_z_to_t_withinwallet() { } } - #[test] +#[test] fn test_multi_t() { const AMOUNT: u64 = 5000000; const AMOUNT_SENT1: u64 = 20000; @@ -1368,7 +1412,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-1); + assert_eq!(zaddr2_note.account, 2); assert_eq!(zaddr2_note.is_change, false); assert_eq!(zaddr2_note.spent, None); assert_eq!(zaddr2_note.unconfirmed_spent, None); @@ -1376,7 +1420,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-1); + assert_eq!(zaddr3_note.account, 3); assert_eq!(zaddr3_note.is_change, false); assert_eq!(zaddr3_note.spent, None); assert_eq!(zaddr3_note.unconfirmed_spent, None); @@ -1737,8 +1781,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(); - let zaddr2 = wallet.add_zaddr(); + let zaddr1 = wallet.add_zaddr(); // This is actually address at index 2 + let zaddr2 = wallet.add_zaddr(); // This is actually address at index 3 let taddr0 = wallet.address_from_sk(&wallet.tkeys.read().unwrap()[0]); let taddr1 = wallet.add_taddr(); @@ -1769,15 +1813,15 @@ fn test_lock_unlock() { { let extsks = wallet.extsks.read().unwrap(); let tkeys = wallet.tkeys.read().unwrap(); - assert_eq!(extsks.len(), 3); + assert_eq!(extsks.len(), 4); // 3 zaddrs + 1 added originally in get_test_wallet() 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[1]).default_address().unwrap().1)); - assert_eq!(zaddr2, encode_payment_address(config.hrp_sapling_address(), &ExtendedFullViewingKey::from(&extsks[2]).default_address().unwrap().1)); + assert_eq!(zaddr2, encode_payment_address(config.hrp_sapling_address(), + &ExtendedFullViewingKey::from(&extsks[3]).default_address().unwrap().1)); assert_eq!(taddr0, wallet.address_from_sk(&tkeys[0])); assert_eq!(taddr1, wallet.address_from_sk(&tkeys[1])); @@ -1805,15 +1849,15 @@ fn test_lock_unlock() { { let extsks = wallet2.extsks.read().unwrap(); let tkeys = wallet2.tkeys.read().unwrap(); - assert_eq!(extsks.len(), 3); + assert_eq!(extsks.len(), 4); 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[1]).default_address().unwrap().1)); - assert_eq!(zaddr2, encode_payment_address(wallet2.config.hrp_sapling_address(), &ExtendedFullViewingKey::from(&extsks[2]).default_address().unwrap().1)); + assert_eq!(zaddr2, encode_payment_address(wallet2.config.hrp_sapling_address(), + &ExtendedFullViewingKey::from(&extsks[3]).default_address().unwrap().1)); assert_eq!(taddr0, wallet2.address_from_sk(&tkeys[0])); assert_eq!(taddr1, wallet2.address_from_sk(&tkeys[1])); @@ -1996,7 +2040,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, 1); + assert_eq!(zaddr2_note.account, 2); assert_eq!(zaddr2_note.is_change, false); assert_eq!(zaddr2_note.spent, None); assert_eq!(zaddr2_note.unconfirmed_spent, None);