From e5f3bca2b4c90dc429e8919c0f37a6eaabd73bc0 Mon Sep 17 00:00:00 2001
From: Steven Roose <steven@stevenroose.org>
Date: Sun, 25 Oct 2020 18:27:45 +0000
Subject: [PATCH] Add serde_utils module to fix JSON serialization

---
 src/lib.rs                  |   3 +
 src/serde_utils.rs          | 107 ++++++++++++++++++++++++++++++++++++
 src/util/psbt/map/global.rs |   3 +
 src/util/psbt/map/input.rs  |   4 ++
 src/util/psbt/map/output.rs |   3 +
 src/util/psbt/mod.rs        |  94 +++++++++++++++++++++++++++++++
 src/util/psbt/raw.rs        |   4 ++
 7 files changed, 218 insertions(+)
 create mode 100644 src/serde_utils.rs

diff --git a/src/lib.rs b/src/lib.rs
index 5f61648e..7007cf9e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -57,6 +57,9 @@ compile_error!("rust-bitcoin cannot be used on 16-bit architectures");
 mod test_macros;
 #[macro_use]
 mod internal_macros;
+#[cfg(feature = "serde")]
+mod serde_utils;
+
 #[macro_use]
 pub mod network;
 pub mod blockdata;
diff --git a/src/serde_utils.rs b/src/serde_utils.rs
new file mode 100644
index 00000000..0c0b243d
--- /dev/null
+++ b/src/serde_utils.rs
@@ -0,0 +1,107 @@
+
+//! Module for special serde serializations.
+
+pub mod btreemap {
+    //! Module for serialization of BTreeMaps because serde_json will
+    //! not serialize hashmaps with non-string keys be default.
+    #![allow(missing_docs)]
+
+    // NOTE: This module can be exactly copied to use with HashMap.
+
+    use ::std::collections::BTreeMap;
+    use serde;
+
+    pub fn serialize<S, T, U>(v: &BTreeMap<T, U>, s: S)
+        -> Result<S::Ok, S::Error> where
+        S: serde::Serializer,
+        T: serde::Serialize + ::std::hash::Hash + Eq + Ord,
+        U: serde::Serialize,
+    {
+        use serde::ser::SerializeSeq;
+
+        let mut seq = s.serialize_seq(Some(v.len()))?;
+        for pair in v.iter() {
+            seq.serialize_element(&pair)?;
+        }
+        seq.end()
+    }
+
+    pub fn deserialize<'de, D, T, U>(d: D)
+        -> Result<BTreeMap<T, U>, D::Error> where
+        D: serde::Deserializer<'de>,
+        T: serde::Deserialize<'de> + ::std::hash::Hash + Eq + Ord,
+        U: serde::Deserialize<'de>,
+    {
+        use ::std::marker::PhantomData;
+
+        struct Visitor<T, U>(PhantomData<(T, U)>);
+        impl<'de, T, U> serde::de::Visitor<'de> for Visitor<T, U> where
+            T: serde::Deserialize<'de> + ::std::hash::Hash + Eq + Ord,
+            U: serde::Deserialize<'de>,
+        {
+            type Value = BTreeMap<T, U>;
+
+            fn expecting(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
+                write!(f, "a sequence of pairs")
+            }
+
+            fn visit_seq<A: serde::de::SeqAccess<'de>>(self, mut a: A)
+                -> Result<Self::Value, A::Error>
+            {
+                let mut ret = BTreeMap::new();
+                while let Some((key, value)) = a.next_element()? {
+                    ret.insert(key, value);
+                }
+                Ok(ret)
+            }
+        }
+
+        d.deserialize_seq(Visitor(PhantomData))
+    }
+}
+
+pub mod hex_bytes {
+    //! Module for serialization of byte arrays as hex strings.
+    #![allow(missing_docs)]
+
+    use hashes::hex::{FromHex, ToHex};
+    use serde;
+
+    pub fn serialize<T, S>(bytes: &T, serializer: S) -> Result<S::Ok, S::Error>
+        where T: AsRef<[u8]>, S: serde::Serializer
+    {
+        serializer.serialize_str(&bytes.as_ref().to_hex())
+    }
+
+    pub fn deserialize<'de, D, B>(d: D) -> Result<B, D::Error>
+        where D: serde::Deserializer<'de>, B: FromHex,
+    {
+        struct Visitor<B>(::std::marker::PhantomData<B>);
+
+        impl<'de, B: FromHex> serde::de::Visitor<'de> for Visitor<B> {
+            type Value = B;
+
+            fn expecting(&self, formatter: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
+                formatter.write_str("an ASCII hex string")
+            }
+
+            fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
+                where E: serde::de::Error,
+            {
+                if let Ok(hex) = ::std::str::from_utf8(v) {
+                    FromHex::from_hex(hex).map_err(E::custom)
+                } else {
+                    return Err(E::invalid_value(serde::de::Unexpected::Bytes(v), &self));
+                }
+            }
+
+            fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
+                where E: serde::de::Error,
+            {
+                FromHex::from_hex(v).map_err(E::custom)
+            }
+        }
+
+        d.deserialize_str(Visitor(::std::marker::PhantomData))
+    }
+}
diff --git a/src/util/psbt/map/global.rs b/src/util/psbt/map/global.rs
index 1897e49b..aa5d6254 100644
--- a/src/util/psbt/map/global.rs
+++ b/src/util/psbt/map/global.rs
@@ -46,10 +46,13 @@ pub struct Global {
     pub version: u32,
     /// A global map from extended public keys to the used key fingerprint and
     /// derivation path as defined by BIP 32
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap"))]
     pub xpub: BTreeMap<ExtendedPubKey, KeySource>,
     /// Global proprietary key-value pairs.
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap"))]
     pub proprietary: BTreeMap<raw::ProprietaryKey, Vec<u8>>,
     /// Unknown global key-value pairs.
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap"))]
     pub unknown: BTreeMap<raw::Key, Vec<u8>>,
 }
 
