From 12bcc64f277f642ece03c25653e726f2276f0d51 Mon Sep 17 00:00:00 2001
From: Andrew Chow <achow101-github@achow101.com>
Date: Wed, 27 Jun 2018 16:56:30 -0700
Subject: [PATCH] Add pubkeys and whether input was witness to SignatureData

Stores pubkeys in SignatureData and retrieves them when using GetPubKey().

Stores whether the signatures in a SignatureData are for a witness input.
---
 src/script/sign.cpp | 17 ++++++++++++++---
 src/script/sign.h   |  2 ++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/script/sign.cpp b/src/script/sign.cpp
index afbcb22d1..b62fe2e92 100644
--- a/src/script/sign.cpp
+++ b/src/script/sign.cpp
@@ -49,9 +49,10 @@ static bool GetCScript(const SigningProvider& provider, const SignatureData& sig
     return false;
 }
 
-static bool GetPubKey(const SigningProvider& provider, const SignatureData& sigdata, const CKeyID& address, CPubKey& pubkey)
+static bool GetPubKey(const SigningProvider& provider, SignatureData& sigdata, const CKeyID& address, CPubKey& pubkey)
 {
     if (provider.GetPubKey(address, pubkey)) {
+        sigdata.misc_pubkeys.emplace(pubkey.GetID(), pubkey);
         return true;
     }
     // Look for pubkey in all partial sigs
@@ -60,6 +61,12 @@ static bool GetPubKey(const SigningProvider& provider, const SignatureData& sigd
         pubkey = it->second.first;
         return true;
     }
+    // Look for pubkey in pubkey list
+    const auto& pk_it = sigdata.misc_pubkeys.find(address);
+    if (pk_it != sigdata.misc_pubkeys.end()) {
+        pubkey = pk_it->second;
+        return true;
+    }
     return false;
 }
 
@@ -70,9 +77,9 @@ static bool CreateSig(const BaseSignatureCreator& creator, SignatureData& sigdat
         sig_out = it->second.second;
         return true;
     }
+    CPubKey pubkey;
+    GetPubKey(provider, sigdata, keyid, pubkey);
     if (creator.CreateSig(provider, sig_out, keyid, scriptcode, sigversion)) {
-        CPubKey pubkey;
-        GetPubKey(provider, sigdata, keyid, pubkey);
         auto i = sigdata.signatures.emplace(keyid, SigPair(pubkey, sig_out));
         assert(i.second);
         return true;
@@ -200,6 +207,7 @@ bool ProduceSignature(const SigningProvider& provider, const BaseSignatureCreato
         txnouttype subType;
         solved = solved && SignStep(provider, creator, witnessscript, result, subType, SigVersion::WITNESS_V0, sigdata);
         sigdata.scriptWitness.stack = result;
+        sigdata.witness = true;
         result.clear();
     }
     else if (solved && whichType == TX_WITNESS_V0_SCRIPTHASH)
@@ -210,7 +218,10 @@ bool ProduceSignature(const SigningProvider& provider, const BaseSignatureCreato
         solved = solved && SignStep(provider, creator, witnessscript, result, subType, SigVersion::WITNESS_V0, sigdata) && subType != TX_SCRIPTHASH && subType != TX_WITNESS_V0_SCRIPTHASH && subType != TX_WITNESS_V0_KEYHASH;
         result.push_back(std::vector<unsigned char>(witnessscript.begin(), witnessscript.end()));
         sigdata.scriptWitness.stack = result;
+        sigdata.witness = true;
         result.clear();
+    } else if (solved && whichType == TX_WITNESS_UNKNOWN) {
+        sigdata.witness = true;
     }
 
     if (P2SH) {
diff --git a/src/script/sign.h b/src/script/sign.h
index d1cf91861..5384f9753 100644
--- a/src/script/sign.h
+++ b/src/script/sign.h
@@ -66,11 +66,13 @@ typedef std::pair<CPubKey, std::vector<unsigned char>> SigPair;
 // in order to construct final scriptSigs and scriptWitnesses.
 struct SignatureData {
     bool complete = false; ///< Stores whether the scriptSig and scriptWitness are complete
+    bool witness = false; ///< Stores whether the input this SigData corresponds to is a witness input
     CScript scriptSig; ///< The scriptSig of an input. Contains complete signatures or the traditional partial signatures format
     CScript redeem_script; ///< The redeemScript (if any) for the input
     CScript witness_script; ///< The witnessScript (if any) for the input. witnessScripts are used in P2WSH outputs.
     CScriptWitness scriptWitness; ///< The scriptWitness of an input. Contains complete signatures or the traditional partial signatures format. scriptWitness is part of a transaction input per BIP 144.
     std::map<CKeyID, SigPair> signatures; ///< BIP 174 style partial signatures for the input. May contain all signatures necessary for producing a final scriptSig or scriptWitness.
+    std::map<CKeyID, CPubKey> misc_pubkeys;
 
     SignatureData() {}
     explicit SignatureData(const CScript& script) : scriptSig(script) {}