feat(wallet): add functions to lock and unlock utxos

crates/wallet/src/wallet/mod.rs

@@ -58,6 +58,7 @@ use miniscript::{
 };
 
 use bdk_chain::tx_graph::CalculateFeeError;
+use chain::collections::HashSet;
 
 mod changeset;
 pub mod coin_selection;
@@ -116,6 +117,7 @@ pub struct Wallet {
     change_signers: Arc<SignersContainer>,
     chain: LocalChain,
     indexed_graph: IndexedTxGraph<ConfirmationBlockTime, KeychainTxOutIndex<KeychainKind>>,
+    locked_unspent: HashSet<OutPoint>,
     stage: ChangeSet,
     network: Network,
     secp: SecpCtx,
@@ -308,7 +310,7 @@ impl Wallet {
     /// received on the external keychain (including change), and without a change keychain
     /// BDK lacks enough information to distinguish between change and outside payments.
     ///
-    /// Additionally because this wallet has no internal (change) keychain, all methods that
+    /// Additionally, because this wallet has no internal (change) keychain, all methods that
     /// require a [`KeychainKind`] as input, e.g. [`reveal_next_address`] should only be called
     /// using the [`External`] variant. In most cases passing [`Internal`] is treated as the
     /// equivalent of [`External`] but this behavior must not be relied on.
@@ -434,6 +436,7 @@ impl Wallet {
             network,
             chain,
             indexed_graph,
+            locked_unspent: Default::default(),
             stage,
             secp,
         })
@@ -625,6 +628,7 @@ impl Wallet {
             change_signers,
             chain,
             indexed_graph,
+            locked_unspent: Default::default(),
             stage,
             network,
             secp,
@@ -834,6 +838,27 @@ impl Wallet {
             .map(|((k, i), full_txo)| new_local_utxo(k, i, full_txo))
     }
 
+    /// Lock unspent output
+    ///
+    /// The wallet's locked unspent outputs are automatically added to the [`TxBuilder`]
+    /// unspendable [`TxOut`]s. See [`TxBuilder::add_unspendable`].
+    ///
+    /// Returns true if the given [`TxOut`] is unspent and not already locked.
+    pub fn lock_unspent(&mut self, outpoint: OutPoint) -> bool {
+        // if outpoint is unspent and not already locked insert it in locked_unspent
+        if self.list_unspent().any(|utxo| utxo.outpoint == outpoint) {
+            return self.locked_unspent.insert(outpoint);
+        }
+        false
+    }
+
+    /// Unlock unspent output
+    ///
+    /// returns true if the given [`TxOut`] was locked.
+    pub fn unlock_unspent(&mut self, outpoint: OutPoint) -> bool {
+        self.locked_unspent.remove(&outpoint)
+    }
+
     /// List all relevant outputs (includes both spent and unspent, confirmed and unconfirmed).
     ///
     /// To list only unspent outputs (UTXOs), use [`Wallet::list_unspent`] instead.
@@ -1214,7 +1239,10 @@ impl Wallet {
 
     /// Start building a transaction.
     ///
-    /// This returns a blank [`TxBuilder`] from which you can specify the parameters for the transaction.
+    /// This returns a [`TxBuilder`] from which you can specify the parameters for the transaction.
+    ///
+    /// Locked unspent [`TxOut`] are automatically added to the unspendable set. See [`Wallet::lock_unspent`]
+    /// and [`TxBuilder::add_unspendable`].
     ///
     /// ## Example
     ///
@@ -1238,12 +1266,16 @@ impl Wallet {
     /// // sign and broadcast ...
     /// # Ok::<(), anyhow::Error>(())
     /// ```
-    ///
-    /// [`TxBuilder`]: crate::TxBuilder
+
     pub fn build_tx(&mut self) -> TxBuilder<'_, DefaultCoinSelectionAlgorithm> {
+        let params = TxParams {
+            unspendable: self.locked_unspent.clone(),
+            ..Default::default()
+        };
+
         TxBuilder {
             wallet: alloc::rc::Rc::new(core::cell::RefCell::new(self)),
-            params: TxParams::default(),
+            params,
             coin_selection: DefaultCoinSelectionAlgorithm::default(),
         }
     }
@@ -1585,6 +1617,9 @@ impl Wallet {
     /// *replace by fee* (RBF). If the transaction can be fee bumped then it returns a [`TxBuilder`]
     /// pre-populated with the inputs and outputs of the original transaction.
     ///
+    /// Locked unspent [`TxOut`] are automatically added to the unspendable set. See [`Wallet::lock_unspent`]
+    /// and [`TxBuilder::add_unspendable`].
+    ///
     /// ## Example
     ///
     /// ```no_run

crates/wallet/tests/wallet.rs

@@ -13,7 +13,7 @@ use bdk_wallet::error::CreateTxError;
 use bdk_wallet::psbt::PsbtUtils;
 use bdk_wallet::signer::{SignOptions, SignerError};
 use bdk_wallet::tx_builder::AddForeignUtxoError;
-use bdk_wallet::{AddressInfo, Balance, ChangeSet, Wallet, WalletPersister, WalletTx};
+use bdk_wallet::{AddressInfo, Balance, ChangeSet, TxOrdering, Wallet, WalletPersister, WalletTx};
 use bdk_wallet::{KeychainKind, LoadError, LoadMismatch, LoadWithPersistError};
 use bitcoin::constants::ChainHash;
 use bitcoin::hashes::Hash;
@@ -4309,3 +4309,45 @@ fn test_transactions_sort_by() {
         .collect();
     assert_eq!([None, Some(2000), Some(1000)], conf_heights.as_slice());
 }
+
+#[test]
+fn test_locked_unlocked_utxo() {
+    // create a wallet with 2 utxos
+    let (mut wallet, _txid) = get_funded_wallet_wpkh();
+    receive_output(&mut wallet, 25_000, ChainPosition::Unconfirmed(0));
+    let unspent = wallet.list_unspent().collect::<Vec<_>>();
+    assert_eq!(unspent.len(), 2);
+
+    // get a drain-to address and fee_rate
+    let spk = wallet
+        .next_unused_address(KeychainKind::External)
+        .script_pubkey();
+    let fee_rate = FeeRate::from_sat_per_vb_unchecked(1);
+
+    // lock utxo 0 and verify it is NOT included in drain all utxo tx
+    wallet.lock_unspent(unspent[0].outpoint);
+
+    // verify locking an already locked utxo returns false
+    assert!(!wallet.lock_unspent(unspent[0].outpoint));
+
+    // verify locked utxo is not spent
+    let mut builder = wallet.build_tx();
+    builder.drain_to(spk.clone()).drain_wallet().ordering(TxOrdering::Untouched).fee_rate(fee_rate);
+    let tx_inputs = builder.finish().unwrap().unsigned_tx.input;
+    assert_eq!(tx_inputs.len(), 1);
+    assert_eq!(tx_inputs[0].previous_output, unspent[1].outpoint);
+
+    // unlock utxo 0 and verify it IS included in drain all utxo tx
+    wallet.unlock_unspent(unspent[0].outpoint);
+
+    // verify unlocking an already unlocked utxo returns false
+    assert!(!wallet.unlock_unspent(unspent[0].outpoint));
+
+    // verify all utxos are spent
+    let mut builder = wallet.build_tx();
+    builder.drain_to(spk).drain_wallet().ordering(TxOrdering::Untouched).fee_rate(fee_rate);
+    let tx_inputs = builder.finish().unwrap().unsigned_tx.input;
+    assert_eq!(tx_inputs.len(), 2);
+    assert_eq!(tx_inputs[0].previous_output, unspent[0].outpoint);
+    assert_eq!(tx_inputs[1].previous_output, unspent[1].outpoint);
+}