diff --git a/src/util/psbt/map/mod.rs b/src/util/psbt/map/mod.rs index e0ec96ebb2..3b25bd21f4 100644 --- a/src/util/psbt/map/mod.rs +++ b/src/util/psbt/map/mod.rs @@ -24,7 +24,7 @@ mod input; mod output; pub use self::input::{Input, PsbtSighashType}; -pub use self::output::{Output, TapTree}; +pub use self::output::{Output, TapTree, IncompleteTapTree}; /// A trait that describes a PSBT key-value map. pub(super) trait Map { diff --git a/src/util/psbt/map/output.rs b/src/util/psbt/map/output.rs index 53edd2d54e..c49eb760a3 100644 --- a/src/util/psbt/map/output.rs +++ b/src/util/psbt/map/output.rs @@ -79,6 +79,39 @@ pub struct Output { pub unknown: BTreeMap>, } +/// Error happening when [`TapTree`] is constructed from a [`TaprootBuilder`] +/// having hidden branches or not being finalized. +#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)] +pub enum IncompleteTapTree { + /// Indicates an attempt to construct a tap tree from a builder containing incomplete branches. + NotFinalized(TaprootBuilder), + /// Indicates an attempt to construct a tap tree from a builder containing hidden parts. + HiddenParts(TaprootBuilder), +} + +impl IncompleteTapTree { + /// Converts error into the original incomplete [`TaprootBuilder`] instance. + pub fn into_builder(self) -> TaprootBuilder { + match self { + IncompleteTapTree::NotFinalized(builder) | + IncompleteTapTree::HiddenParts(builder) => builder + } + } +} + +impl core::fmt::Display for IncompleteTapTree { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str(match self { + IncompleteTapTree::NotFinalized(_) => "an attempt to construct a tap tree from a builder containing incomplete branches.", + IncompleteTapTree::HiddenParts(_) => "an attempt to construct a tap tree from a builder containing hidden parts.", + }) + } +} + +#[cfg(feature = "std")] +#[cfg_attr(docsrs, doc(cfg(feature = "std")))] +impl ::std::error::Error for IncompleteTapTree {} + /// Taproot Tree representing a finalized [`TaprootBuilder`] (a complete binary tree). #[derive(Clone, Debug)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] @@ -104,13 +137,16 @@ impl TapTree { /// Converts a [`TaprootBuilder`] into a tree if it is complete binary tree. /// - /// # Return - /// A `TapTree` iff the `inner` builder is complete, otherwise return the inner as `Err`. - pub fn from_inner(inner: TaprootBuilder) -> Result { - if inner.is_complete() { - Ok(TapTree(inner)) + /// # Returns + /// A [`TapTree`] iff the `inner` builder is complete, otherwise return [`IncompleteTapTree`] + /// error with the content of incomplete builder `inner` instance. + pub fn from_inner(inner: TaprootBuilder) -> Result { + if !inner.is_finalized() { + Err(IncompleteTapTree::NotFinalized(inner)) + } else if inner.has_hidden_nodes() { + Err(IncompleteTapTree::HiddenParts(inner)) } else { - Err(inner) + Ok(TapTree(inner)) } } diff --git a/src/util/psbt/mod.rs b/src/util/psbt/mod.rs index 989542c2a2..4cf84ff698 100644 --- a/src/util/psbt/mod.rs +++ b/src/util/psbt/mod.rs @@ -41,7 +41,7 @@ mod macros; pub mod serialize; mod map; -pub use self::map::{Input, Output, TapTree, PsbtSighashType}; +pub use self::map::{Input, Output, TapTree, PsbtSighashType, IncompleteTapTree}; use self::map::Map; use util::bip32::{ExtendedPubKey, KeySource}; diff --git a/src/util/psbt/serialize.rs b/src/util/psbt/serialize.rs index 7657e5393a..438693f902 100644 --- a/src/util/psbt/serialize.rs +++ b/src/util/psbt/serialize.rs @@ -355,7 +355,7 @@ impl Deserialize for TapTree { builder = builder.add_leaf_with_ver(*depth, script, leaf_version) .map_err(|_| encode::Error::ParseFailed("Tree not in DFS order"))?; } - if builder.is_complete() { + if builder.is_finalized() || !builder.has_hidden_nodes() { Ok(TapTree(builder)) } else { Err(encode::Error::ParseFailed("Incomplete taproot Tree")) @@ -370,8 +370,41 @@ fn key_source_len(key_source: &KeySource) -> usize { #[cfg(test)] mod tests { + use hashes::hex::FromHex; use super::*; + // Composes tree matching a given depth map, filled with dumb script leafs, + // each of which consists of a single push-int op code, with int value + // increased for each consecutive leaf. + pub fn compose_taproot_builder<'map>(opcode: u8, depth_map: impl IntoIterator) -> TaprootBuilder { + let mut val = opcode; + let mut builder = TaprootBuilder::new(); + for depth in depth_map { + let script = Script::from_hex(&format!("{:02x}", val)).unwrap(); + builder = builder.add_leaf(*depth, script).unwrap(); + let (new_val, _) = val.overflowing_add(1); + val = new_val; + } + builder + } + + #[test] + fn taptree_hidden() { + let mut builder = compose_taproot_builder(0x51, &[2, 2, 2]); + builder = builder.add_leaf_with_ver(3, Script::from_hex("b9").unwrap(), LeafVersion::from_consensus(0xC2).unwrap()).unwrap(); + builder = builder.add_hidden_node(3, sha256::Hash::default()).unwrap(); + assert!(TapTree::from_inner(builder.clone()).is_err()); + } + + #[test] + fn taptree_roundtrip() { + let mut builder = compose_taproot_builder(0x51, &[2, 2, 2, 3]); + builder = builder.add_leaf_with_ver(3, Script::from_hex("b9").unwrap(), LeafVersion::from_consensus(0xC2).unwrap()).unwrap(); + let tree = TapTree::from_inner(builder).unwrap(); + let tree_prime = TapTree::deserialize(&tree.serialize()).unwrap(); + assert_eq!(tree, tree_prime); + } + #[test] fn can_deserialize_non_standard_psbt_sighash_type() { let non_standard_sighash = [222u8, 0u8, 0u8, 0u8]; // 32 byte value. diff --git a/src/util/taproot.rs b/src/util/taproot.rs index b7880402f2..325feaa65a 100644 --- a/src/util/taproot.rs +++ b/src/util/taproot.rs @@ -456,16 +456,28 @@ impl TaprootBuilder { /// Adds a hidden/omitted node at `depth` to the builder. Errors if the leaves are not provided /// in DFS walk order. The depth of the root node is 0. - pub fn add_hidden(self, depth: u8, hash: sha256::Hash) -> Result { - let node = NodeInfo::new_hidden(hash); + pub fn add_hidden_node(self, depth: u8, hash: sha256::Hash) -> Result { + let node = NodeInfo::new_hidden_node(hash); self.insert(node, depth) } - /// Checks if the builder is a complete tree. - pub fn is_complete(&self) -> bool { + /// Checks if the builder has finalized building a tree. + pub fn is_finalized(&self) -> bool { self.branch.len() == 1 && self.branch[0].is_some() } + /// Checks if the builder has hidden nodes. + pub fn has_hidden_nodes(&self) -> bool { + for node in &self.branch { + if let Some(node) = node { + if node.has_hidden_nodes { + return true + } + } + } + false + } + /// Creates a [`TaprootSpendInfo`] with the given internal key. pub fn finalize( mut self, @@ -549,14 +561,17 @@ pub struct NodeInfo { pub(crate) hash: sha256::Hash, /// Information about leaves inside this node. pub(crate) leaves: Vec, + /// Tracks information on hidden nodes below this node. + pub(crate) has_hidden_nodes: bool, } impl NodeInfo { /// Creates a new [`NodeInfo`] with omitted/hidden info. - pub fn new_hidden(hash: sha256::Hash) -> Self { + pub fn new_hidden_node(hash: sha256::Hash) -> Self { Self { hash: hash, leaves: vec![], + has_hidden_nodes: true } } @@ -566,6 +581,7 @@ impl NodeInfo { Self { hash: leaf.hash(), leaves: vec![leaf], + has_hidden_nodes: false, } } @@ -584,6 +600,7 @@ impl NodeInfo { Ok(Self { hash: sha256::Hash::from_inner(hash.into_inner()), leaves: all_leaves, + has_hidden_nodes: a.has_hidden_nodes || b.has_hidden_nodes }) } }