From 9b4b450e88d48f5a8e7dd511bb637b3a856c9f86 Mon Sep 17 00:00:00 2001
From: benthecarman <benthecarman@live.com>
Date: Mon, 18 Dec 2023 13:50:16 -0600
Subject: [PATCH] Change Signer::sign_tx_input to Signer::sign_psbt_input

With the previous API, signing taproot inputs comes with a lot of
problems. Taproot signatures commit to the previous txo of all the
inputs so to be able to construct a signature, you need to have all of
them available. The previous API would require the user to look up the
outputs themself which does not really work for mobile clients and it is
a bad practice. This changes it to a PSBT so we can give all the
necessary info to the user directly.
---
 bitcoin-rpc-provider/src/lib.rs     |  37 +++++++--
 dlc-manager/src/channel_updater.rs  |  12 ++-
 dlc-manager/src/contract_updater.rs | 118 ++++++++++++++++------------
 dlc-manager/src/lib.rs              |   7 +-
 dlc-manager/src/manager.rs          |   4 +-
 mocks/src/mock_wallet.rs            |  11 ++-
 simple-wallet/src/lib.rs            |  45 +++++++++--
 7 files changed, 153 insertions(+), 81 deletions(-)

diff --git a/bitcoin-rpc-provider/src/lib.rs b/bitcoin-rpc-provider/src/lib.rs
index b8153ba5..413b2903 100644
--- a/bitcoin-rpc-provider/src/lib.rs
+++ b/bitcoin-rpc-provider/src/lib.rs
@@ -7,6 +7,7 @@ use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
 use bitcoin::consensus::encode::Error as EncodeError;
