PSBT: proprietary keys type system according to BIP 174

src/util/psbt/error.rs

@@ -16,6 +16,7 @@ use std::error;
 use std::fmt;
 
 use blockdata::transaction::Transaction;
+use consensus::encode;
 use util::psbt::raw;
 
 use hashes;
@@ -38,6 +39,8 @@ pub enum Error {
     InvalidSeparator,
     /// Known keys must be according to spec.
     InvalidKey(raw::Key),
+    /// Non-proprietary key type found when proprietary key was expected
+    InvalidProprietaryKey,
     /// Keys within key-value map should never be duplicated.
     DuplicateKey(raw::Key),
     /// The scriptSigs for the unsigned transaction must be empty.
@@ -71,12 +74,15 @@ pub enum Error {
     },
     /// Data inconsistency/conflicting data during merge procedure
     MergeConflict(String),
+    /// Serialization error in bitcoin consensus-encoded structures
+    ConsensusEncoding,
 }
 
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match *self {
             Error::InvalidKey(ref rkey) => write!(f, "invalid key: {}", rkey),
+            Error::InvalidProprietaryKey => write!(f, "non-proprietary key type found when proprietary key was expected"),
             Error::DuplicateKey(ref rkey) => write!(f, "duplicate key: {}", rkey),
             Error::UnexpectedUnsignedTx { expected: ref e, actual: ref a } => write!(f, "different unsigned transaction: expected {}, actual {}", e.txid(), a.txid()),
             Error::NonStandardSigHashType(ref sht) => write!(f, "non-standard sighash type: {}", sht),
@@ -94,6 +100,7 @@ impl fmt::Display for Error {
                 write!(f, "Preimage {:?} does not match {:?} hash {:?}", preimage, hash_type, hash )
             }
             Error::MergeConflict(ref s) => { write!(f, "Merge conflict: {}", s) }
+            Error::ConsensusEncoding => f.write_str("bitcoin consensus encoding error"),
         }
     }
 }
@@ -106,3 +113,12 @@ impl From<hashes::Error> for Error {
         Error::HashParseError(e)
     }
 }
+
+impl From<encode::Error> for Error {
+    fn from(err: encode::Error) -> Self {
+        match err {
+            encode::Error::Psbt(err) => err,
+            _ => Error::ConsensusEncoding,
+        }
+    }
+}

src/util/psbt/map/global.rs

@@ -47,7 +47,7 @@ pub struct Global {
     /// derivation path as defined by BIP 32
     pub xpub: BTreeMap<ExtendedPubKey, KeySource>,
     /// Global proprietary key-value pairs.
-    pub proprietary: BTreeMap<raw::Key, Vec<u8>>,
+    pub proprietary: BTreeMap<raw::ProprietaryKey, Vec<u8>>,
     /// Unknown global key-value pairs.
     pub unknown: BTreeMap<raw::Key, Vec<u8>>,
 }
