From f259595c5abd2e59da6ac515f9d30953dbf21d06 Mon Sep 17 00:00:00 2001
From: Gavin Andresen <gavinandresen@gmail.com>
Date: Wed, 22 Jul 2015 16:20:06 -0400
Subject: [PATCH] Allow precise tracking of validation sigops / bytes hashed

Adds a BlockValidationResourceTracker class that is passed to
CheckInputs() / CScriptCheck() to keep track of the exact number
of signature operations required to validate a block, and the
exact number of bytes hashed to compute signature hashes.

Also extends CHashWriter to keep track of number of bytes hashed.

Intended to be the starting point for consensus rule changes
to replace the incorrect, messy sigop counting and size-limiting
rules currently in place.
---
 src/hash.h                     |  7 ++++-
 src/main.cpp                   | 25 ++++++++++++-----
 src/main.h                     | 50 +++++++++++++++++++++++++++++++---
 src/miner.cpp                  |  3 +-
 src/script/interpreter.cpp     | 12 ++++++--
 src/script/interpreter.h       | 14 +++++++---
 src/test/script_P2SH_tests.cpp |  2 +-
 7 files changed, 92 insertions(+), 21 deletions(-)

diff --git a/src/hash.h b/src/hash.h
index 0771555623975..6139d7c244fe9 100644
--- a/src/hash.h
+++ b/src/hash.h
@@ -123,15 +123,17 @@ class CHashWriter
 {
 private:
     CHash256 ctx;
+    size_t nBytesHashed;
 
 public:
     int nType;
     int nVersion;
 
-    CHashWriter(int nTypeIn, int nVersionIn) : nType(nTypeIn), nVersion(nVersionIn) {}
+    CHashWriter(int nTypeIn, int nVersionIn) : nBytesHashed(0), nType(nTypeIn), nVersion(nVersionIn) {}
 
     CHashWriter& write(const char *pch, size_t size) {
         ctx.Write((const unsigned char*)pch, size);
+        nBytesHashed += size;
         return (*this);
     }
 
@@ -141,6 +143,9 @@ class CHashWriter
         ctx.Finalize((unsigned char*)&result);
         return result;
     }
