Merge rust-bitcoin#3601: Fix bug in witness stack getters

1ed3fc9 Add unit test for tapscript function (Tobin C. Harding)
195615c Fix bug in witness stack getters (Tobin C. Harding)

Pull request description:

  In rust-bitcoin#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: rust-bitcoin#3598

ACKs for top commit:
  sanket1729:
    ACK 1ed3fc9
  apoelstra:
    ACK 1ed3fc9; successfully ran local tests

Tree-SHA512: fdaa254a05e58a5a8800ac3316ee1dd2e6da83910fa4860156683ba27c594be19549f22cdf7f170142767733240eba2b31f3a31c10626a1ba5c21409e62b9ec5

Author: apoelstra

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.
@@ -314,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");

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)?;