Implement de/ser of PSBT's Taproot fields

1 files changed, 257 insertions, 0 deletions

diff --git a/src/psbt.h b/src/psbt.h
index b8f3e70eee..a143a99988 100644
--- a/src/psbt.h
+++ b/src/psbt.h
@@ -40,12 +40,21 @@ static constexpr uint8_t PSBT_IN_RIPEMD160 = 0x0A;
 static constexpr uint8_t PSBT_IN_SHA256 = 0x0B;
 static constexpr uint8_t PSBT_IN_HASH160 = 0x0C;
 static constexpr uint8_t PSBT_IN_HASH256 = 0x0D;
+static constexpr uint8_t PSBT_IN_TAP_KEY_SIG = 0x13;
+static constexpr uint8_t PSBT_IN_TAP_SCRIPT_SIG = 0x14;
+static constexpr uint8_t PSBT_IN_TAP_LEAF_SCRIPT = 0x15;
+static constexpr uint8_t PSBT_IN_TAP_BIP32_DERIVATION = 0x16;
+static constexpr uint8_t PSBT_IN_TAP_INTERNAL_KEY = 0x17;
+static constexpr uint8_t PSBT_IN_TAP_MERKLE_ROOT = 0x18;
 static constexpr uint8_t PSBT_IN_PROPRIETARY = 0xFC;
 
 // Output types
 static constexpr uint8_t PSBT_OUT_REDEEMSCRIPT = 0x00;
 static constexpr uint8_t PSBT_OUT_WITNESSSCRIPT = 0x01;
 static constexpr uint8_t PSBT_OUT_BIP32_DERIVATION = 0x02;
+static constexpr uint8_t PSBT_OUT_TAP_INTERNAL_KEY = 0x05;
+static constexpr uint8_t PSBT_OUT_TAP_TREE = 0x06;
+static constexpr uint8_t PSBT_OUT_TAP_BIP32_DERIVATION = 0x07;
 static constexpr uint8_t PSBT_OUT_PROPRIETARY = 0xFC;
 
 // The separator is 0x00. Reading this in means that the unserializer can interpret it
@@ -193,6 +202,15 @@ struct PSBTInput
     std::map<uint256, std::vector<unsigned char>> sha256_preimages;
     std::map<uint160, std::vector<unsigned char>> hash160_preimages;
     std::map<uint256, std::vector<unsigned char>> hash256_preimages;
+
+    // Taproot fields
+    std::vector<unsigned char> m_tap_key_sig;
+    std::map<std::pair<XOnlyPubKey, uint256>, std::vector<unsigned char>> m_tap_script_sigs;
+    std::map<std::pair<CScript, int>, std::set<std::vector<unsigned char>, ShortestVectorFirstComparator>> m_tap_scripts;
+    std::map<XOnlyPubKey, std::pair<std::set<uint256>, KeyOriginInfo>> m_tap_bip32_paths;
+    XOnlyPubKey m_tap_internal_key;
+    uint256 m_tap_merkle_root;
+
     std::map<std::vector<unsigned char>, std::vector<unsigned char>> unknown;
     std::set<PSBTProprietary> m_proprietary;
     std::optional<int> sighash_type;
@@ -267,6 +285,53 @@ struct PSBTInput
                 SerializeToVector(s, CompactSizeWriter(PSBT_IN_HASH256), Span{hash});
                 s << preimage;
             }
+
+            // Write taproot key sig
+            if (!m_tap_key_sig.empty()) {
+                SerializeToVector(s, PSBT_IN_TAP_KEY_SIG);
+                s << m_tap_key_sig;
+            }
+
+            // Write taproot script sigs
+            for (const auto& [pubkey_leaf, sig] : m_tap_script_sigs) {
+                const auto& [xonly, leaf_hash] = pubkey_leaf;
+                SerializeToVector(s, PSBT_IN_TAP_SCRIPT_SIG, xonly, leaf_hash);
+                s << sig;
+            }
+
+            // Write taproot leaf scripts
+            for (const auto& [leaf, control_blocks] : m_tap_scripts) {
+                const auto& [script, leaf_ver] = leaf;
+                for (const auto& control_block : control_blocks) {
+                    SerializeToVector(s, PSBT_IN_TAP_LEAF_SCRIPT, Span{control_block});
+                    std::vector<unsigned char> value_v(script.begin(), script.end());
+                    value_v.push_back((uint8_t)leaf_ver);
+                    s << value_v;
+                }
+            }
+
+            // Write taproot bip32 keypaths
+            for (const auto& [xonly, leaf_origin] : m_tap_bip32_paths) {
+                const auto& [leaf_hashes, origin] = leaf_origin;
+                SerializeToVector(s, PSBT_IN_TAP_BIP32_DERIVATION, xonly);
+                std::vector<unsigned char> value;
+                CVectorWriter s_value(s.GetType(), s.GetVersion(), value, 0);
+                s_value << leaf_hashes;
+                SerializeKeyOrigin(s_value, origin);
+                s << value;
+            }
+
+            // Write taproot internal key
+            if (!m_tap_internal_key.IsNull()) {
+                SerializeToVector(s, PSBT_IN_TAP_INTERNAL_KEY);
+                s << ToByteVector(m_tap_internal_key);
+            }
+
+            // Write taproot merkle root
+            if (!m_tap_merkle_root.IsNull()) {
+                SerializeToVector(s, PSBT_IN_TAP_MERKLE_ROOT);
+                SerializeToVector(s, m_tap_merkle_root);
+            }
         }
 
         // Write script sig
