Provide a default CoinSelectionSource implementation via a new trait

Author: wpaulino

lightning/src/events/bump_transaction.rs

@@ -23,6 +23,7 @@ use crate::ln::chan_utils::{
 };
 use crate::events::Event;
 use crate::prelude::HashMap;
+use crate::sync::Mutex;
 use crate::util::logger::Logger;
 use crate::util::ser::Writeable;
 
@@ -308,7 +309,8 @@ pub struct CoinSelection {
 
 /// An abstraction over a bitcoin wallet that can perform coin selection over a set of UTXOs and can
 /// sign for them. The coin selection method aims to mimic Bitcoin Core's `fundrawtransaction` RPC,
-/// which most wallets should be able to satisfy.
+/// which most wallets should be able to satisfy. Otherwise, consider implementing [`WalletSource`],
+/// which can provide a default implementation of this trait when used with [`Wallet`].
 pub trait CoinSelectionSource {
 	/// Performs coin selection of a set of UTXOs, with at least 1 confirmation each, that are
 	/// available to spend. Implementations are free to pick their coin selection algorithm of
@@ -336,6 +338,141 @@ pub trait CoinSelectionSource {
 	fn sign_tx(&self, tx: &mut Transaction) -> Result<(), ()>;
 }
 
+/// An alternative to [`CoinSelectionSource`] that can be implemented and used along [`Wallet`] to
+/// provide a default implementation to [`CoinSelectionSource`].
+pub trait WalletSource {
+	/// Returns all UTXOs, with at least 1 confirmation each, that are available to spend.
+	fn list_confirmed_utxos(&self) -> Result<Vec<Utxo>, ()>;
+	/// Returns a script to use for change above dust resulting from a successful coin selection
+	/// attempt.
+	fn change_script(&self) -> Result<Script, ()>;
+	/// Signs and provides the full witness for all inputs within the transaction known to the
+	/// wallet (i.e., any provided via [`WalletSource::list_confirmed_utxos`]).
+	fn sign_tx(&self, tx: &mut Transaction) -> Result<(), ()>;
+}
+
+/// A wrapper over [`WalletSource`] that implements [`CoinSelection`] by preferring UTXOs that would
+/// avoid conflicting double spends. If not enough UTXOs are available to do so, conflicting double
+/// spends may happen.
+pub struct Wallet<W: Deref> where W::Target: WalletSource {
+	source: W,
+	// TODO: Do we care about cleaning this up once the UTXOs have a confirmed spend? We can do so
+	// by checking whether any UTXOs that exist in the map are no longer returned in
+	// `list_confirmed_utxos`.
+	locked_utxos: Mutex<HashMap<OutPoint, BumpId>>,
+}
+
+impl<W: Deref> Wallet<W> where W::Target: WalletSource {
+	/// Returns a new instance backed by the given [`WalletSource`] that serves as an implementation
+	/// of [`CoinSelectionSource`].
+	pub fn new(source: W) -> Self {
+		Self { source, locked_utxos: Mutex::new(HashMap::new()) }
+	}
+
+	/// Performs coin selection on the set of UTXOs obtained from
+	/// [`WalletSource::list_confirmed_utxos`]. Its algorithm can be described as "smallest
+	/// above-dust-after-spend first", with a slight twist: we may skip UTXOs that are above dust
+	/// after spending them at the target feerate if `force_conflicting_utxo_spend` is unset to
+	/// avoid producing conflicting transactions.
+	fn select_confirmed_utxos_internal(
+		&self, bump_id: BumpId, force_conflicting_utxo_spend: bool,
+		target_feerate_sat_per_1000_weight: u32, preexisting_tx_weight: u64, target_amount: u64,
+	) -> Result<CoinSelection, ()> {
+		let mut utxos = self.source.list_confirmed_utxos()?;
+		utxos.sort_unstable_by_key(|utxo| utxo.output.value);
+
+		let mut selected_amount = 0;
+		let mut fee_amount = preexisting_tx_weight * target_feerate_sat_per_1000_weight as u64;
+		let selected_utxos = {
+			let mut locked_utxos = self.locked_utxos.lock().unwrap();
+			let selected_utxos = utxos.into_iter().scan(
+				(&mut selected_amount, &mut fee_amount), |(selected_amount, fee_amount), utxo| {
+					if let Some(utxo_bump_id) = locked_utxos.get(&utxo.outpoint) {
+						if *utxo_bump_id != bump_id && !force_conflicting_utxo_spend {
+							return None;
+						}
+					}
+					let need_more_inputs = **selected_amount < target_amount + **fee_amount;
+					if need_more_inputs {
+						let fee_to_spend_utxo = target_feerate_sat_per_1000_weight as u64 *
+							((41 * WITNESS_SCALE_FACTOR) as u64 + utxo.witness_weight);
+						let utxo_value_after_fee = utxo.output.value.saturating_sub(fee_to_spend_utxo);
+						if utxo_value_after_fee > 0 {
+							**selected_amount += utxo.output.value;
+							**fee_amount += fee_to_spend_utxo;
+							Some(utxo)
+						} else {
+							None
+						}
+					} else {
+						None
+					}
+				}
+			).collect::<Vec<_>>();
+			let need_more_inputs = selected_amount < target_amount + fee_amount;
+			if need_more_inputs {
+				return Err(());
+			}
+			for utxo in &selected_utxos {
+				locked_utxos.insert(utxo.outpoint, bump_id);
+			}
+			selected_utxos
+		};
+
+		let remaining_amount = selected_amount - target_amount - fee_amount;
+		let change_script = self.source.change_script()?;
+		let change_output_fee = target_feerate_sat_per_1000_weight as u64
+			* (8 + change_script.consensus_encode(&mut sink()).unwrap() as u64);
+		let change_output_amount = remaining_amount.saturating_sub(change_output_fee);
+		let change_output = if change_output_amount < change_script.dust_value().to_sat() {
+			None
+		} else {
+			Some(TxOut { script_pubkey: change_script, value: change_output_amount })
+		};
+
+		Ok(CoinSelection {
+			confirmed_utxos: selected_utxos,
+			change_output,
+		})
+	}
+}
+
+impl<W: Deref> CoinSelectionSource for Wallet<W> where W::Target: WalletSource {
+	fn select_confirmed_utxos(
+		&self, bump_id: BumpId, must_spend: Vec<Input>, must_pay_to: Vec<TxOut>,
+		target_feerate_sat_per_1000_weight: u32,
+	) -> Result<CoinSelection, ()> {
+		// TODO: Use fee estimation utils when we upgrade to bitcoin v0.30.0.
+		let base_tx_weight = 4 /* version */ + 1 /* input count */ + 1 /* output count */ + 4 /* locktime */;
+		let total_input_weight = must_spend.len() *
+			(32 /* txid */ + 4 /* vout */ + 4 /* sequence */ + 1 /* script sig */);
+		let total_output_weight: usize = must_pay_to.iter().map(|output|
+			8 /* value */ + 1 /* script len */ + output.script_pubkey.len()
+		).sum();
+		let total_non_witness_weight = base_tx_weight + total_input_weight + total_output_weight;
+		let total_witness_weight: u64 = must_spend.iter().map(|input| input.witness_weight).sum();
+
+		let preexisting_tx_weight = 2 /* segwit marker & flag */ + total_witness_weight +
+			(total_non_witness_weight * WITNESS_SCALE_FACTOR) as u64;
+		let target_amount = must_pay_to.iter().map(|output| output.value).sum();
+		let do_coin_selection = |force_conflicting_utxo_spend: bool| {
+			self.select_confirmed_utxos_internal(
+				bump_id, force_conflicting_utxo_spend, target_feerate_sat_per_1000_weight,
+				preexisting_tx_weight, target_amount,
+			)
+		};
+		do_coin_selection(false).or_else(|_| do_coin_selection(true))
+	}
+
+	fn change_script(&self) -> Result<Script, ()> {
+		self.source.change_script()
+	}
+
+	fn sign_tx(&self, tx: &mut Transaction) -> Result<(), ()> {
+		self.source.sign_tx(tx)
+	}
+}
+
 /// A handler for [`Event::BumpTransaction`] events that sources confirmed UTXOs from a
 /// [`CoinSelectionSource`] to fee bump transactions via Child-Pays-For-Parent (CPFP) or
 /// Replace-By-Fee (RBF).
@@ -402,7 +539,6 @@ where
 		// input, without caring where the change goes, we use an output just above dust backed by
 		// the wallet's change script. If the wallet ends up producing its own change output when
 		// funding the transaction, we'll join them into one, saving the user a few satoshis.
-		// TODO: Prevent change address inflation.
 		let change_script = self.utxo_source.change_script()?;
 		let dust_change_output = TxOut {
 			value: change_script.dust_value().to_sat(),