+use bitcoin::psbt::PartiallySignedTransaction;
 use bitcoin::secp256k1::rand::thread_rng;
 use bitcoin::secp256k1::{PublicKey, SecretKey};
 use bitcoin::{
@@ -185,14 +186,32 @@ impl Signer for BitcoinCoreProvider {
         Ok(pk.inner)
     }
 
-    fn sign_tx_input(
+    fn sign_psbt_input(
         &self,
-        tx: &mut Transaction,
+        psbt: &mut PartiallySignedTransaction,
         input_index: usize,
-        tx_out: &TxOut,
-        redeem_script: Option<Script>,
     ) -> Result<(), ManagerError> {
-        let outpoint = &tx.input[input_index].previous_output;
+        let outpoint = &psbt.unsigned_tx.input[input_index].previous_output;
+        let tx_out = if let Some(input) = psbt.inputs.get(input_index) {
+            if let Some(wit_utxo) = &input.witness_utxo {
+                Ok(wit_utxo.clone())
+            } else if let Some(in_tx) = &input.non_witness_utxo {
+                Ok(in_tx.output[outpoint.vout as usize].clone())
+            } else {
+                Err(ManagerError::InvalidParameters(
+                    "No TxOut for PSBT input".to_string(),
+                ))
+            }
+        } else {
+            Err(ManagerError::InvalidParameters(
+                "No TxOut for PSBT input".to_string(),
+            ))
+        }?;
+
+        let redeem_script = psbt
+            .inputs
+            .get(input_index)
+            .and_then(|i| i.redeem_script.clone());
 
         let input = json::SignRawTransactionInput {
             txid: outpoint.txid,
@@ -206,13 +225,15 @@ impl Signer for BitcoinCoreProvider {
             .client
             .lock()
             .unwrap()
-            .sign_raw_transaction_with_wallet(&*tx, Some(&[input]), None)
+            .sign_raw_transaction_with_wallet(&psbt.unsigned_tx, Some(&[input]), None)
             .map_err(rpc_err_to_manager_err)?;
         let signed_tx = Transaction::consensus_decode(&mut sign_result.hex.as_slice())
             .map_err(enc_err_to_manager_err)?;
 
-        tx.input[input_index].script_sig = signed_tx.input[input_index].script_sig.clone();
-        tx.input[input_index].witness = signed_tx.input[input_index].witness.clone();
+        psbt.inputs[input_index].final_script_sig =
+            Some(signed_tx.input[input_index].script_sig.clone());
+        psbt.inputs[input_index].final_script_witness =
+            Some(signed_tx.input[input_index].witness.clone());
 
         Ok(())
     }
diff --git a/dlc-manager/src/channel_updater.rs b/dlc-manager/src/channel_updater.rs
index dfd9cf74..9afdaf15 100644
--- a/dlc-manager/src/channel_updater.rs
+++ b/dlc-manager/src/channel_updater.rs
@@ -435,7 +435,7 @@ pub fn verify_signed_channel<S: Deref>(
     accepted_contract: &AcceptedContract,
     sign_channel: &SignChannel,
     signer: &S,
-) -> Result<(SignedChannel, SignedContract), Error>
+) -> Result<(SignedChannel, SignedContract, Transaction), Error>
 where
     S::Target: Signer,
 {
@@ -458,7 +458,7 @@ where
 
     let cet_adaptor_signatures: Vec<_> = (&sign_channel.cet_adaptor_signatures).into();
 
-    let (signed_contract, fund_tx) = verify_signed_contract_internal(
+    let (signed_contract, signed_fund_tx) = verify_signed_contract_internal(
         secp,
         accepted_contract,
         &sign_channel.refund_signature,
@@ -493,7 +493,11 @@ where
             is_offer: false,
         },
         update_idx: INITIAL_UPDATE_NUMBER,
-        fund_tx,
+        fund_tx: signed_contract
+            .accepted_contract
+            .dlc_transactions
+            .fund
+            .clone(),
         fund_script_pubkey: accepted_contract
             .dlc_transactions
             .funding_script_pubkey
@@ -507,7 +511,7 @@ where
             .fee_rate_per_vb,
     };
 
-    Ok((signed_channel, signed_contract))
+    Ok((signed_channel, signed_contract, signed_fund_tx))
 }
 
 /// Creates a [`SettleOffer`] message from the given [`SignedChannel`] and parameters,
diff --git a/dlc-manager/src/contract_updater.rs b/dlc-manager/src/contract_updater.rs
index 4fefb8c2..4cf9fd03 100644
--- a/dlc-manager/src/contract_updater.rs
+++ b/dlc-manager/src/contract_updater.rs
@@ -2,6 +2,7 @@
 
 use std::ops::Deref;
 
+use bitcoin::psbt::PartiallySignedTransaction;
 use bitcoin::{consensus::Decodable, Script, Transaction, Witness};
 use dlc::{DlcTransactions, PartyParams};
 use dlc_messages::{
@@ -284,6 +285,31 @@ where
     Ok((signed_contract, signed_msg))
 }
 
+fn populate_psbt(
+    psbt: &mut PartiallySignedTransaction,
+    all_funding_inputs: &[&FundingInputInfo],
+) -> Result<(), Error> {
+    // add witness utxo to fund_psbt for all inputs
+    for (input_index, x) in all_funding_inputs.iter().enumerate() {
+        let tx = Transaction::consensus_decode(&mut x.funding_input.prev_tx.as_slice()).map_err(
+            |_| {
+                Error::InvalidParameters(
+                    "Could not decode funding input previous tx parameter".to_string(),
+                )
+            },
+        )?;
+        let vout = x.funding_input.prev_tx_vout;
+        let tx_out = tx.output.get(vout as usize).ok_or_else(|| {
+            Error::InvalidParameters(format!("Previous tx output not found at index {}", vout))
+        })?;
+
+        psbt.inputs[input_index].witness_utxo = Some(tx_out.clone());
+        psbt.inputs[input_index].redeem_script = Some(x.funding_input.redeem_script.clone());
+    }
+
+    Ok(())
+}
+
 pub(crate) fn verify_accepted_and_sign_contract_internal<S: Deref>(
     secp: &Secp256k1<All>,
     offered_contract: &OfferedContract,
@@ -309,7 +335,8 @@ where
         funding_script_pubkey,
     } = dlc_transactions;
 
-    let mut fund = fund.clone();
+    let mut fund_psbt = PartiallySignedTransaction::from_unsigned_tx(fund.clone())
+        .map_err(|_| Error::InvalidState("Tried to create PSBT from signed tx".to_string()))?;
     let mut cets = cets.clone();
 
     let input_script_pubkey = input_script_pubkey.unwrap_or_else(|| funding_script_pubkey.clone());
@@ -392,43 +419,42 @@ where
         own_signatures.extend(sigs);
     }
 
-    let mut input_serial_ids: Vec<_> = offered_contract
+    // get all funding inputs
+    let mut all_funding_inputs = offered_contract
         .funding_inputs_info
         .iter()
-        .map(|x| x.funding_input.input_serial_id)
-        .chain(accept_params.inputs.iter().map(|x| x.serial_id))
-        .collect();
-    input_serial_ids.sort_unstable();
+        .chain(funding_inputs_info.iter())
+        .collect::<Vec<_>>();
+    // sort by serial id
+    all_funding_inputs.sort_by_key(|x| x.funding_input.input_serial_id);
+
+    populate_psbt(&mut fund_psbt, &all_funding_inputs)?;
 
     // Vec<Witness>
     let witnesses: Vec<Witness> = offered_contract
         .funding_inputs_info
         .iter()
         .map(|x| {
-            let input_index = input_serial_ids
+            let input_index = all_funding_inputs
                 .iter()
-                .position(|y| y == &x.funding_input.input_serial_id)
+                .position(|y| y.funding_input == x.funding_input)
                 .ok_or_else(|| {
                     Error::InvalidState(format!(
                         "Could not find input for serial id {}",
                         x.funding_input.input_serial_id
                     ))
                 })?;
-            let tx = Transaction::consensus_decode(&mut x.funding_input.prev_tx.as_slice())
-                .map_err(|_| {
-                    Error::InvalidParameters(
-                        "Could not decode funding input previous tx parameter".to_string(),
-                    )
-                })?;
-            let vout = x.funding_input.prev_tx_vout;
-            let tx_out = tx.output.get(vout as usize).ok_or_else(|| {
-                Error::InvalidParameters(format!("Previous tx output not found at index {}", vout))
-            })?;
 
-            // pass wallet instead of privkeys
-            signer.sign_tx_input(&mut fund, input_index, tx_out, None)?;
+            signer.sign_psbt_input(&mut fund_psbt, input_index)?;
 
-            Ok(fund.input[input_index].witness.clone())
+            let witness = fund_psbt.inputs[input_index]
+                .final_script_witness
+                .clone()
+                .ok_or(Error::InvalidParameters(
+                    "No witness from signing psbt input".to_string(),
+                ))?;
+
+            Ok(witness)
         })
         .collect::<Result<Vec<_>, Error>>()?;
 
@@ -445,8 +471,6 @@ where
         })
         .collect::<Result<Vec<_>, Error>>()?;
 
