Taproot psbt impl BIP 371

sanket1729 · sanket1729 · commit d22e0149ad7c · 2021-12-28T20:40:58.000+05:30
diff --git a/src/util/psbt/map/input.rs b/src/util/psbt/map/input.rs
@@ -28,6 +28,9 @@ use util::psbt::raw;
 use util::psbt::serialize::Deserialize;
 use util::psbt::{Error, error};
 
+use schnorr;
+use util::taproot::{ControlBlock, LeafVersion, TapLeafHash, TapBranchHash};
+
 /// Type: Non-Witness UTXO PSBT_IN_NON_WITNESS_UTXO = 0x00
 const PSBT_IN_NON_WITNESS_UTXO: u8 = 0x00;
 /// Type: Witness UTXO PSBT_IN_WITNESS_UTXO = 0x01
@@ -54,6 +57,18 @@ const PSBT_IN_SHA256: u8 = 0x0b;
 const PSBT_IN_HASH160: u8 = 0x0c;
 /// Type: HASH256 preimage PSBT_IN_HASH256 = 0x0d
 const PSBT_IN_HASH256: u8 = 0x0d;
+/// Type: Schnorr Signature in Key Spend PSBT_IN_TAP_KEY_SIG = 0x13
+const PSBT_IN_TAP_KEY_SIG: u8 = 0x13;
+/// Type: Schnorr Signature in Script Spend PSBT_IN_TAP_SCRIPT_SIG = 0x14
+const PSBT_IN_TAP_SCRIPT_SIG: u8 = 0x14;
+/// Type: Taproot Leaf Script PSBT_IN_TAP_LEAF_SCRIPT = 0x14
+const PSBT_IN_TAP_LEAF_SCRIPT: u8 = 0x15;
+/// Type: Taproot Key BIP 32 Derivation Path PSBT_IN_TAP_BIP32_DERIVATION = 0x16
+const PSBT_IN_TAP_BIP32_DERIVATION : u8 = 0x16;
+/// Type: Taproot Internal Key PSBT_IN_TAP_INTERNAL_KEY = 0x17
+const PSBT_IN_TAP_INTERNAL_KEY : u8 = 0x17;
+/// Type: Taproot Merkle Root PSBT_IN_TAP_MERKLE_ROOT = 0x18
+const PSBT_IN_TAP_MERKLE_ROOT : u8 = 0x18;
 /// Type: Proprietary Use Type PSBT_IN_PROPRIETARY = 0xFC
 const PSBT_IN_PROPRIETARY: u8 = 0xFC;
 
