From c0f4225b32299ebeb6ac9b54198ffb8e0117a6ee Mon Sep 17 00:00:00 2001
From: Christopher Jeffrey <chjjeffrey@gmail.com>
Date: Mon, 19 Sep 2016 04:41:29 -0700
Subject: [PATCH] tx/mempool: fix mutated witnesses in rejects filter.

---
 lib/mempool/mempool.js    |  20 ++++--
 lib/net/pool.js           |   8 ++-
 lib/primitives/tx.js      | 132 ++++++++++++++++++++++++++++++++++++++
 lib/protocol/constants.js |   6 +-
 4 files changed, 159 insertions(+), 7 deletions(-)

diff --git a/lib/mempool/mempool.js b/lib/mempool/mempool.js
index 185f9d17..d1bc9237 100644
--- a/lib/mempool/mempool.js
+++ b/lib/mempool/mempool.js
@@ -815,11 +815,21 @@ Mempool.prototype.verify = function verify(entry, callback) {
         0));
     }
 
-    if (self.requireStandard && !tx.hasStandardInputs()) {
-      return callback(new VerifyError(tx,
-        'nonstandard',
-        'bad-txns-nonstandard-inputs',
-        0));
+    if (self.requireStandard) {
+      if (!tx.hasStandardInputs()) {
+        return callback(new VerifyError(tx,
+          'nonstandard',
+          'bad-txns-nonstandard-inputs',
+          0));
+      }
+      // if (self.chain.state.hasWitness()) {
+      //   if (!tx.hasStandardWitness()) {
+      //     return callback(new VerifyError(tx,
+      //       'nonstandard',
+      //       'bad-witness-nonstandard',
+      //       0));
+      //   }
+      // }
     }
 
     if (tx.getSigopsWeight(flags) > constants.tx.MAX_SIGOPS_WEIGHT) {
diff --git a/lib/net/pool.js b/lib/net/pool.js
index a62adfa6..0d821588 100644
--- a/lib/net/pool.js
+++ b/lib/net/pool.js
@@ -1359,7 +1359,13 @@ Pool.prototype._handleTX = function _handleTX(tx, peer, callback) {
       if (err.type === 'VerifyError') {
         if (err.score !== -1)
           peer.reject(tx, err.code, err.reason, err.score);
-        self.rejectsFilter.add(tx.hash());
+
+        // Once we test it more, we can do:
+        // if (err.reason !== 'bad-witness-nonstandard')
+
+        if (!tx.hasWitness())
+          self.rejectsFilter.add(tx.hash());
+
         return callback(err);
       }
       return callback(err);
diff --git a/lib/primitives/tx.js b/lib/primitives/tx.js
index d034cec5..74234be5 100644
--- a/lib/primitives/tx.js
+++ b/lib/primitives/tx.js
@@ -1359,6 +1359,138 @@ TX.prototype.hasStandardInputs = function hasStandardInputs() {
   return true;
 };
 
+/**
+ * Perform contextual checks to verify coin and witness standardness.
+ * @see IsBadWitness()
+ * @returns {Boolean}
+ */
+
+TX.prototype.hasStandardWitness = function hasStandardWitness(strict) {
+  var i, j, input, prev, hash, redeem, m, n;
+
+  if (this.isCoinbase())
+    return true;
+
+  for (i = 0; i < this.inputs.length; i++) {
+    input = this.inputs[i];
+
+    if (!input.coin)
+      continue;
+
+    prev = input.coin.script;
+
+    if (prev.isScripthash()) {
+      prev = input.script.getRedeem();
+      if (!prev)
+        return false;
+    }
+
+    if (!prev.isProgram()) {
+      if (input.witness.length !== 0)
+        return false;
+      continue;
+    }
+
+    if (prev.isWitnessPubkeyhash()) {
+      if (input.witness.length !== 2)
+        return false;
+
+      if (input.witness.get(0).length > 73)
+        return false;
+
+      hash = crypto.hash160(input.witness.get(1));
+
+      if (!utils.equal(hash, prev.get(1)))
+        return false;
+
+      continue;
+    }
+
+    if (!prev.isWitnessScripthash()) {
+      // Unknown program type,
+      // let it through for now.
+      continue;
+    }
+
+    if (input.witness.length === 0)
+      return false;
+
+    // Based on Johnson Lau's calculations:
+    if (input.witness.length - 1 > 604)
+      return false;
+
+    if (strict) {
+      if (input.witness.length - 1 > constants.script.MAX_P2WSH_STACK)
+        return false;
+    }
+
+    for (j = 0; j < input.witness.length; j++) {
+      if (input.witness.get(j).length > constants.script.MAX_PUSH)
+        return false;
+
+      if (strict) {
+        if (input.witness.get(j).length > constants.script.MAX_P2WSH_PUSH)
+          return false;
+      }
+    }
+
+    redeem = input.witness.get(input.witness.length - 1);
+
+    if (redeem.length > constants.script.MAX_SIZE)
+      return false;
+
+    if (strict) {
+      if (redeem.length > constants.script.MAX_P2WSH_SIZE)
+        return false;
+    }
+
+    hash = crypto.sha256(redeem);
+
+    if (!utils.equal(hash, prev.get(1)))
+      return false;
+
+    redeem = new bcoin.script(redeem);
+
+    if (redeem.isPubkey()) {
+      if (input.witness.length - 1 !== 1)
+        return false;
+
+      if (input.witness.get(0).length > 73)
+        return false;
+
+      continue;
+    }
+
+    if (redeem.isPubkeyhash()) {
+      if (input.witness.length - 1 !== 2)
+        return false;
+
+      if (input.witness.get(0).length > 73)
+        return false;
+
+      continue;
+    }
+
+    if (redeem.isMultisig()) {
+      m = redeem.getSmall(0);
+      n = redeem.getSmall(redeem.length - 2);
+
+      if (input.witness.length - 1 !== m + 1)
+        return false;
+
+      if (input.witness.get(0).length !== 0)
+        return false;
+
+      for (j = 1; j < input.witness.length - 1; j++) {
+        if (input.witness.get(i).length > 73)
+          return false;
+      }
+    }
+  }
+
+  return true;
+};
+
 /**
  * Perform contextual checks to verify input, output,
  * and fee values, as well as coinbase spend maturity
diff --git a/lib/protocol/constants.js b/lib/protocol/constants.js
index c30e5b52..e44622be 100644
--- a/lib/protocol/constants.js
+++ b/lib/protocol/constants.js
@@ -441,7 +441,11 @@ exports.script = {
   MAX_MULTISIG_PUBKEYS: 20,
   MAX_SCRIPTHASH_SIGOPS: 15,
   MAX_OP_RETURN_BYTES: 83,
-  MAX_OP_RETURN: 80
+  MAX_OP_RETURN: 80,
+  BYTES_PER_SIGOP: 20,
+  MAX_P2WSH_STACK: 100,
+  MAX_P2WSH_PUSH: 80,
+  MAX_P2WSH_SIZE: 3600
 };
 
 /**