From 0422beb9bdfcd7d1cc095ab681a760d8bc7a1918 Mon Sep 17 00:00:00 2001
From: Andrew Chow <achow101-github@achow101.com>
Date: Tue, 3 Jul 2018 17:18:52 -0700
Subject: [PATCH] Make SignatureData able to store signatures and scripts

In addition to having the scriptSig and scriptWitness, have SignatureData
also be able to store just the signatures (pubkeys mapped to sigs) and
scripts (script ids mapped to scripts).

Also have DataFromTransaction be able to extract signatures and scripts
from the scriptSig and scriptWitness of an input to put them in SignatureData.

Adds a new SignatureChecker which takes a SignatureData and puts pubkeys
and signatures into it when it successfully verifies a signature.

Adds a new field in SignatureData which stores whether the SignatureData
was complete. This allows us to also update the scriptSig and
scriptWitness to the final one when updating a SignatureData with another
one.
---
 src/bench/verify_script.cpp    |   1 +
 src/bitcoin-tx.cpp             |   4 +-
 src/rpc/rawtransaction.cpp     |   6 +-
 src/script/sign.cpp            | 130 +++++++++++++++++++++++++++------
 src/script/sign.h              |  17 ++++-
 src/test/transaction_tests.cpp |   8 +-
 6 files changed, 129 insertions(+), 37 deletions(-)

diff --git a/src/bench/verify_script.cpp b/src/bench/verify_script.cpp
index 4100519d48..ae60588c2d 100644
--- a/src/bench/verify_script.cpp
+++ b/src/bench/verify_script.cpp
@@ -9,6 +9,7 @@
 #endif
 #include <script/script.h>
 #include <script/sign.h>
+#include <script/standard.h>
 #include <streams.h>
 
 #include <array>
diff --git a/src/bitcoin-tx.cpp b/src/bitcoin-tx.cpp
index e6eb723cf4..1fc18b1e46 100644
--- a/src/bitcoin-tx.cpp
+++ b/src/bitcoin-tx.cpp
@@ -645,13 +645,13 @@ static void MutateTxSign(CMutableTransaction& tx, const std::string& flagStr)
         const CScript& prevPubKey = coin.out.scriptPubKey;
         const CAmount& amount = coin.out.nValue;
 
-        SignatureData sigdata;
+        SignatureData sigdata = DataFromTransaction(mergedTx, i, coin.out);
         // Only sign SIGHASH_SINGLE if there's a corresponding output:
         if (!fHashSingle || (i < mergedTx.vout.size()))
             ProduceSignature(keystore, MutableTransactionSignatureCreator(&mergedTx, i, amount, nHashType), prevPubKey, sigdata);
 
         // ... and merge in other signatures:
-        sigdata = CombineSignatures(prevPubKey, MutableTransactionSignatureChecker(&mergedTx, i, amount), sigdata, DataFromTransaction(txv, i));
+        sigdata = CombineSignatures(prevPubKey, MutableTransactionSignatureChecker(&mergedTx, i, amount), sigdata, DataFromTransaction(txv, i, coin.out));
         UpdateInput(txin, sigdata);
     }
 
diff --git a/src/rpc/rawtransaction.cpp b/src/rpc/rawtransaction.cpp
index 3b3f43edea..550fffb326 100644
--- a/src/rpc/rawtransaction.cpp
+++ b/src/rpc/rawtransaction.cpp
@@ -744,7 +744,7 @@ static UniValue combinerawtransaction(const JSONRPCRequest& request)
         // ... and merge in other signatures:
         for (const CMutableTransaction& txv : txVariants) {
             if (txv.vin.size() > i) {
-                sigdata = CombineSignatures(prevPubKey, TransactionSignatureChecker(&txConst, i, amount), sigdata, DataFromTransaction(txv, i));
+                sigdata = CombineSignatures(prevPubKey, TransactionSignatureChecker(&txConst, i, amount), sigdata, DataFromTransaction(txv, i, coin.out));
             }
         }
 
