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] 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 0844161cd..e0a8d298a 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 95737c7ca..4eb025bf4 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)?;