psbt: Implement un/ser of musig2 fields

src/psbt.h

@@ -50,6 +50,9 @@ 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_MUSIG2_PARTICIPANT_PUBKEYS = 0x1a;
+static constexpr uint8_t PSBT_IN_MUSIG2_PUB_NONCE = 0x1b;
+static constexpr uint8_t PSBT_IN_MUSIG2_PARTIAL_SIG = 0x1c;
 static constexpr uint8_t PSBT_IN_PROPRIETARY = 0xFC;
 
 // Output types
@@ -59,6 +62,7 @@ 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_MUSIG2_PARTICIPANT_PUBKEYS = 0x08;
 static constexpr uint8_t PSBT_OUT_PROPRIETARY = 0xFC;
 
 // The separator is 0x00. Reading this in means that the unserializer can interpret it
@@ -193,6 +197,49 @@ void SerializeHDKeypaths(Stream& s, const std::map<CPubKey, KeyOriginInfo>& hd_k
     }
 }
 
+// Deserialize a PSBT_{IN/OUT}_MUSIG2_PARTICIPANT_PUBKEYS field
+template<typename Stream>
+void DeserializeMuSig2ParticipantPubkeys(Stream& s, SpanReader& skey, std::map<CPubKey, std::vector<CPubKey>>& out, std::string context)
+{
+    std::array<unsigned char, CPubKey::COMPRESSED_SIZE> agg_pubkey_bytes;
+    skey >> std::as_writable_bytes(std::span{agg_pubkey_bytes});
+    CPubKey agg_pubkey(agg_pubkey_bytes);
+
+    std::vector<CPubKey> participants;
+    std::vector<unsigned char> val;
+    s >> val;
+    SpanReader s_val{val};
+    while (s_val.size() >= CPubKey::COMPRESSED_SIZE) {
+        std::array<unsigned char, CPubKey::COMPRESSED_SIZE> part_pubkey_bytes;
+        s_val >> std::as_writable_bytes(std::span{part_pubkey_bytes});
+        participants.emplace_back(std::span{part_pubkey_bytes});
+    }
+    if (!s_val.empty()) {
+        throw std::ios_base::failure(context + " musig2 participants pubkeys value size is not a multiple of 33");
+    }
+
+    out.emplace(agg_pubkey, participants);
+}
+
+// Deserialize the MuSig2 participant identifiers from PSBT_MUSIG2_{PUBNONCE/PARTIAL_SIG} fields
+// Both fields contain the same data after the type byte - aggregate pubkey | participant pubkey | leaf script hash
+template<typename Stream>
+void DeserializeMuSig2ParticipantDataIdentifier(Stream& skey, CPubKey& agg_pub, CPubKey& part_pub, uint256& leaf_hash)
+{
+    leaf_hash.SetNull();
+
+    std::array<unsigned char, CPubKey::COMPRESSED_SIZE> part_pubkey_bytes;
+    std::array<unsigned char, CPubKey::COMPRESSED_SIZE> agg_pubkey_bytes;
+
+    skey >> std::as_writable_bytes(std::span{part_pubkey_bytes}) >> std::as_writable_bytes(std::span{agg_pubkey_bytes});
+    agg_pub.Set(agg_pubkey_bytes.begin(), agg_pubkey_bytes.end());
+    part_pub.Set(part_pubkey_bytes.begin(), part_pubkey_bytes.end());
+
+    if (!skey.empty()) {
+        skey >> leaf_hash;
+    }
+}
+
 /** A structure for PSBTs which contain per-input information */
 struct PSBTInput
 {
@@ -217,6 +264,13 @@ struct PSBTInput
     XOnlyPubKey m_tap_internal_key;
     uint256 m_tap_merkle_root;
 
+    // MuSig2 fields
+    std::map<CPubKey, std::vector<CPubKey>> m_musig2_participants;
+    // Key is the aggregate pubkey and the script leaf hash, value is a map of participant pubkey to pubnonce
+    std::map<std::pair<CPubKey, uint256>, std::map<CPubKey, std::vector<uint8_t>>> m_musig2_pubnonces;
+    // Key is the aggregate pubkey and the script leaf hash, value is a map of participant pubkey to partial_sig
+    std::map<std::pair<CPubKey, uint256>, std::map<CPubKey, uint256>> m_musig2_partial_sigs;
+
     std::map<std::vector<unsigned char>, std::vector<unsigned char>> unknown;
     std::set<PSBTProprietary> m_proprietary;
     std::optional<int> sighash_type;
@@ -337,6 +391,43 @@ struct PSBTInput
                 SerializeToVector(s, PSBT_IN_TAP_MERKLE_ROOT);
                 SerializeToVector(s, m_tap_merkle_root);
             }
