Merge #929: Fix TapTree hidden branches bug

c036b0d Unit test for failing TapTree on builder containing hidden nodes. (Dr Maxim Orlovsky)
7771531 Prevent TapTree from hidden parts (Dr Maxim Orlovsky)
b0f3992 Rename TaprootBuilder::is_complete into is_finalized (Dr Maxim Orlovsky)
efa800f Make TapTree::from_inner return a proper error type (Dr Maxim Orlovsky)
e24c6e2 TapTree serialization roundtrip unit test (Dr Maxim Orlovsky)
56adfa4 TaprootBuilder::has_hidden_nodes method (Dr Maxim Orlovsky)
e69701e Rename taproot `*_hidden` API into `*_hidden_nodes` (Dr Maxim Orlovsky)
6add0dd Track information about hidden leaves in taproot NodeInfo (Dr Maxim Orlovsky)

Pull request description:

  Closes #928

ACKs for top commit:
  sanket1729:
    ACK c036b0d. Reviewed the range diff
  apoelstra:
    ACK c036b0d

Tree-SHA512: 3a8193e6d6dd985da30a2094d1111471b5971f422525870003b77b6ac47cd4ad6e718d46a6d86bbb5e92e5253ac53804badf67edd98bbccbdc11e6383c675663

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 {

src/util/psbt/map/output.rs

@@ -79,6 +79,39 @@ pub struct Output {
     pub unknown: BTreeMap<raw::Key, Vec<u8>>,
 }
 
+/// 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<Self, TaprootBuilder> {
-        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<Self, IncompleteTapTree> {
+        if !inner.is_finalized() {
+            Err(IncompleteTapTree::NotFinalized(inner))
+        } else if inner.has_hidden_nodes() {
+            Err(IncompleteTapTree::HiddenParts(inner))
         } else {
-            Err(inner)
+            Ok(TapTree(inner))
         }
     }
 

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};

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<Item = &'map u8>) -> 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.

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<Self, TaprootBuilderError> {
-        let node = NodeInfo::new_hidden(hash);
+    pub fn add_hidden_node(self, depth: u8, hash: sha256::Hash) -> Result<Self, TaprootBuilderError> {
+        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<C: secp256k1::Verification>(
         mut self,
@@ -549,14 +561,17 @@ pub struct NodeInfo {
     pub(crate) hash: sha256::Hash,
     /// Information about leaves inside this node.
     pub(crate) leaves: Vec<LeafInfo>,
+    /// 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
         })
     }
 }