@@ -85,9 +85,9 @@ impl Map for Global {
 
         match raw_key.type_value {
             PSBT_GLOBAL_UNSIGNED_TX => return Err(Error::DuplicateKey(raw_key).into()),
-            PSBT_GLOBAL_PROPRIETARY => match self.proprietary.entry(raw_key) {
+            PSBT_GLOBAL_PROPRIETARY => match self.proprietary.entry(raw::ProprietaryKey::from_key(raw_key.clone())?) {
                 Entry::Vacant(empty_key) => {empty_key.insert(raw_value);},
-                Entry::Occupied(k) => return Err(Error::DuplicateKey(k.key().clone()).into()),
+                Entry::Occupied(_) => return Err(Error::DuplicateKey(raw_key).into()),
             }
             _ => match self.unknown.entry(raw_key) {
                 Entry::Vacant(empty_key) => {empty_key.insert(raw_value);},
@@ -146,7 +146,7 @@ impl Map for Global {
 
         for (key, value) in self.proprietary.iter() {
             rv.push(raw::Pair {
-                key: key.clone(),
+                key: key.to_key(),
                 value: value.clone(),
             });
         }
@@ -232,7 +232,7 @@ impl Decodable for Global {
         let mut version: Option<u32> = None;
         let mut unknowns: BTreeMap<raw::Key, Vec<u8>> = Default::default();
         let mut xpub_map: BTreeMap<ExtendedPubKey, (Fingerprint, DerivationPath)> = Default::default();
-        let mut proprietary: BTreeMap<raw::Key, Vec<u8>> = Default::default();
+        let mut proprietary: BTreeMap<raw::ProprietaryKey, Vec<u8>> = Default::default();
 
         loop {
             match raw::Pair::consensus_decode(&mut d) {
@@ -317,9 +317,9 @@ impl Decodable for Global {
                                 return Err(Error::InvalidKey(pair.key).into())
                             }
                         }
-                        PSBT_GLOBAL_PROPRIETARY => match proprietary.entry(pair.key) {
+                        PSBT_GLOBAL_PROPRIETARY => match proprietary.entry(raw::ProprietaryKey::from_key(pair.key.clone())?) {
                             Entry::Vacant(empty_key) => {empty_key.insert(pair.value);},
-                            Entry::Occupied(k) => return Err(Error::DuplicateKey(k.key().clone()).into()),
+                            Entry::Occupied(_) => return Err(Error::DuplicateKey(pair.key).into()),
                         }
                         _ => match unknowns.entry(pair.key) {
                             Entry::Vacant(empty_key) => {empty_key.insert(pair.value);},

src/util/psbt/map/input.rs

@@ -96,7 +96,7 @@ pub struct Input {
     /// HAS256 hash to preimage map
     pub hash256_preimages: BTreeMap<sha256d::Hash, Vec<u8>>,
     /// Proprietary key-value pairs for this input.
-    pub proprietary: BTreeMap<raw::Key, Vec<u8>>,
+    pub proprietary: BTreeMap<raw::ProprietaryKey, Vec<u8>>,
     /// Unknown key-value pairs for this input.
     pub unknown: BTreeMap<raw::Key, Vec<u8>>,
 }
@@ -173,9 +173,9 @@ impl Map for Input {
             PSBT_IN_HASH256 => {
                 psbt_insert_hash_pair(&mut self.hash256_preimages, raw_key, raw_value, error::PsbtHash::Hash256)?;
             }
-            PSBT_IN_PROPRIETARY => match self.proprietary.entry(raw_key) {
+            PSBT_IN_PROPRIETARY => match self.proprietary.entry(raw::ProprietaryKey::from_key(raw_key.clone())?) {
                 ::std::collections::btree_map::Entry::Vacant(empty_key) => {empty_key.insert(raw_value);},
-                ::std::collections::btree_map::Entry::Occupied(k) => return Err(Error::DuplicateKey(k.key().clone()).into()),
+                ::std::collections::btree_map::Entry::Occupied(_) => return Err(Error::DuplicateKey(raw_key).into()),
             }
             _ => match self.unknown.entry(raw_key) {
                 ::std::collections::btree_map::Entry::Vacant(empty_key) => {
@@ -247,7 +247,7 @@ impl Map for Input {
 
         for (key, value) in self.proprietary.iter() {
             rv.push(raw::Pair {
-                key: key.clone(),
+                key: key.to_key(),
                 value: value.clone(),
             });
         }

src/util/psbt/map/output.rs

@@ -45,7 +45,7 @@ pub struct Output {
     /// corresponding master key fingerprints and derivation paths.
     pub bip32_derivation: BTreeMap<PublicKey, KeySource>,
     /// Proprietary key-value pairs for this output.
-    pub proprietary: BTreeMap<raw::Key, Vec<u8>>,
+    pub proprietary: BTreeMap<raw::ProprietaryKey, Vec<u8>>,
     /// Unknown key-value pairs for this output.
     pub unknown: BTreeMap<raw::Key, Vec<u8>>,
 }
@@ -76,9 +76,9 @@ impl Map for Output {
                     self.bip32_derivation <= <raw_key: PublicKey>|<raw_value: KeySource>
                 }
             }
-            PSBT_OUT_PROPRIETARY => match self.proprietary.entry(raw_key) {
+            PSBT_OUT_PROPRIETARY => match self.proprietary.entry(raw::ProprietaryKey::from_key(raw_key.clone())?) {
                 Entry::Vacant(empty_key) => {empty_key.insert(raw_value);},
-                Entry::Occupied(k) => return Err(Error::DuplicateKey(k.key().clone()).into()),
+                Entry::Occupied(_) => return Err(Error::DuplicateKey(raw_key.clone()).into()),
             }
             _ => match self.unknown.entry(raw_key) {
                     Entry::Vacant(empty_key) => {empty_key.insert(raw_value);},
@@ -106,7 +106,7 @@ impl Map for Output {
 
         for (key, value) in self.proprietary.iter() {
             rv.push(raw::Pair {
-                key: key.clone(),
+                key: key.to_key(),
                 value: value.clone(),
             });
         }

src/util/psbt/raw.rs

@@ -19,7 +19,7 @@
 
 use std::{fmt, io};
 
-use consensus::encode::{self, Decodable, Encodable, VarInt, MAX_VEC_SIZE};
+use consensus::encode::{self, ReadExt, WriteExt, Decodable, Encodable, VarInt, serialize, deserialize, MAX_VEC_SIZE};
 use hashes::hex::ToHex;
 use util::psbt::Error;
 
@@ -43,6 +43,23 @@ pub struct Pair {
 }
 serde_struct_impl!(Pair, key, value);
 
+/// Default implementation for proprietary key subtyping
+pub type ProprietaryType = u8;
+
+/// Proprietary keys (i.e. keys starting with 0xFC byte) with their internal
+/// structure according to BIP 174.
+#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
+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
+    pub prefix: Vec<u8>,
+    /// Custom proprietary subtype
+    pub subtype: Subtype,
+    /// Additional key bytes (like serialized public key data etc)
+    pub key: Vec<u8>,
+}
+serde_struct_impl!(ProprietaryKey, prefix, subtype, key);
+
 impl fmt::Display for Key {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(
@@ -121,3 +138,47 @@ impl Decodable for Pair {
         })
     }
 }
+
+impl<Subtype> Encodable for ProprietaryKey<Subtype> where Subtype: Copy + From<u8> + Into<u8> {
+    fn consensus_encode<W: io::Write>(&self, mut e: W) -> Result<usize, encode::Error> {
+        let mut len = self.prefix.consensus_encode(&mut e)? + 1;
+        e.emit_u8(self.subtype.into())?;
+        len += e.write(&self.key)?;
+        Ok(len)
+    }
+}
+
+impl<Subtype> Decodable for ProprietaryKey<Subtype> where Subtype: Copy + From<u8> + Into<u8> {
+    fn consensus_decode<D: io::Read>(mut d: D) -> Result<Self, encode::Error> {
+        let prefix = Vec::<u8>::consensus_decode(&mut d)?;
+        let mut key = vec![];
+        let subtype = Subtype::from(d.read_u8()?);
+        d.read_to_end(&mut key)?;
+
+        Ok(ProprietaryKey {
+            prefix,
+            subtype,
+            key
+        })
+    }
+}
+
+impl<Subtype> ProprietaryKey<Subtype> where Subtype: Copy + From<u8> + Into<u8> {
+    /// Constructs [ProprietaryKey] from [Key]; returns
+    /// [Error::InvalidProprietaryKey] if `key` do not starts with 0xFC byte
+    pub fn from_key(key: Key) -> Result<Self, Error> {
+        if key.type_value != 0xFC {
+            return Err(Error::InvalidProprietaryKey)
+        }
+
+        Ok(deserialize(&key.key)?)
+    }
+
+    /// Constructs full [Key] corresponding to this proprietary key type
+    pub fn to_key(&self) -> Key {
+        Key {
+            type_value: 0xFC,
+            key: serialize(self)
+        }
+    }
+}