diff --git a/src/wallet/mod.rs b/src/wallet/mod.rs index f693ecfc7..fe0660b6e 100644 --- a/src/wallet/mod.rs +++ b/src/wallet/mod.rs @@ -331,6 +331,13 @@ impl Wallet { .is_some() } + /// Finds how the wallet derived the script pubkey `spk`. + /// + /// Will only return `Some(_)` if the wallet has given out the spk. + pub fn derivation_of_spk(&self, spk: &Script) -> Option<(KeychainKind, u32)> { + self.keychain_tracker.txout_index.index_of_spk(spk).copied() + } + /// Return the list of unspent outputs of this wallet /// /// Note that this method only operates on the internal database, which first needs to be @@ -1405,6 +1412,22 @@ impl Wallet { self.keychain_tracker.txout_index.next_index(&keychain).0 } + /// Informs the wallet that you no longer intend to broadcast a tx that was built from it. + /// + /// This frees up the change address used when creating the tx for use in future transactions. + /// + // TODO: Make this free up reserved utxos when that's implemented + pub fn cancel_tx(&mut self, tx: &Transaction) { + let txout_index = &mut self.keychain_tracker.txout_index; + for txout in &tx.output { + if let Some(&(keychain, index)) = txout_index.index_of_spk(&txout.script_pubkey) { + // NOTE: unmark_used will **not** make something unused if it has actually been used + // by a tx in the tracker. It only removes the superficial marking. + txout_index.unmark_used(&keychain, index); + } + } + } + fn map_keychain(&self, keychain: KeychainKind) -> KeychainKind { if keychain == KeychainKind::Internal && self.public_descriptor(KeychainKind::Internal).is_none() diff --git a/tests/common.rs b/tests/common.rs index 2ba9942f7..de9467032 100644 --- a/tests/common.rs +++ b/tests/common.rs @@ -5,8 +5,11 @@ use bitcoin::hashes::Hash; use bitcoin::{BlockHash, Network, Transaction, TxOut}; /// Return a fake wallet that appears to be funded for testing. -pub fn get_funded_wallet(descriptor: &str) -> (Wallet, bitcoin::Txid) { - let mut wallet = Wallet::new_no_persist(descriptor, None, Network::Regtest).unwrap(); +pub fn get_funded_wallet_with_change( + descriptor: &str, + change: Option<&str>, +) -> (Wallet, bitcoin::Txid) { + let mut wallet = Wallet::new_no_persist(descriptor, change, Network::Regtest).unwrap(); let address = wallet.get_address(AddressIndex::New).address; let tx = Transaction { @@ -38,6 +41,10 @@ pub fn get_funded_wallet(descriptor: &str) -> (Wallet, bitcoin::Txid) { (wallet, tx.txid()) } +pub fn get_funded_wallet(descriptor: &str) -> (Wallet, bitcoin::Txid) { + get_funded_wallet_with_change(descriptor, None) +} + pub fn get_test_wpkh() -> &'static str { "wpkh(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW)" } diff --git a/tests/wallet.rs b/tests/wallet.rs index 533c21bfb..9e2d06852 100644 --- a/tests/wallet.rs +++ b/tests/wallet.rs @@ -3240,3 +3240,72 @@ fn test_taproot_load_descriptor_duplicated_keys() { "bcrt1pvysh4nmh85ysrkpwtrr8q8gdadhgdejpy6f9v424a8v9htjxjhyqw9c5s5" ); } + +#[test] +/// The wallet should re-use previously allocated change addresses when the tx using them is cancelled +fn test_tx_cancellation() { + macro_rules! new_tx { + ($wallet:expr) => {{ + let addr = Address::from_str("2N4eQYCbKUHCCTUjBJeHcJp9ok6J2GZsTDt").unwrap(); + let mut builder = $wallet.build_tx(); + builder.add_recipient(addr.script_pubkey(), 10_000); + + let (psbt, _) = builder.finish().unwrap(); + + psbt + }}; + } + + let (mut wallet, _) = + get_funded_wallet_with_change(get_test_wpkh(), Some(get_test_tr_single_sig_xprv())); + + let psbt1 = new_tx!(wallet); + let change_derivation_1 = psbt1 + .unsigned_tx + .output + .iter() + .find_map(|txout| wallet.derivation_of_spk(&txout.script_pubkey)) + .unwrap(); + assert_eq!(change_derivation_1, (KeychainKind::Internal, 0)); + + let psbt2 = new_tx!(wallet); + + let change_derivation_2 = psbt2 + .unsigned_tx + .output + .iter() + .find_map(|txout| wallet.derivation_of_spk(&txout.script_pubkey)) + .unwrap(); + assert_eq!(change_derivation_2, (KeychainKind::Internal, 1)); + + wallet.cancel_tx(&psbt1.extract_tx()); + + let psbt3 = new_tx!(wallet); + let change_derivation_3 = psbt3 + .unsigned_tx + .output + .iter() + .find_map(|txout| wallet.derivation_of_spk(&txout.script_pubkey)) + .unwrap(); + assert_eq!(change_derivation_3, (KeychainKind::Internal, 0)); + + let psbt3 = new_tx!(wallet); + let change_derivation_3 = psbt3 + .unsigned_tx + .output + .iter() + .find_map(|txout| wallet.derivation_of_spk(&txout.script_pubkey)) + .unwrap(); + assert_eq!(change_derivation_3, (KeychainKind::Internal, 2)); + + wallet.cancel_tx(&psbt3.extract_tx()); + + let psbt3 = new_tx!(wallet); + let change_derivation_4 = psbt3 + .unsigned_tx + .output + .iter() + .find_map(|txout| wallet.derivation_of_spk(&txout.script_pubkey)) + .unwrap(); + assert_eq!(change_derivation_4, (KeychainKind::Internal, 2)); +}