P2WSH Signer fix + tests (#1300)

* P2WSH Signer fix + tests

* Respond to review

core-test/src/test/scala/org/bitcoins/core/wallet/signer/SignerTest.scala

@@ -1,19 +1,51 @@
 package org.bitcoins.core.wallet.signer
 
-import org.bitcoins.core.crypto.ECDigitalSignature
+import org.bitcoins.core.crypto.{
-import org.bitcoins.core.currency.Satoshis
+  BaseTxSigComponent,
+  ECDigitalSignature,
+  WitnessTxSigComponentP2SH,
+  WitnessTxSigComponentRaw
+}
+import org.bitcoins.core.currency.{CurrencyUnits, Satoshis}
+import org.bitcoins.core.number.UInt32
+import org.bitcoins.core.policy.Policy
 import org.bitcoins.core.protocol.script.{
+  CLTVScriptPubKey,
+  CSVScriptPubKey,
+  ConditionalScriptPubKey,
+  EmptyScriptPubKey,
   EmptyScriptWitness,
+  MultiSignatureScriptPubKey,
+  NonStandardScriptPubKey,
+  P2PKHScriptPubKey,
+  P2PKScriptPubKey,
+  P2PKWithTimeoutScriptPubKey,
+  P2SHScriptPubKey,
+  P2SHScriptSignature,
   P2WPKHWitnessV0,
   P2WSHWitnessV0,
-  WitnessScriptPubKey
+  UnassignedWitnessScriptPubKey,
+  WitnessCommitment,
+  WitnessScriptPubKey,
+  WitnessScriptPubKeyV0
 }
-import org.bitcoins.core.protocol.transaction.WitnessTransaction
+import org.bitcoins.core.protocol.transaction.{
+  EmptyWitness,
+  Transaction,
+  TransactionInput,
+  TransactionOutput,
+  WitnessTransaction
+}
+import org.bitcoins.core.psbt.InputPSBTRecord.PartialSignature
+import org.bitcoins.core.psbt.PSBT
+import org.bitcoins.core.script.PreExecutionScriptProgram
+import org.bitcoins.core.script.interpreter.ScriptInterpreter
 import org.bitcoins.core.wallet.builder.BitcoinTxBuilder
 import org.bitcoins.core.wallet.fee.SatoshisPerVirtualByte
 import org.bitcoins.core.wallet.utxo.{
   P2WPKHV0SpendingInfo,
   P2WSHV0SpendingInfoFull,
+  UTXOSpendingInfo,
   UTXOSpendingInfoSingle,
   UnassignedSegwitNativeUTXOSpendingInfo
 }
@@ -146,4 +178,128 @@ class SignerTest extends BitcoinSAsyncTest {
         }
     }
   }
