From dd792cea62b1a69fa9e1fc04a46fd2b2310e7675 Mon Sep 17 00:00:00 2001
From: Andrew Chow <achow101-github@achow101.com>
Date: Mon, 4 Jan 2021 16:07:25 -0500
Subject: [PATCH] Implement PSBTv2 fields de/ser

---
 src/psbt.h      | 219 +++++++++++++++++++++++++++++++++++++++++++++++-
 src/serialize.h |  13 +++
 2 files changed, 228 insertions(+), 4 deletions(-)

diff --git a/src/psbt.h b/src/psbt.h
index 11401e2de9..8d91eff7e4 100644
--- a/src/psbt.h
+++ b/src/psbt.h
@@ -14,6 +14,8 @@
 #include <script/sign.h>
 #include <script/signingprovider.h>
 
+#include <bitset>
+
 // Magic bytes
 static constexpr uint8_t PSBT_MAGIC_BYTES[5] = {'p', 's', 'b', 't', 0xff};
 
@@ -179,6 +181,11 @@ struct PSBTInput
     CScriptWitness final_script_witness;
     std::map<CPubKey, KeyOriginInfo> hd_keypaths;
     std::map<CKeyID, SigPair> partial_sigs;
+    uint256 prev_txid;
+    Optional<uint32_t> prev_out;
+    Optional<uint32_t> sequence;
+    Optional<uint32_t> time_locktime;
+    Optional<uint32_t> height_locktime;
     std::map<std::vector<unsigned char>, std::vector<unsigned char>> unknown;
     std::set<PSBTProprietary> m_proprietary;
     int sighash_type = 0;
@@ -242,6 +249,28 @@ struct PSBTInput
             SerializeToVector(s, final_script_witness.stack);
         }
 
+        // Write prev txid, vout, sequence, and lock times
+        if (!prev_txid.IsNull()) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_IN_PREVIOUS_TXID));
+            SerializeToVector(s, prev_txid);
+        }
+        if (prev_out != nullopt) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_IN_OUTPUT_INDEX));
+            SerializeToVector(s, *prev_out);
+        }
+        if (sequence != nullopt) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_IN_SEQUENCE));
+            SerializeToVector(s, *sequence);
+        }
+        if (time_locktime != nullopt) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_IN_REQUIRED_TIME_LOCKTIME));
+            SerializeToVector(s, *time_locktime);
+        }
+        if (height_locktime != nullopt) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_IN_REQUIRED_HEIGHT_LOCKTIME));
+            SerializeToVector(s, *height_locktime);
+        }
+
         // Write proprietary things
         for (const auto& entry : m_proprietary) {
             s << entry.key;
@@ -379,6 +408,64 @@ struct PSBTInput
                     UnserializeFromVector(s, final_script_witness.stack);
                     break;
                 }
