bark-bitcoin-ext 0.1.2

use std::borrow::BorrowMut;
use std::sync::Arc;

use bdk_wallet::{AddressInfo, TxBuilder, Wallet};
use bdk_wallet::chain::BlockId;
use bdk_wallet::coin_selection::InsufficientFunds;
use bdk_wallet::error::CreateTxError;
use bitcoin::consensus::encode::serialize_hex;
use bitcoin::{Amount, BlockHash, FeeRate, OutPoint, Transaction, TxOut, Txid, Weight, Witness};
use bitcoin::psbt::{ExtractTxError, Input};
use log::{debug, trace};

use crate::{BlockHeight, TransactionExt};
use crate::cpfp::MakeCpfpFees;
use crate::fee::FEE_ANCHOR_SPEND_WEIGHT;

/// Balance categorized by our recursive trust model.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct TrustedBalance {
	/// Funds in UTXOs we trust (confirmed or all-ours unconfirmed chains).
	pub trusted: Amount,
	/// Funds in UTXOs we don't trust.
	pub untrusted: Amount,
}

impl TrustedBalance {
	pub fn total(&self) -> Amount {
		self.trusted + self.untrusted
	}
}

/// The [bdk_wallet::KeychainKind] that is always used, because we only use a single keychain.
pub const KEYCHAIN: bdk_wallet::KeychainKind = bdk_wallet::KeychainKind::External;


/// An extension trait for [TxBuilder].
pub trait TxBuilderExt<'a, A>: BorrowMut<TxBuilder<'a, A>> {
	/// Add an input to the tx that spends a fee anchor.
	fn add_fee_anchor_spend(&mut self, anchor: OutPoint, output: &TxOut)
	where
		A: bdk_wallet::coin_selection::CoinSelectionAlgorithm,
	{
		let psbt_in = Input {
			witness_utxo: Some(output.clone()),
			final_script_witness: Some(Witness::new()),
			..Default::default()
		};
		self.borrow_mut().add_foreign_utxo(anchor, psbt_in, FEE_ANCHOR_SPEND_WEIGHT)
			.expect("adding foreign utxo");
	}
}
impl<'a, A> TxBuilderExt<'a, A> for TxBuilder<'a, A> {}

#[derive(Debug, thiserror::Error)]
pub enum CpfpInternalError {
	#[error("{0}")]
	General(String),
	#[error("Unable to construct transaction: {0}")]
	Create(CreateTxError),
	#[error("Unable to extract the final transaction after signing the PSBT: {0}")]
	Extract(ExtractTxError),
	#[error("Failed to determine the weight/fee when creating a P2A CPFP")]
	Fee(),
	#[error("Unable to finalize CPFP transaction: {0}")]
	FinalizeError(String),
	#[error("You need more confirmations on your on-chain funds: {0}")]
	InsufficientConfirmedFunds(InsufficientFunds),
	#[error("Transaction has no fee anchor: {0}")]
	NoFeeAnchor(Txid),
	#[allow(deprecated)]
	#[error("Unable to sign transaction: {0}")]
	Signer(bdk_wallet::signer::SignerError),
}

/// An extension trait for [Wallet].
pub trait WalletExt: BorrowMut<Wallet> {
	/// Peek into the next address.
	fn peek_next_address(&self) -> AddressInfo {
		self.borrow().peek_address(KEYCHAIN, self.borrow().next_derivation_index(KEYCHAIN))
	}

	/// Returns an iterator for each unconfirmed transaction in the wallet.
	fn unconfirmed_txids(&self) -> impl Iterator<Item = Txid> {
		self.borrow().transactions().filter_map(|tx| {
			if tx.chain_position.is_unconfirmed() {
				Some(tx.tx_node.txid)
			} else {
				None
			}
		})
	}

	/// Returns an iterator for each unconfirmed transaction in the wallet, useful for syncing
	/// with bitcoin core.
	fn unconfirmed_txs(&self) -> impl Iterator<Item = Arc<Transaction>> {
		self.borrow().transactions().filter_map(|tx| {
			if tx.chain_position.is_unconfirmed() {
				Some(tx.tx_node.tx.clone())
			} else {
				None
			}
		})
	}

