From 124cbe4b673f4a437d93a6b28032c37d6ac4b90f Mon Sep 17 00:00:00 2001
From: Chris Stewart <stewart.chris1234@gmail.com>
Date: Thu, 23 Jun 2022 12:12:38 -0500
Subject: [PATCH] Add annex to `TaprootKeyPath` (#4416)

* Add annex to TaprootKeyPath

* Add TaprootWitness.annexHashOpt

* Fix duplicate method
---
 .../protocol/script/TaprootWitnessTest.scala  |  9 +++
 .../core/protocol/script/ScriptWitness.scala  | 58 +++++++++++++------
 2 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/core-test/src/test/scala/org/bitcoins/core/protocol/script/TaprootWitnessTest.scala b/core-test/src/test/scala/org/bitcoins/core/protocol/script/TaprootWitnessTest.scala
index eae39a0004..905eb09e01 100644
--- a/core-test/src/test/scala/org/bitcoins/core/protocol/script/TaprootWitnessTest.scala
+++ b/core-test/src/test/scala/org/bitcoins/core/protocol/script/TaprootWitnessTest.scala
@@ -91,4 +91,13 @@ class TaprootWitnessTest extends BitcoinSUnitTest {
     assert(taprootScriptPath.annexHashOpt.get.hex == expectedAnnexHash)
   }
 
+  it must "construct a taproot keypath witness with an annex" in {
+    val vec = Vector(
+      "c2bdc23435c7bbdce741081181eecd31865f7d94fad6c49c8b1f4619aad72b83354530dbc9446243ff81e0dac2e77b2d437b9d53d279b535a23fb8c599454b3e02",
+      "50ba")
+    val stack = vec.map(ByteVector.fromValidHex(_))
+    val tr = TaprootWitness.fromStack(stack.reverse)
+    assert(tr.isInstanceOf[TaprootKeyPath])
+    assert(tr.annexOpt == Some(stack.last))
+  }
 }
diff --git a/core/src/main/scala/org/bitcoins/core/protocol/script/ScriptWitness.scala b/core/src/main/scala/org/bitcoins/core/protocol/script/ScriptWitness.scala
index 8165098ae2..c23ced7f1f 100644
--- a/core/src/main/scala/org/bitcoins/core/protocol/script/ScriptWitness.scala
+++ b/core/src/main/scala/org/bitcoins/core/protocol/script/ScriptWitness.scala
@@ -205,20 +205,37 @@ object ScriptWitness extends Factory[ScriptWitness] {
 
 sealed trait TaprootWitness extends ScriptWitness {
   override def bytes: ByteVector = RawScriptWitnessParser.write(this)
+
+  def annexOpt: Option[ByteVector]
+
+  /** As per bip341
+    *  the SHA256 of (compact_size(size of annex) || annex), where annex includes the mandatory 0x50 prefix.
+    *  @see https://github.com/bitcoin/bips/blob/master/bip-0341.mediawiki#signature-validation-rules
+    */
+  def annexHashOpt: Option[Sha256Digest] = {
+    annexOpt.map { annex =>
+      val cmpct = CompactSizeUInt.calc(annex)
+      CryptoUtil.sha256(cmpct.bytes ++ annex)
+    }
+  }
 }
 
 object TaprootWitness {
 
   def fromStack(stack: Vector[ByteVector]): TaprootWitness = {
-    if (stack.length == 1) TaprootKeyPath.fromStack(stack)
-    else TaprootScriptPath(stack)
+    val hasAnnex = TaprootScriptPath.hasAnnex(stack)
+
+    if ((hasAnnex && stack.length == 2) || stack.length == 1) {
+      TaprootKeyPath.fromStack(stack)
+    } else TaprootScriptPath(stack)
   }
 }
 
 /** Spending a taproot output via the key path spend */
 case class TaprootKeyPath(
     signature: SchnorrDigitalSignature,
-    hashType: HashType)
+    hashType: HashType,
+    annexOpt: Option[ByteVector])
     extends TaprootWitness {
   override val stack: Vector[ByteVector] = Vector(signature.bytes)
 }
@@ -226,21 +243,35 @@ case class TaprootKeyPath(
 object TaprootKeyPath {
 
   def fromStack(vec: Vector[ByteVector]): TaprootKeyPath = {
+    val hasAnnex = TaprootScriptPath.hasAnnex(vec)
     require(
-      vec.length == 1,
-      s"Taproot keypath can only have one stack element, got=${vec.length}")
+      vec.length == 1 || (hasAnnex && vec.length == 2),
+      s"Taproot keypath can only have at most 2 stack elements, got=${vec.length}")
 
-    val sigBytes = vec.head
+    val annexOpt = {
+      if (hasAnnex) {
+        Some(vec.head)
+      } else {
+        None
+      }
+    }
+    val sigBytes = {
+      if (hasAnnex) {
+        vec(1)
+      } else {
+        vec.head
+      }
+    }
 
     val keyPath = if (sigBytes.length == 64) {
       //means SIGHASH_ALL is implicitly encoded
       //see: https://github.com/bitcoin/bips/blob/master/bip-0341.mediawiki#Common_signature_message
       val sig = SchnorrDigitalSignature.fromBytes(sigBytes)
-      TaprootKeyPath(sig, HashType.sigHashAll)
+      TaprootKeyPath(sig, HashType.sigHashAll, annexOpt)
     } else if (sigBytes.length == 65) {
       val sig = SchnorrDigitalSignature.fromBytes(sigBytes.dropRight(1))
       val hashType = HashType.fromByte(sigBytes.last)
-      TaprootKeyPath(sig, hashType)
+      TaprootKeyPath(sig, hashType, annexOpt)
     } else {
       sys.error(
         s"Unknown sig bytes length, should be 64 or 65, got=${sigBytes.length}")
@@ -300,17 +331,6 @@ case class TaprootScriptPath(stack: Vector[ByteVector]) extends TaprootWitness {
   /** Let p = c[1:33] and let P = lift_x(int(p)) where lift_x and [:] are defined as in BIP340. Fail if this point is not on the curve.
     */
   def p: XOnlyPubKey = controlBlock.p
-
-  /** As per bip341
-    *  the SHA256 of (compact_size(size of annex) || annex), where annex includes the mandatory 0x50 prefix.
-    *  @see https://github.com/bitcoin/bips/blob/master/bip-0341.mediawiki#signature-validation-rules
-    */
-  def annexHashOpt: Option[Sha256Digest] = {
-    annexOpt.map { annex =>
-      val cmpct = CompactSizeUInt.calc(annex)
-      CryptoUtil.sha256(cmpct.bytes ++ annex)
-    }
-  }
 }
 
 object TaprootScriptPath {