Add Wallet::cancel_tx

To allow you to re-use change addresses from transactions that get cancelled.

diff --git a/src/wallet/mod.rs b/src/wallet/mod.rs
index f693ecfc785fef9b0439c73e44727d9b513dc1e6..fe0660b6e7c4d047a9fec9634778c188c085f034 100644 (file)
--- a/src/wallet/mod.rs
+++ b/src/wallet/mod.rs
@@ -331,6 +331,13 @@ impl<D> Wallet<D> {
             .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<D> Wallet<D> {
         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 2ba9942f71976ca20ce8bd0f762bf7c0f938214b..de94670321465d8372dd3557cb1ad63f5e6981eb 100644 (file)
--- 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 533c21bfba45224369acc30076d887590004f4b4..9e2d06852c4beca30e536fd277bb9bdd9e02dfc5 100644 (file)
--- 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));
+}