diff --git a/wallet/src/wallet/changeset.rs b/wallet/src/wallet/changeset.rs index ebfdb9fb..620f69b3 100644 --- a/wallet/src/wallet/changeset.rs +++ b/wallet/src/wallet/changeset.rs @@ -1,12 +1,15 @@ use bdk_chain::{ indexed_tx_graph, keychain_txout, local_chain, tx_graph, ConfirmationBlockTime, Merge, }; +use bitcoin::{OutPoint, Txid}; use miniscript::{Descriptor, DescriptorPublicKey}; use serde::{Deserialize, Serialize}; type IndexedTxGraphChangeSet = indexed_tx_graph::ChangeSet; +use crate::locked_outpoints; + /// A change set for [`Wallet`] /// /// ## Definition @@ -114,6 +117,8 @@ pub struct ChangeSet { pub tx_graph: tx_graph::ChangeSet, /// Changes to [`KeychainTxOutIndex`](keychain_txout::KeychainTxOutIndex). pub indexer: keychain_txout::ChangeSet, + /// Changes to locked outpoints. + pub locked_outpoints: locked_outpoints::ChangeSet, } impl Merge for ChangeSet { @@ -142,6 +147,9 @@ impl Merge for ChangeSet { self.network = other.network; } + // merge locked outpoints + self.locked_outpoints.merge(other.locked_outpoints); + Merge::merge(&mut self.local_chain, other.local_chain); Merge::merge(&mut self.tx_graph, other.tx_graph); Merge::merge(&mut self.indexer, other.indexer); @@ -154,6 +162,7 @@ impl Merge for ChangeSet { && self.local_chain.is_empty() && self.tx_graph.is_empty() && self.indexer.is_empty() + && self.locked_outpoints.is_empty() } } @@ -163,6 +172,8 @@ impl ChangeSet { pub const WALLET_SCHEMA_NAME: &'static str = "bdk_wallet"; /// Name of table to store wallet descriptors and network. pub const WALLET_TABLE_NAME: &'static str = "bdk_wallet"; + /// Name of table to store wallet locked outpoints. + pub const WALLET_OUTPOINT_LOCK_TABLE_NAME: &'static str = "bdk_wallet_locked_outpoints"; /// Get v0 sqlite [ChangeSet] schema pub fn schema_v0() -> alloc::string::String { @@ -177,12 +188,24 @@ impl ChangeSet { ) } + /// Get v1 sqlite [`ChangeSet`] schema. Schema v1 adds a table for locked outpoints. + pub fn schema_v1() -> alloc::string::String { + format!( + "CREATE TABLE {} ( \ + txid TEXT NOT NULL, \ + vout INTEGER NOT NULL, \ + PRIMARY KEY(txid, vout) \ + ) STRICT;", + Self::WALLET_OUTPOINT_LOCK_TABLE_NAME, + ) + } + /// Initialize sqlite tables for wallet tables. pub fn init_sqlite_tables(db_tx: &chain::rusqlite::Transaction) -> chain::rusqlite::Result<()> { crate::rusqlite_impl::migrate_schema( db_tx, Self::WALLET_SCHEMA_NAME, - &[&Self::schema_v0()], + &[&Self::schema_v0(), &Self::schema_v1()], )?; bdk_chain::local_chain::ChangeSet::init_sqlite_tables(db_tx)?; @@ -220,6 +243,24 @@ impl ChangeSet { changeset.network = network.map(Impl::into_inner); } + // Select locked outpoints. + let mut stmt = db_tx.prepare(&format!( + "SELECT txid, vout FROM {}", + Self::WALLET_OUTPOINT_LOCK_TABLE_NAME, + ))?; + let rows = stmt.query_map([], |row| { + Ok(( + row.get::<_, Impl>("txid")?, + row.get::<_, u32>("vout")?, + )) + })?; + let locked_outpoints = &mut changeset.locked_outpoints.locked_outpoints; + for row in rows { + let (Impl(txid), vout) = row?; + let outpoint = OutPoint::new(txid, vout); + locked_outpoints.insert(outpoint, true); + } + changeset.local_chain = local_chain::ChangeSet::from_sqlite(db_tx)?; changeset.tx_graph = tx_graph::ChangeSet::<_>::from_sqlite(db_tx)?; changeset.indexer = keychain_txout::ChangeSet::from_sqlite(db_tx)?; @@ -268,6 +309,31 @@ impl ChangeSet { })?; } + // Insert or delete locked outpoints. + let mut insert_stmt = db_tx.prepare_cached(&format!( + "REPLACE INTO {}(txid, vout) VALUES(:txid, :vout)", + Self::WALLET_OUTPOINT_LOCK_TABLE_NAME + ))?; + let mut delete_stmt = db_tx.prepare_cached(&format!( + "DELETE FROM {} WHERE txid=:txid AND vout=:vout", + Self::WALLET_OUTPOINT_LOCK_TABLE_NAME, + ))?; + let locked_outpoints = &self.locked_outpoints.locked_outpoints; + for (&outpoint, &is_locked) in locked_outpoints.iter() { + let OutPoint { txid, vout } = outpoint; + if is_locked { + insert_stmt.execute(named_params! { + ":txid": Impl(txid), + ":vout": vout, + })?; + } else { + delete_stmt.execute(named_params! { + ":txid": Impl(txid), + ":vout": vout, + })?; + } + } + self.local_chain.persist_to_sqlite(db_tx)?; self.tx_graph.persist_to_sqlite(db_tx)?; self.indexer.persist_to_sqlite(db_tx)?; @@ -311,3 +377,12 @@ impl From for ChangeSet { } } } + +impl From for ChangeSet { + fn from(locked_outpoints: locked_outpoints::ChangeSet) -> Self { + Self { + locked_outpoints, + ..Default::default() + } + } +} diff --git a/wallet/src/wallet/locked_outpoints.rs b/wallet/src/wallet/locked_outpoints.rs new file mode 100644 index 00000000..c70542ab --- /dev/null +++ b/wallet/src/wallet/locked_outpoints.rs @@ -0,0 +1,26 @@ +//! Module containing the locked outpoints change set. + +use bdk_chain::Merge; +use bitcoin::OutPoint; +use serde::{Deserialize, Serialize}; + +use crate::collections::BTreeMap; + +/// Represents changes to locked outpoints. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct ChangeSet { + /// The lock status of an outpoint, `true == is_locked`. + pub locked_outpoints: BTreeMap, +} + +impl Merge for ChangeSet { + fn merge(&mut self, other: Self) { + // Extend self with other. Any entries in `self` that share the same + // outpoint are overwritten. + self.locked_outpoints.extend(other.locked_outpoints); + } + + fn is_empty(&self) -> bool { + self.locked_outpoints.is_empty() + } +} diff --git a/wallet/src/wallet/mod.rs b/wallet/src/wallet/mod.rs index 2b4fbd4d..0398561f 100644 --- a/wallet/src/wallet/mod.rs +++ b/wallet/src/wallet/mod.rs @@ -53,6 +53,7 @@ mod changeset; pub mod coin_selection; pub mod error; pub mod export; +pub mod locked_outpoints; mod params; mod persisted; pub mod signer; @@ -109,6 +110,7 @@ pub struct Wallet { stage: ChangeSet, network: Network, secp: SecpCtx, + locked_outpoints: BTreeMap, } /// An update to [`Wallet`]. @@ -449,6 +451,7 @@ impl Wallet { let change_descriptor = index.get_descriptor(KeychainKind::Internal).cloned(); let indexed_graph = IndexedTxGraph::new(index); let indexed_graph_changeset = indexed_graph.initial_changeset(); + let locked_outpoints = BTreeMap::new(); let stage = ChangeSet { descriptor, @@ -457,6 +460,7 @@ impl Wallet { tx_graph: indexed_graph_changeset.tx_graph, indexer: indexed_graph_changeset.indexer, network: Some(network), + ..Default::default() }; Ok(Wallet { @@ -467,6 +471,7 @@ impl Wallet { indexed_graph, stage, secp, + locked_outpoints, }) } @@ -654,6 +659,21 @@ impl Wallet { indexed_graph.apply_changeset(changeset.indexer.into()); indexed_graph.apply_changeset(changeset.tx_graph.into()); + // Apply locked outpoints + let locked_outpoints = changeset.locked_outpoints.locked_outpoints; + let locked_outpoints = locked_outpoints + .into_iter() + .map(|(outpoint, is_locked)| { + ( + outpoint, + UtxoLock { + outpoint, + is_locked, + }, + ) + }) + .collect(); + let stage = ChangeSet::default(); Ok(Some(Wallet { @@ -664,6 +684,7 @@ impl Wallet { stage, network, secp, + locked_outpoints, })) } @@ -2108,6 +2129,8 @@ impl Wallet { CanonicalizationParams::default(), self.indexed_graph.index.outpoints().iter().cloned(), ) + // Filter out locked outpoints + .filter(|(_, txo)| !self.is_outpoint_locked(txo.outpoint)) // only create LocalOutput if UTxO is mature .filter_map(move |((k, i), full_txo)| { full_txo @@ -2376,6 +2399,82 @@ impl Wallet { &self.chain } + /// Get a reference to the locked outpoints. + pub fn locked_outpoints(&self) -> &BTreeMap { + &self.locked_outpoints + } + + /// List unspent outpoints that are currently locked. + pub fn list_locked_unspent(&self) -> impl Iterator + '_ { + self.list_unspent() + .filter(|output| self.is_outpoint_locked(output.outpoint)) + .map(|output| output.outpoint) + } + + /// Whether the `outpoint` is locked. See [`Wallet::lock_outpoint`] for more. + pub fn is_outpoint_locked(&self, outpoint: OutPoint) -> bool { + self.locked_outpoints + .get(&outpoint) + .map_or(false, |u| u.is_locked) + } + + /// Lock a wallet output identified by the given `outpoint`. + /// + /// A locked UTXO will not be selected as an input to fund a transaction. This is useful + /// for excluding or reserving candidate inputs during transaction creation. + /// + /// **You must persist the staged change for the lock status to be persistent**. To unlock a + /// previously locked outpoint, see [`Wallet::unlock_outpoint`]. + pub fn lock_outpoint(&mut self, outpoint: OutPoint) { + use crate::collections::btree_map; + let lock_value = true; + let mut changeset = locked_outpoints::ChangeSet::default(); + + // If the lock status changed, update the entry and record the change + // in the changeset. + match self.locked_outpoints.entry(outpoint) { + btree_map::Entry::Occupied(mut e) => { + let utxo = e.get_mut(); + if !utxo.is_locked { + utxo.is_locked = lock_value; + changeset.locked_outpoints.insert(outpoint, lock_value); + } + } + btree_map::Entry::Vacant(e) => { + e.insert(UtxoLock { + outpoint, + is_locked: lock_value, + }); + changeset.locked_outpoints.insert(outpoint, lock_value); + } + }; + + self.stage.merge(changeset.into()); + } + + /// Unlock the wallet output of the specified `outpoint`. + /// + /// **You must persist the staged changes.** + pub fn unlock_outpoint(&mut self, outpoint: OutPoint) { + use crate::collections::btree_map; + let lock_value = false; + let mut changeset = locked_outpoints::ChangeSet::default(); + + match self.locked_outpoints.entry(outpoint) { + btree_map::Entry::Vacant(..) => {} + btree_map::Entry::Occupied(mut e) => { + // If the utxo is currently locked, update the lock value and stage + // the change. + let utxo = e.get_mut(); + if utxo.is_locked { + utxo.is_locked = lock_value; + changeset.locked_outpoints.insert(outpoint, lock_value); + self.stage.merge(changeset.into()); + } + } + } + } + /// Introduces a `block` of `height` to the wallet, and tries to connect it to the /// `prev_blockhash` of the block's header. /// @@ -2579,6 +2678,20 @@ impl AsRef> for Wallet { } } +/// Records the lock status of a wallet UTXO. +/// +/// Only a single [`UtxoLock`] may be assigned to a particular outpoint at a time. Refer to the docs +/// for [`Wallet::lock_outpoint`]. Note that the lock status is user-defined and does not take +/// into account any timelocks encoded by the descriptor. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +#[non_exhaustive] +pub struct UtxoLock { + /// Outpoint. + pub outpoint: OutPoint, + /// Whether the outpoint is locked. + pub is_locked: bool, +} + /// Deterministically generate a unique name given the descriptors defining the wallet /// /// Compatible with [`wallet_name_from_descriptor`] diff --git a/wallet/tests/persisted_wallet.rs b/wallet/tests/persisted_wallet.rs index 997575c8..7ca23830 100644 --- a/wallet/tests/persisted_wallet.rs +++ b/wallet/tests/persisted_wallet.rs @@ -6,7 +6,9 @@ use assert_matches::assert_matches; use bdk_chain::{ keychain_txout::DEFAULT_LOOKAHEAD, ChainPosition, ConfirmationBlockTime, DescriptorExt, }; +use bdk_wallet::coin_selection::InsufficientFunds; use bdk_wallet::descriptor::IntoWalletDescriptor; +use bdk_wallet::error::CreateTxError; use bdk_wallet::test_utils::*; use bdk_wallet::{ ChangeSet, KeychainKind, LoadError, LoadMismatch, LoadWithPersistError, Wallet, WalletPersister, @@ -358,3 +360,88 @@ fn single_descriptor_wallet_persist_and_recover() { "single descriptor wallet should refuse change descriptor param" ); } + +#[test] +fn test_lock_outpoint_persist() -> anyhow::Result<()> { + use bdk_chain::rusqlite; + let mut conn = rusqlite::Connection::open_in_memory()?; + + let (desc, change_desc) = get_test_tr_single_sig_xprv_and_change_desc(); + let mut wallet = Wallet::create(desc, change_desc) + .network(Network::Signet) + .create_wallet(&mut conn)?; + + // Receive coins. + let mut outpoints = vec![]; + for i in 0..3 { + let op = receive_output(&mut wallet, Amount::from_sat(10_000), ReceiveTo::Mempool(i)); + outpoints.push(op); + } + + // Test: lock outpoints + let unspent = wallet.list_unspent().collect::>(); + assert!(!unspent.is_empty()); + for utxo in unspent { + wallet.lock_outpoint(utxo.outpoint); + assert!( + wallet.is_outpoint_locked(utxo.outpoint), + "Expect outpoint is locked" + ); + } + wallet.persist(&mut conn)?; + + // Test: The lock value is persistent + { + wallet = Wallet::load() + .load_wallet(&mut conn)? + .expect("wallet is persisted"); + + let unspent = wallet.list_unspent().collect::>(); + assert!(!unspent.is_empty()); + for utxo in unspent { + assert!( + wallet.is_outpoint_locked(utxo.outpoint), + "Expect recover lock value" + ); + } + let locked_unspent = wallet.list_locked_unspent().collect::>(); + assert_eq!(locked_unspent, outpoints); + + // Test: Locked outpoints are excluded from coin selection + let addr = wallet.next_unused_address(KeychainKind::External).address; + let mut tx_builder = wallet.build_tx(); + tx_builder.add_recipient(addr, Amount::from_sat(10_000)); + let res = tx_builder.finish(); + assert!( + matches!( + res, + Err(CreateTxError::CoinSelection(InsufficientFunds { + available: Amount::ZERO, + .. + })), + ), + "Locked outpoints should not be selected", + ); + } + + // Test: Unlock outpoints + { + wallet = Wallet::load() + .load_wallet(&mut conn)? + .expect("wallet is persisted"); + + let unspent = wallet.list_unspent().collect::>(); + for utxo in &unspent { + wallet.unlock_outpoint(utxo.outpoint); + assert!( + !wallet.is_outpoint_locked(utxo.outpoint), + "Expect outpoint is not locked" + ); + } + assert!(!wallet.locked_outpoints().values().any(|u| u.is_locked)); + assert!(wallet.list_locked_unspent().next().is_none()); + wallet.persist(&mut conn)?; + } + + Ok(()) +}