From 382c8f9e4f91bb06b41938c7d9d1acaa02d9f398 Mon Sep 17 00:00:00 2001
From: sanket1729 <sanket1729@gmail.com>
Date: Fri, 14 Jan 2022 04:06:12 +0530
Subject: [PATCH] Introduce PsbtSigHashType

---
 src/util/psbt/map/input.rs | 54 +++++++++++++++++++++++++++++++++++---
 src/util/psbt/map/mod.rs   |  2 +-
 src/util/psbt/mod.rs       |  7 ++---
 src/util/psbt/serialize.rs | 17 ++++--------
 src/util/sighash.rs        |  4 +--
 5 files changed, 63 insertions(+), 21 deletions(-)

diff --git a/src/util/psbt/map/input.rs b/src/util/psbt/map/input.rs
index 970e7849..bc6772b7 100644
--- a/src/util/psbt/map/input.rs
+++ b/src/util/psbt/map/input.rs
@@ -19,7 +19,7 @@ use ::{EcdsaSig, io};
 use secp256k1;
 use blockdata::script::Script;
 use blockdata::witness::Witness;
-use blockdata::transaction::{EcdsaSigHashType, Transaction, TxOut};
+use blockdata::transaction::{Transaction, TxOut, NonStandardSigHashType};
 use consensus::encode;
 use hashes::{self, hash160, ripemd160, sha256, sha256d};
 use secp256k1::XOnlyPublicKey;
@@ -32,6 +32,8 @@ use util::psbt::{Error, error};
 
 use ::{SchnorrSig};
 use util::taproot::{ControlBlock, LeafVersion, TapLeafHash, TapBranchHash};
+use util::sighash;
+use {EcdsaSigHashType, SchnorrSigHashType};
 
 /// Type: Non-Witness UTXO PSBT_IN_NON_WITNESS_UTXO = 0x00
 const PSBT_IN_NON_WITNESS_UTXO: u8 = 0x00;