+
+  def inputIndex(spendingInfo: UTXOSpendingInfo, tx: Transaction): Int = {
+    tx.inputs.zipWithIndex
+      .find(_._1.previousOutput == spendingInfo.outPoint) match {
+      case Some((_, index)) => index
+      case None =>
+        throw new IllegalArgumentException(
+          "Transaction did not contain expected input.")
+    }
+  }
+
+  def createProgram(
+      tx: Transaction,
+      idx: Int,
+      utxo: UTXOSpendingInfo): PreExecutionScriptProgram = {
+    val output = utxo.output
+
+    val spk = output.scriptPubKey
+
+    val amount = output.value
+
+    val txSigComponent = spk match {
+      case witSPK: WitnessScriptPubKeyV0 =>
+        val o = TransactionOutput(amount, witSPK)
+        WitnessTxSigComponentRaw(tx.asInstanceOf[WitnessTransaction],
+                                 UInt32(idx),
+                                 o,
+                                 Policy.standardFlags)
+      case _: UnassignedWitnessScriptPubKey => ???
+      case x @ (_: P2PKScriptPubKey | _: P2PKHScriptPubKey |
+          _: P2PKWithTimeoutScriptPubKey | _: MultiSignatureScriptPubKey |
+          _: WitnessCommitment | _: CSVScriptPubKey | _: CLTVScriptPubKey |
+          _: ConditionalScriptPubKey | _: NonStandardScriptPubKey |
+          EmptyScriptPubKey) =>
+        val o = TransactionOutput(CurrencyUnits.zero, x)
+        BaseTxSigComponent(tx, UInt32(idx), o, Policy.standardFlags)
+
+      case _: P2SHScriptPubKey =>
+        val p2shScriptSig =
+          tx.inputs(idx).scriptSignature.asInstanceOf[P2SHScriptSignature]
+        p2shScriptSig.redeemScript match {
+
+          case _: WitnessScriptPubKey =>
+            WitnessTxSigComponentP2SH(transaction =
+                                        tx.asInstanceOf[WitnessTransaction],
+                                      inputIndex = UInt32(idx),
+                                      output = output,
+                                      flags = Policy.standardFlags)
+
+          case _ =>
+            BaseTxSigComponent(tx, UInt32(idx), output, Policy.standardFlags)
+        }
+    }
+
+    PreExecutionScriptProgram(txSigComponent)
+  }
+
+  def verifyScripts(
+      tx: Transaction,
+      utxos: Vector[UTXOSpendingInfo]): Boolean = {
+    val programs: Vector[PreExecutionScriptProgram] =
+      tx.inputs.zipWithIndex.toVector.map {
+        case (input: TransactionInput, idx: Int) =>
+          val utxo = utxos.find(_.outPoint == input.previousOutput).get
+          createProgram(tx, idx, utxo)
+      }
+    ScriptInterpreter.runAllVerify(programs)
+  }
+
+  it must "sign p2wsh inputs correctly when provided no witness data" in {
+    forAllAsync(CreditingTxGen.inputsAndOutputs(CreditingTxGen.p2wshOutputs),
+                ScriptGenerators.scriptPubKey,
+                ChainParamsGenerator.bitcoinNetworkParams) {
+      case ((creditingTxsInfos, destinations), changeSPK, network) =>
+        val fee = SatoshisPerVirtualByte(Satoshis(100))
+
+        for {
+          builder <- BitcoinTxBuilder(destinations,
+                                      creditingTxsInfos,
+                                      fee,
+                                      changeSPK._1,
+                                      network)
+          unsignedTx <- builder.unsignedTx
+
+          singleSigs: Vector[Vector[PartialSignature]] <- {
+            val singleInfosVec: Vector[Vector[UTXOSpendingInfoSingle]] =
+              creditingTxsInfos.toVector.map(_.toSingles)
+            val sigVecFs = singleInfosVec.map { singleInfos =>
+              val sigFs = singleInfos.map { singleInfo =>
+                val wtx =
+                  WitnessTransaction(unsignedTx.version,
+                                     unsignedTx.inputs,
+                                     unsignedTx.outputs,
+                                     unsignedTx.lockTime,
+                                     EmptyWitness.fromInputs(unsignedTx.inputs))
+                BitcoinSignerSingle.signSingle(singleInfo,
+                                               wtx,
+                                               isDummySignature = false)
+
+              }
+
+              Future.sequence(sigFs)
+            }
+
+            Future.sequence(sigVecFs)
+          }
+        } yield {
+
+          val psbt =
+            creditingTxsInfos.foldLeft(PSBT.fromUnsignedTx(unsignedTx)) {
+              (psbt, spendInfo) =>
+                val idx = inputIndex(spendInfo, unsignedTx)
+                psbt
+                  .addWitnessUTXOToInput(spendInfo.output, idx)
+                  .addScriptWitnessToInput(spendInfo.scriptWitnessOpt.get, idx)
+                  .addSignatures(singleSigs(idx), idx)
+            }
+
+          val signedTx = psbt.finalizePSBT.get.extractTransactionAndValidate
+
+          assert(verifyScripts(signedTx.get, creditingTxsInfos.toVector))
+        }
+    }
+  }
 }

core/src/main/scala/org/bitcoins/core/psbt/PSBT.scala

@@ -196,6 +196,30 @@ case class PSBT(
 
   }
 
