From 34f59e8a600f2006219285cc49d13cedd0dc2f31 Mon Sep 17 00:00:00 2001
From: Christopher Jeffrey <chjjeffrey@gmail.com>
Date: Thu, 18 Aug 2016 20:03:58 -0700
Subject: [PATCH] mtx: refactor.

---
 lib/bcoin/keyring.js |   4 +-
 lib/bcoin/mtx.js     | 393 ++++++++++++++++++++++++-------------------
 lib/bcoin/script.js  |  23 +++
 lib/bcoin/wallet.js  |   4 +-
 4 files changed, 249 insertions(+), 175 deletions(-)

diff --git a/lib/bcoin/keyring.js b/lib/bcoin/keyring.js
index efccb1bb..0d34ad32 100644
--- a/lib/bcoin/keyring.js
+++ b/lib/bcoin/keyring.js
@@ -623,7 +623,7 @@ KeyRing.prototype.getRedeem = function(hash) {
  */
 
 KeyRing.prototype.scriptInputs = function scriptInputs(tx) {
-  return tx.template(this.publicKey, this.script, this.program);
+  return tx.template(this.publicKey, this.script);
 };
 
 /**
@@ -635,7 +635,7 @@ KeyRing.prototype.scriptInputs = function scriptInputs(tx) {
  */
 
 KeyRing.prototype.sign = function sign(tx, key) {
-  return tx.sign(key, this.script, this.program);
+  return tx.sign(key, this.script);
 };
 
 /**
diff --git a/lib/bcoin/mtx.js b/lib/bcoin/mtx.js
index 4ff80089..03746e7e 100644
--- a/lib/bcoin/mtx.js
+++ b/lib/bcoin/mtx.js
@@ -211,37 +211,16 @@ MTX.prototype.addOutput = function addOutput(options, value) {
  * Build input script (or witness) templates (with
  * OP_0 in place of signatures).
  * @param {Number} index - Input index.
- * @param {KeyRing} ring - Address used to build. The address
- * must be able to redeem the coin.
+ * @param {Buffer} key - Public key.
+ * @param {Script} script - Redeem script.
  * @returns {Boolean} Whether the script was able to be built.
- * @throws on unavailable coins.
  */
 
-function getRedeem(hash, script, program) {
-  if (program) {
-    if (utils.equal(program.hash160(), hash))
-      return program;
-  }
+MTX.prototype.scriptInput = function scriptInput(index, key, script) {
+  var input = this.inputs[index];
+  var prev, redeem;
 
-  if (script) {
-    if (utils.equal(script.hash160(), hash))
-      return script;
-
-    if (utils.equal(script.sha256(), hash))
-      return script;
-  }
-};
-
-
-MTX.prototype.buildInput = function buildInput(index, key, script, program) {
-  var input, prev, redeem;
-
-  if (key.getPublicKey)
-    key = key.getPublicKey();
-
-  // Get the input
-  input = this.inputs[index];
-  assert(input);
+  assert(input, 'Input does not exist.');
 
   // We should have previous outputs by now.
   if (!input.coin)
@@ -249,8 +228,13 @@ MTX.prototype.buildInput = function buildInput(index, key, script, program) {
 
   // Don't bother with any below calculation
   // if the output is already templated.
-  if (input.script.length !== 0 || input.witness.length !== 0)
+  if (input.script.length !== 0
+      || input.witness.length !== 0) {
     return true;
+  }
+
+  if (key.getPublicKey)
+    key = key.getPublicKey();
 
   // Get the previous output's script
   prev = input.coin.script;
@@ -259,7 +243,7 @@ MTX.prototype.buildInput = function buildInput(index, key, script, program) {
   // with segwit: figuring out where the redeem script and witness
   // redeem scripts go.
   if (prev.isScripthash()) {
-    redeem = getRedeem(prev.get(1), script, program);
+    redeem = this._getRedeem(prev.get(1), key, script);
 
     if (!redeem)
       return false;
@@ -268,22 +252,31 @@ MTX.prototype.buildInput = function buildInput(index, key, script, program) {
     if (redeem.isProgram()) {
       // P2WSH nested within pay-to-scripthash.
       if (redeem.isWitnessScripthash()) {
-        prev = getRedeem(redeem.get(1), script, program);
+        prev = this._getRedeem(redeem.get(1), key, script);
+
         if (!prev)
           return false;
-        this.scriptVector(prev, input.witness, key);
+
+        if (!this.scriptVector(prev, input.witness, key))
+          return false;
+
         input.witness.push(prev.toRaw());
         input.script.push(redeem.toRaw());
         input.script.compile();
+
         return true;
       }
 
       // P2WPKH nested within pay-to-scripthash.
       if (redeem.isWitnessPubkeyhash()) {
         prev = Script.fromPubkeyhash(utils.hash160(key));
-        this.scriptVector(prev, input.witness, key);
+
+        if (!this.scriptVector(prev, input.witness, key))
+          return false;
+
         input.script.push(redeem.toRaw());
         input.script.compile();
+
         return true;
       }
 
@@ -292,9 +285,12 @@ MTX.prototype.buildInput = function buildInput(index, key, script, program) {
     }
 
     // Regular P2SH.
-    this.scriptVector(redeem, input.script, key);
+    if (!this.scriptVector(redeem, input.script, key))
+      return false;
+
     input.script.push(redeem.toRaw());
     input.script.compile();
+
     return true;
   }
 
@@ -302,22 +298,29 @@ MTX.prototype.buildInput = function buildInput(index, key, script, program) {
   if (prev.isProgram()) {
     // Bare P2WSH.
     if (prev.isWitnessScripthash()) {
-      redeem = getRedeem(prev.get(1), script, program);
+      redeem = this._getRedeem(prev.get(1), key, script);
 
       if (!redeem)
         return false;
 
-      this.scriptVector(redeem, input.witness, key);
+      if (!this.scriptVector(redeem, input.witness, key))
+        return false;
+
       input.witness.push(redeem.toRaw());
       input.script.compile();
+
       return true;
     }
 
     // Bare P2WPKH.
     if (prev.isWitnessPubkeyhash()) {
       prev = Script.fromPubkeyhash(prev.get(1));
-      this.scriptVector(prev, input.witness, key);
+
+      if (!this.scriptVector(prev, input.witness, key))
+        return false;
+
       input.script.compile();
+
       return true;
     }
 
@@ -326,10 +329,18 @@ MTX.prototype.buildInput = function buildInput(index, key, script, program) {
   }
 
   // Wow, a normal output! Praise be to Jengus and Gord.
-  this.scriptVector(prev, input.script, key);
-  return true;
+  return this.scriptVector(prev, input.script, key);
 };
 
+/**
+ * Build script for a single vector
+ * based on a previous script.
+ * @param {Script} prev
+ * @param {Witness|Script} vector
+ * @param {Buffer} key
+ * @return {Boolean}
+ */
+
 MTX.prototype.scriptVector = function scriptVector(prev, vector, key) {
   var i, n;
 
@@ -338,10 +349,6 @@ MTX.prototype.scriptVector = function scriptVector(prev, vector, key) {
     if (!utils.equal(prev.get(1), key))
       return false;
 
-    // Already has a script template (at least)
-    if (vector.length !== 0)
-      return true;
-
     vector.set(0, opcodes.OP_0);
 
     return true;
@@ -352,10 +359,6 @@ MTX.prototype.scriptVector = function scriptVector(prev, vector, key) {
     if (!utils.equal(prev.get(2), utils.hash160(key)))
       return false;
 
-    // Already has a script template (at least)
-    if (vector.length !== 0)
-      return true;
-
     vector.set(0, opcodes.OP_0);
     vector.set(1, key);
 
@@ -367,10 +370,6 @@ MTX.prototype.scriptVector = function scriptVector(prev, vector, key) {
     if (prev.indexOf(key) === -1)
       return false;
 
-    // Already has a script template (at least)
-    if (vector.length !== 0)
-      return true;
-
     // Technically we should create m signature slots,
     // but we create n signature slots so we can order
     // the signatures properly.
@@ -386,35 +385,61 @@ MTX.prototype.scriptVector = function scriptVector(prev, vector, key) {
     return true;
   }
 
-  if (prev.indexOf(key) === -1)
-    return false;
+  return false;
+};
 
-  // Already has a script template (at least)
-  if (vector.length !== 0)
-    return true;
+/**
+ * Calculate a redeem script based on hash.
+ * Test against passed in redeem script.
+ * @private
+ * @param {Buffer} hash - 32 or 20 byte hash.
+ * @param {Buffer} key - Public key.
+ * @param {Script} script - Known redeem script.
+ * @returns {Script|null}
+ */
 
-  // Likely a non-standard scripthash multisig
-  // input. Determine n value by counting keys.
-  // Also, only allow nonstandard types for
-  // scripthash.
-  vector.set(0, opcodes.OP_0);
+MTX.prototype._getRedeem = function getRedeem(hash, key, script) {
+  var program;
 
-  // Fill script with `n` signature slots.
-  for (i = 0; i < prev.length; i++) {
-    if (Script.isKey(prev.get(i)))
-      vector.set(i + 1, opcodes.OP_0);
+  if (!key)
+    return;
+
+  switch (hash.length) {
+    case 20:
+      program = bcoin.script.fromProgram(0, utils.hash160(key));
+
+      if (utils.equal(program.hash160(), hash))
+        return program;
+
+      if (!script)
+        return;
+
+      program = script.forWitness();
+
+      if (utils.equal(program.hash160(), hash))
+        return program;
+
+      if (utils.equal(script.hash160(), hash))
+        return script;
+
+      break;
+    case 32:
+      if (!script)
+        return;
+
+      if (utils.equal(script.sha256(), hash))
+        return script;
+
+      break;
   }
 };
 
 /**
  * Sign an input.
  * @param {Number} index - Index of input being signed.
- * @param {KeyRing} ring - Address used to sign. The address
- * must be able to redeem the coin.
  * @param {HDPrivateKey|KeyPair|Buffer} key - Private key.
  * @param {SighashType} type
  * @returns {Boolean} Whether the input was able to be signed.
- * @throws on unavailable coins.
  */
 
 MTX.prototype.signInput = function signInput(index, key, type) {
@@ -423,15 +448,15 @@ MTX.prototype.signInput = function signInput(index, key, type) {
   var redeem = false;
   var prev, vector, sig, result;
 
-  assert(input);
-
-  if (key.getPrivateKey)
-    key = key.getPrivateKey();
+  assert(input, 'Input does not exist.');
 
   // We should have previous outputs by now.
   if (!input.coin)
     return false;
 
+  if (key.getPrivateKey)
+    key = key.getPrivateKey();
+
   // Get the previous output's script
   prev = input.coin.script;
   vector = input.script;
@@ -465,7 +490,7 @@ MTX.prototype.signInput = function signInput(index, key, type) {
   }
 
   // Create our signature.
-  sig = this.createSignature(index, prev, key, type, version);
+  sig = this.signature(index, prev, key, type, version);
 
   if (redeem) {
     redeem = vector.pop();
@@ -478,9 +503,19 @@ MTX.prototype.signInput = function signInput(index, key, type) {
   return this.signVector(prev, vector, sig, key);
 };
 
+/**
+ * Add a signature to a vector
+ * based on a previous script.
+ * @param {Script} prev
+ * @param {Witness|Script} vector
+ * @param {Buffer} sig
+ * @param {Buffer} key
+ * @return {Boolean}
+ */
+
 MTX.prototype.signVector = function signVector(prev, vector, sig, key) {
   var pub = bcoin.ec.publicKeyCreate(key, true);
-  var keys, i, m, n, signatures, keyIndex;
+  var i, m, n, keys, keyIndex, total;
 
   // P2PK
   if (prev.isPubkey()) {
@@ -534,102 +569,109 @@ MTX.prototype.signVector = function signVector(prev, vector, sig, key) {
 
     // Grab `n` value (number of keys).
     n = prev.getSmall(prev.length - 2);
-  } else {
-    keys = [];
 
-    for (i = 0; i < prev.length; i++) {
-      if (Script.isKey(prev.get(i)))
-        keys.push(prev.get(i));
+    if (vector.getSmall(0) !== 0)
+      throw new Error('Input has not been templated.');
+
+    // Too many signature slots. Abort.
+    if (vector.length - 1 > n)
+      return false;
+
+    // Count the number of current signatures.
+    total = 0;
+    for (i = 1; i < vector.length; i++) {
+      if (Script.isSignature(vector.get(i)))
+        total++;
     }
 
-    // We don't know what m is, so
-    // we can never finalize the signatures.
-    m = keys.length;
-    n = keys.length;
-  }
+    // Signatures are already finalized.
+    if (total === m && vector.length - 1 === m)
+      return true;
 
-  if (vector.getSmall(0) !== 0)
-    throw new Error('Input has not been templated.');
+    // Add some signature slots for us to use if
+    // there was for some reason not enough.
+    while (vector.length - 1 < n)
+      vector.push(opcodes.OP_0);
 
-  // Too many signature slots. Abort.
-  if (vector.length - 1 > n)
-    return false;
+    // Find the key index so we can place
+    // the signature in the same index.
+    keyIndex = utils.indexOf(keys, pub);
 
-  // Count the number of current signatures.
-  signatures = 0;
-  for (i = 1; i < vector.length; i++) {
-    if (Script.isSignature(vector.get(i)))
-      signatures++;
-  }
+    // Our public key is not in the prev_out
+    // script. We tried to sign a transaction
+    // that is not redeemable by us.
+    if (keyIndex === -1)
+      return false;
+
+    // Offset key index by one to turn it into
+    // "sig index". Accounts for OP_0 byte at
+    // the start.
+    keyIndex++;
+
+    // Add our signature to the correct slot
+    // and increment the total number of
+    // signatures.
+    if (keyIndex < vector.length && total < m) {
+      if (vector.getSmall(keyIndex) === 0) {
+        vector.set(keyIndex, sig);
+        total++;
+      }
+    }
+
+    // All signatures added. Finalize.
+    if (total >= m) {
+      // Remove empty slots left over.
+      for (i = vector.length - 1; i >= 1; i--) {
+        if (vector.getSmall(i) === 0)
+          vector.remove(i);
+      }
+
+      // Remove signatures which are not required.
+      // This should never happen.
+      while (total > m) {
+        vector.pop();
+        total--;
+      }
+
+      // Sanity checks.
+      assert(total === m);
+      assert(vector.length - 1 === m);
+    }
+
+    vector.compile();
+
+    if (total !== m)
+      return false;
 
-  // Signatures are already finalized.
-  if (signatures === m && vector.length - 1 === m)
     return true;
-
-  // Add some signature slots for us to use if
-  // there was for some reason not enough.
-  while (vector.length - 1 < n)
-    vector.push(opcodes.OP_0);
-
-  // Find the key index so we can place
-  // the signature in the same index.
-  keyIndex = utils.indexOf(keys, pub);
-
-  // Our public key is not in the prev_out
-  // script. We tried to sign a transaction
-  // that is not redeemable by us.
-  if (keyIndex === -1)
-    return false;
-
-  // Offset key index by one to turn it into
-  // "sig index". Accounts for OP_0 byte at
-  // the start.
-  keyIndex++;
-
-  // Add our signature to the correct slot
-  // and increment the total number of
-  // signatures.
-  if (keyIndex < vector.length && signatures < m) {
-    if (vector.getSmall(keyIndex) === 0) {
-      vector.set(keyIndex, sig);
-      signatures++;
-    }
   }
 
-  // All signatures added. Finalize.
-  if (signatures >= m) {
-    // Remove empty slots left over.
-    for (i = vector.length - 1; i >= 1; i--) {
-      if (vector.getSmall(i) === 0)
-        vector.remove(i);
-    }
-
-    // Remove signatures which are not required.
-    // This should never happen.
-    while (signatures > m) {
-      vector.pop();
-      signatures--;
-    }
-
-    // Sanity checks.
-    assert(signatures === m);
-    assert(vector.length - 1 === m);
-  }
-
-  vector.compile();
-
-  return signatures === m;
+  return false;
 };
 
-MTX.prototype.combineMultisig = function combineMultisig(index, prev, version, script, signature) {
+/**
+ * Combine and sort multisig signatures for script.
+ * Mimics bitcoind's behavior.
+ * @param {Number} index
+ * @param {Script} prev
+ * @param {Witness|Script} vector
+ * @param {Number} version
+ * @param {Buffer} data
+ * @return {Boolean}
+ */
+
+MTX.prototype.combine = function combine(index, prev, vector, version, data) {
   var m = prev.getSmall(0);
-  var sigs = [signature];
+  var sigs = [];
   var map = {};
-  var result;
+  var result = false;
   var i, j, sig, type, msg, key, pub, res;
 
-  for (i = 1; i < script.length; i++) {
-    sig = script.get(i);
+  if (data)
+    sigs.push(data);
+
+  for (i = 1; i < vector.length; i++) {
+    sig = vector.get(i);
     if (Script.isSignature(sig))
       sigs.push(sig);
   }
@@ -651,27 +693,29 @@ MTX.prototype.combineMultisig = function combineMultisig(index, prev, version, s
 
       if (res) {
         map[pub] = sig;
-        if (utils.equal(sig, signature))
+        if (utils.equal(sig, data))
           result = true;
         break;
       }
     }
   }
 
-  script.clear();
-  script.push(opcodes.OP_0);
+  vector.clear();
+  vector.push(opcodes.OP_0);
+
   for (i = 1; i < prev.length - 2; i++) {
     key = prev.get(i);
     pub = key.toString('hex');
     sig = map[pub];
     if (sig)
-      script.push(sig);
+      vector.push(sig);
   }
 
-  while (script.length - 1 < m)
-    script.push(opcodes.OP_0);
+  while (vector.length - 1 < m)
+    vector.push(opcodes.OP_0);
+
+  vector.compile();
 
-  script.compile();
   return result;
 };
 
@@ -686,7 +730,7 @@ MTX.prototype.combineMultisig = function combineMultisig(index, prev, version, s
  * @returns {Buffer} Signature in DER format.
  */
 
-MTX.prototype.createSignature = function createSignature(index, prev, key, type, version) {
+MTX.prototype.signature = function signature(index, prev, key, type, version) {
   var hash;
 
   if (type == null)
@@ -703,7 +747,6 @@ MTX.prototype.createSignature = function createSignature(index, prev, key, type,
   return Script.sign(hash, key, type);
 };
 
-
 /**
  * Test whether the transaction is fully-signed.
  * @returns {Boolean}
@@ -754,10 +797,16 @@ MTX.prototype.isSigned = function isSigned() {
     if (prev.isPubkey()) {
       if (!Script.isSignature(vector.get(0)))
         return false;
-    } else if (prev.isPubkeyhash()) {
+      continue;
+    }
+
+    if (prev.isPubkeyhash()) {
       if (!Script.isSignature(vector.get(0)))
         return false;
-    } else if (prev.isMultisig()) {
+      continue;
+    }
+
+    if (prev.isMultisig()) {
       // Grab `m` value (number of required sigs).
       m = prev.getSmall(0);
 
@@ -771,9 +820,11 @@ MTX.prototype.isSigned = function isSigned() {
       // of required signatures.
       if (len - 1 !== m)
         return false;
-    } else {
-      return false;
+
+      continue;
     }
+
+    return false;
   }
 
   return true;
@@ -790,7 +841,7 @@ MTX.prototype.isSigned = function isSigned() {
  * @throws on unavailable coins.
  */
 
-MTX.prototype.template = function template(key, script, program, type) {
+MTX.prototype.template = function template(key, script) {
   var total = 0;
   var i;
 
@@ -799,7 +850,7 @@ MTX.prototype.template = function template(key, script, program, type) {
 
   for (i = 0; i < this.inputs.length; i++) {
     // Build script for input
-    if (!this.buildInput(i, key, script, program))
+    if (!this.scriptInput(i, key, script))
       continue;
     total++;
   }
@@ -818,7 +869,7 @@ MTX.prototype.template = function template(key, script, program, type) {
  * @throws on unavailable coins.
  */
 
-MTX.prototype.sign = function sign(key, script, program, type) {
+MTX.prototype.sign = function sign(key, script, type) {
   var total = 0;
   var i, pub;
 
@@ -829,7 +880,7 @@ MTX.prototype.sign = function sign(key, script, program, type) {
 
   for (i = 0; i < this.inputs.length; i++) {
     // Build script for input
-    if (!this.buildInput(i, pub, script, program))
+    if (!this.scriptInput(i, pub, script))
       continue;
 
     // Sign input
@@ -1465,7 +1516,7 @@ MTX.isMTX = function isMTX(obj) {
   return obj
     && Array.isArray(obj.inputs)
     && typeof obj.locktime === 'number'
-    && typeof obj.buildInput === 'function';
+    && typeof obj.scriptInput === 'function';
 };
 
 /*
diff --git a/lib/bcoin/script.js b/lib/bcoin/script.js
index bd989685..8bed9b5d 100644
--- a/lib/bcoin/script.js
+++ b/lib/bcoin/script.js
@@ -3085,6 +3085,29 @@ Script.prototype.toProgram = function toProgram() {
   return new Program(version, data);
 };
 
+/**
+ * Get the script to the equivalent witness
+ * program (mimics bitcoind's scriptForWitness).
+ * @returns {Program|null}
+ */
+
+Script.prototype.forWitness = function() {
+  var hash;
+
+  if (this.isProgram())
+    return this;
+
+  if (this.isPubkey()) {
+    hash = utils.hash160(this.get(0));
+    return Script.fromProgram(0, hash);
+  }
+
+  if (this.isPubkeyhash())
+    return Script.fromProgram(0, this.get(2));
+
+  return Script.fromProgram(0, this.sha256());
+};
+
 /**
  * Test whether the output script is
  * a pay-to-witness-pubkeyhash program.
diff --git a/lib/bcoin/wallet.js b/lib/bcoin/wallet.js
index a6171ec5..dc3f2225 100644
--- a/lib/bcoin/wallet.js
+++ b/lib/bcoin/wallet.js
@@ -879,8 +879,8 @@ Wallet.prototype.fund = function fund(tx, options, callback, force) {
           height: self.db.height,
           rate: rate,
           wallet: self,
-          m: self.m,
-          n: self.n
+          m: account.m,
+          n: account.n
         });
       } catch (e) {
         return callback(e);