From 195615c14d2ee3632af54859bb403d6c622c95b7 Mon Sep 17 00:00:00 2001
From: "Tobin C. Harding" <me@tobin.cc>
Date: Mon, 11 Nov 2024 13:09:13 +1100
Subject: [PATCH 1/2] Fix bug in witness stack getters

In #2646 we introduced a bug in the taproot witness stack getter
functions, of which we have three:

- `tapscript`
- `taproot_control_block`
- `taproot_annex`

Each returns `Some` if a possible bytes slice is found (with no other
guarantees).

Use `taproot_annex` combined with getters from `primitives` to implement
the other two getters. This simplifies the code and fixes the bug.

Add an additional getter to `primitives` `Witness::third_from_last`.

Fix: #3598
---
 bitcoin/src/blockdata/witness.rs | 44 +++++++++++++-------------------
 primitives/src/witness.rs        |  9 +++++++
 2 files changed, 27 insertions(+), 26 deletions(-)

diff --git a/bitcoin/src/blockdata/witness.rs b/bitcoin/src/blockdata/witness.rs
index 0844161cd5..e0a8d298aa 100644
--- a/bitcoin/src/blockdata/witness.rs
+++ b/bitcoin/src/blockdata/witness.rs
@@ -147,19 +147,15 @@ crate::internal_macros::define_extension_trait! {
         ///
         /// See [`Script::is_p2tr`] to check whether this is actually a Taproot witness.
         fn tapscript(&self) -> Option<&Script> {
-            self.last().and_then(|last| {
-                // From BIP341:
-                // If there are at least two witness elements, and the first byte of
-                // the last element is 0x50, this last element is called annex a
-                // and is removed from the witness stack.
-                if self.len() >= 3 && last.first() == Some(&TAPROOT_ANNEX_PREFIX) {
-                    self.nth(self.len() - 3).map(Script::from_bytes)
-                } else if self.len() >= 2 {
-                    self.nth(self.len() - 2).map(Script::from_bytes)
-                } else {
-                    None
-                }
-            })
+            if self.is_empty() {
+                return None;
+            }
+
+            if self.taproot_annex().is_some() {
+                self.third_to_last().map(Script::from_bytes)
+            } else {
+                self.second_to_last().map(Script::from_bytes)
+            }
         }
 
         /// Get the taproot control block following BIP341 rules.
@@ -170,19 +166,15 @@ crate::internal_macros::define_extension_trait! {
         ///
         /// See [`Script::is_p2tr`] to check whether this is actually a Taproot witness.
         fn taproot_control_block(&self) -> Option<&[u8]> {
-            self.last().and_then(|last| {
-                // From BIP341:
-                // If there are at least two witness elements, and the first byte of
-                // the last element is 0x50, this last element is called annex a
-                // and is removed from the witness stack.
-                if self.len() >= 3 && last.first() == Some(&TAPROOT_ANNEX_PREFIX) {
-                    self.nth(self.len() - 2)
-                } else if self.len() >= 2 {
-                    Some(last)
-                } else {
-                    None
-                }
-            })
+            if self.is_empty() {
+                return None;
+            }
+
+            if self.taproot_annex().is_some() {
+                self.second_to_last()
+            } else {
+                self.last()
+            }
         }
 
         /// Get the taproot annex following BIP341 rules.
diff --git a/primitives/src/witness.rs b/primitives/src/witness.rs
index 95737c7ca0..4eb025bf4d 100644
--- a/primitives/src/witness.rs
+++ b/primitives/src/witness.rs
@@ -192,6 +192,15 @@ impl Witness {
         }
     }
 
+    /// Returns the third-to-last element in the witness, if any.
+    pub fn third_to_last(&self) -> Option<&[u8]> {
+        if self.witness_elements <= 2 {
+            None
+        } else {
+            self.nth(self.witness_elements - 3)
+        }
+    }
+
     /// Return the nth element in the witness, if any
     pub fn nth(&self, index: usize) -> Option<&[u8]> {
         let pos = decode_cursor(&self.content, self.indices_start, index)?;

From 1ed3fc9270c6e46f42f447bf1c63cfa64dafa5e8 Mon Sep 17 00:00:00 2001
From: "Tobin C. Harding" <me@tobin.cc>
Date: Wed, 13 Nov 2024 10:53:03 +1100
Subject: [PATCH 2/2] Add unit test for tapscript function

Add a unit test that verifies we correctly remove the annex when getting
the tapscript from the witness stack.

Unit test written by Casey, pulled out of #3599.

Co-developed-by: Casey Rodarmor <casey@rodarmor.com>
---
 bitcoin/src/blockdata/witness.rs | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/bitcoin/src/blockdata/witness.rs b/bitcoin/src/blockdata/witness.rs
index e0a8d298aa..318baaf508 100644
--- a/bitcoin/src/blockdata/witness.rs
+++ b/bitcoin/src/blockdata/witness.rs
@@ -306,6 +306,26 @@ mod test {
         assert_eq!(witness_annex.tapscript(), Some(Script::from_bytes(&tapscript[..])));
     }
 
+    #[test]
+    fn test_get_tapscript_from_keypath() {
+        let signature = hex!("deadbeef");
+        // annex starting with 0x50 causes the branching logic.
+        let annex = hex!("50");
+
+        let witness_vec = vec![signature.clone()];
+        let witness_vec_annex = vec![signature.clone(), annex];
+
+        let witness_serialized: Vec<u8> = serialize(&witness_vec);
+        let witness_serialized_annex: Vec<u8> = serialize(&witness_vec_annex);
+
+        let witness = deserialize::<Witness>(&witness_serialized[..]).unwrap();
+        let witness_annex = deserialize::<Witness>(&witness_serialized_annex[..]).unwrap();
+
+        // With or without annex, no tapscript should be returned.
+        assert_eq!(witness.tapscript(), None);
+        assert_eq!(witness_annex.tapscript(), None);
+    }
+
     #[test]
     fn test_get_control_block() {
         let tapscript = hex!("deadbeef");