+                case PSBT_IN_PREVIOUS_TXID:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, previous txid is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Previous txid key is more than one byte type");
+                    }
+                    UnserializeFromVector(s, prev_txid);
+                    break;
+                }
+                case PSBT_IN_OUTPUT_INDEX:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, previous output's index is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Previous output's index is more than one byte type");
+                    }
+                    uint32_t v;
+                    UnserializeFromVector(s, v);
+                    prev_out = v;
+                    break;
+                }
+                case PSBT_IN_SEQUENCE:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, sequence is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Sequence key is more than one byte type");
+                    }
+                    uint32_t v;
+                    UnserializeFromVector(s, v);
+                    sequence = v;
+                    break;
+                }
+                case PSBT_IN_REQUIRED_TIME_LOCKTIME:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, required time based locktime is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Required time based locktime is more than one byte type");
+                    }
+                    uint32_t v;
+                    UnserializeFromVector(s, v);
+                    time_locktime = v;
+                    break;
+                }
+                case PSBT_IN_REQUIRED_HEIGHT_LOCKTIME:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, required height based locktime is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Required height based locktime is more than one byte type");
+                    }
+                    uint32_t v;
+                    UnserializeFromVector(s, v);
+                    height_locktime = v;
+                    break;
+                }
                 case PSBT_IN_PROPRIETARY:
                 {
                     PSBTProprietary this_prop;
@@ -423,6 +510,8 @@ struct PSBTOutput
     CScript redeem_script;
     CScript witness_script;
     std::map<CPubKey, KeyOriginInfo> hd_keypaths;
+    Optional<CAmount> amount;
+    CScript script;
     std::map<std::vector<unsigned char>, std::vector<unsigned char>> unknown;
     std::set<PSBTProprietary> m_proprietary;
 
@@ -449,6 +538,16 @@ struct PSBTOutput
         // Write any hd keypaths
         SerializeHDKeypaths(s, hd_keypaths, CompactSizeWriter(PSBT_OUT_BIP32_DERIVATION));
 
+        // Write amount and spk
+        if (amount != nullopt) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_OUT_AMOUNT));
+            SerializeToVector(s, *amount);
+        }
+        if (!script.empty()) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_OUT_SCRIPT));
+            s << script;
+        }
+
         // Write proprietary things
         for (const auto& entry : m_proprietary) {
             s << entry.key;
@@ -515,6 +614,28 @@ struct PSBTOutput
                     DeserializeHDKeypaths(s, key, hd_keypaths);
                     break;
                 }
+                case PSBT_OUT_AMOUNT:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, output amount is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Output amount key is more than one byte type");
+                    }
+                    CAmount v;
+                    UnserializeFromVector(s, v);
+                    amount = v;
+                    break;
+                }
+                case PSBT_OUT_SCRIPT:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, output script is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Output script key is more than one byte type");
+                    }
+                    s >> script;
+                    break;
+                }
                 case PSBT_OUT_PROPRIETARY:
                 {
                     PSBTProprietary this_prop;
@@ -561,6 +682,9 @@ struct PartiallySignedTransaction
     // We use a set of CExtPubKey in the event that there happens to be the same KeyOriginInfos for different CExtPubKeys
     // Note that this map swaps the key and values from the serialization
     std::map<KeyOriginInfo, std::set<CExtPubKey>> m_xpubs;
+    Optional<int32_t> tx_version;
+    Optional<uint32_t> fallback_locktime;
+    Optional<std::bitset<8>> m_tx_modifiable;
     std::vector<PSBTInput> inputs;
     std::vector<PSBTOutput> outputs;
     std::map<std::vector<unsigned char>, std::vector<unsigned char>> unknown;
@@ -611,6 +735,26 @@ struct PartiallySignedTransaction
             }
         }
 
+        // Write PSBTv2 tx version, locktime, counts, etc.
+        if (tx_version != nullopt) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_GLOBAL_TX_VERSION));
+            SerializeToVector(s, *tx_version);
+        }
+        if (fallback_locktime != nullopt) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_GLOBAL_FALLBACK_LOCKTIME));
+            SerializeToVector(s, *fallback_locktime);
+        }
+        if (m_version != nullopt && *m_version >= 2) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_GLOBAL_INPUT_COUNT));
+            SerializeToVector(s, CompactSizeWriter(inputs.size()));
+            SerializeToVector(s, CompactSizeWriter(PSBT_GLOBAL_OUTPUT_COUNT));
+            SerializeToVector(s, CompactSizeWriter(outputs.size()));
+        }
+        if (m_tx_modifiable != nullopt) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_GLOBAL_TX_MODIFIABLE));
+            SerializeToVector(s, static_cast<uint8_t>(m_tx_modifiable->to_ulong()));
+        }
+
         // PSBT version
         if (GetVersion() > 0) {
             SerializeToVector(s, CompactSizeWriter(PSBT_GLOBAL_VERSION));
@@ -660,6 +804,10 @@ struct PartiallySignedTransaction
 
         // Read global data
         bool found_sep = false;
+        uint64_t input_count = 0;
+        uint64_t output_count = 0;
+        bool found_input_count = false;
+        bool found_output_count = false;
         while(!s.empty()) {
             // Read
             std::vector<unsigned char> key;
@@ -697,6 +845,69 @@ struct PartiallySignedTransaction
                             throw std::ios_base::failure("Unsigned tx does not have empty scriptSigs and scriptWitnesses.");
                         }
                     }