	/// Check whether a UTXO can be trusted for spending.
	///
	/// Delegates to [WalletExt::is_trusted_tx] on the creating transaction.
	fn is_trusted_utxo(&self, outpoint: OutPoint, min_confs: u32) -> bool {
		self.is_trusted_tx(outpoint.txid, min_confs)
	}

	/// Check whether a transaction can be trusted to confirm on chain.
	///
	/// A transaction is trusted if:
	/// - `min_confs` is 0 (unconditionally trusted), or
	/// - it has at least `min_confs` confirmations, or
	/// - all its inputs are ours and their creating transactions are also
	///   trusted (checked recursively).
	///
	/// To keep the check cheap, at most 100 ancestor transactions are visited.
	/// If the budget is exhausted the transaction is considered untrusted.
	fn is_trusted_tx(&self, txid: Txid, min_confs: u32) -> bool {
		let w = self.borrow();
		let tip = w.latest_checkpoint().height();
		let mut budget = 100u32;
		is_trusted_tx_inner(w, txid, min_confs, tip, &mut budget)
	}

	/// Compute the wallet balance using our recursive trust model.
	fn trusted_balance(&self, min_confs: u32) -> TrustedBalance {
		let mut trusted = Amount::ZERO;
		let mut untrusted = Amount::ZERO;
		for utxo in self.borrow().list_unspent() {
			if self.is_trusted_utxo(utxo.outpoint, min_confs) {
				trusted += utxo.txout.value;
			} else {
				untrusted += utxo.txout.value;
			}
		}
		TrustedBalance { trusted, untrusted }
	}

	/// Return all UTXOs that are untrusted.
	///
	/// Delegates to [WalletExt::is_trusted_utxo] for each UTXO.
	fn untrusted_utxos(&self, min_confs: u32) -> Vec<OutPoint> {
		self.borrow().list_unspent()
			.filter(|utxo| !self.is_trusted_utxo(utxo.outpoint, min_confs))
			.map(|utxo| utxo.outpoint)
			.collect()
	}

	/// Insert a checkpoint into the wallet.
	///
	/// It's advised to use this only when recovering a wallet with a birthday.
	fn set_checkpoint(&mut self, height: u32, hash: BlockHash) {
		let checkpoint = BlockId { height, hash };
		let wallet = self.borrow_mut();
		wallet.apply_update(bdk_wallet::Update {
			chain: Some(wallet.latest_checkpoint().insert(checkpoint)),
			..Default::default()
		}).expect("should work, might fail if tip is genesis");
	}