@@ -92,7 +94,7 @@ pub struct Input {
     pub partial_sigs: BTreeMap<secp256k1::PublicKey, EcdsaSig>,
     /// The sighash type to be used for this input. Signatures for this input
     /// must use the sighash type.
-    pub sighash_type: Option<EcdsaSigHashType>,
+    pub sighash_type: Option<PsbtSigHashType>,
     /// The redeem script for this input.
     pub redeem_script: Option<Script>,
     /// The witness script for this input.
@@ -143,6 +145,52 @@ pub struct Input {
     pub unknown: BTreeMap<raw::Key, Vec<u8>>,
 }
 
+
+/// A Signature hash type for the corresponding input. As of taproot upgrade, the signature hash
+/// type can be either [`EcdsaSigHashType`] or [`SchnorrSigHashType`] but it is not possible to know
+/// directly which signature hash type the user is dealing with. Therefore, the user is responsible
+/// for converting to/from [`PsbtSigHashType`] from/to the desired signature hash type they need.
+#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct PsbtSigHashType {
+    pub (in ::util::psbt) inner: u32,
+}
+
+impl From<EcdsaSigHashType> for PsbtSigHashType {
+    fn from(ecdsa_hash_ty: EcdsaSigHashType) -> Self {
+        PsbtSigHashType {inner: ecdsa_hash_ty as u32}
+    }
+}
+
+impl From<SchnorrSigHashType> for PsbtSigHashType {
+    fn from(schnorr_hash_ty: SchnorrSigHashType) -> Self {
+        PsbtSigHashType {inner: schnorr_hash_ty as u32}
+    }
+}
+
+impl PsbtSigHashType {
+    /// Returns the [`EcdsaSigHashType`] if the [`PsbtSigHashType`] can be
+    /// converted to one.
+    pub fn ecdsa_hash_ty(self) -> Result<EcdsaSigHashType, NonStandardSigHashType> {
+        EcdsaSigHashType::from_u32_standard(self.inner)
+    }
+
+    /// Returns the [`SchnorrSigHashType`] if the [`PsbtSigHashType`] can be
+    /// converted to one.
+    pub fn schnorr_hash_ty(self) -> Result<SchnorrSigHashType, sighash::Error> {
+        if self.inner > 0xffu32 {
+            Err(sighash::Error::InvalidSigHashType(self.inner))
+        } else {
+            SchnorrSigHashType::from_u8(self.inner as u8)
+        }
+    }
+
+    /// Obtains the inner sighash byte from this [`PsbtSigHashType`].
+    pub fn inner(self) -> u32 {
+        self.inner
+    }
+}
+
 impl Map for Input {
     fn insert_pair(&mut self, pair: raw::Pair) -> Result<(), encode::Error> {
         let raw::Pair {
@@ -168,7 +216,7 @@ impl Map for Input {
             }
             PSBT_IN_SIGHASH_TYPE => {
                 impl_psbt_insert_pair! {
-                    self.sighash_type <= <raw_key: _>|<raw_value: EcdsaSigHashType>
+                    self.sighash_type <= <raw_key: _>|<raw_value: PsbtSigHashType>
                 }
             }
             PSBT_IN_REDEEM_SCRIPT => {
diff --git a/src/util/psbt/map/mod.rs b/src/util/psbt/map/mod.rs
index a40de9d2..ef0c7963 100644
--- a/src/util/psbt/map/mod.rs
+++ b/src/util/psbt/map/mod.rs
@@ -53,6 +53,6 @@ mod global;
 mod input;
 mod output;
 
-pub use self::input::Input;
+pub use self::input::{Input, PsbtSigHashType};
 pub use self::output::Output;
 pub use self::output::TapTree;
diff --git a/src/util/psbt/mod.rs b/src/util/psbt/mod.rs
index d8512624..4dd3e8d2 100644
--- a/src/util/psbt/mod.rs
+++ b/src/util/psbt/mod.rs
@@ -406,6 +406,7 @@ mod tests {
         //! Create a full PSBT value with various fields filled and make sure it can be JSONized.
         use hashes::sha256d;
         use util::psbt::map::Input;
+        use EcdsaSigHashType;
 
         // create some values to use in the PSBT
         let tx = Transaction {
@@ -469,7 +470,7 @@ mod tests {
                     value: 190303501938,
                     script_pubkey: hex_script!("a914339725ba21efd62ac753a9bcd067d6c7a6a39d0587"),
                 }),
-                sighash_type: Some("SIGHASH_SINGLE|SIGHASH_ANYONECANPAY".parse().unwrap()),
+                sighash_type: Some("SIGHASH_SINGLE|SIGHASH_ANYONECANPAY".parse::<EcdsaSigHashType>().unwrap().into()),
                 redeem_script: Some(vec![0x51].into()),
                 witness_script: None,
                 partial_sigs: vec![(
@@ -732,8 +733,8 @@ mod tests {
                     .is_p2pkh()
             );
             assert_eq!(
-                (&psbt.inputs[0].sighash_type).as_ref().unwrap(),
-                &EcdsaSigHashType::All
+                (&psbt.inputs[0].sighash_type).as_ref().unwrap().ecdsa_hash_ty().unwrap(),
+                EcdsaSigHashType::All
             );
         }
 
diff --git a/src/util/psbt/serialize.rs b/src/util/psbt/serialize.rs
index 9a5d2141..791c52b9 100644
--- a/src/util/psbt/serialize.rs
+++ b/src/util/psbt/serialize.rs
@@ -30,10 +30,9 @@ use secp256k1::{self, XOnlyPublicKey};
 use util::bip32::{ChildNumber, Fingerprint, KeySource};
 use hashes::{hash160, ripemd160, sha256, sha256d, Hash};
 use util::ecdsa::EcdsaSig;
-use util::psbt;
 use util::taproot::{TapBranchHash, TapLeafHash, ControlBlock, LeafVersion};
 use schnorr;
-use super::map::TapTree;
+use super::map::{TapTree, PsbtSigHashType};
 
 use util::taproot::TaprootBuilder;
 use util::sighash::SchnorrSigHashType;
@@ -160,22 +159,16 @@ impl Deserialize for Vec<u8> {
     }
 }
 
-impl Serialize for EcdsaSigHashType {
+impl Serialize for PsbtSigHashType {
     fn serialize(&self) -> Vec<u8> {
-        serialize(&self.as_u32())
+        serialize(&self.inner())
     }
 }
 
-impl Deserialize for EcdsaSigHashType {
+impl Deserialize for PsbtSigHashType {
     fn deserialize(bytes: &[u8]) -> Result<Self, encode::Error> {
         let raw: u32 = encode::deserialize(bytes)?;
-        let rv: EcdsaSigHashType = EcdsaSigHashType::from_u32_consensus(raw);
-
-        if rv.as_u32() == raw {
-            Ok(rv)
-        } else {
-            Err(psbt::Error::NonStandardSigHashType(raw).into())
-        }
+        Ok(PsbtSigHashType { inner: raw })
     }
 }
 
diff --git a/src/util/sighash.rs b/src/util/sighash.rs
index b878949e..44b7c140 100644
--- a/src/util/sighash.rs
+++ b/src/util/sighash.rs
@@ -167,7 +167,7 @@ pub enum Error {
     WrongAnnex,
 
     /// Invalid Sighash type
-    InvalidSigHashType(u8),
+    InvalidSigHashType(u32),
 }
 
 impl fmt::Display for Error {
@@ -286,7 +286,7 @@ impl SchnorrSigHashType {
             0x81 => Ok(SchnorrSigHashType::AllPlusAnyoneCanPay),
             0x82 => Ok(SchnorrSigHashType::NonePlusAnyoneCanPay),
             0x83 => Ok(SchnorrSigHashType::SinglePlusAnyoneCanPay),
-            x => Err(Error::InvalidSigHashType(x)),
+            x => Err(Error::InvalidSigHashType(x as u32)),
         }
     }
 }