From 8f3aba0ebb6831f161700ae771940034c16c6423 Mon Sep 17 00:00:00 2001
From: Pieter Wuille <pieter.wuille@gmail.com>
Date: Wed, 6 Jan 2016 01:15:24 +0100
Subject: [PATCH] SigOp counting in witnesses

---
 src/main.cpp               | 23 ++++++++++++++++++-
 src/script/interpreter.cpp | 47 ++++++++++++++++++++++++++++++++++++++
 src/script/interpreter.h   |  2 ++
 3 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/src/main.cpp b/src/main.cpp
index 2ff14f9771efc..ae49417423090 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -739,6 +739,19 @@ unsigned int GetP2SHSigOpCount(const CTransaction& tx, const CCoinsViewCache& in
     return nSigOps;
 }
 
+unsigned int GetWitnessSigOpCount(const CTransaction& tx, const CCoinsViewCache& inputs, int flags)
+{
+    if (tx.IsCoinBase())
+        return 0;
+
+    unsigned int nSigOps = 0;
+    for (unsigned int i = 0; i < tx.vin.size(); i++)
+    {
+        const CTxOut &prevout = inputs.GetOutputFor(tx.vin[i]);
+        nSigOps += CountWitnessSigOps(tx.vin[i].scriptSig, prevout.scriptPubKey, i < tx.wit.vtxinwit.size() ? &tx.wit.vtxinwit[i].scriptWitness : NULL, flags);
+    }
+    return nSigOps;
+}
 
 
 
@@ -945,6 +958,11 @@ bool AcceptToMemoryPoolWorker(CTxMemPool& pool, CValidationState &state, const C
 
         unsigned int nSigOps = GetLegacySigOpCount(tx);
         nSigOps += GetP2SHSigOpCount(tx, view);
+        nSigOps += (GetWitnessSigOpCount(tx, view, STANDARD_SCRIPT_VERIFY_FLAGS) + 3) / 4;
+
+        if (nSigOps > MAX_STANDARD_TX_SIGOPS)
+            return state.DoS(0, false, REJECT_NONSTANDARD, "bad-txns-too-many-sigops", false,
+                strprintf("%d > %d", nSigOps, MAX_STANDARD_TX_SIGOPS));
 
         CAmount nValueOut = tx.GetValueOut();
         CAmount nFees = nValueIn-nValueOut;
@@ -2084,6 +2102,7 @@ bool ConnectBlock(const CBlock& block, CValidationState& state, CBlockIndex* pin
     CAmount nFees = 0;
     int nInputs = 0;
     unsigned int nSigOps = 0;
+    unsigned int nWitSigOps = 0;
     CDiskTxPos pos(pindex->GetBlockPos(), GetSizeOfCompactSize(block.vtx.size()));
     std::vector<std::pair<uint256, CDiskTxPos> > vPos;
     vPos.reserve(block.vtx.size());
@@ -2104,13 +2123,15 @@ bool ConnectBlock(const CBlock& block, CValidationState& state, CBlockIndex* pin
                 return state.DoS(100, error("ConnectBlock(): inputs missing/spent"),
                                  REJECT_INVALID, "bad-txns-inputs-missingorspent");
 
+            nWitSigOps += GetWitnessSigOpCount(tx, view, flags);
+
             if (fStrictPayToScriptHash)
             {
                 // Add in sigops done by pay-to-script-hash inputs;
                 // this is to prevent a "rogue miner" from creating
                 // an incredibly-expensive-to-validate block.
                 nSigOps += GetP2SHSigOpCount(tx, view);
-                if (nSigOps > MAX_BLOCK_SIGOPS)
+                if (nSigOps + (nWitSigOps + 3) / 4 > MAX_BLOCK_SIGOPS)
                     return state.DoS(100, error("ConnectBlock(): too many sigops"),
                                      REJECT_INVALID, "bad-blk-sigops");
             }
diff --git a/src/script/interpreter.cpp b/src/script/interpreter.cpp
index d64c7444e53c3..776f49274cebe 100644
--- a/src/script/interpreter.cpp
+++ b/src/script/interpreter.cpp
@@ -1369,3 +1369,50 @@ bool VerifyScript(const CScript& scriptSig, const CScript& scriptPubKey, const C
 
     return set_success(serror);
 }
+
+size_t static WitnessSigOps(int witversion, const std::vector<unsigned char>& witprogram, const CScriptWitness& witness, int flags)
+{
+    if (witversion == 0) {
+        CScript script(witprogram.begin(), witprogram.end());
+        return script.GetSigOpCount(true);
+    }
+
+    if (witversion == 1 && witness.stack.size() > 0) {
+        CScript subscript(witness.stack.back().begin(), witness.stack.back().end());
+        return subscript.GetSigOpCount(true);
+    }
+
+    // Future flags may be implemented here.
+    return 0;
+}
+
+size_t CountWitnessSigOps(const CScript& scriptSig, const CScript& scriptPubKey, const CScriptWitness* witness, unsigned int flags)
+{
+    static const CScriptWitness witnessEmpty;
+
+    if ((flags & SCRIPT_VERIFY_WITNESS) == 0) {
+        return 0;
+    }
+    assert((flags & SCRIPT_VERIFY_P2SH) != 0);
+
+    int witnessversion;
+    std::vector<unsigned char> witnessprogram;
+    if (scriptPubKey.IsWitnessProgram(witnessversion, witnessprogram)) {
+        return WitnessSigOps(witnessversion, witnessprogram, witness ? *witness : witnessEmpty, flags);
+    }
+
+    if (scriptPubKey.IsPayToScriptHash() && scriptSig.IsPushOnly()) {
+        CScript::const_iterator pc = scriptSig.begin();
+        vector<unsigned char> data;
+        while (pc < scriptSig.end()) {
+            opcodetype opcode;
+            scriptSig.GetOp(pc, opcode, data);
+        }
+        CScript subscript(data.begin(), data.end());
+        if (subscript.IsWitnessProgram(witnessversion, witnessprogram)) {
+            return WitnessSigOps(witnessversion, witnessprogram, witness ? *witness : witnessEmpty, flags);
+        }
+    }
+
+    return 0;
+}
diff --git a/src/script/interpreter.h b/src/script/interpreter.h
index 67125fd927d89..6b72804408631 100644
--- a/src/script/interpreter.h
+++ b/src/script/interpreter.h
@@ -139,4 +139,6 @@ class MutableTransactionSignatureChecker : public TransactionSignatureChecker
 bool EvalScript(std::vector<std::vector<unsigned char> >& stack, const CScript& script, unsigned int flags, const BaseSignatureChecker& checker, int sigversion, ScriptError* error = NULL);
 bool VerifyScript(const CScript& scriptSig, const CScript& scriptPubKey, const CScriptWitness* witness, unsigned int flags, const BaseSignatureChecker& checker, ScriptError* serror = NULL);
 
+size_t CountWitnessSigOps(const CScript& scriptSig, const CScript& scriptPubKey, const CScriptWitness* witness, unsigned int flags);
+
 #endif // BITCOIN_SCRIPT_INTERPRETER_H