-    input_serial_ids.sort_unstable();
-
     let offer_refund_signature = dlc::util::get_raw_sig_for_tx_input(
         secp,
         refund,
@@ -457,7 +481,7 @@ where
     )?;
 
     let dlc_transactions = DlcTransactions {
-        fund,
+        fund: fund.clone(),
         cets,
         refund: refund.clone(),
         funding_script_pubkey: funding_script_pubkey.clone(),
@@ -565,24 +589,29 @@ where
         )?;
     }
 
-    let mut input_serials: Vec<_> = offered_contract
+    let fund_tx = &accepted_contract.dlc_transactions.fund;
+    let mut fund_psbt = PartiallySignedTransaction::from_unsigned_tx(fund_tx.clone())
+        .map_err(|_| Error::InvalidState("Tried to create PSBT from signed tx".to_string()))?;
+
+    // get all funding inputs
+    let mut all_funding_inputs = offered_contract
         .funding_inputs_info
         .iter()
         .chain(accepted_contract.funding_inputs.iter())
-        .map(|x| x.funding_input.input_serial_id)
-        .collect();
-    input_serials.sort_unstable();
+        .collect::<Vec<_>>();
+    // sort by serial id
+    all_funding_inputs.sort_by_key(|x| x.funding_input.input_serial_id);
 
-    let mut fund_tx = accepted_contract.dlc_transactions.fund.clone();
+    populate_psbt(&mut fund_psbt, &all_funding_inputs)?;
 
     for (funding_input, funding_signatures) in offered_contract
         .funding_inputs_info
         .iter()
         .zip(funding_signatures.funding_signatures.iter())
     {
-        let input_index = input_serials
+        let input_index = all_funding_inputs
             .iter()
-            .position(|x| x == &funding_input.funding_input.input_serial_id)
+            .position(|x| x.funding_input == funding_input.funding_input)
             .ok_or_else(|| {
                 Error::InvalidState(format!(
                     "Could not find input for serial id {}",
@@ -590,38 +619,27 @@ where
                 ))
             })?;
 
-        fund_tx.input[input_index].witness = Witness::from_vec(
+        fund_psbt.inputs[input_index].final_script_witness = Some(Witness::from_vec(
             funding_signatures
                 .witness_elements
                 .iter()
                 .map(|x| x.witness.clone())
                 .collect(),
-        );
+        ));
     }
 