	fn make_signed_p2a_cpfp(
		&mut self,
		tx: &Transaction,
		fees: MakeCpfpFees,
	) -> Result<Transaction, CpfpInternalError> {
		let wallet = self.borrow_mut();
		let (outpoint, txout) = tx.fee_anchor()
			.ok_or_else(|| CpfpInternalError::NoFeeAnchor(tx.compute_txid()))?;

		// Since BDK doesn't support adding extra weight for fees, we have to loop to achieve the
		// effective fee rate and potential minimum fee we need.
		let p2a_weight = tx.weight();
		let extra_fee_needed = p2a_weight * fees.effective();

		// Since BDK doesn't allow tx without recipients, we add a drain output.
		let change_addr = wallet.reveal_next_address(KEYCHAIN);
		let dust_limit = change_addr.address.script_pubkey().minimal_non_dust();

		// We will loop, constructing the transaction and signing it until we exceed the effective
		// fee rate and meet any minimum fee requirements
		let mut spend_weight = Weight::ZERO;
		let mut fee_needed = extra_fee_needed;
		for i in 0..100 {
			// We need to account for a particularly annoying BDK bug when using foreign UTXOs when
			// BDK tries to use the P2A value to pay the fees. If the P2A has a value of 420 sats
			// and the absolute fee is 200 sats, this will produce a 220 sat change output which
			// results in a coin selection error. Ideally, BDK would pull in an extra UTXO to ensure
			// the change output is more than the dust limit; however, this seems to be an edge case
			// with experimental foreign UTXOs.
			if fee_needed < txout.value {
				if txout.value - fee_needed < dust_limit {
					fee_needed = txout.value + Amount::ONE_SAT;
				}
			}

			let mut b = wallet.build_tx();
			b.only_witness_utxo();
			b.exclude_unconfirmed();
			b.version(3); // for 1p1c package relay, all inputs must be confirmed
			b.add_fee_anchor_spend(outpoint, txout);
			b.drain_to(change_addr.address.script_pubkey());
			b.fee_absolute(fee_needed);

			// Attempt to create and sign the transaction
			let mut psbt = b.finish().map_err(|e| match e {
				CreateTxError::CoinSelection(e) => CpfpInternalError::InsufficientConfirmedFunds(e),
				_ => CpfpInternalError::Create(e),
			})?;
			#[allow(deprecated)]
			let opts = bdk_wallet::SignOptions {
				trust_witness_utxo: true,
				..Default::default()
			};
			let finalized = wallet.sign(&mut psbt, opts)
				.map_err(|e| CpfpInternalError::Signer(e))?;
			if !finalized {
				return Err(CpfpInternalError::FinalizeError("finalization failed".into()));
			}
			let tx = psbt.extract_tx()
				.map_err(|e| CpfpInternalError::Extract(e))?;
			let anchor_weight = FEE_ANCHOR_SPEND_WEIGHT.to_wu();
			assert!(tx.input.iter().any(|i| i.witness.size() as u64 == anchor_weight),
				"Missing anchor spend, tx is {}", serialize_hex(&tx),
			);

			// We can finally check the fees and weight
			let tx_weight = tx.weight();
			let total_weight = tx_weight + p2a_weight;
			if tx_weight != spend_weight {
				// Since the weight changed, we can drop the transaction and recalculate the
				// required fee amount.
				wallet.cancel_tx(&tx);
				spend_weight = tx_weight;
				fee_needed = match fees {
					MakeCpfpFees::Effective(fr) => total_weight * fr,
					MakeCpfpFees::Rbf { min_effective_fee_rate, current_package_fee } => {
						// RBF requires that you spend at least the total fee of every
						// unconfirmed ancestor and the transaction you want to replace,
						// then you must add mintxrelayfee * package_vbytes on top.
						let min_tx_relay_fee = FeeRate::from_sat_per_vb(1).unwrap();
						let min_package_fee = current_package_fee +
							p2a_weight * min_tx_relay_fee +
							tx_weight * min_tx_relay_fee;

						// This is the fee we want to pay based on the given minimum effective fee
						// rate. It's possible that the desired fee is lower than the minimum
						// package fee if the currently broadcast child transaction is bigger than
						// the transaction we just produced.
						let desired_fee = total_weight * min_effective_fee_rate;
						if desired_fee < min_package_fee {
							debug!("Using a minimum fee of {} instead of the desired fee of {} for RBF",
								min_package_fee, desired_fee,
							);
							min_package_fee
						} else {
							trace!("Attempting to use the desired fee of {} for CPFP RBF",
								desired_fee,
							);
							desired_fee
						}
					}
				}
			} else {
				debug!("Created P2A CPFP with weight {} and fee {} in {} iterations",
					total_weight, fee_needed, i,
				);
				return Ok(tx);
			}
		}
		Err(CpfpInternalError::General("Reached max iterations".into()))
	}
}

impl WalletExt for Wallet {}

fn is_trusted_tx_inner(
	w: &Wallet, txid: Txid, min_confs: u32, tip: BlockHeight, budget: &mut u32,
) -> bool {
	if *budget == 0 {
		return false;
	}
	*budget -= 1;

	let Some(tx) = w.get_tx(txid) else {
		return false;
	};

	// Trust transactions with enough confirmations (unconfirmed = 0 confs).
	let nb_confs = match tx.chain_position.confirmation_height_upper_bound() {
		Some(h) => tip.saturating_sub(h) + 1,
		None => 0,
	};
	if nb_confs >= min_confs {
		return true;
	}

	// Recursively check that all inputs are ours and come from trusted transactions.
	tx.tx_node.tx.input.iter().all(|input| {
		let prev = input.previous_output;
		let Some(prev_tx) = w.get_tx(prev.txid) else {
			return false;
		};
		let Some(txout) = prev_tx.tx_node.tx.output.get(prev.vout as usize) else {
			return false;
		};
		w.is_mine(txout.script_pubkey.clone())
			&& is_trusted_tx_inner(w, prev.txid, min_confs, tip, budget)
	})
}