+  /**
+    * Adds the TransactionOutput to the indexed InputPSBTMap to the WitnessUTXO field
+    * @param output TransactionOutput to add to PSBT
+    * @param index index of the InputPSBTMap to add tx to
+    * @return PSBT with added tx
+    */
+  def addWitnessUTXOToInput(output: TransactionOutput, index: Int): PSBT = {
+    require(WitnessScriptPubKey.isWitnessScriptPubKey(output.scriptPubKey.asm),
+            s"Given output was not a Witness UTXO: $output")
+    require(
+      index < inputMaps.size,
+      s"index must be less than the number of input maps present in the psbt, $index >= ${inputMaps.size}")
+
+    val inputMap = inputMaps(index)
+    require(!inputMap.isFinalized,
+            s"Cannot update an InputPSBTMap that is finalized, index: $index")
+
+    val elements = inputMap.filterRecords(WitnessUTXOKeyId) :+ WitnessUTXO(
+      output)
+    val newInputMaps = inputMaps.updated(index, InputPSBTMap(elements))
+
+    PSBT(globalMap, newInputMaps, outputMaps)
+  }
+
   /**
     * Adds script to the indexed InputPSBTMap to either the RedeemScript
     * or WitnessScript field depending on the script and available information in the PSBT
@@ -494,6 +518,32 @@ case class PSBT(
     PSBT(globalMap, newInputMaps, outputMaps)
   }
 
+  /** Adds all the PartialSignatures to the input map at the given index */
+  def addSignatures(
+      partialSignatures: Vector[PartialSignature],
+      inputIndex: Int): PSBT = {
+    require(
+      inputIndex < inputMaps.size,
+      s"index must be less than the number of input maps present in the psbt, $inputIndex >= ${inputMaps.size}")
+    require(
+      !inputMaps(inputIndex).isFinalized,
+      s"Cannot update an InputPSBTMap that is finalized, index: $inputIndex")
+    val intersect =
+      inputMaps(inputIndex).partialSignatures.intersect(partialSignatures)
+    val allSigs = inputMaps(inputIndex).partialSignatures ++ partialSignatures
+
+    require(
+      allSigs.groupBy(_.pubKey).values.forall(_.length == 1),
+      s"Cannot add differing signatures for associated public keys, got: $intersect"
+    )
+    val newElements = inputMaps(inputIndex).elements ++ partialSignatures
+
+    val newInputMaps =
+      inputMaps.updated(inputIndex, InputPSBTMap(newElements))
+
+    PSBT(globalMap, newInputMaps, outputMaps)
+  }
+
   /**
     * Extracts the serialized from the serialized, fully signed transaction from
     * this PSBT and validates the script signatures using the ScriptInterpreter.

core/src/main/scala/org/bitcoins/core/wallet/signer/Signer.scala

@@ -60,21 +60,14 @@ sealed abstract class SignerUtils {
 
     spendingInfo.output.scriptPubKey match {
       case _: WitnessScriptPubKey =>
-        val wtx = unsignedTx match {
+        val wtx = {
-          case btx: BaseTransaction =>
+          val noWitnessWtx = WitnessTransaction.toWitnessTx(unsignedTx)
-            val transactionWitnessOpt =
+          spendingInfo.scriptWitnessOpt match {
-              spendingInfo.scriptWitnessOpt.map(scriptWit =>
+            case None =>
-                TransactionWitness(Vector(scriptWit)))
+              noWitnessWtx
-            val transactionWitness =
+            case Some(scriptWitness) =>
-              transactionWitnessOpt.getOrElse(
+              noWitnessWtx.updateWitness(index.toInt, scriptWitness)
-                EmptyWitness.fromInputs(btx.inputs))
+          }
-
-            WitnessTransaction(btx.version,
-                               btx.inputs,
-                               btx.outputs,
-                               btx.lockTime,
-                               transactionWitness)
-          case wtx: WitnessTransaction => wtx
         }
 
         WitnessTxSigComponent(wtx, index, spendingInfo.output, flags)
@@ -800,7 +793,10 @@ sealed abstract class P2WSHSignerSingle
       isDummySignature: Boolean,
       spendingInfoToSatisfy: P2WSHV0SpendingInfoSingle)(
       implicit ec: ExecutionContext): Future[PartialSignature] = {
-    val wtx = WitnessTransaction.toWitnessTx(unsignedTx)
+    val idx = inputIndex(spendingInfo, unsignedTx)
+    val wtx = WitnessTransaction
+      .toWitnessTx(unsignedTx)
+      .updateWitness(idx.toInt, spendingInfoToSatisfy.scriptWitness)
 
     BitcoinSignerSingle.signSingle(spendingInfo,
                                    wtx,
@@ -825,7 +821,9 @@ sealed abstract class P2WSHSigner
     } else {
       val (_, output, inputIndex, _) = relevantInfo(spendingInfo, unsignedTx)
 
-      val wtx = WitnessTransaction.toWitnessTx(unsignedTx)
+      val wtx = WitnessTransaction
+        .toWitnessTx(unsignedTx)
+        .updateWitness(inputIndex.toInt, spendingInfoToSatisfy.scriptWitness)
 
       val signedSigComponentF = BitcoinSigner.sign(
         spendingInfo,