From 4446da6e8a632506905fe36e74616f47aa1fa9d6 Mon Sep 17 00:00:00 2001 From: Steve Myers Date: Thu, 31 Oct 2024 19:57:52 -0500 Subject: [PATCH] feat(wallet): add functions to lock and unlock utxos --- crates/wallet/src/wallet/mod.rs | 45 +++++++++++++++++++++++++++++---- crates/wallet/tests/wallet.rs | 44 +++++++++++++++++++++++++++++++- 2 files changed, 83 insertions(+), 6 deletions(-) diff --git a/crates/wallet/src/wallet/mod.rs b/crates/wallet/src/wallet/mod.rs index 58b504dc9..0d1087294 100644 --- a/crates/wallet/src/wallet/mod.rs +++ b/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, chain: LocalChain, indexed_graph: IndexedTxGraph>, + locked_unspent: HashSet, 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 diff --git a/crates/wallet/tests/wallet.rs b/crates/wallet/tests/wallet.rs index d51af3352..d70be0193 100644 --- a/crates/wallet/tests/wallet.rs +++ b/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::>(); + 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); +}