+
+            // Write MuSig2 Participants
+            for (const auto& [agg_pubkey, part_pubs] : m_musig2_participants) {
+                SerializeToVector(s, CompactSizeWriter(PSBT_IN_MUSIG2_PARTICIPANT_PUBKEYS), std::span{agg_pubkey});
+                std::vector<unsigned char> value;
+                VectorWriter s_value{value, 0};
+                for (auto& pk : part_pubs) {
+                    s_value << std::span{pk};
+                }
+                s << value;
+            }
+
+            // Write MuSig2 pubnonces
+            for (const auto& [agg_pubkey_leaf_hash, pubnonces] : m_musig2_pubnonces) {
+                const auto& [agg_pubkey, leaf_hash] = agg_pubkey_leaf_hash;
+                for (const auto& [part_pubkey, pubnonce] : pubnonces) {
+                    if (leaf_hash.IsNull()) {
+                        SerializeToVector(s, CompactSizeWriter(PSBT_IN_MUSIG2_PUB_NONCE), std::span{part_pubkey}, std::span{agg_pubkey});
+                    } else {
+                        SerializeToVector(s, CompactSizeWriter(PSBT_IN_MUSIG2_PUB_NONCE), std::span{part_pubkey}, std::span{agg_pubkey}, leaf_hash);
+                    }
+                    s << pubnonce;
+                }
+            }
+
+            // Write MuSig2 partial signatures
+            for (const auto& [agg_pubkey_leaf_hash, psigs] : m_musig2_partial_sigs) {
+                const auto& [agg_pubkey, leaf_hash] = agg_pubkey_leaf_hash;
+                for (const auto& [pubkey, psig] : psigs) {
+                    if (leaf_hash.IsNull()) {
+                        SerializeToVector(s, CompactSizeWriter(PSBT_IN_MUSIG2_PARTIAL_SIG), std::span{pubkey}, std::span{agg_pubkey});
+                    } else {
+                        SerializeToVector(s, CompactSizeWriter(PSBT_IN_MUSIG2_PARTIAL_SIG), std::span{pubkey}, std::span{agg_pubkey}, leaf_hash);
+                    }
+                    SerializeToVector(s, psig);
+                }
+            }
         }
 
         // Write script sig
@@ -672,6 +763,53 @@ struct PSBTInput
                     UnserializeFromVector(s, m_tap_merkle_root);
                     break;
                 }
