From 82a30fade70a2a95c2bbeac4aa06dafda600479d Mon Sep 17 00:00:00 2001
From: Andrew Chow <achow101-github@achow101.com>
Date: Mon, 10 Feb 2020 19:50:59 -0500
Subject: [PATCH] Move key and script filling and signing from
 CWallet::FillPSBT to ScriptPubKeyMan::FillPSBT

Instead of fetching a SigningProvider from ScriptPubKeyMan in order
to fill and sign the keys and scripts for a PSBT, just pass that
PSBT to a new FillPSBT function that does all that for us.
---
 src/wallet/scriptpubkeyman.cpp        | 42 ++++++++++++++++++++++
 src/wallet/scriptpubkeyman.h          |  5 +++
 src/wallet/test/psbt_wallet_tests.cpp |  4 +--
 src/wallet/wallet.cpp                 | 51 ++++++++++++++++++---------
 4 files changed, 82 insertions(+), 20 deletions(-)

diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
index b110275654..cfb8184c3b 100644
--- a/src/wallet/scriptpubkeyman.cpp
+++ b/src/wallet/scriptpubkeyman.cpp
@@ -511,6 +511,48 @@ bool LegacyScriptPubKeyMan::SignTransaction(CMutableTransaction& tx, const std::
     return ::SignTransaction(tx, this, coins, sighash, input_errors);
 }
 
