Skip to content

Commit

Permalink
feat(wallet): add functions to lock and unlock utxos
Browse files Browse the repository at this point in the history
  • Loading branch information
notmandatory committed Nov 1, 2024
1 parent 493638e commit 423b4a8
Show file tree
Hide file tree
Showing 2 changed files with 165 additions and 6 deletions.
46 changes: 41 additions & 5 deletions crates/wallet/src/wallet/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ use miniscript::{
};

use bdk_chain::tx_graph::CalculateFeeError;
use chain::collections::HashSet;

mod changeset;
pub mod coin_selection;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -434,6 +436,7 @@ impl Wallet {
network,
chain,
indexed_graph,
locked_unspent: Default::default(),
stage,
secp,
})
Expand Down Expand Up @@ -625,6 +628,7 @@ impl Wallet {
change_signers,
chain,
indexed_graph,
locked_unspent: Default::default(),
stage,
network,
secp,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
///
Expand All @@ -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(),
}
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -1744,6 +1779,7 @@ impl Wallet {
absolute: fee,
rate: fee_rate,
}),
unspendable: self.locked_unspent.clone(),
..Default::default()
};

Expand Down
125 changes: 124 additions & 1 deletion crates/wallet/tests/wallet.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -4309,3 +4309,126 @@ 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);
}

#[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::<Vec<_>>();
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);
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);
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);
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);
}

0 comments on commit 423b4a8

Please sign in to comment.