+                case PSBT_IN_MUSIG2_PARTICIPANT_PUBKEYS:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, input participant pubkeys for an aggregate key already provided");
+                    } else if (key.size() != CPubKey::COMPRESSED_SIZE + 1) {
+                        throw std::ios_base::failure("Input musig2 participants pubkeys aggregate key is not 34 bytes");
+                    }
+                    DeserializeMuSig2ParticipantPubkeys(s, skey, m_musig2_participants, std::string{"Input"});
+                    break;
+                }
+                case PSBT_IN_MUSIG2_PUB_NONCE:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, input musig2 pubnonce already provided");
+                    } else if (key.size() != 2 * CPubKey::COMPRESSED_SIZE + 1 && key.size() != 2 * CPubKey::COMPRESSED_SIZE + CSHA256::OUTPUT_SIZE + 1) {
+                        throw std::ios_base::failure("Input musig2 pubnonce key is not expected size of 67 or 99 bytes");
+                    }
+                    CPubKey agg_pub, part_pub;
+                    uint256 leaf_hash;
+                    DeserializeMuSig2ParticipantDataIdentifier(skey, agg_pub, part_pub, leaf_hash);
+
+                    std::vector<uint8_t> pubnonce;
+                    s >> pubnonce;
+                    if (pubnonce.size() != 66) {
+                        throw std::ios_base::failure("Input musig2 pubnonce value is not 66 bytes");
+                    }
+
+                    m_musig2_pubnonces[std::make_pair(agg_pub, leaf_hash)].emplace(part_pub, pubnonce);
+                    break;
+                }
+                case PSBT_IN_MUSIG2_PARTIAL_SIG:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, input musig2 partial sig already provided");
+                    } else if (key.size() != 2 * CPubKey::COMPRESSED_SIZE + 1 && key.size() != 2 * CPubKey::COMPRESSED_SIZE + CSHA256::OUTPUT_SIZE + 1) {
+                        throw std::ios_base::failure("Input musig2 partial sig key is not expected size of 67 or 99 bytes");
+                    }
+                    CPubKey agg_pub, part_pub;
+                    uint256 leaf_hash;
+                    DeserializeMuSig2ParticipantDataIdentifier(skey, agg_pub, part_pub, leaf_hash);
+
+                    uint256 partial_sig;
+                    UnserializeFromVector(s, partial_sig);
+
+                    m_musig2_partial_sigs[std::make_pair(agg_pub, leaf_hash)].emplace(part_pub, partial_sig);
+                    break;
+                }
                 case PSBT_IN_PROPRIETARY:
                 {
                     PSBTProprietary this_prop;
@@ -719,6 +857,7 @@ struct PSBTOutput
     XOnlyPubKey m_tap_internal_key;
     std::vector<std::tuple<uint8_t, uint8_t, std::vector<unsigned char>>> m_tap_tree;
     std::map<XOnlyPubKey, std::pair<std::set<uint256>, KeyOriginInfo>> m_tap_bip32_paths;
+    std::map<CPubKey, std::vector<CPubKey>> m_musig2_participants;
     std::map<std::vector<unsigned char>, std::vector<unsigned char>> unknown;
     std::set<PSBTProprietary> m_proprietary;
 
@@ -781,6 +920,17 @@ struct PSBTOutput
             s << value;
         }
 
+        // Write MuSig2 Participants
+        for (const auto& [agg_pubkey, part_pubs] : m_musig2_participants) {
+            SerializeToVector(s, CompactSizeWriter(PSBT_OUT_MUSIG2_PARTICIPANT_PUBKEYS), std::span{agg_pubkey});
+            std::vector<unsigned char> value;
+            VectorWriter s_value{value, 0};
+            for (auto& pk : part_pubs) {
+                s_value << std::span{pk};
+            }
+            s << value;
+        }
+
         // Write unknown things
         for (auto& entry : unknown) {
             s << entry.first;
@@ -907,6 +1057,16 @@ struct PSBTOutput
                     m_tap_bip32_paths.emplace(xonly, std::make_pair(leaf_hashes, DeserializeKeyOrigin(s, origin_len)));
                     break;
                 }
+                case PSBT_OUT_MUSIG2_PARTICIPANT_PUBKEYS:
+                {
+                    if (!key_lookup.emplace(key).second) {
+                        throw std::ios_base::failure("Duplicate Key, output participant pubkeys for an aggregate key already provided");
+                    } else if (key.size() != CPubKey::COMPRESSED_SIZE + 1) {
+                        throw std::ios_base::failure("Output musig2 participants pubkeys aggregate key is not 34 bytes");
+                    }
+                    DeserializeMuSig2ParticipantPubkeys(s, skey, m_musig2_participants, std::string{"Output"});
+                    break;
+                }
                 case PSBT_OUT_PROPRIETARY:
                 {
                     PSBTProprietary this_prop;