@@ -503,6 +568,103 @@ struct PSBTInput
                     hash256_preimages.emplace(hash, std::move(preimage));
                     break;
                 }
+                case PSBT_IN_TAP_KEY_SIG:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, input Taproot key signature already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Input Taproot key signature key is more than one byte type");
+                    }
+                    s >> m_tap_key_sig;
+                    if (m_tap_key_sig.size() < 64) {
+                        throw std::ios_base::failure("Input Taproot key path signature is shorter than 64 bytes");
+                    } else if (m_tap_key_sig.size() > 65) {
+                        throw std::ios_base::failure("Input Taproot key path signature is longer than 65 bytes");
+                    }
+                    break;
+                }
+                case PSBT_IN_TAP_SCRIPT_SIG:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, input Taproot script signature already provided");
+                    } else if (key.size() != 65) {
+                        throw std::ios_base::failure("Input Taproot script signature key is not 65 bytes");
+                    }
+                    SpanReader s_key(s.GetType(), s.GetVersion(), Span{key}.subspan(1));
+                    XOnlyPubKey xonly;
+                    uint256 hash;
+                    s_key >> xonly;
+                    s_key >> hash;
+                    std::vector<unsigned char> sig;
+                    s >> sig;
+                    if (sig.size() < 64) {
+                        throw std::ios_base::failure("Input Taproot script path signature is shorter than 64 bytes");
+                    } else if (sig.size() > 65) {
+                        throw std::ios_base::failure("Input Taproot script path signature is longer than 65 bytes");
+                    }
+                    m_tap_script_sigs.emplace(std::make_pair(xonly, hash), sig);
+                    break;
+                }
+                case PSBT_IN_TAP_LEAF_SCRIPT:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, input Taproot leaf script already provided");
+                    } else if (key.size() < 34) {
+                        throw std::ios_base::failure("Taproot leaf script key is not at least 34 bytes");
+                    } else if ((key.size() - 2) % 32 != 0) {
+                        throw std::ios_base::failure("Input Taproot leaf script key's control block size is not valid");
+                    }
+                    std::vector<unsigned char> script_v;
+                    s >> script_v;
+                    if (script_v.empty()) {
+                        throw std::ios_base::failure("Input Taproot leaf script must be at least 1 byte");
+                    }
+                    uint8_t leaf_ver = script_v.back();
+                    script_v.pop_back();
+                    const auto leaf_script = std::make_pair(CScript(script_v.begin(), script_v.end()), (int)leaf_ver);
+                    m_tap_scripts[leaf_script].insert(std::vector<unsigned char>(key.begin() + 1, key.end()));
+                    break;
+                }
+                case PSBT_IN_TAP_BIP32_DERIVATION:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, input Taproot BIP32 keypath already provided");
+                    } else if (key.size() != 33) {
+                        throw std::ios_base::failure("Input Taproot BIP32 keypath key is not at 33 bytes");
+                    }
+                    SpanReader s_key(s.GetType(), s.GetVersion(), Span{key}.subspan(1));
+                    XOnlyPubKey xonly;
+                    s_key >> xonly;
+                    std::set<uint256> leaf_hashes;
+                    uint64_t value_len = ReadCompactSize(s);
+                    size_t before_hashes = s.size();
+                    s >> leaf_hashes;
+                    size_t after_hashes = s.size();
+                    size_t hashes_len = before_hashes - after_hashes;
+                    size_t origin_len = value_len - hashes_len;
+                    m_tap_bip32_paths.emplace(xonly, std::make_pair(leaf_hashes, DeserializeKeyOrigin(s, origin_len)));
+                    break;
+                }
+                case PSBT_IN_TAP_INTERNAL_KEY:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, input Taproot internal key already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Input Taproot internal key key is more than one byte type");
+                    }
+                    UnserializeFromVector(s, m_tap_internal_key);
+                    break;
+                }
+                case PSBT_IN_TAP_MERKLE_ROOT:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, input Taproot merkle root already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Input Taproot merkle root key is more than one byte type");
+                    }
+                    UnserializeFromVector(s, m_tap_merkle_root);
+                    break;
+                }
                 case PSBT_IN_PROPRIETARY:
                 {
                     PSBTProprietary this_prop;
@@ -547,6 +709,9 @@ struct PSBTOutput
     CScript redeem_script;
     CScript witness_script;
     std::map<CPubKey, KeyOriginInfo> hd_keypaths;
+    XOnlyPubKey m_tap_internal_key;
+    std::optional<TaprootBuilder> m_tap_tree;
+    std::map<XOnlyPubKey, std::pair<std::set<uint256>, KeyOriginInfo>> m_tap_bip32_paths;
     std::map<std::vector<unsigned char>, std::vector<unsigned char>> unknown;
     std::set<PSBTProprietary> m_proprietary;
 
@@ -579,6 +744,40 @@ struct PSBTOutput
             s << entry.value;
         }
 