-    for funding_input_info in &accepted_contract.funding_inputs {
-        let input_index = input_serials
+    for funding_input in &accepted_contract.funding_inputs {
+        let input_index = all_funding_inputs
             .iter()
-            .position(|x| x == &funding_input_info.funding_input.input_serial_id)
+            .position(|x| x.funding_input == funding_input.funding_input)
             .ok_or_else(|| {
                 Error::InvalidState(format!(
                     "Could not find input for serial id {}",
-                    funding_input_info.funding_input.input_serial_id,
+                    funding_input.funding_input.input_serial_id
                 ))
             })?;
-        let tx =
-            Transaction::consensus_decode(&mut funding_input_info.funding_input.prev_tx.as_slice())
-                .map_err(|_| {
-                    Error::InvalidParameters(
-                        "Could not decode funding input previous tx parameter".to_string(),
-                    )
-                })?;
-        let vout = funding_input_info.funding_input.prev_tx_vout;
-        let tx_out = tx.output.get(vout as usize).ok_or_else(|| {
-            Error::InvalidParameters(format!("Previous tx output not found at index {}", vout))
-        })?;
 
-        signer.sign_tx_input(&mut fund_tx, input_index, tx_out, None)?;
+        signer.sign_psbt_input(&mut fund_psbt, input_index)?;
     }
 
     let signed_contract = SignedContract {
@@ -632,7 +650,7 @@ where
         channel_id,
     };
 
-    Ok((signed_contract, fund_tx))
+    Ok((signed_contract, fund_psbt.extract_tx()))
 }
 
 /// Signs and return the CET that can be used to close the given contract.
diff --git a/dlc-manager/src/lib.rs b/dlc-manager/src/lib.rs
index 0ffc4154..52b31310 100644
--- a/dlc-manager/src/lib.rs
+++ b/dlc-manager/src/lib.rs
@@ -36,6 +36,7 @@ pub mod manager;
 pub mod payout_curve;
 mod utils;
 
+use bitcoin::psbt::PartiallySignedTransaction;
 use bitcoin::{Address, Block, OutPoint, Script, Transaction, TxOut, Txid};
 use chain_monitor::ChainMonitor;
 use channel::offered_channel::OfferedChannel;
