Add fee_amount() and fee_rate() functions to PsbtUtils trait

Author: notmandatory

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<TxOut>;
+
+    /// 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<u64>;
+
+    /// 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<FeeRate>;
 }
 
 impl PsbtUtils for Psbt {
@@ -37,15 +48,39 @@ impl PsbtUtils for Psbt {
             None
         }
     }
+
+    fn fee_amount(&self) -> Option<u64> {
+        let tx = &self.unsigned_tx;
+        let utxos: Option<Vec<TxOut>> = (0..tx.input.len()).map(|i| self.get_utxo_for(i)).collect();
+
+        utxos.map(|inputs| {
+            let input_amount = inputs.iter().fold(0, |acc, i| acc + i.value);
+
+            let output_amount = &self
+                .unsigned_tx
+                .output
+                .iter()
+                .fold(0, |acc, o| acc + o.value);
+
+            input_amount - output_amount
+        })
+    }
+
+    fn fee_rate(&self) -> Option<FeeRate> {
+        let fee_amount = self.fee_amount();
+        let weight = self.clone().extract_tx().weight();
+        fee_amount.map(|fee| FeeRate::from_wu(fee, weight))
+    }
 }
 
 #[cfg(test)]
 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 +153,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());
+    }
 }