+        // Write taproot internal key
+        if (!m_tap_internal_key.IsNull()) {
+            SerializeToVector(s, PSBT_OUT_TAP_INTERNAL_KEY);
+            s << ToByteVector(m_tap_internal_key);
+        }
+
+        // Write taproot tree
+        if (m_tap_tree.has_value()) {
+            SerializeToVector(s, PSBT_OUT_TAP_TREE);
+            std::vector<unsigned char> value;
+            CVectorWriter s_value(s.GetType(), s.GetVersion(), value, 0);
+            const auto& tuples = m_tap_tree->GetTreeTuples();
+            for (const auto& tuple : tuples) {
+                uint8_t depth = std::get<0>(tuple);
+                uint8_t leaf_ver = std::get<1>(tuple);
+                CScript script = std::get<2>(tuple);
+                s_value << depth;
+                s_value << leaf_ver;
+                s_value << script;
+            }
+            s << value;
+        }
+
+        // Write taproot bip32 keypaths
+        for (const auto& [xonly, leaf] : m_tap_bip32_paths) {
+            const auto& [leaf_hashes, origin] = leaf;
+            SerializeToVector(s, PSBT_OUT_TAP_BIP32_DERIVATION, xonly);
+            std::vector<unsigned char> value;
+            CVectorWriter s_value(s.GetType(), s.GetVersion(), value, 0);
+            s_value << leaf_hashes;
+            SerializeKeyOrigin(s_value, origin);
+            s << value;
+        }
+
         // Write unknown things
         for (auto& entry : unknown) {
             s << entry.first;
@@ -639,6 +838,59 @@ struct PSBTOutput
                     DeserializeHDKeypaths(s, key, hd_keypaths);
                     break;
                 }
+                case PSBT_OUT_TAP_INTERNAL_KEY:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, output Taproot internal key already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Output Taproot internal key key is more than one byte type");
+                    }
+                    UnserializeFromVector(s, m_tap_internal_key);
+                    break;
+                }
+                case PSBT_OUT_TAP_TREE:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, output Taproot tree already provided");
+                    } else if (key.size() != 1) {
+                        throw std::ios_base::failure("Output Taproot tree key is more than one byte type");
+                    }
+                    m_tap_tree.emplace();
+                    std::vector<unsigned char> tree_v;
+                    s >> tree_v;
+                    SpanReader s_tree(s.GetType(), s.GetVersion(), tree_v);
+                    while (!s_tree.empty()) {
+                        uint8_t depth;
+                        uint8_t leaf_ver;
+                        CScript script;
+                        s_tree >> depth;
+                        s_tree >> leaf_ver;
+                        s_tree >> script;
+                        m_tap_tree->Add((int)depth, script, (int)leaf_ver, true /* track */);
+                    }
+                    if (!m_tap_tree->IsComplete()) {
+                        throw std::ios_base::failure("Output Taproot tree is malformed");
+                    }
+                    break;
+                }
+                case PSBT_OUT_TAP_BIP32_DERIVATION:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, output Taproot BIP32 keypath already provided");
+                    } else if (key.size() != 33) {
+                        throw std::ios_base::failure("Output Taproot BIP32 keypath key is not at 33 bytes");
+                    }
+                    XOnlyPubKey xonly(uint256({key.begin() + 1, key.begin() + 33}));
+                    std::set<uint256> leaf_hashes;
+                    uint64_t value_len = ReadCompactSize(s);
+                    size_t before_hashes = s.size();
+                    s >> leaf_hashes;
+                    size_t after_hashes = s.size();
+                    size_t hashes_len = before_hashes - after_hashes;
+                    size_t origin_len = value_len - hashes_len;
+                    m_tap_bip32_paths.emplace(xonly, std::make_pair(leaf_hashes, DeserializeKeyOrigin(s, origin_len)));
+                    break;
+                }
                 case PSBT_OUT_PROPRIETARY:
                 {
                     PSBTProprietary this_prop;
@@ -667,6 +919,11 @@ struct PSBTOutput
             }
         }
 
+        // Finalize m_tap_tree so that all of the computed things are computed
+        if (m_tap_tree.has_value() && m_tap_tree->IsComplete() && m_tap_internal_key.IsFullyValid()) {
+            m_tap_tree->Finalize(m_tap_internal_key);
+        }
+
         if (!found_sep) {
             throw std::ios_base::failure("Separator is missing at the end of an output map");
         }