@@ -104,6 +119,21 @@ pub struct Input {
     /// HAS256 hash to preimage map
     #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap_byte_values"))]
     pub hash256_preimages: BTreeMap<sha256d::Hash, Vec<u8>>,
+    /// Serialized schnorr signature with sighash type for key spend
+    pub tap_key_sig: Option<schnorr::SchnorrSig>,
+    /// Map of <xonlypubkey>|<leafhash> with signature
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap_as_seq"))]
+    pub tap_script_sigs: BTreeMap<(schnorr::PublicKey, TapLeafHash), schnorr::SchnorrSig>,
+    /// Map of Control blocks to Script version pair
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap_as_seq"))]
+    pub tap_scripts: BTreeMap<ControlBlock, (Script, LeafVersion)>,
+    /// Map of tap root x only keys to origin info and leaf hashes contained in it
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap_as_seq"))]
+    pub tap_key_origins: BTreeMap<schnorr::PublicKey, (Vec<TapLeafHash>, KeySource)>,
+    /// Taproot Internal key
+    pub tap_internal_key : Option<schnorr::PublicKey>,
+    /// Taproot Merkle root
+    pub tap_merkle_root : Option<TapBranchHash>,
     /// Proprietary key-value pairs for this input.
     #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap_as_seq_byte_values"))]
     pub proprietary: BTreeMap<raw::ProprietaryKey, Vec<u8>>,
@@ -177,6 +207,36 @@ impl Map for Input {
             PSBT_IN_HASH256 => {
                 psbt_insert_hash_pair(&mut self.hash256_preimages, raw_key, raw_value, error::PsbtHash::Hash256)?;
             }
+            PSBT_IN_TAP_KEY_SIG => {
+                impl_psbt_insert_pair! {
+                    self.tap_key_sig <= <raw_key: _>|<raw_value: schnorr::SchnorrSig>
+                }
+            }
+            PSBT_IN_TAP_SCRIPT_SIG => {
+                impl_psbt_insert_pair! {
+                    self.tap_script_sigs <= <raw_key: (schnorr::PublicKey, TapLeafHash)>|<raw_value: schnorr::SchnorrSig>
+                }
+            }
+            PSBT_IN_TAP_LEAF_SCRIPT=> {
+                impl_psbt_insert_pair! {
+                    self.tap_scripts <= <raw_key: ControlBlock>|< raw_value: (Script, LeafVersion)>
+                }
+            }
+            PSBT_IN_TAP_BIP32_DERIVATION => {
+                impl_psbt_insert_pair! {
+                    self.tap_key_origins <= <raw_key: schnorr::PublicKey>|< raw_value: (Vec<TapLeafHash>, KeySource)>
+                }
+            }
+            PSBT_IN_TAP_INTERNAL_KEY => {
+                impl_psbt_insert_pair! {
+                    self.tap_internal_key <= <raw_key: _>|< raw_value: schnorr::PublicKey>
+                }
+            }
+            PSBT_IN_TAP_MERKLE_ROOT => {
+                impl_psbt_insert_pair! {
+                    self.tap_merkle_root <= <raw_key: _>|< raw_value: TapBranchHash>
+                }
+            }
             PSBT_IN_PROPRIETARY => match self.proprietary.entry(raw::ProprietaryKey::from_key(raw_key.clone())?) {
                 btree_map::Entry::Vacant(empty_key) => {empty_key.insert(raw_value);},
                 btree_map::Entry::Occupied(_) => return Err(Error::DuplicateKey(raw_key).into()),
@@ -249,6 +309,30 @@ impl Map for Input {
             rv.push(self.hash256_preimages as <PSBT_IN_HASH256, sha256d::Hash>|<Vec<u8>>)
         }
 
+        impl_psbt_get_pair! {
+            rv.push(self.tap_key_sig as <PSBT_IN_TAP_KEY_SIG, _>|<Vec<u8>>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.tap_script_sigs as <PSBT_IN_TAP_SCRIPT_SIG, (schnorr::PublicKey, TapLeafHash)>|<Vec<u8>>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.tap_scripts as <PSBT_IN_TAP_LEAF_SCRIPT, ControlBlock>|<(Script, LeafVersion)>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.tap_key_origins as <PSBT_IN_TAP_BIP32_DERIVATION,
+                schnorr::PublicKey>|<(Vec<TapLeafHash>, KeySource)>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.tap_internal_key as <PSBT_IN_TAP_INTERNAL_KEY, _>|<schnorr::PublicKey>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.tap_merkle_root as <PSBT_IN_TAP_MERKLE_ROOT, _>|<TapBranchHash>)
+        }
         for (key, value) in self.proprietary.iter() {
             rv.push(raw::Pair {
                 key: key.to_key(),
@@ -280,13 +364,19 @@ impl Map for Input {
         self.sha256_preimages.extend(other.sha256_preimages);
         self.hash160_preimages.extend(other.hash160_preimages);
         self.hash256_preimages.extend(other.hash256_preimages);
+        self.tap_script_sigs.extend(other.tap_script_sigs);
+        self.tap_scripts.extend(other.tap_scripts);
+        self.tap_key_origins.extend(other.tap_key_origins);
         self.proprietary.extend(other.proprietary);
         self.unknown.extend(other.unknown);
 
         merge!(redeem_script, self, other);
         merge!(witness_script, self, other);
         merge!(final_script_sig, self, other);
         merge!(final_script_witness, self, other);
+        merge!(tap_key_sig, self, other);
+        merge!(tap_internal_key, self, other);
+        merge!(tap_merkle_root, self, other);
 
         Ok(())
     }
diff --git a/src/util/psbt/map/mod.rs b/src/util/psbt/map/mod.rs
@@ -55,3 +55,4 @@ mod output;
 
 pub use self::input::Input;
 pub use self::output::Output;
+pub use self::output::TapTree;
diff --git a/src/util/psbt/map/output.rs b/src/util/psbt/map/output.rs
@@ -25,12 +25,23 @@ use util::psbt::map::Map;
 use util::psbt::raw;
 use util::psbt::Error;
 
+use schnorr;
+use util::taproot::TapLeafHash;
+
+use util::taproot::{NodeInfo, TaprootBuilder};
+
 /// Type: Redeem Script PSBT_OUT_REDEEM_SCRIPT = 0x00
 const PSBT_OUT_REDEEM_SCRIPT: u8 = 0x00;
 /// Type: Witness Script PSBT_OUT_WITNESS_SCRIPT = 0x01
 const PSBT_OUT_WITNESS_SCRIPT: u8 = 0x01;
 /// Type: BIP 32 Derivation Path PSBT_OUT_BIP32_DERIVATION = 0x02
 const PSBT_OUT_BIP32_DERIVATION: u8 = 0x02;
+/// Type: Taproot Internal Key PSBT_OUT_TAP_INTERNAL_KEY = 0x05
+const PSBT_OUT_TAP_INTERNAL_KEY: u8 = 0x05;
+/// Type: Taproot Tree PSBT_OUT_TAP_TREE = 0x06
+const PSBT_OUT_TAP_TREE: u8 = 0x06;
+/// Type: Taproot Key BIP 32 Derivation Path PSBT_OUT_TAP_BIP32_DERIVATION = 0x07
+const PSBT_OUT_TAP_BIP32_DERIVATION: u8 = 0x07;
 /// Type: Proprietary Use Type PSBT_IN_PROPRIETARY = 0xFC
 const PSBT_OUT_PROPRIETARY: u8 = 0xFC;
 
@@ -47,14 +58,67 @@ pub struct Output {
     /// corresponding master key fingerprints and derivation paths.
     #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap_as_seq"))]
     pub bip32_derivation: BTreeMap<PublicKey, KeySource>,
+    /// The internal pubkey
+    pub tap_internal_key: Option<schnorr::PublicKey>,
+    /// Taproot Output tree
+    pub tap_tree: Option<TapTree>,
+    /// Map of tap root x only keys to origin info and leaf hashes contained in it
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap_as_seq"))]
+    pub tap_key_origins: BTreeMap<schnorr::PublicKey, (Vec<TapLeafHash>, KeySource)>,
     /// Proprietary key-value pairs for this output.
-    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap_as_seq_byte_values"))]
+    #[cfg_attr(
+        feature = "serde",
+        serde(with = "::serde_utils::btreemap_as_seq_byte_values")
+    )]
     pub proprietary: BTreeMap<raw::ProprietaryKey, Vec<u8>>,
     /// Unknown key-value pairs for this output.
-    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap_as_seq_byte_values"))]
+    #[cfg_attr(
+        feature = "serde",
+        serde(with = "::serde_utils::btreemap_as_seq_byte_values")
+    )]
     pub unknown: BTreeMap<raw::Key, Vec<u8>>,
 }
 
+/// Taproot Tree representing a finalized [`TaprootBuilder`] (a complete binary tree)
+#[derive(Clone, Debug)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct TapTree(pub(crate) TaprootBuilder);
+
+impl PartialEq for TapTree {
+    fn eq(&self, other: &Self) -> bool {
+        self.node_info().hash.eq(&other.node_info().hash)
+    }
+}
+
+impl Eq for TapTree {}
+
+impl TapTree {
+    // get the inner node info as the builder is finalized
+    fn node_info(&self) -> &NodeInfo {
+        // The builder algorithm invariant guarantees that is_complete builder
+        // have only 1 element in branch and that is not None.
+        // We make sure that we only allow is_complete builders via the from_inner
+        // constructor
+        self.0.branch()[0].as_ref().expect("from_inner only parses is_complete builders")
+    }
+
+    /// Convert a [`TaprootBuilder`] into a tree if it is complete binary tree.
+    /// Returns the inner as Err if it is not a complete tree
+    pub fn from_inner(inner: TaprootBuilder) -> Result<Self, TaprootBuilder> {
+        if inner.is_complete() {
+            Ok(TapTree(inner))
+        } else {
+            Err(inner)
+        }
+    }
+
+    /// Convert self into builder [`TaprootBuilder`]. The builder is guaranteed to
+    /// be finalized.
+    pub fn into_inner(self) -> TaprootBuilder {
+        self.0
+    }
+}
+
 impl Map for Output {
     fn insert_pair(&mut self, pair: raw::Pair) -> Result<(), encode::Error> {
         let raw::Pair {
@@ -82,10 +146,29 @@ impl Map for Output {
                 btree_map::Entry::Vacant(empty_key) => {empty_key.insert(raw_value);},
                 btree_map::Entry::Occupied(_) => return Err(Error::DuplicateKey(raw_key).into()),
             }
-            _ => match self.unknown.entry(raw_key) {
-                btree_map::Entry::Vacant(empty_key) => {empty_key.insert(raw_value);},
-                btree_map::Entry::Occupied(k) => return Err(Error::DuplicateKey(k.key().clone()).into()),
+            PSBT_OUT_TAP_INTERNAL_KEY => {
+                impl_psbt_insert_pair! {
+                    self.tap_internal_key <= <raw_key: _>|<raw_value: schnorr::PublicKey>
+                }
+            }
+            PSBT_OUT_TAP_TREE => {
+                impl_psbt_insert_pair! {
+                    self.tap_tree <= <raw_key: _>|<raw_value: TapTree>
+                }
+            }
+            PSBT_OUT_TAP_BIP32_DERIVATION => {
+                impl_psbt_insert_pair! {
+                    self.tap_key_origins <= <raw_key: schnorr::PublicKey>|< raw_value: (Vec<TapLeafHash>, KeySource)>
+                }
             }
+            _ => match self.unknown.entry(raw_key) {
+                btree_map::Entry::Vacant(empty_key) => {
+                    empty_key.insert(raw_value);
+                }
+                btree_map::Entry::Occupied(k) => {
+                    return Err(Error::DuplicateKey(k.key().clone()).into())
+                }
+            },
         }
 
         Ok(())
@@ -106,6 +189,19 @@ impl Map for Output {
             rv.push(self.bip32_derivation as <PSBT_OUT_BIP32_DERIVATION, PublicKey>|<KeySource>)
         }
 
+        impl_psbt_get_pair! {
+            rv.push(self.tap_internal_key as <PSBT_OUT_TAP_INTERNAL_KEY, _>|<schnorr::PublicKey>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.tap_tree as <PSBT_OUT_TAP_TREE, _>|<TaprootBuilder>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.tap_key_origins as <PSBT_OUT_TAP_BIP32_DERIVATION,
+                    schnorr::PublicKey>|<(Vec<TapLeafHash>, KeySource)>)
+        }
+
         for (key, value) in self.proprietary.iter() {
             rv.push(raw::Pair {
                 key: key.to_key(),
@@ -127,9 +223,12 @@ impl Map for Output {
         self.bip32_derivation.extend(other.bip32_derivation);
         self.proprietary.extend(other.proprietary);
         self.unknown.extend(other.unknown);
+        self.tap_key_origins.extend(other.tap_key_origins);
 
         merge!(redeem_script, self, other);
         merge!(witness_script, self, other);
+        merge!(tap_internal_key, self, other);
+        merge!(tap_tree, self, other);
 
         Ok(())
     }
diff --git a/src/util/psbt/mod.rs b/src/util/psbt/mod.rs
@@ -39,7 +39,7 @@ mod macros;
 pub mod serialize;
 
 mod map;
-pub use self::map::{Map, Input, Output};
+pub use self::map::{Map, Input, Output, TapTree};
 
 use util::bip32::{ExtendedPubKey, KeySource};
 
diff --git a/src/util/psbt/serialize.rs b/src/util/psbt/serialize.rs
diff --git a/src/util/taproot.rs b/src/util/taproot.rs

Original file line number	Diff line number	Diff line change
`@@ -55,3 +55,4 @@ mod output;`
`55`	`55`
`56`	`56`	`pub use self::input::Input;`
`57`	`57`	`pub use self::output::Output;`
	`58`	`+pub use self::output::TapTree;`