From 638f8eda65ef0faf4c3b1d87610818d00153ea86 Mon Sep 17 00:00:00 2001
From: Jeremy Rubin <j@rubin.io>
Date: Mon, 15 Aug 2022 18:37:58 -0700
Subject: [PATCH] Partial fix: fail deserialization of Taproot PSBTs with non
 BIP-174 compliant PSBTs

---
 src/util/psbt/error.rs      |  8 +++++++-
 src/util/psbt/map/output.rs |  6 ++++++
 src/util/psbt/serialize.rs  | 15 ++++++++++++++-
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/util/psbt/error.rs b/src/util/psbt/error.rs
index 7965e0e612..6816f1b92c 100644
--- a/src/util/psbt/error.rs
+++ b/src/util/psbt/error.rs
@@ -7,6 +7,7 @@ use core::fmt;
 use crate::blockdata::transaction::Transaction;
 use crate::consensus::encode;
 use crate::util::psbt::raw;
+use crate::util::psbt::map::IncompleteTapTree;
 
 use crate::hashes;
 use crate::util::bip32::ExtendedPubKey;
@@ -73,6 +74,9 @@ pub enum Error {
     CombineInconsistentKeySources(Box<ExtendedPubKey>),
     /// Serialization error in bitcoin consensus-encoded structures
     ConsensusEncoding,
+    /// Incomplete TapTree Error
+    // N.B. No need to Box as it's just a vector internally.
+    IncompleteTapTree(IncompleteTapTree)
 }
 
 impl fmt::Display for Error {
@@ -100,6 +104,7 @@ impl fmt::Display for Error {
             },
             Error::CombineInconsistentKeySources(ref s) => { write!(f, "combine conflict: {}", s) },
             Error::ConsensusEncoding => f.write_str("bitcoin consensus or BIP-174 encoding error"),
+            Error::IncompleteTapTree(ref b) => write!(f, "{}", b),
         }
     }
 }
@@ -127,7 +132,8 @@ impl std::error::Error for Error {
             | NonStandardSighashType(_)
             | InvalidPreimageHashPair{ .. }
             | CombineInconsistentKeySources(_)
-            | ConsensusEncoding => None,
+            | ConsensusEncoding
+            | IncompleteTapTree(_) => None,
         }
     }
 }
diff --git a/src/util/psbt/map/output.rs b/src/util/psbt/map/output.rs
index 2545c1e937..d402a7e517 100644
--- a/src/util/psbt/map/output.rs
+++ b/src/util/psbt/map/output.rs
@@ -99,6 +99,12 @@ impl core::fmt::Display for IncompleteTapTree {
     }
 }
 
+impl From<IncompleteTapTree> for crate::util::psbt::Error {
+    fn from(i: IncompleteTapTree) -> Self {
+        crate::util::psbt::Error::IncompleteTapTree(i)
+    }
+}
+
 #[cfg(feature = "std")]
 #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
 impl std::error::Error for IncompleteTapTree {
diff --git a/src/util/psbt/serialize.rs b/src/util/psbt/serialize.rs
index c10b98f286..11046f5746 100644
--- a/src/util/psbt/serialize.rs
+++ b/src/util/psbt/serialize.rs
@@ -7,6 +7,8 @@
 //!
 
 use crate::prelude::*;
+use core;
+use core::convert::TryFrom;
 
 use crate::io;
 
@@ -344,7 +346,7 @@ impl Deserialize for TapTree {
                 .map_err(|_| encode::Error::ParseFailed("Tree not in DFS order"))?;
         }
         if builder.is_finalizable() && !builder.has_hidden_nodes() {
-            Ok(TapTree(builder))
+            Ok(TapTree::try_from(builder).map_err(psbt::Error::from)?)
         } else {
             Err(encode::Error::ParseFailed("Incomplete taproot Tree"))
         }
@@ -395,6 +397,17 @@ mod tests {
         assert_eq!(tree, tree_prime);
     }
 
+    #[cfg(feature = "base64")]
+    #[test]
+    fn taptree_empty() {
+        use core::str::FromStr;
+        use crate::psbt::PartiallySignedTransaction;
+        let test_data_containing_empty =
+            "cHNidP8BAIkCAAAAAWlZ9F5/nRjE0quO8bE8cfKIiVTKHa/Zu6PKwWfpZDmpAAAAAAD9////Anz3nHYAAAAAIlEg7T17MEFOm47gN7FhZB5g4HjviGW4YxZWReJr1c/w0I+AlpgAAAAAACJRIDspRhxBQj5+aJw088IczFSbyRKB7wofy/SQXnErwCcQAAAAAAABASsAlDV3AAAAACJRIBMUosKNoVNU5tEW2Q/tylLRzabMGIprODHM0LbB8h1/IRaXdf3KRBBZMc8epg7uuU3pr42i3YMFnheNDW10F0cpnxkAHzYIzlYAAIABAACAAAAAgAAAAABRAAAAARcgl3X9ykQQWTHPHqYO7rlN6a+Not2DBZ4XjQ1tdBdHKZ8AAQUgT+NhtL213EJgJ51h6FiZYhPgUv48YQKACnnMYpOZPMEBBgAhB0/jYbS9tdxCYCedYehYmWIT4FL+PGECgAp5zGKTmTzBGQAfNgjOVgAAgAEAAIAAAACAAQAAADEBAAAAAQUgrUXPD+fhaOeY7GlxJDZ0pAgUZ0reW+A3wmxMwBAIt14BBgAhB61Fzw/n4WjnmOxpcSQ2dKQIFGdK3lvgN8JsTMAQCLdeGQAfNgjOVgAAgAEAAIAAAACAAAAAAFIAAAAA";
+        let psbt = PartiallySignedTransaction::from_str(test_data_containing_empty);
+        assert!(matches!(psbt, Err(psbt::PsbtParseError::PsbtEncoding(encode::Error::ParseFailed("Incomplete taproot Tree")))));
+    }
+
     #[test]
     fn can_deserialize_non_standard_psbt_sighash_type() {
         let non_standard_sighash = [222u8, 0u8, 0u8, 0u8]; // 32 byte value.