From badb0f2a7763f425ef359d855687ec6e1ae6e8a4 Mon Sep 17 00:00:00 2001
From: Carl Dong <accounts@carldong.me>
Date: Sat, 8 Sep 2018 21:30:38 -0700
Subject: [PATCH] Add PSBT input data key-value map type

- Implement psbt::Map trait for psbt::Input
- Add (en)decoding logic for psbt::Input

- Implement PSBT (de)serialization trait for relevant psbt::Input types
---
 src/util/psbt/map/input.rs | 183 +++++++++++++++++++++++++++++++++++++
 src/util/psbt/map/mod.rs   |   2 +
 src/util/psbt/mod.rs       |   2 +-
 src/util/psbt/serialize.rs |  36 +++++++-
 4 files changed, 221 insertions(+), 2 deletions(-)
 create mode 100644 src/util/psbt/map/input.rs

diff --git a/src/util/psbt/map/input.rs b/src/util/psbt/map/input.rs
new file mode 100644
index 00000000..b53fcacd
--- /dev/null
+++ b/src/util/psbt/map/input.rs
@@ -0,0 +1,183 @@
+use std::collections::HashMap;
+
+use blockdata::script::Script;
+use blockdata::transaction::{SigHashType, Transaction, TxOut};
+use consensus::encode;
+use util::bip32::{DerivationPath, Fingerprint};
+use util::key::PublicKey;
+use util::psbt;
+use util::psbt::map::Map;
+use util::psbt::raw;
+use util::psbt::Error;
+
+/// A key-value map for an input of the corresponding index in the unsigned
+/// transaction.
+#[derive(Clone, Default, Debug, PartialEq)]
+pub struct Input {
+    /// The non-witness transaction this input spends from. Should only be
+    /// [std::option::Option::Some] for inputs which spend non-segwit outputs or
+    /// if it is unknown whether an input spends a segwit output.
+    pub non_witness_utxo: Option<Transaction>,
+    /// The transaction output this input spends from. Should only be
+    /// [std::option::Option::Some] for inputs which spend segwit outputs,
+    /// including P2SH embedded ones.
+    pub witness_utxo: Option<TxOut>,
+    /// A map from public keys to their corresponding signature as would be
+    /// pushed to the stack from a scriptSig or witness.
+    pub partial_sigs: HashMap<PublicKey, Vec<u8>>,
+    /// The sighash type to be used for this input. Signatures for this input
+    /// must use the sighash type.
+    pub sighash_type: Option<SigHashType>,
+    /// The redeem script for this input.
+    pub redeem_script: Option<Script>,
+    /// The witness script for this input.
+    pub witness_script: Option<Script>,
+    /// A map from public keys needed to sign this input to their corresponding
+    /// master key fingerprints and derivation paths.
+    pub hd_keypaths: HashMap<PublicKey, (Fingerprint, DerivationPath)>,
+    /// The finalized, fully-constructed scriptSig with signatures and any other
+    /// scripts necessary for this input to pass validation.
+    pub final_script_sig: Option<Script>,
+    /// The finalized, fully-constructed scriptWitness with signatures and any
+    /// other scripts necessary for this input to pass validation.
+    pub final_script_witness: Option<Vec<Vec<u8>>>,
+    /// Unknown key-value pairs for this input.
+    pub unknown: HashMap<raw::Key, Vec<u8>>,
+}
+
+impl Map for Input {
+    fn insert_pair(&mut self, pair: raw::Pair) -> Result<(), encode::Error> {
+        let raw::Pair {
+            key: raw_key,
+            value: raw_value,
+        } = pair;
+
+        match raw_key.type_value {
+            0u8 => {
+                impl_psbt_insert_pair! {
+                    self.non_witness_utxo <= <raw_key: _>|<raw_value: Transaction>
+                }
+            }
+            1u8 => {
+                impl_psbt_insert_pair! {
+                    self.witness_utxo <= <raw_key: _>|<raw_value: TxOut>
+                }
+            }
+            3u8 => {
+                impl_psbt_insert_pair! {
+                    self.sighash_type <= <raw_key: _>|<raw_value: SigHashType>
+                }
+            }
+            4u8 => {
+                impl_psbt_insert_pair! {
+                    self.redeem_script <= <raw_key: _>|<raw_value: Script>
+                }
+            }
+            5u8 => {
+                impl_psbt_insert_pair! {
+                    self.witness_script <= <raw_key: _>|<raw_value: Script>
+                }
+            }
+            7u8 => {
+                impl_psbt_insert_pair! {
+                    self.final_script_sig <= <raw_key: _>|<raw_value: Script>
+                }
+            }
+            8u8 => {
+                impl_psbt_insert_pair! {
+                    self.final_script_witness <= <raw_key: _>|<raw_value: Vec<Vec<u8>>>
+                }
+            }
+            2u8 => {
+                impl_psbt_insert_pair! {
+                    self.partial_sigs <= <raw_key: PublicKey>|<raw_value: Vec<u8>>
+                }
+            }
+            6u8 => {
+                impl_psbt_insert_pair! {
+                    self.hd_keypaths <= <raw_key: PublicKey>|<raw_value: (Fingerprint, DerivationPath)>
+                }
+            }
+            _ => {
+                if self.unknown.contains_key(&raw_key) {
+                    return Err(Error::DuplicateKey(raw_key).into());
+                } else {
+                    self.unknown.insert(raw_key, raw_value);
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    fn get_pairs(&self) -> Result<Vec<raw::Pair>, encode::Error> {
+        let mut rv: Vec<raw::Pair> = Default::default();
+
+        impl_psbt_get_pair! {
+            rv.push(self.non_witness_utxo as <0u8, _>|<Transaction>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.witness_utxo as <1u8, _>|<TxOut>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.partial_sigs as <2u8, PublicKey>|<Vec<u8>>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.sighash_type as <3u8, _>|<SigHashType>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.redeem_script as <4u8, _>|<Script>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.witness_script as <5u8, _>|<Script>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.hd_keypaths as <6u8, PublicKey>|<(Fingerprint, DerivationPath)>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.final_script_sig as <7u8, _>|<Script>)
+        }
+
+        impl_psbt_get_pair! {
+            rv.push(self.final_script_witness as <8u8, _>|<Script>)
+        }
+
+        for (key, value) in self.unknown.iter() {
+            rv.push(raw::Pair {
+                key: key.clone(),
+                value: value.clone(),
+            });
+        }
+
+        Ok(rv)
+    }
+
+    fn merge(&mut self, other: Self) -> Result<(), psbt::Error> {
+        merge!(non_witness_utxo, self, other);
+
+        if let (&None, Some(witness_utxo)) = (&self.witness_utxo, other.witness_utxo) {
+            self.witness_utxo = Some(witness_utxo);
+            self.non_witness_utxo = None; // Clear out any non-witness UTXO when we set a witness one
+        }
+
+        self.partial_sigs.extend(other.partial_sigs);
+        self.hd_keypaths.extend(other.hd_keypaths);
+        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);
+
+        Ok(())
+    }
+}
+
+impl_psbtmap_consensus_enc_dec_oding!(Input);
diff --git a/src/util/psbt/map/mod.rs b/src/util/psbt/map/mod.rs
index 6ad915ce..b585b60c 100644
--- a/src/util/psbt/map/mod.rs
+++ b/src/util/psbt/map/mod.rs
@@ -16,7 +16,9 @@ pub trait Map {
 
 // place at end to pick up macros
 mod global;
+mod input;
 mod output;
 
 pub use self::global::Global;
+pub use self::input::Input;
 pub use self::output::Output;
diff --git a/src/util/psbt/mod.rs b/src/util/psbt/mod.rs
index d0a30222..0b8fbbd1 100644
--- a/src/util/psbt/mod.rs
+++ b/src/util/psbt/mod.rs
@@ -15,7 +15,7 @@ mod macros;
 pub mod serialize;
 
 mod map;
-pub use self::map::{Map, Global, Output};
+pub use self::map::{Map, Global, Input, Output};
 
 #[cfg(test)]
 mod tests {
diff --git a/src/util/psbt/serialize.rs b/src/util/psbt/serialize.rs
index 953b0f3e..0f8bd6b7 100644
--- a/src/util/psbt/serialize.rs
+++ b/src/util/psbt/serialize.rs
@@ -6,7 +6,7 @@
 use std::io::{self, Cursor};
 
 use blockdata::script::Script;
-use blockdata::transaction::Transaction;
+use blockdata::transaction::{SigHashType, Transaction, TxOut};
 use consensus::encode::{self, serialize, Decodable};
 use util::bip32::{ChildNumber, DerivationPath, Fingerprint};
 use util::key::PublicKey;
@@ -26,6 +26,8 @@ pub trait Deserialize: Sized {
 }
 
 impl_psbt_de_serialize!(Transaction);
+impl_psbt_de_serialize!(TxOut);
+impl_psbt_de_serialize!(Vec<Vec<u8>>); // scriptWitness
 
 impl Serialize for Script {
     fn serialize(&self) -> Vec<u8> {
@@ -94,3 +96,35 @@ impl Deserialize for (Fingerprint, DerivationPath) {
         Ok((fprint, dpath.into()))
     }
 }
+
+// partial sigs
+impl Serialize for Vec<u8> {
+    fn serialize(&self) -> Vec<u8> {
+        self.clone()
+    }
+}
+
+impl Deserialize for Vec<u8> {
+    fn deserialize(bytes: &[u8]) -> Result<Self, encode::Error> {
+        Ok(bytes.to_vec())
+    }
+}
+
+impl Serialize for SigHashType {
+    fn serialize(&self) -> Vec<u8> {
+        serialize(&self.as_u32())
+    }
+}
+
+impl Deserialize for SigHashType {
+    fn deserialize(bytes: &[u8]) -> Result<Self, encode::Error> {
+        let raw: u32 = encode::deserialize(bytes)?;
+        let rv: SigHashType = SigHashType::from_u32(raw);
+
+        if rv.as_u32() == raw {
+            Ok(rv)
+        } else {
+            Err(psbt::Error::NonStandardSigHashType(raw).into())
+        }
+    }
+}