@@ -875,12 +875,12 @@ UniValue SignTransaction(CMutableTransaction& mtx, const UniValue& prevTxsUnival
         const CScript& prevPubKey = coin.out.scriptPubKey;
         const CAmount& amount = coin.out.nValue;
 
-        SignatureData sigdata;
+        SignatureData sigdata = DataFromTransaction(mtx, i, coin.out);
         // Only sign SIGHASH_SINGLE if there's a corresponding output:
         if (!fHashSingle || (i < mtx.vout.size())) {
             ProduceSignature(*keystore, MutableTransactionSignatureCreator(&mtx, i, amount, nHashType), prevPubKey, sigdata);
         }
-        sigdata = CombineSignatures(prevPubKey, TransactionSignatureChecker(&txConst, i, amount), sigdata, DataFromTransaction(mtx, i));
+        sigdata = CombineSignatures(prevPubKey, TransactionSignatureChecker(&txConst, i, amount), sigdata, DataFromTransaction(mtx, i, coin.out));
 
         UpdateInput(txin, sigdata);
 
diff --git a/src/script/sign.cpp b/src/script/sign.cpp
index 05fde3b549..cb63c2159e 100644
--- a/src/script/sign.cpp
+++ b/src/script/sign.cpp
@@ -126,6 +126,8 @@ static CScript PushAll(const std::vector<valtype>& values)
 
 bool ProduceSignature(const SigningProvider& provider, const BaseSignatureCreator& creator, const CScript& fromPubKey, SignatureData& sigdata)
 {
+    if (sigdata.complete) return true;
+
     std::vector<valtype> result;
     txnouttype whichType;
     bool solved = SignStep(provider, creator, fromPubKey, result, whichType, SigVersion::BASE);
@@ -168,15 +170,117 @@ bool ProduceSignature(const SigningProvider& provider, const BaseSignatureCreato
     sigdata.scriptSig = PushAll(result);
 
     // Test solution
-    return solved && VerifyScript(sigdata.scriptSig, fromPubKey, &sigdata.scriptWitness, STANDARD_SCRIPT_VERIFY_FLAGS, creator.Checker());
+    sigdata.complete = solved && VerifyScript(sigdata.scriptSig, fromPubKey, &sigdata.scriptWitness, STANDARD_SCRIPT_VERIFY_FLAGS, creator.Checker());
+    return sigdata.complete;
+}
+
+class SignatureExtractorChecker final : public BaseSignatureChecker
+{
+private:
+    SignatureData& sigdata;
+    BaseSignatureChecker& checker;
+
+public:
+    SignatureExtractorChecker(SignatureData& sigdata, BaseSignatureChecker& checker) : sigdata(sigdata), checker(checker) {}
+    bool CheckSig(const std::vector<unsigned char>& scriptSig, const std::vector<unsigned char>& vchPubKey, const CScript& scriptCode, SigVersion sigversion) const override;
+};
+
+bool SignatureExtractorChecker::CheckSig(const std::vector<unsigned char>& scriptSig, const std::vector<unsigned char>& vchPubKey, const CScript& scriptCode, SigVersion sigversion) const
+{
+    if (checker.CheckSig(scriptSig, vchPubKey, scriptCode, sigversion)) {
+        CPubKey pubkey(vchPubKey);
+        sigdata.signatures.emplace(pubkey.GetID(), SigPair(pubkey, scriptSig));
+        return true;
+    }
+    return false;
+}
+
+namespace
+{
+struct Stacks
+{
+    std::vector<valtype> script;
+    std::vector<valtype> witness;
+
+    Stacks() {}
+    explicit Stacks(const std::vector<valtype>& scriptSigStack_) : script(scriptSigStack_), witness() {}
+    explicit Stacks(const SignatureData& data) : witness(data.scriptWitness.stack) {
+        EvalScript(script, data.scriptSig, SCRIPT_VERIFY_STRICTENC, BaseSignatureChecker(), SigVersion::BASE);
+    }
+
+    SignatureData Output() const {
+        SignatureData result;
+        result.scriptSig = PushAll(script);
+        result.scriptWitness.stack = witness;
+        return result;
+    }
+};
 }
 
-SignatureData DataFromTransaction(const CMutableTransaction& tx, unsigned int nIn)
+// Extracts signatures and scripts from incomplete scriptSigs. Please do not extend this, use PSBT instead
+SignatureData DataFromTransaction(const CMutableTransaction& tx, unsigned int nIn, const CTxOut& txout)
 {
     SignatureData data;
     assert(tx.vin.size() > nIn);
     data.scriptSig = tx.vin[nIn].scriptSig;
     data.scriptWitness = tx.vin[nIn].scriptWitness;
+    Stacks stack(data);
+
+    // Get signatures
+    MutableTransactionSignatureChecker tx_checker(&tx, nIn, txout.nValue);
+    SignatureExtractorChecker extractor_checker(data, tx_checker);
+    if (VerifyScript(data.scriptSig, txout.scriptPubKey, &data.scriptWitness, STANDARD_SCRIPT_VERIFY_FLAGS, extractor_checker)) {
+        data.complete = true;
+        return data;
+    }
+
+    // Get scripts
+    txnouttype script_type;
+    std::vector<std::vector<unsigned char>> solutions;
+    Solver(txout.scriptPubKey, script_type, solutions);
+    SigVersion sigversion = SigVersion::BASE;
+    CScript next_script = txout.scriptPubKey;
+
+    if (script_type == TX_SCRIPTHASH && !stack.script.empty() && !stack.script.back().empty()) {
+        // Get the redeemScript
+        CScript redeem_script(stack.script.back().begin(), stack.script.back().end());
+        data.redeem_script = redeem_script;
+        next_script = std::move(redeem_script);
+
+        // Get redeemScript type
+        Solver(next_script, script_type, solutions);
+        stack.script.pop_back();
+    }
+    if (script_type == TX_WITNESS_V0_SCRIPTHASH && !stack.witness.empty() && !stack.witness.back().empty()) {
+        // Get the witnessScript
+        CScript witness_script(stack.witness.back().begin(), stack.witness.back().end());
+        data.witness_script = witness_script;
+        next_script = std::move(witness_script);
+
+        // Get witnessScript type
+        Solver(next_script, script_type, solutions);
+        stack.witness.pop_back();
+        stack.script = std::move(stack.witness);
+        stack.witness.clear();
+        sigversion = SigVersion::WITNESS_V0;
+    }
+    if (script_type == TX_MULTISIG && !stack.script.empty()) {
+        // Build a map of pubkey -> signature by matching sigs to pubkeys:
+        assert(solutions.size() > 1);
+        unsigned int num_pubkeys = solutions.size()-2;
+        unsigned int last_success_key = 0;
+        for (const valtype& sig : stack.script) {
+            for (unsigned int i = last_success_key; i < num_pubkeys; ++i) {
+                const valtype& pubkey = solutions[i+1];
+                // We either have a signature for this pubkey, or we have found a signature and it is valid
+                if (data.signatures.count(CPubKey(pubkey).GetID()) || extractor_checker.CheckSig(sig, pubkey, next_script, sigversion)) {
+                    last_success_key = i + 1;
+                    break;
+                }
+            }
+        }
+    }
+
     return data;
 }
 
@@ -263,28 +367,6 @@ static std::vector<valtype> CombineMultisig(const CScript& scriptPubKey, const B
     return result;
 }
 
-namespace
-{
-struct Stacks
-{
-    std::vector<valtype> script;
-    std::vector<valtype> witness;
-
-    Stacks() {}
-    explicit Stacks(const std::vector<valtype>& scriptSigStack_) : script(scriptSigStack_), witness() {}
-    explicit Stacks(const SignatureData& data) : witness(data.scriptWitness.stack) {
-        EvalScript(script, data.scriptSig, SCRIPT_VERIFY_STRICTENC, BaseSignatureChecker(), SigVersion::BASE);
-    }
-
-    SignatureData Output() const {
-        SignatureData result;
-        result.scriptSig = PushAll(script);
-        result.scriptWitness.stack = witness;
-        return result;
-    }
-};
-}
-
 static Stacks CombineSignatures(const CScript& scriptPubKey, const BaseSignatureChecker& checker,
                                  const txnouttype txType, const std::vector<valtype>& vSolutions,
                                  Stacks sigs1, Stacks sigs2, SigVersion sigversion)
diff --git a/src/script/sign.h b/src/script/sign.h
index 8ef0306bfe..7dc4bedb06 100644
--- a/src/script/sign.h
+++ b/src/script/sign.h
@@ -53,9 +53,18 @@ public:
 /** A signature creator that just produces 72-byte empty signatures. */
 extern const BaseSignatureCreator& DUMMY_SIGNATURE_CREATOR;
 
+typedef std::pair<CPubKey, std::vector<unsigned char>> SigPair;
+
+// This struct contains information from a transaction input and also contains signatures for that input.
+// The information contained here can be used to create a signature and is also filled by ProduceSignature
+// in order to construct final scriptSigs and scriptWitnesses.
 struct SignatureData {
-    CScript scriptSig;
-    CScriptWitness scriptWitness;
+    bool complete = false; ///< Stores whether the scriptSig and scriptWitness are complete
+    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.
 
     SignatureData() {}
     explicit SignatureData(const CScript& script) : scriptSig(script) {}
@@ -71,8 +80,8 @@ bool SignSignature(const SigningProvider &provider, const CTransaction& txFrom,
 /** Combine two script signatures using a generic signature checker, intelligently, possibly with OP_0 placeholders. */
 SignatureData CombineSignatures(const CScript& scriptPubKey, const BaseSignatureChecker& checker, const SignatureData& scriptSig1, const SignatureData& scriptSig2);
 
-/** Extract signature data from a transaction, and insert it. */
-SignatureData DataFromTransaction(const CMutableTransaction& tx, unsigned int nIn);
+/** Extract signature data from a transaction input, and insert it. */
+SignatureData DataFromTransaction(const CMutableTransaction& tx, unsigned int nIn, const CTxOut& txout);
 void UpdateInput(CTxIn& input, const SignatureData& data);
 
 /* Check whether we know how to sign for an output like this, assuming we
diff --git a/src/test/transaction_tests.cpp b/src/test/transaction_tests.cpp
index 65c5b8ea1d..9ec3fd35f0 100644
--- a/src/test/transaction_tests.cpp
+++ b/src/test/transaction_tests.cpp
@@ -629,7 +629,7 @@ BOOST_AUTO_TEST_CASE(test_witness)
     CreateCreditAndSpend(keystore2, scriptMulti, output2, input2, false);
     CheckWithFlag(output2, input2, 0, false);
     BOOST_CHECK(*output1 == *output2);
-    UpdateInput(input1.vin[0], CombineSignatures(output1->vout[0].scriptPubKey, MutableTransactionSignatureChecker(&input1, 0, output1->vout[0].nValue), DataFromTransaction(input1, 0), DataFromTransaction(input2, 0)));
+    UpdateInput(input1.vin[0], CombineSignatures(output1->vout[0].scriptPubKey, MutableTransactionSignatureChecker(&input1, 0, output1->vout[0].nValue), DataFromTransaction(input1, 0, output1->vout[0]), DataFromTransaction(input2, 0, output1->vout[0])));
     CheckWithFlag(output1, input1, STANDARD_SCRIPT_VERIFY_FLAGS, true);
 
     // P2SH 2-of-2 multisig
@@ -640,7 +640,7 @@ BOOST_AUTO_TEST_CASE(test_witness)
     CheckWithFlag(output2, input2, 0, true);
     CheckWithFlag(output2, input2, SCRIPT_VERIFY_P2SH, false);
     BOOST_CHECK(*output1 == *output2);
-    UpdateInput(input1.vin[0], CombineSignatures(output1->vout[0].scriptPubKey, MutableTransactionSignatureChecker(&input1, 0, output1->vout[0].nValue), DataFromTransaction(input1, 0), DataFromTransaction(input2, 0)));
+    UpdateInput(input1.vin[0], CombineSignatures(output1->vout[0].scriptPubKey, MutableTransactionSignatureChecker(&input1, 0, output1->vout[0].nValue), DataFromTransaction(input1, 0, output1->vout[0]), DataFromTransaction(input2, 0, output1->vout[0])));
     CheckWithFlag(output1, input1, SCRIPT_VERIFY_P2SH, true);
     CheckWithFlag(output1, input1, STANDARD_SCRIPT_VERIFY_FLAGS, true);
 
@@ -652,7 +652,7 @@ BOOST_AUTO_TEST_CASE(test_witness)
     CheckWithFlag(output2, input2, 0, true);
     CheckWithFlag(output2, input2, SCRIPT_VERIFY_P2SH | SCRIPT_VERIFY_WITNESS, false);
     BOOST_CHECK(*output1 == *output2);
-    UpdateInput(input1.vin[0], CombineSignatures(output1->vout[0].scriptPubKey, MutableTransactionSignatureChecker(&input1, 0, output1->vout[0].nValue), DataFromTransaction(input1, 0), DataFromTransaction(input2, 0)));
+    UpdateInput(input1.vin[0], CombineSignatures(output1->vout[0].scriptPubKey, MutableTransactionSignatureChecker(&input1, 0, output1->vout[0].nValue), DataFromTransaction(input1, 0, output1->vout[0]), DataFromTransaction(input2, 0, output1->vout[0])));
     CheckWithFlag(output1, input1, SCRIPT_VERIFY_P2SH | SCRIPT_VERIFY_WITNESS, true);
     CheckWithFlag(output1, input1, STANDARD_SCRIPT_VERIFY_FLAGS, true);
 
@@ -664,7 +664,7 @@ BOOST_AUTO_TEST_CASE(test_witness)
     CheckWithFlag(output2, input2, SCRIPT_VERIFY_P2SH, true);
     CheckWithFlag(output2, input2, SCRIPT_VERIFY_P2SH | SCRIPT_VERIFY_WITNESS, false);
     BOOST_CHECK(*output1 == *output2);
-    UpdateInput(input1.vin[0], CombineSignatures(output1->vout[0].scriptPubKey, MutableTransactionSignatureChecker(&input1, 0, output1->vout[0].nValue), DataFromTransaction(input1, 0), DataFromTransaction(input2, 0)));
+    UpdateInput(input1.vin[0], CombineSignatures(output1->vout[0].scriptPubKey, MutableTransactionSignatureChecker(&input1, 0, output1->vout[0].nValue), DataFromTransaction(input1, 0, output1->vout[0]), DataFromTransaction(input2, 0, output1->vout[0])));
     CheckWithFlag(output1, input1, SCRIPT_VERIFY_P2SH | SCRIPT_VERIFY_WITNESS, true);
     CheckWithFlag(output1, input1, STANDARD_SCRIPT_VERIFY_FLAGS, true);
 }