+TransactionError LegacyScriptPubKeyMan::FillPSBT(PartiallySignedTransaction& psbtx, int sighash_type, bool sign, bool bip32derivs) const
+{
+    for (unsigned int i = 0; i < psbtx.tx->vin.size(); ++i) {
+        const CTxIn& txin = psbtx.tx->vin[i];
+        PSBTInput& input = psbtx.inputs.at(i);
+
+        if (PSBTInputSigned(input)) {
+            continue;
+        }
+
+        // Verify input looks sane. This will check that we have at most one uxto, witness or non-witness.
+        if (!input.IsSane()) {
+            return TransactionError::INVALID_PSBT;
+        }
+
+        // Get the Sighash type
+        if (sign && input.sighash_type > 0 && input.sighash_type != sighash_type) {
+            return TransactionError::SIGHASH_MISMATCH;
+        }
+
+        // Check non_witness_utxo has specified prevout
+        if (input.non_witness_utxo) {
+            if (txin.prevout.n >= input.non_witness_utxo->vout.size()) {
+                return TransactionError::MISSING_INPUTS;
+            }
+        } else if (input.witness_utxo.IsNull()) {
+            // There's no UTXO so we can just skip this now
+            continue;
+        }
+        SignatureData sigdata;
+        input.FillSignatureData(sigdata);
+        SignPSBTInput(HidingSigningProvider(this, !sign, !bip32derivs), psbtx, i, sighash_type);
+    }
+
+    // Fill in the bip32 keypaths and redeemscripts for the outputs so that hardware wallets can identify change
+    for (unsigned int i = 0; i < psbtx.tx->vout.size(); ++i) {
+        UpdatePSBTOutput(HidingSigningProvider(this, true, !bip32derivs), psbtx, i);
+    }
+
+    return TransactionError::OK;
+}
+
 const CKeyMetadata* LegacyScriptPubKeyMan::GetMetadata(const CTxDestination& dest) const
 {
     LOCK(cs_KeyStore);
diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h
index df4ea3ad7b..1628001a74 100644
--- a/src/wallet/scriptpubkeyman.h
+++ b/src/wallet/scriptpubkeyman.h
@@ -5,8 +5,10 @@
 #ifndef BITCOIN_WALLET_SCRIPTPUBKEYMAN_H
 #define BITCOIN_WALLET_SCRIPTPUBKEYMAN_H
 
+#include <psbt.h>
 #include <script/signingprovider.h>
 #include <script/standard.h>
+#include <util/error.h>
 #include <wallet/crypter.h>
 #include <wallet/ismine.h>
 #include <wallet/walletdb.h>
@@ -212,6 +214,8 @@ public:
 
     /** Creates new signatures and adds them to the transaction. Returns whether all inputs were signed */
     virtual bool SignTransaction(CMutableTransaction& tx, const std::map<COutPoint, Coin>& coins, int sighash, std::map<int, std::string>& input_errors) const { return false; }
+    /** Adds script and derivation path information to a PSBT, and optionally signs it. */
+    virtual TransactionError FillPSBT(PartiallySignedTransaction& psbt, int sighash_type = 1 /* SIGHASH_ALL */, bool sign = true, bool bip32derivs = false) const { return TransactionError::INVALID_PSBT; }
 
     virtual uint256 GetID() const { return uint256(); }
 
@@ -354,6 +358,7 @@ public:
     bool CanProvide(const CScript& script, SignatureData& sigdata) override;
 
     bool SignTransaction(CMutableTransaction& tx, const std::map<COutPoint, Coin>& coins, int sighash, std::map<int, std::string>& input_errors) const override;
+    TransactionError FillPSBT(PartiallySignedTransaction& psbt, int sighash_type = 1 /* SIGHASH_ALL */, bool sign = true, bool bip32derivs = false) const override;
 
     uint256 GetID() const override;
 
diff --git a/src/wallet/test/psbt_wallet_tests.cpp b/src/wallet/test/psbt_wallet_tests.cpp
index 008272a69b..8b7b7af21d 100644
--- a/src/wallet/test/psbt_wallet_tests.cpp
+++ b/src/wallet/test/psbt_wallet_tests.cpp
@@ -73,9 +73,7 @@ BOOST_AUTO_TEST_CASE(psbt_updater_test)
 
     // Try to sign the mutated input
     SignatureData sigdata;
-    psbtx.inputs[0].FillSignatureData(sigdata);
-    const std::unique_ptr<SigningProvider> provider = m_wallet.GetSigningProvider(ws1, sigdata);
-    BOOST_CHECK(!SignPSBTInput(*provider, psbtx, 0, SIGHASH_ALL));
+    BOOST_CHECK(spk_man->FillPSBT(psbtx, SIGHASH_ALL, true, true) != TransactionError::OK);
 }
 
 BOOST_AUTO_TEST_CASE(parse_hd_keypath)
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 099a79ed64..de09e60014 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -2481,7 +2481,6 @@ TransactionError CWallet::FillPSBT(PartiallySignedTransaction& psbtx, bool& comp
 {
     LOCK(cs_wallet);
     // Get all of the previous transactions
-    complete = true;
     for (unsigned int i = 0; i < psbtx.tx->vin.size(); ++i) {
         const CTxIn& txin = psbtx.tx->vin[i];
         PSBTInput& input = psbtx.inputs.at(i);
@@ -2506,13 +2505,22 @@ TransactionError CWallet::FillPSBT(PartiallySignedTransaction& psbtx, bool& comp
                 input.non_witness_utxo = wtx.tx;
             }
         }
+    }
 
-        // Get the Sighash type
-        if (sign && input.sighash_type > 0 && input.sighash_type != sighash_type) {
-            return TransactionError::SIGHASH_MISMATCH;
+    // Fill in information from ScriptPubKeyMans
+    // Because each ScriptPubKeyMan may be able to fill more than one input, we need to keep track of each ScriptPubKeyMan that has filled this psbt.
+    // Each iteration, we may fill more inputs than the input that is specified in that iteration.
+    // We assume that each input is filled by only one ScriptPubKeyMan
+    std::set<uint256> visited_spk_mans;
+    for (unsigned int i = 0; i < psbtx.tx->vin.size(); ++i) {
+        const CTxIn& txin = psbtx.tx->vin[i];
+        PSBTInput& input = psbtx.inputs.at(i);
+
+        if (PSBTInputSigned(input)) {
+            continue;
         }
 
-        // Get the scriptPubKey to know which SigningProvider to use
+        // Get the scriptPubKey to know which ScriptPubKeyMan to use
         CScript script;
         if (!input.witness_utxo.IsNull()) {
             script = input.witness_utxo.scriptPubKey;
@@ -2523,27 +2531,36 @@ TransactionError CWallet::FillPSBT(PartiallySignedTransaction& psbtx, bool& comp
             script = input.non_witness_utxo->vout[txin.prevout.n].scriptPubKey;
         } else {
             // There's no UTXO so we can just skip this now
-            complete = false;
             continue;
         }
         SignatureData sigdata;
         input.FillSignatureData(sigdata);
-        std::unique_ptr<SigningProvider> provider = GetSigningProvider(script, sigdata);
-        if (!provider) {
-            complete = false;
+        std::set<ScriptPubKeyMan*> spk_mans = GetScriptPubKeyMans(script, sigdata);
+        if (spk_mans.size() == 0) {
             continue;
         }
 
-        complete &= SignPSBTInput(HidingSigningProvider(provider.get(), !sign, !bip32derivs), psbtx, i, sighash_type);
+        for (auto& spk_man : spk_mans) {
+            // If we've already been signed by this spk_man, skip it
+            if (visited_spk_mans.count(spk_man->GetID()) > 0) {
+                continue;
+            }
+
+            // Fill in the information from the spk_man
+            TransactionError res = spk_man->FillPSBT(psbtx, sighash_type, sign, bip32derivs);
+            if (res != TransactionError::OK) {
+                return res;
+            }
+
+            // Add this spk_man to visited_spk_mans so we can skip it later
+            visited_spk_mans.insert(spk_man->GetID());
+        }
     }
 
-    // Fill in the bip32 keypaths and redeemscripts for the outputs so that hardware wallets can identify change
-    for (unsigned int i = 0; i < psbtx.tx->vout.size(); ++i) {
-        const CTxOut& out = psbtx.tx->vout.at(i);
-        std::unique_ptr<SigningProvider> provider = GetSigningProvider(out.scriptPubKey);
-        if (provider) {
-            UpdatePSBTOutput(HidingSigningProvider(provider.get(), true, !bip32derivs), psbtx, i);
-        }
+    // Complete if every input is now signed
+    complete = true;
+    for (const auto& input : psbtx.inputs) {
+        complete &= PSBTInputSigned(input);
     }
 
     return TransactionError::OK;