From 7cb5f6bb003729291a38cb071a70ee36021036d2 Mon Sep 17 00:00:00 2001
From: Andrew Chow <achow101-github@achow101.com>
Date: Wed, 16 Oct 2019 14:48:29 -0400
Subject: [PATCH] Sign PSBTs with peg-in data

---
 src/wallet/psbtwallet.cpp         | 22 ++++++++++++++++++++++
 test/functional/feature_fedpeg.py |  5 +++++
 2 files changed, 27 insertions(+)

diff --git a/src/wallet/psbtwallet.cpp b/src/wallet/psbtwallet.cpp
index 1e89485967..d261c15d6f 100644
--- a/src/wallet/psbtwallet.cpp
+++ b/src/wallet/psbtwallet.cpp
@@ -3,6 +3,7 @@
 // file COPYING or http://www.opensource.org/licenses/mit-license.php.
 
 #include <confidential_validation.h>
+#include <pegins.h>
 #include <wallet/psbtwallet.h>
 
 TransactionError FillPSBTInputsData(const CWallet* pwallet, PartiallySignedTransaction& psbtx, bool bip32derivs)
@@ -95,6 +96,27 @@ TransactionError SignPSBT(const CWallet* pwallet, PartiallySignedTransaction& ps
         }
     }
 
+    // Stuff in the peg-in data
+    for (unsigned int i = 0; i < tx.vin.size(); ++i) {
+        PSBTInput& input = psbtx.inputs[i];
+        if (input.value && input.peg_in_tx.which() != 0 && input.txout_proof.which() != 0 && !input.claim_script.empty() && !input.genesis_hash.IsNull()) {
+            CScriptWitness pegin_witness;
+            if (Params().GetConsensus().ParentChainHasPow()) {
+                const Sidechain::Bitcoin::CTransactionRef& btc_peg_in_tx = boost::get<Sidechain::Bitcoin::CTransactionRef>(input.peg_in_tx);
+                const Sidechain::Bitcoin::CMerkleBlock& btc_txout_proof = boost::get<Sidechain::Bitcoin::CMerkleBlock>(input.txout_proof);
+                pegin_witness = CreatePeginWitness(*input.value, input.asset, input.genesis_hash, input.claim_script, btc_peg_in_tx, btc_txout_proof);
+            } else {
+                const CTransactionRef& elem_peg_in_tx = boost::get<CTransactionRef>(input.peg_in_tx);
+                const CMerkleBlock& elem_txout_proof = boost::get<CMerkleBlock>(input.txout_proof);
+                pegin_witness = CreatePeginWitness(*input.value, input.asset, input.genesis_hash, input.claim_script, elem_peg_in_tx, elem_txout_proof);
+            }
+            tx.vin[i].m_is_pegin = true;
+            tx.witness.vtxinwit[i].m_pegin_witness = pegin_witness;
+            // Set the witness utxo
+            input.witness_utxo = GetPeginOutputFromWitness(tx.witness.vtxinwit[i].m_pegin_witness);
+        }
+    }
+
     // This is a convenience/usability check -- it's not invalid to sign an unbalanced transaction, but it's easy to shoot yourself in the foot.
     if (!imbalance_ok) {
         // Get UTXOs for all inputs, to check that amounts balance before signing.
diff --git a/test/functional/feature_fedpeg.py b/test/functional/feature_fedpeg.py
index 9e44e9e083..f32edc10e6 100755
--- a/test/functional/feature_fedpeg.py
+++ b/test/functional/feature_fedpeg.py
@@ -311,6 +311,11 @@ class FedPegTest(BitcoinTestFramework):
         assert 'pegin_genesis_hash' not in decoded_psbt['inputs'][0]
         merged_pegin_psbt = sidechain.combinepsbt([pegin_psbt, merge_pegin_psbt])
         assert_equal(pegin_psbt, merged_pegin_psbt)
+        # Now sign the psbt
+        signed_psbt = sidechain.walletsignpsbt(pegin_psbt)
+        # Finalize and extract and compare
+        fin_psbt = sidechain.finalizepsbt(signed_psbt['psbt'])
+        assert_equal(fin_psbt, signed_pegin)
 
         sample_pegin_struct = FromHex(CTransaction(), signed_pegin["hex"])
         # Round-trip peg-in transaction using python serialization