From 2f1609b9189d877810a80d40a4c8e6653f440434 Mon Sep 17 00:00:00 2001
From: junderw <junderwood@bitcoinbank.co.jp>
Date: Thu, 11 Jul 2019 11:28:09 +0900
Subject: [PATCH] Fix: P2WPKH was signing with nonWitnessUtxo

---
 src/psbt.js             | 23 +++++++++++++----------
 test/fixtures/psbt.json |  2 +-
 ts_src/psbt.ts          | 25 +++++++++++++++----------
 3 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/src/psbt.js b/src/psbt.js
index c7e8b55..752714f 100644
--- a/src/psbt.js
+++ b/src/psbt.js
@@ -404,6 +404,7 @@ const isP2MS = isPaymentFactory(payments.p2ms);
 const isP2PK = isPaymentFactory(payments.p2pk);
 const isP2PKH = isPaymentFactory(payments.p2pkh);
 const isP2WPKH = isPaymentFactory(payments.p2wpkh);
+const isP2WSHScript = isPaymentFactory(payments.p2wsh);
 function check32Bit(num) {
   if (
     typeof num !== 'number' ||
@@ -611,19 +612,16 @@ function getHashForSig(inputIndex, input, cache, sighashTypes) {
       // If a redeemScript is provided, the scriptPubKey must be for that redeemScript
       checkRedeemScript(inputIndex, prevout.script, input.redeemScript);
       script = input.redeemScript;
-      hash = unsignedTx.hashForSignature(
-        inputIndex,
-        input.redeemScript,
-        sighashType,
-      );
     } else {
       script = prevout.script;
-      hash = unsignedTx.hashForSignature(
-        inputIndex,
-        prevout.script,
-        sighashType,
+    }
+    if (isP2WPKH(script) || isP2WSHScript(script)) {
+      throw new Error(
+        `Input #${inputIndex} has nonWitnessUtxo but segwit script: ` +
+          `${script.toString('hex')}`,
       );
     }
+    hash = unsignedTx.hashForSignature(inputIndex, script, sighashType);
   } else if (input.witnessUtxo) {
     let _script; // so we don't shadow the `let script` above
     if (input.redeemScript) {
@@ -647,7 +645,7 @@ function getHashForSig(inputIndex, input, cache, sighashTypes) {
         sighashType,
       );
       script = _script;
-    } else {
+    } else if (isP2WSHScript(_script)) {
       if (!input.witnessScript)
         throw new Error('Segwit input needs witnessScript if not P2WPKH');
       checkWitnessScript(inputIndex, _script, input.witnessScript);
@@ -659,6 +657,11 @@ function getHashForSig(inputIndex, input, cache, sighashTypes) {
       );
       // want to make sure the script we return is the actual meaningful script
       script = input.witnessScript;
+    } else {
+      throw new Error(
+        `Input #${inputIndex} has witnessUtxo but non-segwit script: ` +
+          `${_script.toString('hex')}`,
+      );
     }
   } else {
     throw new Error('Need a Utxo input item for signing');
diff --git a/test/fixtures/psbt.json b/test/fixtures/psbt.json
index 61ab114..fbad821 100644
--- a/test/fixtures/psbt.json
+++ b/test/fixtures/psbt.json
@@ -121,7 +121,7 @@
     "failSignChecks": [
       {
         "description": "A Witness UTXO is provided for a non-witness input",
-        "errorMessage": "Segwit input needs witnessScript if not P2WPKH",
+        "errorMessage": "Input #0 has witnessUtxo but non-segwit script",
         "psbt": "cHNidP8BAKACAAAAAqsJSaCMWvfEm4IS9Bfi8Vqz9cM9zxU4IagTn4d6W3vkAAAAAAD+////qwlJoIxa98SbghL0F+LxWrP1wz3PFTghqBOfh3pbe+QBAAAAAP7///8CYDvqCwAAAAAZdqkUdopAu9dAy+gdmI5x3ipNXHE5ax2IrI4kAAAAAAAAGXapFG9GILVT+glechue4O/p+gOcykWXiKwAAAAAAAEBItPf9QUAAAAAGXapFNSO0xELlAFMsRS9Mtb00GbcdCVriKwAAQEgAOH1BQAAAAAXqRQ1RebjO4MsRwUPJNPuuTycA5SLx4cBBBYAFIXRNTfy4mVAWjTbr6nj3aAfuCMIACICAurVlmh8qAYEPtw94RbN8p1eklfBls0FXPaYyNAr8k6ZELSmumcAAACAAAAAgAIAAIAAIgIDlPYr6d8ZlSxVh3aK63aYBhrSxKJciU9H2MFitNchPQUQtKa6ZwAAAIABAACAAgAAgAA=",
         "inputToCheck": 0
       },
diff --git a/ts_src/psbt.ts b/ts_src/psbt.ts
index 8fb1e97..928ca04 100644
--- a/ts_src/psbt.ts
+++ b/ts_src/psbt.ts
@@ -509,6 +509,7 @@ const isP2MS = isPaymentFactory(payments.p2ms);
 const isP2PK = isPaymentFactory(payments.p2pk);
 const isP2PKH = isPaymentFactory(payments.p2pkh);
 const isP2WPKH = isPaymentFactory(payments.p2wpkh);
+const isP2WSHScript = isPaymentFactory(payments.p2wsh);
 
 function check32Bit(num: number): void {
   if (
@@ -764,19 +765,18 @@ function getHashForSig(
       // If a redeemScript is provided, the scriptPubKey must be for that redeemScript
       checkRedeemScript(inputIndex, prevout.script, input.redeemScript);
       script = input.redeemScript;
-      hash = unsignedTx.hashForSignature(
-        inputIndex,
-        input.redeemScript,
-        sighashType,
-      );
     } else {
       script = prevout.script;
-      hash = unsignedTx.hashForSignature(
-        inputIndex,
-        prevout.script,
-        sighashType,
+    }
+
+    if (isP2WPKH(script) || isP2WSHScript(script)) {
+      throw new Error(
+        `Input #${inputIndex} has nonWitnessUtxo but segwit script: ` +
+          `${script.toString('hex')}`,
       );
     }
+
+    hash = unsignedTx.hashForSignature(inputIndex, script, sighashType);
   } else if (input.witnessUtxo) {
     let _script: Buffer; // so we don't shadow the `let script` above
     if (input.redeemScript) {
@@ -800,7 +800,7 @@ function getHashForSig(
         sighashType,
       );
       script = _script;
-    } else {
+    } else if (isP2WSHScript(_script)) {
       if (!input.witnessScript)
         throw new Error('Segwit input needs witnessScript if not P2WPKH');
       checkWitnessScript(inputIndex, _script, input.witnessScript);
@@ -812,6 +812,11 @@ function getHashForSig(
       );
       // want to make sure the script we return is the actual meaningful script
       script = input.witnessScript;
+    } else {
+      throw new Error(
+        `Input #${inputIndex} has witnessUtxo but non-segwit script: ` +
+          `${_script.toString('hex')}`,
+      );
     }
   } else {
     throw new Error('Need a Utxo input item for signing');