+                    // Set the input and output counts
+                    input_count = tx->vin.size();
+                    output_count = tx->vout.size();
+                    break;
+                }
+                case PSBT_GLOBAL_TX_VERSION:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, global transaction version is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Global transaction version key is more than one byte type");
+                    }
+                    uint32_t v;
+                    UnserializeFromVector(s, v);
+                    tx_version = v;
+                    break;
+                }
+                case PSBT_GLOBAL_FALLBACK_LOCKTIME:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, global fallback locktime is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Global fallback locktime key is more than one byte type");
+                    }
+                    uint32_t v;
+                    UnserializeFromVector(s, v);
+                    fallback_locktime = v;
+                    break;
+                }
+                case PSBT_GLOBAL_INPUT_COUNT:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, global input count is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Global input count key is more than one byte type");
+                    }
+                    CompactSizeReader reader(input_count);
+                    UnserializeFromVector(s, reader);
+                    found_input_count = true;
+                    break;
+                }
+                case PSBT_GLOBAL_OUTPUT_COUNT:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, global output count is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Global output count key is more than one byte type");
+                    }
+                    CompactSizeReader reader(output_count);
+                    UnserializeFromVector(s, reader);
+                    found_output_count = true;
+                    break;
+                }
+                case PSBT_GLOBAL_TX_MODIFIABLE:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, tx modifiable flags is already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Global tx modifiable flags key is more than one byte type");
+                    }
+                    uint8_t tx_mod;
+                    UnserializeFromVector(s, tx_mod);
+                    m_tx_modifiable.emplace(tx_mod);
                     break;
                 }
                 case PSBT_GLOBAL_XPUB:
@@ -782,7 +993,7 @@ struct PartiallySignedTransaction
 
         // Read input data
         unsigned int i = 0;
-        while (!s.empty() && i < tx->vin.size()) {
+        while (!s.empty() && i < input_count) {
             PSBTInput input;
             s >> input;
             inputs.push_back(input);
@@ -794,20 +1005,20 @@ struct PartiallySignedTransaction
             ++i;
         }
         // Make sure that the number of inputs matches the number of inputs in the transaction
-        if (inputs.size() != tx->vin.size()) {
+        if (inputs.size() != input_count) {
             throw std::ios_base::failure("Inputs provided does not match the number of inputs in transaction.");
         }
 
         // Read output data
         i = 0;
-        while (!s.empty() && i < tx->vout.size()) {
+        while (!s.empty() && i < output_count) {
             PSBTOutput output;
             s >> output;
             outputs.push_back(output);
             ++i;
         }
         // Make sure that the number of outputs matches the number of outputs in the transaction
-        if (outputs.size() != tx->vout.size()) {
+        if (outputs.size() != output_count) {
             throw std::ios_base::failure("Outputs provided does not match the number of outputs in transaction.");
         }
     }
diff --git a/src/serialize.h b/src/serialize.h
index e6142648e8..f9c6fbeefd 100644
--- a/src/serialize.h
+++ b/src/serialize.h
@@ -563,6 +563,19 @@ struct CompactSizeFormatter
     }
 };
 
+class CompactSizeReader
+{
+protected:
+    uint64_t& n;
+public:
+    explicit CompactSizeReader(uint64_t& n_in) : n(n_in) {}
+
+    template<typename Stream>
+    void Unserialize(Stream &s) const {
+        n = ReadCompactSize<Stream>(s);
+    }
+};
+
 class CompactSizeWriter
 {
 protected: