From d841f2a88f6af94655e042d6411a1ebaf5d4d923 Mon Sep 17 00:00:00 2001
From: benthecarman <benthecarman@live.com>
Date: Wed, 15 Nov 2023 13:27:58 -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     |  40 ++++++++--
 dlc-manager/src/contract_updater.rs | 116 ++++++++++++++++------------
 dlc-manager/src/lib.rs              |   7 +-
 mocks/src/mock_wallet.rs            |  11 ++-
 simple-wallet/src/lib.rs            |  45 +++++++++--
 5 files changed, 144 insertions(+), 75 deletions(-)

diff --git a/bitcoin-rpc-provider/src/lib.rs b/bitcoin-rpc-provider/src/lib.rs
index b8153ba5..bbee24c9 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,35 @@ 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[psbt.unsigned_tx.input[input_index].previous_output.vout as usize]
+                        .clone(),
+                )
+            } else {
+                Err(ManagerError::InvalidParameters(
+                    "No TxOut for PSBT inout".to_string(),
+                ))
+            }
+        } else {
+            Err(ManagerError::InvalidParameters(
+                "No TxOut for PSBT inout".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 +228,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/contract_updater.rs b/dlc-manager/src/contract_updater.rs
index 4fefb8c2..4557c16a 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,7 @@ where
         funding_script_pubkey,
     } = dlc_transactions;
 
-    let mut fund = fund.clone();
+    let mut fund_psbt = PartiallySignedTransaction::from_unsigned_tx(fund.clone()).unwrap();
     let mut cets = cets.clone();
 
     let input_script_pubkey = input_script_pubkey.unwrap_or_else(|| funding_script_pubkey.clone());
@@ -392,43 +418,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 +470,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 +480,7 @@ where
     )?;
 
     let dlc_transactions = DlcTransactions {
-        fund,
+        fund: fund.clone(),
         cets,
         refund: refund.clone(),
         funding_script_pubkey: funding_script_pubkey.clone(),
@@ -565,24 +588,28 @@ where
         )?;
     }
 
-    let mut input_serials: Vec<_> = offered_contract
+    let fund_tx = accepted_contract.dlc_transactions.fund.clone();
+    let mut fund_psbt = PartiallySignedTransaction::from_unsigned_tx(fund_tx).unwrap();
+
+    // 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 +617,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 +648,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/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(())
     }