milksad
/
rust-bitcoin-unsafe-fast
14
0
Fork 0
Code Issues Pull Requests Activity

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
This commit is contained in:
Martin Habovstiak 2025-02-21 14:59:55 +01:00
parent d1171381bc
commit e810ecff7c
1 changed files with 92 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

119
bitcoin/src/blockdata/witness.rs
View File

@ -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<Self> {
// 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 {}