diff --git a/src/util/psbt/map/input.rs b/src/util/psbt/map/input.rs
index 0e79fc2c..112eae9a 100644
--- a/src/util/psbt/map/input.rs
+++ b/src/util/psbt/map/input.rs
@@ -70,6 +70,7 @@ pub struct Input {
     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.
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap"))]
     pub partial_sigs: BTreeMap<PublicKey, Vec<u8>>,
     /// The sighash type to be used for this input. Signatures for this input
     /// must use the sighash type.
@@ -80,6 +81,7 @@ pub struct 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.
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap"))]
     pub bip32_derivation: BTreeMap<PublicKey, KeySource>,
     /// The finalized, fully-constructed scriptSig with signatures and any other
     /// scripts necessary for this input to pass validation.
@@ -97,8 +99,10 @@ pub struct Input {
     /// HAS256 hash to preimage map
     pub hash256_preimages: BTreeMap<sha256d::Hash, Vec<u8>>,
     /// Proprietary key-value pairs for this input.
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap"))]
     pub proprietary: BTreeMap<raw::ProprietaryKey, Vec<u8>>,
     /// Unknown key-value pairs for this input.
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap"))]
     pub unknown: BTreeMap<raw::Key, Vec<u8>>,
 }
 
diff --git a/src/util/psbt/map/output.rs b/src/util/psbt/map/output.rs
index 2577ed87..34861ba5 100644
--- a/src/util/psbt/map/output.rs
+++ b/src/util/psbt/map/output.rs
@@ -44,10 +44,13 @@ pub struct Output {
     pub witness_script: Option<Script>,
     /// A map from public keys needed to spend this output to their
     /// corresponding master key fingerprints and derivation paths.
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap"))]
     pub bip32_derivation: BTreeMap<PublicKey, KeySource>,
     /// Proprietary key-value pairs for this output.
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap"))]
     pub proprietary: BTreeMap<raw::ProprietaryKey, Vec<u8>>,
     /// Unknown key-value pairs for this output.
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::btreemap"))]
     pub unknown: BTreeMap<raw::Key, Vec<u8>>,
 }
 
diff --git a/src/util/psbt/mod.rs b/src/util/psbt/mod.rs
index 7de1109c..feadf195 100644
--- a/src/util/psbt/mod.rs
+++ b/src/util/psbt/mod.rs
@@ -315,6 +315,100 @@ mod tests {
         assert_eq!(hex, serialize_hex(&psbt));
     }
 
+    #[cfg(feature = "serde")]
+    #[test]
+    fn test_serde_psbt() {
+        //! Create a full PSBT value with various fields filled and make sure it can be JSONized.
+        use util::psbt::map::Input;
+
+        // create some values to use in the PSBT
+        let tx = Transaction {
+            version: 1,
+            lock_time: 0,
+            input: vec![TxIn {
+                previous_output: OutPoint {
+                    txid: Txid::from_hex("e567952fb6cc33857f392efa3a46c995a28f69cca4bb1b37e0204dab1ec7a389").unwrap(),
+                    vout: 1,
+                },
+                script_sig: hex_script!("160014be18d152a9b012039daf3da7de4f53349eecb985"),
+                sequence: 4294967295,
+                witness: vec![Vec::from_hex("03d2e15674941bad4a996372cb87e1856d3652606d98562fe39c5e9e7e413f2105").unwrap()],
+            }],
+            output: vec![
+                TxOut {
+                    value: 190303501938,
+                    script_pubkey: hex_script!("a914339725ba21efd62ac753a9bcd067d6c7a6a39d0587"),
+                },
+            ],
+        };
+        let unknown: BTreeMap<raw::Key, Vec<u8>> = vec![(
+            raw::Key { type_value: 1, key: vec![0, 1] },
+            vec![3, 4 ,5],
+        )].into_iter().collect();
+        let key_source = ("deadbeef".parse().unwrap(), "m/0'/1".parse().unwrap());
+        let keypaths: BTreeMap<PublicKey, KeySource> = vec![(
+            "0339880dc92394b7355e3d0439fa283c31de7590812ea011c4245c0674a685e883".parse().unwrap(),
+            key_source.clone(),
+        )].into_iter().collect();
+
+        let proprietary: BTreeMap<raw::ProprietaryKey, Vec<u8>> = vec![(
+            raw::ProprietaryKey {
+                prefix: "prefx".as_bytes().to_vec(),
+                subtype: 42,
+                key: "test_key".as_bytes().to_vec(),
+            },
+            vec![5, 6, 7],
+        )].into_iter().collect();
+
+        let psbt = PartiallySignedTransaction {
+            global: Global {
+                version: 0,
+                xpub: {
+                    let xpub: ExtendedPubKey =
+                        "xpub661MyMwAqRbcGoRVtwfvzZsq2VBJR1LAHfQstHUoxqDorV89vRoMxUZ27kLrraAj6MPi\
+                        QfrDb27gigC1VS1dBXi5jGpxmMeBXEkKkcXUTg4".parse().unwrap();
+                    vec![(xpub, key_source.clone())].into_iter().collect()
+                },
+                unsigned_tx: {
+                    let mut unsigned = tx.clone();
+                    unsigned.input[0].script_sig = Script::new();
+                    unsigned.input[0].witness = Vec::new();
+                    unsigned
+                },
+                proprietary: proprietary.clone(),
+                unknown: unknown.clone(),
+            },
+            inputs: vec![Input {
+                non_witness_utxo: Some(tx),
+                witness_utxo: Some(TxOut {
+                    value: 190303501938,
+                    script_pubkey: hex_script!("a914339725ba21efd62ac753a9bcd067d6c7a6a39d0587"),
+                }),
+                sighash_type: Some("SIGHASH_SINGLE|SIGHASH_ANYONECANPAY".parse().unwrap()),
+                redeem_script: Some(vec![0x51].into()),
+                witness_script: None,
+                partial_sigs: vec![(
+                    "0339880dc92394b7355e3d0439fa283c31de7590812ea011c4245c0674a685e883".parse().unwrap(),
+                    vec![8, 5, 4],
+                )].into_iter().collect(),
+                bip32_derivation: keypaths.clone(),
+                final_script_witness: Some(vec![vec![1, 3], vec![5]]),
+                proprietary: proprietary.clone(),
+                unknown: unknown.clone(),
+                ..Default::default()
+            }],
+            outputs: vec![Output {
+                bip32_derivation: keypaths.clone(),
+                proprietary: proprietary.clone(),
+                unknown: unknown.clone(),
+                ..Default::default()
+            }],
+        };
+        let encoded = ::serde_json::to_string(&psbt).unwrap();
+        let decoded: PartiallySignedTransaction = ::serde_json::from_str(&encoded).unwrap();
+        assert_eq!(psbt, decoded);
+    }
+
     mod bip_vectors {
         use std::collections::BTreeMap;
 
diff --git a/src/util/psbt/raw.rs b/src/util/psbt/raw.rs
index af2e5963..74290b31 100644
--- a/src/util/psbt/raw.rs
+++ b/src/util/psbt/raw.rs
@@ -30,6 +30,7 @@ pub struct Key {
     /// The type of this PSBT key.
     pub type_value: u8,
     /// The key itself in raw byte form.
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::hex_bytes"))]
     pub key: Vec<u8>,
 }
 
@@ -40,6 +41,7 @@ pub struct Pair {
     /// The key of this key-value pair.
     pub key: Key,
     /// The value of this key-value pair in raw byte form.
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::hex_bytes"))]
     pub value: Vec<u8>,
 }
 
@@ -53,10 +55,12 @@ pub type ProprietaryType = u8;
 pub struct ProprietaryKey<Subtype = ProprietaryType> where Subtype: Copy + From<u8> + Into<u8> {
     /// Proprietary type prefix used for grouping together keys under some
     /// application and avoid namespace collision
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::hex_bytes"))]
     pub prefix: Vec<u8>,
     /// Custom proprietary subtype
     pub subtype: Subtype,
     /// Additional key bytes (like serialized public key data etc)
+    #[cfg_attr(feature = "serde", serde(with = "::serde_utils::hex_bytes"))]
     pub key: Vec<u8>,
 }