diff --git a/src/psbt/mod.rs b/src/psbt/mod.rs index e02362a45..f06b5297c 100644 --- a/src/psbt/mod.rs +++ b/src/psbt/mod.rs @@ -9,11 +9,22 @@ // You may not use this file except in accordance with one or both of these // licenses. +use crate::FeeRate; use bitcoin::util::psbt::PartiallySignedTransaction as Psbt; use bitcoin::TxOut; pub trait PsbtUtils { fn get_utxo_for(&self, input_index: usize) -> Option; + + /// The total transaction fee amount, sum of input amounts minus sum of output amounts, in Sats. + /// If the PSBT is missing a TxOut for an input returns None. + fn fee_amount(&self) -> Option; + + /// The transaction's fee rate. This value will only be accurate if calculated AFTER the + /// `PartiallySignedTransaction` is finalized and all witness/signature data is added to the + /// transaction. + /// If the PSBT is missing a TxOut for an input returns None. + fn fee_rate(&self) -> Option; } impl PsbtUtils for Psbt { @@ -37,6 +48,27 @@ impl PsbtUtils for Psbt { None } } + + fn fee_amount(&self) -> Option { + let tx = &self.unsigned_tx; + let utxos: Option> = (0..tx.input.len()).map(|i| self.get_utxo_for(i)).collect(); + + utxos.map(|inputs| { + let input_amount: u64 = inputs.iter().map(|i| i.value).sum(); + let output_amount: u64 = self.unsigned_tx.output.iter().map(|o| o.value).sum(); + input_amount + .checked_sub(output_amount) + .expect("input amount must be greater than output amount") + }) + } + + fn fee_rate(&self) -> Option { + let fee_amount = self.fee_amount(); + fee_amount.map(|fee| { + let weight = self.clone().extract_tx().weight(); + FeeRate::from_wu(fee, weight) + }) + } } #[cfg(test)] @@ -44,8 +76,9 @@ mod test { use crate::bitcoin::TxIn; use crate::psbt::Psbt; use crate::wallet::AddressIndex; + use crate::wallet::AddressIndex::New; use crate::wallet::{get_funded_wallet, test::get_test_wpkh}; - use crate::SignOptions; + use crate::{psbt, FeeRate, SignOptions}; use std::str::FromStr; // from bip 174 @@ -118,4 +151,83 @@ mod test { let _ = wallet.sign(&mut psbt, SignOptions::default()).unwrap(); } + + #[test] + fn test_psbt_fee_rate_with_witness_utxo() { + use psbt::PsbtUtils; + + let expected_fee_rate = 1.2345; + + let (wallet, _, _) = get_funded_wallet("wpkh(tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*)"); + let addr = wallet.get_address(New).unwrap(); + let mut builder = wallet.build_tx(); + builder.drain_to(addr.script_pubkey()).drain_wallet(); + builder.fee_rate(FeeRate::from_sat_per_vb(expected_fee_rate)); + let (mut psbt, _) = builder.finish().unwrap(); + let fee_amount = psbt.fee_amount(); + assert!(fee_amount.is_some()); + + let unfinalized_fee_rate = psbt.fee_rate().unwrap(); + + let finalized = wallet.sign(&mut psbt, Default::default()).unwrap(); + assert!(finalized); + + let finalized_fee_rate = psbt.fee_rate().unwrap(); + assert!(finalized_fee_rate.as_sat_per_vb() >= expected_fee_rate); + assert!(finalized_fee_rate.as_sat_per_vb() < unfinalized_fee_rate.as_sat_per_vb()); + } + + #[test] + fn test_psbt_fee_rate_with_nonwitness_utxo() { + use psbt::PsbtUtils; + + let expected_fee_rate = 1.2345; + + let (wallet, _, _) = get_funded_wallet("pkh(tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*)"); + let addr = wallet.get_address(New).unwrap(); + let mut builder = wallet.build_tx(); + builder.drain_to(addr.script_pubkey()).drain_wallet(); + builder.fee_rate(FeeRate::from_sat_per_vb(expected_fee_rate)); + let (mut psbt, _) = builder.finish().unwrap(); + let fee_amount = psbt.fee_amount(); + assert!(fee_amount.is_some()); + let unfinalized_fee_rate = psbt.fee_rate().unwrap(); + + let finalized = wallet.sign(&mut psbt, Default::default()).unwrap(); + assert!(finalized); + + let finalized_fee_rate = psbt.fee_rate().unwrap(); + assert!(finalized_fee_rate.as_sat_per_vb() >= expected_fee_rate); + assert!(finalized_fee_rate.as_sat_per_vb() < unfinalized_fee_rate.as_sat_per_vb()); + } + + #[test] + fn test_psbt_fee_rate_with_missing_txout() { + use psbt::PsbtUtils; + + let expected_fee_rate = 1.2345; + + let (wpkh_wallet, _, _) = get_funded_wallet("wpkh(tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*)"); + let addr = wpkh_wallet.get_address(New).unwrap(); + let mut builder = wpkh_wallet.build_tx(); + builder.drain_to(addr.script_pubkey()).drain_wallet(); + builder.fee_rate(FeeRate::from_sat_per_vb(expected_fee_rate)); + let (mut wpkh_psbt, _) = builder.finish().unwrap(); + + wpkh_psbt.inputs[0].witness_utxo = None; + wpkh_psbt.inputs[0].non_witness_utxo = None; + assert!(wpkh_psbt.fee_amount().is_none()); + assert!(wpkh_psbt.fee_rate().is_none()); + + let (pkh_wallet, _, _) = get_funded_wallet("pkh(tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*)"); + let addr = pkh_wallet.get_address(New).unwrap(); + let mut builder = pkh_wallet.build_tx(); + builder.drain_to(addr.script_pubkey()).drain_wallet(); + builder.fee_rate(FeeRate::from_sat_per_vb(expected_fee_rate)); + let (mut pkh_psbt, _) = builder.finish().unwrap(); + + pkh_psbt.inputs[0].non_witness_utxo = None; + assert!(pkh_psbt.fee_amount().is_none()); + assert!(pkh_psbt.fee_rate().is_none()); + } }