From a8168c3f81a76165022af3f3aeec82317d8a6bf1 Mon Sep 17 00:00:00 2001
From: Martin Habovstiak <martin.habovstiak@gmail.com>
Date: Fri, 21 Feb 2025 16:33:00 +0100
Subject: [PATCH] Add `taproot_leaf_script` methood to `Witness`

We already have `tapscript` method on `Witness` which is broken because
it doesn't check that the leaf script is a tapscript, however that
behavior might have been intended by some consumers who want to inspect
the script independent of the version. To resolve the confusion, we're
going to add a new method that returns both the leaf script and, to
avoid forgetting version check, also the leaf version.

This doesn't touch the `tapscript` method yet to make backporting of
this commit easier. It's also worth noting that leaf script is often
used together with version. To make passing them around easier it'd be
helpful to use a separate type. Thus this also adds a public POD type
containing the script and the version. In anticipation of if being
usable in different APIs it's also generic over the script type.
Similarly to the `tapscript` method, this also only adds the type and
doesn't change other functions to use it yet. Only the newly added
`taproot_leaf_script` method uses it now.

This is a part of #4073
---
 bitcoin/src/blockdata/witness.rs | 44 +++++++++++++++++++++++++++++++-
 bitcoin/src/taproot/mod.rs       |  9 +++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/bitcoin/src/blockdata/witness.rs b/bitcoin/src/blockdata/witness.rs
index 20efc0b5f..d0ac6f526 100644
--- a/bitcoin/src/blockdata/witness.rs
+++ b/bitcoin/src/blockdata/witness.rs
@@ -13,7 +13,7 @@ use crate::crypto::ecdsa;
 use crate::prelude::Vec;
 #[cfg(doc)]
 use crate::script::ScriptExt as _;
-use crate::taproot::{self, TAPROOT_ANNEX_PREFIX};
+use crate::taproot::{self, LeafScript, LeafVersion, TAPROOT_ANNEX_PREFIX, TAPROOT_CONTROL_BASE_SIZE, TAPROOT_LEAF_MASK};
 use crate::Script;
 
 #[rustfmt::skip]                // Keep public re-exports separate.
@@ -156,6 +156,22 @@ crate::internal_macros::define_extension_trait! {
             }
         }
 
+        /// Returns the leaf script with its version but without the merkle proof.
+        ///
+        /// This does not guarantee that this represents a P2TR [`Witness`]. It
+        /// merely gets the second to last or third to last element depending on
+        /// the first byte of the last element being equal to 0x50 and the associated
+        /// version.
+        fn taproot_leaf_script(&self) -> Option<LeafScript<&Script>> {
+            match P2TrSpend::from_witness(self) {
+                Some(P2TrSpend::Script { leaf_script, control_block, .. }) if control_block.len() >= TAPROOT_CONTROL_BASE_SIZE => {
+                    let version = LeafVersion::from_consensus(control_block[0] & TAPROOT_LEAF_MASK).ok()?;
+                    Some(LeafScript { version, script: leaf_script, })
+                },
+                _ => None,
+            }
+        }
+
         /// Get the taproot control block following BIP341 rules.
         ///
         /// This does not guarantee that this represents a P2TR [`Witness`]. It
@@ -371,6 +387,32 @@ mod test {
         assert_eq!(witness_annex.tapscript(), Some(Script::from_bytes(&tapscript[..])));
     }
 
+    #[test]
+    fn get_taproot_leaf_script() {
+        let tapscript = hex!("deadbeef");
+        let control_block = hex!("c0ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff");
+        // 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, 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();
+
+        let expected_leaf_script = LeafScript {
+            version: LeafVersion::TapScript,
+            script: Script::from_bytes(&tapscript),
+        };
+
+        // With or without annex, the tapscript should be returned.
+        assert_eq!(witness.taproot_leaf_script().unwrap(), expected_leaf_script);
+        assert_eq!(witness_annex.taproot_leaf_script().unwrap(), expected_leaf_script);
+    }
+
     #[test]
     fn get_tapscript_from_keypath() {
         let signature = hex!("deadbeef");
diff --git a/bitcoin/src/taproot/mod.rs b/bitcoin/src/taproot/mod.rs
index 775357f07..69746a1a5 100644
--- a/bitcoin/src/taproot/mod.rs
+++ b/bitcoin/src/taproot/mod.rs
@@ -142,6 +142,15 @@ pub const TAPROOT_CONTROL_BASE_SIZE: usize = 33;
 pub const TAPROOT_CONTROL_MAX_SIZE: usize =
     TAPROOT_CONTROL_BASE_SIZE + TAPROOT_CONTROL_NODE_SIZE * TAPROOT_CONTROL_MAX_NODE_COUNT;
 
+/// The leaf script with its version.
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
+pub struct LeafScript<S> {
+    /// The version of the script.
+    pub version: LeafVersion,
+    /// The script, usually `ScriptBuf` or `&Script`.
+    pub script: S,
+}
+
 // type alias for versioned tap script corresponding Merkle proof
 type ScriptMerkleProofMap = BTreeMap<(ScriptBuf, LeafVersion), BTreeSet<TaprootMerkleBranch>>;