+    size_t GetNumBytesHashed() const {
+        return nBytesHashed;
+    }
 
     template<typename T>
     CHashWriter& operator<<(const T& obj) {
diff --git a/src/main.cpp b/src/main.cpp
index 0cf5d56d76357..3f0d3616b19c4 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -1155,7 +1155,7 @@ bool AcceptToMemoryPool(CTxMemPool& pool, CValidationState &state, const CTransa
 
         // Check against previous transactions
         // This is done last to help prevent CPU exhaustion denial-of-service attacks.
-        if (!CheckInputs(tx, state, view, true, STANDARD_SCRIPT_VERIFY_FLAGS, true))
+        if (!CheckInputs(tx, state, view, true, STANDARD_SCRIPT_VERIFY_FLAGS, true, NULL))
         {
             return error("AcceptToMemoryPool: ConnectInputs failed %s", hash.ToString());
         }
@@ -1169,7 +1169,7 @@ bool AcceptToMemoryPool(CTxMemPool& pool, CValidationState &state, const CTransa
         // There is a similar check in CreateNewBlock() to prevent creating
         // invalid blocks, however allowing such transactions into the mempool
         // can be exploited as a DoS attack.
-        if (!CheckInputs(tx, state, view, true, MANDATORY_SCRIPT_VERIFY_FLAGS, true))
+        if (!CheckInputs(tx, state, view, true, MANDATORY_SCRIPT_VERIFY_FLAGS, true, NULL))
         {
             return error("AcceptToMemoryPool: BUG! PLEASE REPORT THIS! ConnectInputs failed against MANDATORY but not STANDARD flags %s", hash.ToString());
         }
@@ -1508,14 +1508,23 @@ void UpdateCoins(const CTransaction& tx, CValidationState &state, CCoinsViewCach
 }
 
 bool CScriptCheck::operator()() {
+    if (resourceTracker && !resourceTracker->IsWithinLimits())
+        return false; // Don't do any more checks if already past limits
+
     const CScript &scriptSig = ptxTo->vin[nIn].scriptSig;
-    if (!VerifyScript(scriptSig, scriptPubKey, nFlags, CachingTransactionSignatureChecker(ptxTo, nIn, cacheStore), &error)) {
+    CachingTransactionSignatureChecker checker(ptxTo, nIn, cacheStore);
+    if (!VerifyScript(scriptSig, scriptPubKey, nFlags, checker, &error)) {
         return ::error("CScriptCheck(): %s:%d VerifySignature failed: %s", ptxTo->GetHash().ToString(), nIn, ScriptErrorString(error));
     }
+    if (resourceTracker) {
+        if (!resourceTracker->Update(ptxTo->GetHash(), checker.GetNumSigops(), checker.GetBytesHashed()))
+            return ::error("CScriptCheck(): %s:%d sigop and/or sighash byte limit exceeded",
+                           ptxTo->GetHash().ToString(), nIn);
+    }
     return true;
 }
 
-bool CheckInputs(const CTransaction& tx, CValidationState &state, const CCoinsViewCache &inputs, bool fScriptChecks, unsigned int flags, bool cacheStore, std::vector<CScriptCheck> *pvChecks)
+bool CheckInputs(const CTransaction& tx, CValidationState &state, const CCoinsViewCache &inputs, bool fScriptChecks, unsigned int flags, bool cacheStore, BlockValidationResourceTracker* resourceTracker, std::vector<CScriptCheck> *pvChecks)
 {
     if (!tx.IsCoinBase())
     {
@@ -1584,7 +1593,7 @@ bool CheckInputs(const CTransaction& tx, CValidationState &state, const CCoinsVi
                 assert(coins);
 
                 // Verify signature
-                CScriptCheck check(*coins, tx, i, flags, cacheStore);
+                CScriptCheck check(resourceTracker, *coins, tx, i, flags, cacheStore);
                 if (pvChecks) {
                     pvChecks->push_back(CScriptCheck());
                     check.swap(pvChecks->back());
@@ -1596,7 +1605,7 @@ bool CheckInputs(const CTransaction& tx, CValidationState &state, const CCoinsVi
                         // arguments; if so, don't trigger DoS protection to
                         // avoid splitting the network between upgraded and
                         // non-upgraded nodes.
-                        CScriptCheck check(*coins, tx, i,
+                        CScriptCheck check(NULL, *coins, tx, i,
                                 flags & ~STANDARD_NOT_MANDATORY_VERIFY_FLAGS, cacheStore);
                         if (check())
                             return state.Invalid(false, REJECT_NONSTANDARD, strprintf("non-mandatory-script-verify-flag (%s)", ScriptErrorString(check.GetScriptError())));
@@ -1966,6 +1975,7 @@ bool ConnectBlock(const CBlock& block, CValidationState& state, CBlockIndex* pin
 
     CBlockUndo blockundo;
 
+    BlockValidationResourceTracker resourceTracker(std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max());
     CCheckQueueControl<CScriptCheck> control(fScriptChecks && nScriptCheckThreads ? &scriptcheckqueue : NULL);
 
     int64_t nTimeStart = GetTimeMicros();
@@ -2006,7 +2016,8 @@ bool ConnectBlock(const CBlock& block, CValidationState& state, CBlockIndex* pin
             nFees += view.GetValueIn(tx)-tx.GetValueOut();
 
             std::vector<CScriptCheck> vChecks;
-            if (!CheckInputs(tx, state, view, fScriptChecks, flags, false, nScriptCheckThreads ? &vChecks : NULL))
+            if (!CheckInputs(tx, state, view, fScriptChecks, flags, false,
+                             &resourceTracker, nScriptCheckThreads ? &vChecks : NULL))
                 return false;
             control.Add(vChecks);
         }
diff --git a/src/main.h b/src/main.h
index d4a5d9cf5609a..eb5da6b0a5e79 100644
--- a/src/main.h
+++ b/src/main.h
@@ -37,6 +37,7 @@
 
 #include <boost/unordered_map.hpp>
 
+class BlockValidationResourceTracker;
 class CBlockIndex;
 class CBlockTreeDB;
 class CBloomFilter;
@@ -317,7 +318,8 @@ unsigned int GetP2SHSigOpCount(const CTransaction& tx, const CCoinsViewCache& ma
  * instead of being performed inline.
  */
 bool CheckInputs(const CTransaction& tx, CValidationState &state, const CCoinsViewCache &view, bool fScriptChecks,
-                 unsigned int flags, bool cacheStore, std::vector<CScriptCheck> *pvChecks = NULL);
+                 unsigned int flags, bool cacheStore, BlockValidationResourceTracker* resourceTracker,
+                 std::vector<CScriptCheck> *pvChecks = NULL);
 
 /** Apply the effects of this transaction on the UTXO set represented by view */
 void UpdateCoins(const CTransaction& tx, CValidationState &state, CCoinsViewCache &inputs, int nHeight);
@@ -345,6 +347,44 @@ bool IsFinalTx(const CTransaction &tx, int nBlockHeight, int64_t nBlockTime);
  */
 bool CheckFinalTx(const CTransaction &tx, int flags = -1);
 
+/**
+ * Class that keeps track of number of signature operations
+ * and bytes hashed to compute signature hashes.
+ */
+class BlockValidationResourceTracker
+{
+private:
+    mutable CCriticalSection cs;
+    uint64_t nSigops;
+    const uint64_t nMaxSigops;
+    uint64_t nSighashBytes;
+    const uint64_t nMaxSighashBytes;
+
+public:
+    BlockValidationResourceTracker(uint64_t nMaxSigopsIn, uint64_t nMaxSighashBytesIn) :
+                                  nSigops(0), nMaxSigops(nMaxSigopsIn),
+                                  nSighashBytes(0), nMaxSighashBytes(nMaxSighashBytesIn) { }
+
+    bool IsWithinLimits() const {
+        LOCK(cs);
+        return (nSigops <= nMaxSigops && nSighashBytes <= nMaxSighashBytes);
+    }
+    bool Update(const uint256& txid, uint64_t nSigopsIn, uint64_t nSighashBytesIn) {
+        LOCK(cs);
+        nSigops += nSigopsIn;
+        nSighashBytes += nSighashBytesIn;
+        return (nSigops <= nMaxSigops && nSighashBytes <= nMaxSighashBytes);
+    }
+    uint64_t GetSigOps() const {
+        LOCK(cs);
+        return nSigops;
+    }
+    uint64_t GetSighashBytes() const {
+        LOCK(cs);
+        return nSighashBytes;
+    }
+};
+
 /** 
  * Closure representing one script verification
  * Note that this stores references to the spending transaction 
@@ -352,6 +392,7 @@ bool CheckFinalTx(const CTransaction &tx, int flags = -1);
 class CScriptCheck
 {
 private:
+    BlockValidationResourceTracker* resourceTracker;
     CScript scriptPubKey;
     const CTransaction *ptxTo;
     unsigned int nIn;
@@ -360,14 +401,15 @@ class CScriptCheck
     ScriptError error;
 
 public:
-    CScriptCheck(): ptxTo(0), nIn(0), nFlags(0), cacheStore(false), error(SCRIPT_ERR_UNKNOWN_ERROR) {}
-    CScriptCheck(const CCoins& txFromIn, const CTransaction& txToIn, unsigned int nInIn, unsigned int nFlagsIn, bool cacheIn) :
-        scriptPubKey(txFromIn.vout[txToIn.vin[nInIn].prevout.n].scriptPubKey),
+    CScriptCheck(): resourceTracker(NULL), ptxTo(0), nIn(0), nFlags(0), cacheStore(false), error(SCRIPT_ERR_UNKNOWN_ERROR) {}
+    CScriptCheck(BlockValidationResourceTracker* resourceTrackerIn, const CCoins& txFromIn, const CTransaction& txToIn, unsigned int nInIn, unsigned int nFlagsIn, bool cacheIn) :
+        resourceTracker(resourceTrackerIn), scriptPubKey(txFromIn.vout[txToIn.vin[nInIn].prevout.n].scriptPubKey),
         ptxTo(&txToIn), nIn(nInIn), nFlags(nFlagsIn), cacheStore(cacheIn), error(SCRIPT_ERR_UNKNOWN_ERROR) { }
 
     bool operator()();
 
     void swap(CScriptCheck &check) {
+        std::swap(resourceTracker, check.resourceTracker);
         scriptPubKey.swap(check.scriptPubKey);
         std::swap(ptxTo, check.ptxTo);
         std::swap(nIn, check.nIn);
diff --git a/src/miner.cpp b/src/miner.cpp
index 090f2359d4695..929d613c6b9bb 100644
--- a/src/miner.cpp
+++ b/src/miner.cpp
@@ -98,6 +98,7 @@ CBlockTemplate* CreateNewBlock(const CScript& scriptPubKeyIn)
     if(!pblocktemplate.get())
         return NULL;
     CBlock *pblock = &pblocktemplate->block; // pointer for convenience
+    BlockValidationResourceTracker resourceTracker(std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max());
 
     // -regtest only: allow overriding block.nVersion with
     // -blockversion=N to test forking scenarios
@@ -296,7 +297,7 @@ CBlockTemplate* CreateNewBlock(const CScript& scriptPubKeyIn)
             // policy here, but we still have to ensure that the block we
             // create only contains transactions that are valid in new blocks.
             CValidationState state;
-            if (!CheckInputs(tx, state, view, true, MANDATORY_SCRIPT_VERIFY_FLAGS, true))
+            if (!CheckInputs(tx, state, view, true, MANDATORY_SCRIPT_VERIFY_FLAGS, true, &resourceTracker))
                 continue;
 
             UpdateCoins(tx, state, view, nHeight);
diff --git a/src/script/interpreter.cpp b/src/script/interpreter.cpp
index 0b78fdf5a8eed..cd936b3603be4 100644
--- a/src/script/interpreter.cpp
+++ b/src/script/interpreter.cpp
@@ -1075,7 +1075,7 @@ class CTransactionSignatureSerializer {
 
 } // anon namespace
 
-uint256 SignatureHash(const CScript& scriptCode, const CTransaction& txTo, unsigned int nIn, int nHashType)
+uint256 SignatureHash(const CScript& scriptCode, const CTransaction& txTo, unsigned int nIn, int nHashType, size_t* nHashedOut)
 {
     static const uint256 one(uint256S("0000000000000000000000000000000000000000000000000000000000000001"));
     if (nIn >= txTo.vin.size()) {
@@ -1097,6 +1097,8 @@ uint256 SignatureHash(const CScript& scriptCode, const CTransaction& txTo, unsig
     // Serialize and hash
     CHashWriter ss(SER_GETHASH, 0);
     ss << txTmp << nHashType;
+    if (nHashedOut != NULL)
+        *nHashedOut = ss.GetNumBytesHashed();
     return ss.GetHash();
 }
 
@@ -1105,7 +1107,8 @@ bool TransactionSignatureChecker::VerifySignature(const std::vector<unsigned cha
     return pubkey.Verify(sighash, vchSig);
 }
 
-bool TransactionSignatureChecker::CheckSig(const vector<unsigned char>& vchSigIn, const vector<unsigned char>& vchPubKey, const CScript& scriptCode) const
+bool TransactionSignatureChecker::CheckSig(const vector<unsigned char>& vchSigIn, const vector<unsigned char>& vchPubKey,
+                                           const CScript& scriptCode) const
 {
     CPubKey pubkey(vchPubKey);
     if (!pubkey.IsValid())
@@ -1118,7 +1121,10 @@ bool TransactionSignatureChecker::CheckSig(const vector<unsigned char>& vchSigIn
     int nHashType = vchSig.back();
     vchSig.pop_back();
 
-    uint256 sighash = SignatureHash(scriptCode, *txTo, nIn, nHashType);
+    size_t nHashed = 0;
+    uint256 sighash = SignatureHash(scriptCode, *txTo, nIn, nHashType, &nHashed);
+    nBytesHashed += nHashed;
+    ++nSigops;
 
     if (!VerifySignature(vchSig, pubkey, sighash))
         return false;
diff --git a/src/script/interpreter.h b/src/script/interpreter.h
index 35d572f0ad8c1..9230ec55d525f 100644
--- a/src/script/interpreter.h
+++ b/src/script/interpreter.h
@@ -83,7 +83,7 @@ enum
     SCRIPT_VERIFY_CHECKLOCKTIMEVERIFY = (1U << 9),
 };
 
-uint256 SignatureHash(const CScript &scriptCode, const CTransaction& txTo, unsigned int nIn, int nHashType);
+uint256 SignatureHash(const CScript &scriptCode, const CTransaction& txTo, unsigned int nIn, int nHashType, size_t* nHashedOut=NULL);
 
 class BaseSignatureChecker
 {
@@ -106,14 +106,18 @@ class TransactionSignatureChecker : public BaseSignatureChecker
 private:
     const CTransaction* txTo;
     unsigned int nIn;
+    mutable size_t nBytesHashed;
+    mutable size_t nSigops;
 
 protected:
     virtual bool VerifySignature(const std::vector<unsigned char>& vchSig, const CPubKey& vchPubKey, const uint256& sighash) const;
 
 public:
-    TransactionSignatureChecker(const CTransaction* txToIn, unsigned int nInIn) : txTo(txToIn), nIn(nInIn) {}
+    TransactionSignatureChecker(const CTransaction* txToIn, unsigned int nInIn) : txTo(txToIn), nIn(nInIn), nBytesHashed(0), nSigops(0) {}
     bool CheckSig(const std::vector<unsigned char>& scriptSig, const std::vector<unsigned char>& vchPubKey, const CScript& scriptCode) const;
     bool CheckLockTime(const CScriptNum& nLockTime) const;
+    size_t GetBytesHashed() const { return nBytesHashed; }
+    size_t GetNumSigops() const { return nSigops; }
 };
 
 class MutableTransactionSignatureChecker : public TransactionSignatureChecker
@@ -125,7 +129,9 @@ class MutableTransactionSignatureChecker : public TransactionSignatureChecker
     MutableTransactionSignatureChecker(const CMutableTransaction* txToIn, unsigned int nInIn) : TransactionSignatureChecker(&txTo, nInIn), txTo(*txToIn) {}
 };
 
-bool EvalScript(std::vector<std::vector<unsigned char> >& stack, const CScript& script, unsigned int flags, const BaseSignatureChecker& checker, ScriptError* error = NULL);
-bool VerifyScript(const CScript& scriptSig, const CScript& scriptPubKey, unsigned int flags, const BaseSignatureChecker& checker, ScriptError* error = NULL);
+bool EvalScript(std::vector<std::vector<unsigned char> >& stack, const CScript& script, unsigned int flags,
+                const BaseSignatureChecker& checker, ScriptError* error = NULL);
+bool VerifyScript(const CScript& scriptSig, const CScript& scriptPubKey, unsigned int flags,
+                  const BaseSignatureChecker& checker, ScriptError* error = NULL);
 
 #endif // BITCOIN_SCRIPT_INTERPRETER_H
diff --git a/src/test/script_P2SH_tests.cpp b/src/test/script_P2SH_tests.cpp
index c8cfe28729c51..7bc7980ff944d 100644
--- a/src/test/script_P2SH_tests.cpp
+++ b/src/test/script_P2SH_tests.cpp
@@ -115,7 +115,7 @@ BOOST_AUTO_TEST_CASE(sign)
         {
             CScript sigSave = txTo[i].vin[0].scriptSig;
             txTo[i].vin[0].scriptSig = txTo[j].vin[0].scriptSig;
-            bool sigOK = CScriptCheck(CCoins(txFrom, 0), txTo[i], 0, SCRIPT_VERIFY_P2SH | SCRIPT_VERIFY_STRICTENC, false)();
+            bool sigOK = CScriptCheck(NULL, CCoins(txFrom, 0), txTo[i], 0, SCRIPT_VERIFY_P2SH | SCRIPT_VERIFY_STRICTENC, false)();
             if (i == j)
                 BOOST_CHECK_MESSAGE(sigOK, strprintf("VerifySignature %d %d", i, j));
             else