From b0848022ebec44ec9b8abaff6694e5b7329c85c9 Mon Sep 17 00:00:00 2001
From: Steven Roose <steven@stevenroose.org>
Date: Thu, 12 Oct 2023 13:46:24 +0100
Subject: [PATCH] witness: Add Witness::taproot_control_block

---
 bitcoin/src/blockdata/witness.rs | 44 ++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/bitcoin/src/blockdata/witness.rs b/bitcoin/src/blockdata/witness.rs
index e36ffbbb9..5c0044a5f 100644
--- a/bitcoin/src/blockdata/witness.rs
+++ b/bitcoin/src/blockdata/witness.rs
@@ -408,6 +408,29 @@ impl Witness {
             }
         })
     }
+
+    /// Get the taproot control block following BIP341 rules.
+    ///
+    /// This does not guarantee that this represents a P2TR [`Witness`]. It
+    /// merely gets the last or second to last element depending on the first
+    /// byte of the last element being equal to 0x50. See
+    /// [Script::is_p2tr](crate::blockdata::script::Script::is_p2tr) to
+    /// check whether this is actually a Taproot witness.
+    pub 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
+            }
+        })
+    }
 }
 
 impl Index<usize> for Witness {
@@ -705,6 +728,27 @@ mod test {
         assert_eq!(witness_annex.tapscript(), Some(Script::from_bytes(&tapscript[..])));
     }
 
+    #[test]
+    fn test_get_control_block() {
+        let tapscript = hex!("deadbeef");
+        let control_block = hex!("02");
+        // annex starting with 0x50 causes the branching logic.
+        let annex = hex!("50");
+
+        let witness_vec = vec![tapscript.clone(), control_block.clone()];
+        let witness_vec_annex = vec![tapscript.clone(), control_block.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, the tapscript should be returned.
+        assert_eq!(witness.taproot_control_block(), Some(&control_block[..]));
+        assert_eq!(witness_annex.taproot_control_block(), Some(&control_block[..]));
+    }
+
     #[test]
     fn test_tx() {
         const S: &str = "02000000000102b44f26b275b8ad7b81146ba3dbecd081f9c1ea0dc05b97516f56045cfcd3df030100000000ffffffff1cb4749ae827c0b75f3d0a31e63efc8c71b47b5e3634a4c698cd53661cab09170100000000ffffffff020b3a0500000000001976a9143ea74de92762212c96f4dd66c4d72a4deb20b75788ac630500000000000016001493a8dfd1f0b6a600ab01df52b138cda0b82bb7080248304502210084622878c94f4c356ce49c8e33a063ec90f6ee9c0208540888cfab056cd1fca9022014e8dbfdfa46d318c6887afd92dcfa54510e057565e091d64d2ee3a66488f82c0121026e181ffb98ebfe5a64c983073398ea4bcd1548e7b971b4c175346a25a1c12e950247304402203ef00489a0d549114977df2820fab02df75bebb374f5eee9e615107121658cfa02204751f2d1784f8e841bff6d3bcf2396af2f1a5537c0e4397224873fbd3bfbe9cf012102ae6aa498ce2dd204e9180e71b4fb1260fe3d1a95c8025b34e56a9adf5f278af200000000";