@@ -78,12 +79,10 @@ impl Time for SystemTimeProvider {
 /// Provides signing related functionalities.
 pub trait Signer {
     /// Signs a transaction input
-    fn sign_tx_input(
+    fn sign_psbt_input(
         &self,
-        tx: &mut Transaction,
+        psbt: &mut PartiallySignedTransaction,
         input_index: usize,
-        tx_out: &TxOut,
-        redeem_script: Option<Script>,
     ) -> Result<(), Error>;
     /// Get the secret key associated with the provided public key.
     fn get_secret_key_for_pubkey(&self, pubkey: &PublicKey) -> Result<SecretKey, Error>;
diff --git a/dlc-manager/src/manager.rs b/dlc-manager/src/manager.rs
index c21c3c95..8811c1cb 100644
--- a/dlc-manager/src/manager.rs
+++ b/dlc-manager/src/manager.rs
@@ -1390,7 +1390,7 @@ where
             Some(*peer_id)
         )?;
 
-        let (signed_channel, signed_contract) = {
+        let (signed_channel, signed_contract, signed_fund_tx) = {
             let res = verify_signed_channel(
                 &self.secp,
                 &accepted_channel,
@@ -1430,7 +1430,7 @@ where
             unreachable!();
         }
 
-        self.blockchain.send_transaction(&signed_channel.fund_tx)?;
+        self.blockchain.send_transaction(&signed_fund_tx)?;
 
         self.store.upsert_channel(
             Channel::Signed(signed_channel),
diff --git a/mocks/src/mock_wallet.rs b/mocks/src/mock_wallet.rs
index 591acb87..404097f8 100644
--- a/mocks/src/mock_wallet.rs
+++ b/mocks/src/mock_wallet.rs
@@ -1,5 +1,6 @@
 use std::rc::Rc;
 
+use bitcoin::psbt::PartiallySignedTransaction;
 use bitcoin::{Address, PackedLockTime, Script, Transaction, TxOut};
 use dlc_manager::{error::Error, Blockchain, Signer, Utxo, Wallet};
 use secp256k1_zkp::{rand::seq::SliceRandom, SecretKey};
@@ -45,13 +46,11 @@ impl MockWallet {
 }
 
 impl Signer for MockWallet {
-    fn sign_tx_input(
+    fn sign_psbt_input(
         &self,
-        _tx: &mut bitcoin::Transaction,
-        _input_index: usize,
-        _tx_out: &bitcoin::TxOut,
-        _redeem_script: Option<bitcoin::Script>,
-    ) -> Result<(), dlc_manager::error::Error> {
+        _psbt: &mut PartiallySignedTransaction,
+        _idx: usize,
+    ) -> Result<(), Error> {
         Ok(())
     }
 
diff --git a/simple-wallet/src/lib.rs b/simple-wallet/src/lib.rs
index 3717e821..8bf6141c 100644
--- a/simple-wallet/src/lib.rs
+++ b/simple-wallet/src/lib.rs
@@ -5,6 +5,7 @@ use bdk::{
     wallet::coin_selection::{BranchAndBoundCoinSelection, CoinSelectionAlgorithm},
     FeeRate, KeychainKind, LocalUtxo, Utxo as BdkUtxo, WeightedUtxo,
 };
+use bitcoin::psbt::PartiallySignedTransaction;
 use bitcoin::{
     hashes::Hash, Address, Network, PackedLockTime, Script, Sequence, Transaction, TxIn, TxOut,
     Txid, Witness,
@@ -154,10 +155,18 @@ where
         let fee = (weight * fee_rate) / 1000;
         tx.output[0].value -= fee;
 
+        // construct psbt
+        let mut psbt = PartiallySignedTransaction::from_unsigned_tx(tx.clone()).unwrap();
         for (i, utxo) in utxos.iter().enumerate().take(tx.input.len()) {
-            self.sign_tx_input(&mut tx, i, &utxo.tx_out, None)?;
+            psbt.inputs[i].witness_utxo = Some(utxo.tx_out.clone());
         }
 
+        for (i, _) in utxos.iter().enumerate().take(tx.input.len()) {
+            self.sign_psbt_input(&mut psbt, i)?;
+        }
+
+        let tx = psbt.extract_tx();
+
         self.blockchain.send_transaction(&tx)
     }
 }
@@ -167,27 +176,49 @@ where
     B::Target: WalletBlockchainProvider,
     W::Target: WalletStorage,
 {
-    fn sign_tx_input(
+    fn sign_psbt_input(
         &self,
-        tx: &mut bitcoin::Transaction,
+        psbt: &mut PartiallySignedTransaction,
         input_index: usize,
-        tx_out: &bitcoin::TxOut,
-        _: Option<bitcoin::Script>,
-    ) -> Result<()> {
+    ) -> std::result::Result<(), Error> {
+        let tx_out = if let Some(input) = psbt.inputs.get(input_index) {
+            if let Some(wit_utxo) = &input.witness_utxo {
+                Ok(wit_utxo.clone())
+            } else if let Some(in_tx) = &input.non_witness_utxo {
+                Ok(
+                    in_tx.output[psbt.unsigned_tx.input[input_index].previous_output.vout as usize]
+                        .clone(),
+                )
+            } else {
+                Err(Error::InvalidParameters(
+                    "No TxOut for PSBT inout".to_string(),
+                ))
+            }
+        } else {
+            Err(Error::InvalidParameters(
+                "No TxOut for PSBT inout".to_string(),
+            ))
+        }?;
         let address = Address::from_script(&tx_out.script_pubkey, self.network)
             .expect("a valid scriptpubkey");
         let seckey = self
             .storage
             .get_priv_key_for_address(&address)?
             .expect("to have the requested private key");
+
+        let mut tx = psbt.unsigned_tx.clone();
         dlc::util::sign_p2wpkh_input(
             &self.secp_ctx,
             &seckey,
-            tx,
+            &mut tx,
             input_index,
             bitcoin::EcdsaSighashType::All,
             tx_out.value,
         )?;
+
+        let tx_input = tx.input[input_index].clone();
+        psbt.inputs[input_index].final_script_sig = Some(tx_input.script_sig);
+        psbt.inputs[input_index].final_script_witness = Some(tx_input.witness);
         Ok(())
     }