From e882a3a32dccfc009b8443993e9e9bbd5cdb24e9 Mon Sep 17 00:00:00 2001 From: Andrew Chow Date: Wed, 2 Oct 2019 15:49:33 -0400 Subject: [PATCH] Types are compact size uints --- src/psbt.h | 41 +++++++++++++++++++---------------- src/script/sign.h | 2 +- src/serialize.h | 13 +++++++++++ src/test/fuzz/script_sign.cpp | 4 ++-- 4 files changed, 38 insertions(+), 22 deletions(-) diff --git a/src/psbt.h b/src/psbt.h index 7d46334ba7..ef5e10e2c0 100644 --- a/src/psbt.h +++ b/src/psbt.h @@ -68,52 +68,52 @@ struct PSBTInput inline void Serialize(Stream& s) const { // Write the utxo if (non_witness_utxo) { - SerializeToVector(s, PSBT_IN_NON_WITNESS_UTXO); + SerializeToVector(s, CompactSizeWriter(PSBT_IN_NON_WITNESS_UTXO)); OverrideStream os(&s, s.GetType(), s.GetVersion() | SERIALIZE_TRANSACTION_NO_WITNESS); SerializeToVector(os, non_witness_utxo); } if (!witness_utxo.IsNull()) { - SerializeToVector(s, PSBT_IN_WITNESS_UTXO); + SerializeToVector(s, CompactSizeWriter(PSBT_IN_WITNESS_UTXO)); SerializeToVector(s, witness_utxo); } if (final_script_sig.empty() && final_script_witness.IsNull()) { // Write any partial signatures for (auto sig_pair : partial_sigs) { - SerializeToVector(s, PSBT_IN_PARTIAL_SIG, MakeSpan(sig_pair.second.first)); + SerializeToVector(s, CompactSizeWriter(PSBT_IN_PARTIAL_SIG), MakeSpan(sig_pair.second.first)); s << sig_pair.second.second; } // Write the sighash type if (sighash_type > 0) { - SerializeToVector(s, PSBT_IN_SIGHASH); + SerializeToVector(s, CompactSizeWriter(PSBT_IN_SIGHASH)); SerializeToVector(s, sighash_type); } // Write the redeem script if (!redeem_script.empty()) { - SerializeToVector(s, PSBT_IN_REDEEMSCRIPT); + SerializeToVector(s, CompactSizeWriter(PSBT_IN_REDEEMSCRIPT)); s << redeem_script; } // Write the witness script if (!witness_script.empty()) { - SerializeToVector(s, PSBT_IN_WITNESSSCRIPT); + SerializeToVector(s, CompactSizeWriter(PSBT_IN_WITNESSSCRIPT)); s << witness_script; } // Write any hd keypaths - SerializeHDKeypaths(s, hd_keypaths, PSBT_IN_BIP32_DERIVATION); + SerializeHDKeypaths(s, hd_keypaths, CompactSizeWriter(PSBT_IN_BIP32_DERIVATION)); } // Write script sig if (!final_script_sig.empty()) { - SerializeToVector(s, PSBT_IN_SCRIPTSIG); + SerializeToVector(s, CompactSizeWriter(PSBT_IN_SCRIPTSIG)); s << final_script_sig; } // write script witness if (!final_script_witness.IsNull()) { - SerializeToVector(s, PSBT_IN_SCRIPTWITNESS); + SerializeToVector(s, CompactSizeWriter(PSBT_IN_SCRIPTWITNESS)); SerializeToVector(s, final_script_witness.stack); } @@ -146,8 +146,9 @@ struct PSBTInput break; } - // First byte of key is the type - unsigned char type = key[0]; + // Type is compact size uint at beginning of key + VectorReader skey(s.GetType(), s.GetVersion(), key, 0); + uint64_t type = ReadCompactSize(skey); // Do stuff based on type switch(type) { @@ -289,18 +290,18 @@ struct PSBTOutput inline void Serialize(Stream& s) const { // Write the redeem script if (!redeem_script.empty()) { - SerializeToVector(s, PSBT_OUT_REDEEMSCRIPT); + SerializeToVector(s, CompactSizeWriter(PSBT_OUT_REDEEMSCRIPT)); s << redeem_script; } // Write the witness script if (!witness_script.empty()) { - SerializeToVector(s, PSBT_OUT_WITNESSSCRIPT); + SerializeToVector(s, CompactSizeWriter(PSBT_OUT_WITNESSSCRIPT)); s << witness_script; } // Write any hd keypaths - SerializeHDKeypaths(s, hd_keypaths, PSBT_OUT_BIP32_DERIVATION); + SerializeHDKeypaths(s, hd_keypaths, CompactSizeWriter(PSBT_OUT_BIP32_DERIVATION)); // Write unknown things for (auto& entry : unknown) { @@ -331,8 +332,9 @@ struct PSBTOutput break; } - // First byte of key is the type - unsigned char type = key[0]; + // Type is compact size uint at beginning of key + VectorReader skey(s.GetType(), s.GetVersion(), key, 0); + uint64_t type = ReadCompactSize(skey); // Do stuff based on type switch(type) { @@ -419,7 +421,7 @@ struct PartiallySignedTransaction s << PSBT_MAGIC_BYTES; // unsigned tx flag - SerializeToVector(s, PSBT_GLOBAL_UNSIGNED_TX); + SerializeToVector(s, CompactSizeWriter(PSBT_GLOBAL_UNSIGNED_TX)); // Write serialized tx to a stream OverrideStream os(&s, s.GetType(), s.GetVersion() | SERIALIZE_TRANSACTION_NO_WITNESS); @@ -471,8 +473,9 @@ struct PartiallySignedTransaction break; } - // First byte of key is the type - unsigned char type = key[0]; + // Type is compact size uint at beginning of key + VectorReader skey(s.GetType(), s.GetVersion(), key, 0); + uint64_t type = ReadCompactSize(skey); // Do stuff based on type switch(type) { diff --git a/src/script/sign.h b/src/script/sign.h index 4f280af234..7a034980f0 100644 --- a/src/script/sign.h +++ b/src/script/sign.h @@ -135,7 +135,7 @@ void DeserializeHDKeypaths(Stream& s, const std::vector& key, std // Serialize HD keypaths to a stream from a map template -void SerializeHDKeypaths(Stream& s, const std::map& hd_keypaths, uint8_t type) +void SerializeHDKeypaths(Stream& s, const std::map& hd_keypaths, CompactSizeWriter type) { for (auto keypath_pair : hd_keypaths) { if (!keypath_pair.first.IsValid()) { diff --git a/src/serialize.h b/src/serialize.h index 748fdc69a8..e6142648e8 100644 --- a/src/serialize.h +++ b/src/serialize.h @@ -563,6 +563,19 @@ struct CompactSizeFormatter } }; +class CompactSizeWriter +{ +protected: + uint64_t n; +public: + explicit CompactSizeWriter(uint64_t n_in) : n(n_in) { } + + template + void Serialize(Stream &s) const { + WriteCompactSize(s, n); + } +}; + template struct LimitedStringFormatter { diff --git a/src/test/fuzz/script_sign.cpp b/src/test/fuzz/script_sign.cpp index 29fa183ff7..0ebe9dc041 100644 --- a/src/test/fuzz/script_sign.cpp +++ b/src/test/fuzz/script_sign.cpp @@ -42,7 +42,7 @@ void test_one_input(const std::vector& buffer) } catch (const std::ios_base::failure&) { } CDataStream serialized{SER_NETWORK, PROTOCOL_VERSION}; - SerializeHDKeypaths(serialized, hd_keypaths, fuzzed_data_provider.ConsumeIntegral()); + SerializeHDKeypaths(serialized, hd_keypaths, CompactSizeWriter(fuzzed_data_provider.ConsumeIntegral())); } { @@ -60,7 +60,7 @@ void test_one_input(const std::vector& buffer) } CDataStream serialized{SER_NETWORK, PROTOCOL_VERSION}; try { - SerializeHDKeypaths(serialized, hd_keypaths, fuzzed_data_provider.ConsumeIntegral()); + SerializeHDKeypaths(serialized, hd_keypaths, CompactSizeWriter(fuzzed_data_provider.ConsumeIntegral())); } catch (const std::ios_base::failure&) { } std::map deserialized_hd_keypaths;