Implement freeing up reserved utxos and cancel_tx

Author: GideonBature

src/wallet/mod.rs

@@ -109,6 +109,7 @@ pub struct Wallet {
     stage: ChangeSet,
     network: Network,
     secp: SecpCtx,
+    reserved_utxos: HashSet<OutPoint>,
 }
 
 /// An update to [`Wallet`].
@@ -499,6 +500,7 @@ impl Wallet {
             indexed_graph,
             stage,
             secp,
+            reserved_utxos: HashSet::new(),
         })
     }
 
@@ -695,6 +697,7 @@ impl Wallet {
             stage,
             network,
             secp,
+            reserved_utxos: HashSet::new(),
         }))
     }
 
@@ -1650,6 +1653,9 @@ impl Wallet {
 
         let psbt = self.complete_transaction(tx, coin_selection.selected, params)?;
 
+        // Reserve the UTXOs used in this transaction
+        self.reserve_utxos(&psbt.unsigned_tx);
+
         // recording changes to the change keychain
         if let (Excess::Change { .. }, Some((keychain, index))) = (excess, drain_index) {
             if let Some((_, index_changeset)) =
@@ -2093,11 +2099,38 @@ impl Wallet {
             .0
     }
 
-    /// Informs the wallet that you no longer intend to broadcast a tx that was built from it.
+    /// Marks the UTXOs used in the given transaction as reserved.
+    ///
+    /// Reserved UTXOs will be excluded from coin selection until they are unreserved.
+    /// This prevents the same UTXO from being used in multiple unsigned transactions.
+    fn reserve_utxos(&mut self, tx: &Transaction) {
+        for input in &tx.input {
+            self.reserved_utxos.insert(input.previous_output);
+        }
+    }
+
+    /// Frees up UTXOs that were reserved for the given transaction.
+    ///
+    /// This should be called when a transaction is cancelled or when it's confirmed on-chain.
+    /// Returns the number of UTXOs that were unreserved.
+    pub fn unreserve_utxos(&mut self, tx: &Transaction) -> usize {
+        let mut count = 0;
+        for input in &tx.input {
+            if self.reserved_utxos.remove(&input.previous_output) {
+                count += 1;
+            }
+        }
+        count
+    }
+
+    /// Frees up change addresses that were reserved for the given transaction.
+    ///
+    /// This is called internally by [`cancel_tx`] to unreserve change addresses.
+    /// Note: This will **not** unreserve addresses that have actually been used
+    /// by a transaction in the tracker. It only removes superficial markings.
     ///
-    /// 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) {
+    /// [`cancel_tx`]: Self::cancel_tx
+    fn unreserve_change_address(&mut self, tx: &Transaction) {
         let txout_index = &mut self.indexed_graph.index;
         for txout in &tx.output {
             if let Some((keychain, index)) = txout_index.index_of_spk(txout.script_pubkey.clone()) {
@@ -2108,6 +2141,41 @@ impl Wallet {
         }
     }
 
+    /// Informs the wallet that you no longer intend to broadcast a tx that was built from it.
+    ///
+    /// This frees up both the change addresses and the UTXOs (inputs) used when creating the tx
+    /// for use in future transactions. Call this method when you decide not to broadcast a
+    /// transaction that was previously built.
+    ///
+    /// Returns the number of UTXOs that were unreserved.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # use bdk_wallet::*;
+    /// # use bitcoin::*;
+    /// # let mut wallet = doctest_wallet!();
+    /// # let to_address = Address::from_str("2N4eQYCbKUHCCTUjBJeHcJp9ok6J2GZsTDt")
+    /// #     .unwrap()
+    /// #     .assume_checked();
+    /// // Build a transaction
+    /// let mut psbt = {
+    ///     let mut builder = wallet.build_tx();
+    ///     builder.add_recipient(to_address.script_pubkey(), Amount::from_sat(50_000));
+    ///     builder.finish()?
+    /// };
+    ///
+    /// // Later, decide not to broadcast it
+    /// let tx = psbt.unsigned_tx.clone();
+    /// let unreserved_count = wallet.cancel_tx(&tx);
+    /// println!("Unreserved {} UTXOs", unreserved_count);
+    /// # Ok::<(), Box<dyn std::error::Error>>(())
+    /// ```
+    pub fn cancel_tx(&mut self, tx: &Transaction) -> usize {
+        self.unreserve_change_address(tx);
+        self.unreserve_utxos(tx)
+    }
+
     fn get_descriptor_for_txout(&self, txout: &TxOut) -> Option<DerivedDescriptor> {
         let &(keychain, child) = self
             .indexed_graph
@@ -2158,6 +2226,8 @@ impl Wallet {
                 })
                 // only add to optional UTxOs those marked as spendable
                 .filter(|local_output| !params.unspendable.contains(&local_output.outpoint))
+                // exclude reserved UTXOs (those used in other unsigned transactions)
+                .filter(|local_output| !self.reserved_utxos.contains(&local_output.outpoint))
                 // if bumping fees only add to optional UTxOs those confirmed
                 .filter(|local_output| {
                     params.bumping_fee.is_none() || local_output.chain_position.is_confirmed()

tests/wallet.rs

@@ -1,3 +1,4 @@
+use core::panic;
 use std::str::FromStr;
 use std::sync::Arc;
 
@@ -3026,3 +3027,205 @@ fn test_tx_ordering_untouched_preserves_insertion_ordering() {
     // Check vout is sorted by recipient insertion order
     assert!(txouts == vec![400, 300, 500]);
 }
+
+#[test]
+fn test_utxo_reservation_prevents_double_spend() {
+    let (mut wallet, _) = get_funded_wallet_wpkh();
+
+    receive_output_in_latest_block(&mut wallet, Amount::from_sat(30_000));
+    receive_output_in_latest_block(&mut wallet, Amount::from_sat(30_000));
+
+    let addr = wallet.reveal_next_address(KeychainKind::External).address;
+
+    let initial_utxos: Vec<_> = wallet.list_unspent().collect();
+    assert!(
+        initial_utxos.len() >= 2,
+        "Need at least 2 UTXOs for this test"
+    );
+
+    let mut builder1 = wallet.build_tx();
+    builder1.add_recipient(addr.script_pubkey(), Amount::from_sat(25_000));
+    let psbt1 = builder1.finish().expect("should create tx1");
+    let tx1 = &psbt1.unsigned_tx;
+
+    let tx1_inputs: Vec<OutPoint> = tx1
+        .input
+        .iter()
+        .map(|input| input.previous_output)
+        .collect();
+
+    let mut builder2 = wallet.build_tx();
+    builder2.add_recipient(addr.script_pubkey(), Amount::from_sat(25_000));
+    let psbt2 = builder2.finish().expect("should create tx2");
+    let tx2 = &psbt2.unsigned_tx;
+
+    let tx2_inputs: Vec<OutPoint> = tx2
+        .input
+        .iter()
+        .map(|input| input.previous_output)
+        .collect();
+
+    for tx1_input in &tx1_inputs {
+        assert!(
+            !tx2_inputs.contains(tx1_input),
+            "tx2 should not reuse UTXO {:?} from tx1",
+            tx1_input
+        );
+    }
+}
+
+#[test]
+fn test_cancel_tx_unreserves_utxos() {
+    let (mut wallet, _) = get_funded_wallet_wpkh();
+    let addr = wallet.reveal_next_address(KeychainKind::External).address;
+
+    let mut builder1 = wallet.build_tx();
+    builder1.add_recipient(addr.script_pubkey(), Amount::from_sat(25_000));
+    let psbt1 = builder1.finish().expect("should create tx1");
+    let tx1 = psbt1.unsigned_tx.clone();
+
+    let tx1_inputs: Vec<OutPoint> = tx1
+        .input
+        .iter()
+        .map(|input| input.previous_output)
+        .collect();
+
+    wallet.cancel_tx(&tx1);
+
+    let mut builder2 = wallet.build_tx();
+    builder2.add_recipient(addr.script_pubkey(), Amount::from_sat(25_000));
+    let psbt2 = builder2.finish().expect("should create tx2");
+    let tx2 = &psbt2.unsigned_tx;
+
+    let tx2_inputs: Vec<OutPoint> = tx2
+        .input
+        .iter()
+        .map(|input| input.previous_output)
+        .collect();
+
+    // After cancellation, tx2 should be able to use at least some UTXOs from tx1
+    // (coin selection might choose different UTXOs, but they should be available)
+    let shared_inputs: Vec<&OutPoint> = tx1_inputs
+        .iter()
+        .filter(|input| tx2_inputs.contains(input))
+        .collect();
+
+    // We expect tx2 to reuse at least one UTXO from tx1 since they have the same amount
+    assert!(
+        !shared_inputs.is_empty(),
+        "After cancel_tx, tx2 should be able to reuse UTXOs from tx1"
+    );
+}
+
+#[test]
+fn test_unreserve_utxos_returns_correct_count() {
+    let (mut wallet, _) = get_funded_wallet_wpkh();
+    let addr = wallet.reveal_next_address(KeychainKind::External).address;
+
+    let mut builder = wallet.build_tx();
+    builder.add_recipient(addr.script_pubkey(), Amount::from_sat(25_000));
+    let psbt = builder.finish().expect("should create tx");
+    let tx = psbt.unsigned_tx.clone();
+
+    let input_count = tx.input.len();
+
+    let unreserved_count = wallet.unreserve_utxos(&tx);
+
+    // Should unreserve exactly as many UTXOs as there are inputs
+    assert_eq!(
+        unreserved_count, input_count,
+        "Should unreserve {} UTXOs",
+        input_count
+    );
+
+    // Unreserving again should return 0 (UTXOs already unreserved)
+    let second_unreserve_count = wallet.unreserve_utxos(&tx);
+    assert_eq!(
+        second_unreserve_count, 0,
+        "Second unreserve should return 0"
+    );
+}
+
+#[test]
+fn test_multiple_unsigned_transactions() {
+    let (mut wallet, _) = get_funded_wallet_wpkh();
+
+    receive_output_in_latest_block(&mut wallet, Amount::from_sat(30_000));
+    receive_output_in_latest_block(&mut wallet, Amount::from_sat(30_000));
+    receive_output_in_latest_block(&mut wallet, Amount::from_sat(30_000));
+
+    let addr = wallet.reveal_next_address(KeychainKind::External).address;
+
+    let initial_utxos: Vec<_> = wallet.list_unspent().collect();
+    assert!(
+        initial_utxos.len() >= 3,
+        "Need at least 3 UTXOs for this test"
+    );
+
+    let mut transactions = Vec::new();
+    for i in 0..3 {
+        let mut builder = wallet.build_tx();
+        builder.add_recipient(addr.script_pubkey(), Amount::from_sat(20_000 + i * 1_000));
+        let psbt = builder
+            .finish()
+            .unwrap_or_else(|_| panic!("should create tx{}", i + 1));
+        transactions.push(psbt.unsigned_tx.clone());
+    }
+
+    let mut all_inputs: Vec<OutPoint> = Vec::new();
+    for tx in &transactions {
+        for input in &tx.input {
+            all_inputs.push(input.previous_output);
+        }
+    }
+
+    // Verify no UTXO is used twice
+    let mut seen = std::collections::HashSet::new();
+    for input in &all_inputs {
+        assert!(
+            seen.insert(input),
+            "UTXO {:?} was used in multiple transactions",
+            input
+        );
+    }
+
+    // Cancel all transactions and verify we can build new ones
+    for tx in &transactions {
+        wallet.cancel_tx(tx);
+    }
+
+    // After cancelling all, we should be able to build a new transaction
+    let mut builder = wallet.build_tx();
+    builder.add_recipient(addr.script_pubkey(), Amount::from_sat(25_000));
+    let result = builder.finish();
+    assert!(
+        result.is_ok(),
+        "Should be able to build tx after cancelling all previous ones"
+    );
+}
+
+#[test]
+fn test_cancel_tx_unreserves_change_address() {
+    let (mut wallet, _) = get_funded_wallet_wpkh();
+    let addr = wallet.reveal_next_address(KeychainKind::External).address;
+
+    let mut builder = wallet.build_tx();
+    builder.add_recipient(addr.script_pubkey(), Amount::from_sat(10_000));
+    let psbt = builder.finish().expect("should create tx");
+    let tx = psbt.unsigned_tx.clone();
+
+    // Check if transaction has change output (more than 1 output)
+    let has_change = tx.output.len() > 1;
+
+    if has_change {
+        let change_index_before = wallet.next_derivation_index(KeychainKind::Internal);
+        wallet.cancel_tx(&tx);
+
+        let change_index_after = wallet.next_derivation_index(KeychainKind::Internal);
+        // After unreserving, we should be back to the same index
+        assert_eq!(
+            change_index_before, change_index_after,
+            "Change address should be unreserved"
+        );
+    }
+}