From e810ecff7cbaeef3007fe7abf98f0fe383bcd219 Mon Sep 17 00:00:00 2001 From: Martin Habovstiak Date: Fri, 21 Feb 2025 14:59:55 +0100 Subject: [PATCH 1/3] Fix key/script spend detection in `Witness` The `taproot_control_block` did not properly detect whether it deals with script spend or key spend. As a result, if key spend with annex was used it'd return the first element (the signature) as if it was a control block. Further, the conditions identifying which kind of spend it was were repeated multiple times but behaved subtly differently making only `taproot_control_block` buggy but the other places confusing. To resolve these issues this change adds a `P2TrSpend` enum that represents a parsed witness and has a single method doing all the parsing. The other methods can then be trivially implemented by matching on that type. This way only one place needs to be verified and the parsing code is more readable since it uses one big `match` to handle all possibilities. The downside of this is a potential perf impact if the parsing code doesn't get inlined since the common parsing code has to shuffle around data that the caller is not intersted in. I don't think this will be a problem but if it will I suppose it will be solvable (e.g. by using `#[inline(always)]`). The enum also looks somewhat nice and perhaps downstream consumers could make use of it. This change does not expose it yet but is written such that after exposing it the API would be (mostly) idiomatic. Closes #4097 --- bitcoin/src/blockdata/witness.rs | 119 ++++++++++++++++++++++++------- 1 file changed, 92 insertions(+), 27 deletions(-) diff --git a/bitcoin/src/blockdata/witness.rs b/bitcoin/src/blockdata/witness.rs index d6b67a738..d8fa1fa17 100644 --- a/bitcoin/src/blockdata/witness.rs +++ b/bitcoin/src/blockdata/witness.rs @@ -147,14 +147,12 @@ crate::internal_macros::define_extension_trait! { /// /// See [`Script::is_p2tr`] to check whether this is actually a Taproot witness. fn tapscript(&self) -> Option<&Script> { - 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) + match P2TrSpend::from_witness(self) { + // Note: the method is named "tapscript" but historically it was actually returning + // leaf script. This is broken but we now keep the behavior the same to not subtly + // break someone. + Some(P2TrSpend::Script { leaf_script, .. }) => Some(leaf_script), + _ => None, } } @@ -166,14 +164,9 @@ 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]> { - if self.is_empty() { - return None; - } - - if self.taproot_annex().is_some() { - self.second_to_last() - } else { - self.last() + match P2TrSpend::from_witness(self) { + Some(P2TrSpend::Script { control_block, .. }) => Some(control_block), + _ => None, } } @@ -183,17 +176,7 @@ crate::internal_macros::define_extension_trait! { /// /// See [`Script::is_p2tr`] to check whether this is actually a Taproot witness. fn taproot_annex(&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() >= 2 && last.first() == Some(&TAPROOT_ANNEX_PREFIX) { - Some(last) - } else { - None - } - }) + P2TrSpend::from_witness(self)?.annex() } /// Get the p2wsh witness script following BIP141 rules. @@ -206,6 +189,88 @@ crate::internal_macros::define_extension_trait! { } } +/// Represents a possible Taproot spend. +/// +/// Taproot can be spent as key spend or script spend and, depending on which it is, different data +/// is in the witness. This type helps representing that data more cleanly when parsing the witness +/// because there are a lot of conditions that make reasoning hard. It's better to parse it at one +/// place and pass it along. +/// +/// This type is so far private but it could be published eventually. The design is geared towards +/// it but it's not fully finished. +enum P2TrSpend<'a> { + Key { + // This field is technically present in witness in case of key spend but none of our code + // uses it yet. Rather than deleting it, it's kept here commented as documentation and as + // an easy way to add it if anything needs it - by just uncommenting. + // signature: &'a [u8], + annex: Option<&'a [u8]>, + }, + Script { + leaf_script: &'a Script, + control_block: &'a [u8], + annex: Option<&'a [u8]>, + }, +} + +impl<'a> P2TrSpend<'a> { + /// Parses `Witness` to determine what kind of taproot spend this is. + /// + /// Note: this assumes `witness` is a taproot spend. The function cannot figure it out for sure + /// (without knowing the output), so it doesn't attempt to check anything other than what is + /// required for the program to not crash. + /// + /// In other words, if the caller is certain that the witness is a valid p2tr spend (e.g. + /// obtained from Bitcoin Core) then it's OK to unwrap this but not vice versa - `Some` does + /// not imply correctness. + fn from_witness(witness: &'a Witness) -> Option { + // BIP341 says: + // 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. + // + // However here we're not removing anything, so we have to adjust the numbers to account + // for the fact that annex is still there. + match witness.len() { + 0 => None, + 1 => Some(P2TrSpend::Key { /* signature: witness.last().expect("len > 0") ,*/ annex: None }), + 2 if witness.last().expect("len > 0").starts_with(&[TAPROOT_ANNEX_PREFIX]) => { + let spend = P2TrSpend::Key { + // signature: witness.second_to_last().expect("len > 1"), + annex: witness.last(), + }; + Some(spend) + }, + // 2 => this is script spend without annex - same as when there are 3+ elements and the + // last one does NOT start with TAPROOT_ANNEX_PREFIX. This is handled in the catchall + // arm. + 3.. if witness.last().expect("len > 0").starts_with(&[TAPROOT_ANNEX_PREFIX]) => { + let spend = P2TrSpend::Script { + leaf_script: Script::from_bytes(witness.third_to_last().expect("len > 2")), + control_block: witness.second_to_last().expect("len > 1"), + annex: witness.last(), + }; + Some(spend) + }, + _ => { + let spend = P2TrSpend::Script { + leaf_script: Script::from_bytes(witness.second_to_last().expect("len > 1")), + control_block: witness.last().expect("len > 0"), + annex: None, + }; + Some(spend) + }, + } + } + + fn annex(&self) -> Option<&'a [u8]> { + match self { + P2TrSpend::Key { annex, .. } => *annex, + P2TrSpend::Script { annex, .. } => *annex, + } + } +} + mod sealed { pub trait Sealed {} impl Sealed for super::Witness {} From 59f21a291f08af0029691676f86ac95dad04d8c3 Mon Sep 17 00:00:00 2001 From: Martin Habovstiak Date: Fri, 21 Feb 2025 15:34:39 +0100 Subject: [PATCH 2/3] Add a test case checking `taproot_control_block` The previous commit fixed a bug when `taproot_control_block` returned `Some` on key-spends. This adds a test case for it which succeeds when applied after the previous commit and fails if applied before it. --- bitcoin/src/blockdata/witness.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/bitcoin/src/blockdata/witness.rs b/bitcoin/src/blockdata/witness.rs index d8fa1fa17..20efc0b5f 100644 --- a/bitcoin/src/blockdata/witness.rs +++ b/bitcoin/src/blockdata/witness.rs @@ -397,19 +397,24 @@ mod test { let control_block = hex!("02"); // annex starting with 0x50 causes the branching logic. let annex = hex!("50"); + let signature = vec![0xff; 64]; let witness_vec = vec![tapscript.clone(), control_block.clone()]; - let witness_vec_annex = vec![tapscript.clone(), control_block.clone(), annex]; + let witness_vec_annex = vec![tapscript.clone(), control_block.clone(), annex.clone()]; + let witness_vec_key_spend_annex = vec![signature, annex]; let witness_serialized: Vec = serialize(&witness_vec); let witness_serialized_annex: Vec = serialize(&witness_vec_annex); + let witness_serialized_key_spend_annex: Vec = serialize(&witness_vec_key_spend_annex); let witness = deserialize::(&witness_serialized[..]).unwrap(); let witness_annex = deserialize::(&witness_serialized_annex[..]).unwrap(); + let witness_key_spend_annex = deserialize::(&witness_serialized_key_spend_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[..])); + assert!(witness_key_spend_annex.taproot_control_block().is_none()) } #[test] From a8168c3f81a76165022af3f3aeec82317d8a6bf1 Mon Sep 17 00:00:00 2001 From: Martin Habovstiak Date: Fri, 21 Feb 2025 16:33:00 +0100 Subject: [PATCH 3/3] 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> { + 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 = serialize(&witness_vec); + let witness_serialized_annex: Vec = serialize(&witness_vec_annex); + + let witness = deserialize::(&witness_serialized[..]).unwrap(); + let witness_annex = deserialize::(&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 { + /// 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>;