From ae1ea99ed3f26f33ead92d4a91588ae32d9e110b Mon Sep 17 00:00:00 2001
From: Steve Myers <steve@notmandatory.org>
Date: Tue, 18 Oct 2022 15:59:07 -0500
Subject: [PATCH] Add FeeRate struct and fee_amount() and fee_rate() functions
 on PartiallySignedTransaction

---
 src/bdk.udl | 11 +++++++++++
 src/lib.rs  | 48 ++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/src/bdk.udl b/src/bdk.udl
index d696b02..a1e6fe5 100644
--- a/src/bdk.udl
+++ b/src/bdk.udl
@@ -202,6 +202,13 @@ interface Wallet {
   sequence<LocalUtxo> list_unspent();
 };
 
+interface FeeRate {
+  [Name=from_sat_per_vb]
+  constructor(float sat_per_vb);
+
+  float as_sat_per_vb();
+};
+
 interface PartiallySignedTransaction {
   [Throws=BdkError]
   constructor(string psbt_base64);
@@ -214,6 +221,10 @@ interface PartiallySignedTransaction {
 
   [Throws=BdkError]
   PartiallySignedTransaction combine(PartiallySignedTransaction other);
+
+  u64? fee_amount();
+
+  FeeRate? fee_rate();
 };
 
 dictionary TxBuilderResult {
diff --git a/src/lib.rs b/src/lib.rs
index ed98e42..cb4b528 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,7 +3,7 @@ use bdk::bitcoin::hashes::hex::ToHex;
 use bdk::bitcoin::secp256k1::Secp256k1;
 use bdk::bitcoin::util::bip32::DerivationPath as BdkDerivationPath;
 use bdk::bitcoin::util::psbt::serialize::Serialize;
-use bdk::bitcoin::util::psbt::PartiallySignedTransaction;
+use bdk::bitcoin::util::psbt::PartiallySignedTransaction as BdkPartiallySignedTransaction;
 use bdk::bitcoin::Sequence;
 use bdk::bitcoin::{Address as BdkAddress, Network, OutPoint as BdkOutPoint, Txid};
 use bdk::blockchain::any::{AnyBlockchain, AnyBlockchainConfig};
@@ -22,6 +22,7 @@ use bdk::keys::{
     DescriptorSecretKey as BdkDescriptorSecretKey, ExtendedKey, GeneratableKey, GeneratedKey,
 };
 use bdk::miniscript::BareCtx;
+use bdk::psbt::PsbtUtils;
 use bdk::wallet::tx_builder::ChangeSpendPolicy;
 use bdk::wallet::AddressIndex as BdkAddressIndex;
 use bdk::wallet::AddressInfo as BdkAddressInfo;
@@ -347,7 +348,8 @@ pub struct PartiallySignedTransaction {
 
 impl PartiallySignedTransaction {
     fn new(psbt_base64: String) -> Result<Self, BdkError> {
-        let psbt: BdkPartiallySignedTransaction = BdkPartiallySignedTransaction::from_str(&psbt_base64)?;
+        let psbt: BdkPartiallySignedTransaction =
+            BdkPartiallySignedTransaction::from_str(&psbt_base64)?;
         Ok(PartiallySignedTransaction {
             internal: Mutex::new(psbt),
         })
@@ -389,6 +391,20 @@ impl PartiallySignedTransaction {
             internal: Mutex::new(original_psbt),
         }))
     }
+
+    /// 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> {
+        self.internal.lock().unwrap().fee_amount()
+    }
+
+    /// 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<Arc<FeeRate>> {
+        self.internal.lock().unwrap().fee_rate().map(Arc::new)
+    }
 }
 
 /// A Bitcoin wallet.
@@ -1246,4 +1262,32 @@ mod test {
             "e93315d6ce401eb4db803a56232f0ed3e69b053774e6047df54f1bd00e5ea936"
         )
     }
+
+    #[test]
+    fn test_psbt_fee() {
+        let test_wpkh = "wpkh(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW)";
+        let (funded_wallet, _, _) = get_funded_wallet(test_wpkh);
+        let test_wallet = Wallet {
+            wallet_mutex: Mutex::new(funded_wallet),
+        };
+        let drain_to_address = "tb1ql7w62elx9ucw4pj5lgw4l028hmuw80sndtntxt".to_string();
+        let tx_builder = TxBuilder::new()
+            .fee_rate(2.0)
+            .drain_wallet()
+            .drain_to(drain_to_address.clone());
+        //dbg!(&tx_builder);
+        assert!(tx_builder.drain_wallet);
+        assert_eq!(tx_builder.drain_to, Some(drain_to_address));
+
+        let tx_builder_result = tx_builder.finish(&test_wallet).unwrap();
+
+        assert!(tx_builder_result.psbt.fee_rate().is_some());
+        assert_eq!(
+            tx_builder_result.psbt.fee_rate().unwrap().as_sat_per_vb(),
+            2.682927
+        );
+
+        assert!(tx_builder_result.psbt.fee_amount().is_some());
+        assert_eq!(tx_builder_result.psbt.fee_amount().unwrap(), 220);
+    }
 }