From 8f3387e398abebb76026a0dfa7e89728d1e82eb3 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 | 46 ++++++++++-- crates/wallet/tests/wallet.rs | 127 +++++++++++++++++++++++++++++++- 2 files changed, 167 insertions(+), 6 deletions(-) diff --git a/crates/wallet/src/wallet/mod.rs b/crates/wallet/src/wallet/mod.rs index 58b504dc9..9b6f1cc42 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 @@ -1744,6 +1779,7 @@ impl Wallet { absolute: fee, rate: fee_rate, }), + unspendable: self.locked_unspent.clone(), ..Default::default() }; diff --git a/crates/wallet/tests/wallet.rs b/crates/wallet/tests/wallet.rs index d51af3352..bc6a9a2ea 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,128 @@ 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); +} + +#[test] +fn test_bump_fee_with_locked_unlocked_utxo() { + // create a wallet with 2 utxos + let (mut wallet, _txid) = get_funded_wallet_wpkh(); + receive_output_in_latest_block(&mut wallet, 25_000); + receive_output_in_latest_block(&mut wallet, 2000); + let mut unspent = wallet.list_unspent().collect::>(); + unspent.sort_by_key(|utxo| utxo.txout.value); + unspent.reverse(); // now unspent are in largest first order + assert_eq!(unspent.len(), 3); + + // 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 largest utxo + wallet.lock_unspent(unspent[0].outpoint); + + // verify locked largest (50_000 sats) utxo is not spent + let mut builder = wallet + .build_tx() + .coin_selection(LargestFirstCoinSelection::default()); + builder + .set_recipients(vec![(spk.clone(), Amount::from_sat(24_500))]) + .ordering(TxOrdering::Untouched) + .fee_rate(fee_rate); + let mut psbt = builder.finish().unwrap(); + wallet + .finalize_psbt(&mut psbt, SignOptions::default()) + .unwrap(); + + let signed_tx = psbt.extract_tx().unwrap(); + wallet.apply_unconfirmed_txs([(signed_tx.clone(), 100_000)]); + let original_txid = signed_tx.compute_txid(); + let tx_inputs = signed_tx.clone().input; + + // verify unlocked utxo (25_000 sats) used instead of locked largest (50_000 sats) utxo + assert_eq!(tx_inputs.len(), 1); + assert_eq!(tx_inputs[0].previous_output, unspent[1].outpoint); + + // get new bump fee rate: 10 sats/vb + let new_fee_rate = FeeRate::from_sat_per_vb_unchecked(10); + let mut builder = wallet + .build_fee_bump(original_txid) + .unwrap() + .coin_selection(LargestFirstCoinSelection::default()); + builder + .ordering(TxOrdering::Untouched) + .fee_rate(new_fee_rate); + + let tx_inputs = builder.finish().unwrap().unsigned_tx.input; + // verify unlocked utxo used instead of locked largest first utxo + assert_eq!(tx_inputs.len(), 2); + assert_eq!(tx_inputs[0].previous_output, unspent[1].outpoint); + assert_eq!(tx_inputs[1].previous_output, unspent[2].outpoint); + + // confirm locked largest utxo is unlocked + assert!(wallet.unlock_unspent(unspent[0].outpoint)); + let mut builder = wallet + .build_fee_bump(original_txid) + .unwrap() + .coin_selection(LargestFirstCoinSelection::default()); + builder + .ordering(TxOrdering::Untouched) + .fee_rate(new_fee_rate); + + let tx_inputs = builder.finish().unwrap().unsigned_tx.input; + // verify unlocked largest utxo used instead of smaller unlocked utxo + assert_eq!(tx_inputs.len(), 2); + assert_eq!(tx_inputs[0].previous_output, unspent[1].outpoint); + assert_eq!(tx_inputs[1].previous_output, unspent[0].outpoint); +}