From 1a1599f59f711206ea91f025ee48b148e17b01ce Mon Sep 17 00:00:00 2001
From: Cory Fields <cory-nospam-@coryfields.com>
Date: Wed, 8 May 2024 19:57:12 +0000
Subject: [PATCH] Add KeyPair wrapper class

---
 src/key.cpp | 76 +++++++++++++++++++++++++++++++++++++++--------------
 src/key.h   | 12 +++++++++
 2 files changed, 69 insertions(+), 19 deletions(-)

diff --git a/src/key.cpp b/src/key.cpp
index 2bd639629818e5..9c9fe2c14b4e65 100644
--- a/src/key.cpp
+++ b/src/key.cpp
@@ -277,27 +277,11 @@ bool CKey::SignCompact(const uint256 &hash, std::vector<unsigned char>& vchSig)
 
 bool CKey::SignSchnorr(const uint256& hash, Span<unsigned char> sig, const uint256* merkle_root, const uint256& aux) const
 {
-    assert(sig.size() == 64);
-    secp256k1_keypair keypair;
-    if (!secp256k1_keypair_create(secp256k1_context_sign, &keypair, UCharCast(begin()))) return false;
+    KeyPair keypair(*this);
     if (merkle_root) {
-        secp256k1_xonly_pubkey pubkey;
-        if (!secp256k1_keypair_xonly_pub(secp256k1_context_sign, &pubkey, nullptr, &keypair)) return false;
-        unsigned char pubkey_bytes[32];
-        if (!secp256k1_xonly_pubkey_serialize(secp256k1_context_sign, pubkey_bytes, &pubkey)) return false;
-        uint256 tweak = XOnlyPubKey(pubkey_bytes).ComputeTapTweakHash(merkle_root->IsNull() ? nullptr : merkle_root);
-        if (!secp256k1_keypair_xonly_tweak_add(secp256k1_context_static, &keypair, tweak.data())) return false;
-    }
-    bool ret = secp256k1_schnorrsig_sign32(secp256k1_context_sign, sig.data(), hash.data(), &keypair, aux.data());
-    if (ret) {
-        // Additional verification step to prevent using a potentially corrupted signature
-        secp256k1_xonly_pubkey pubkey_verify;
-        ret = secp256k1_keypair_xonly_pub(secp256k1_context_static, &pubkey_verify, nullptr, &keypair);
-        ret &= secp256k1_schnorrsig_verify(secp256k1_context_static, sig.data(), hash.begin(), 32, &pubkey_verify);
+        if(!keypair.ApplyTapTweak(*merkle_root)) return false;
     }
-    if (!ret) memory_cleanse(sig.data(), sig.size());
-    memory_cleanse(&keypair, sizeof(keypair));
-    return ret;
+    return keypair.SignSchnorr(hash, sig, aux);
 }
 
 bool CKey::Load(const CPrivKey &seckey, const CPubKey &vchPubKey, bool fSkipCheck=false) {
@@ -426,6 +410,60 @@ void CExtKey::Decode(const unsigned char code[BIP32_EXTKEY_SIZE]) {
     if ((nDepth == 0 && (nChild != 0 || ReadLE32(vchFingerprint) != 0)) || code[41] != 0) key = CKey();
 }
 
+KeyPair::KeyPair(const CKey& key)
+{
+    m_keydata = make_secure_unique<KeyType>();
+    if (!secp256k1_keypair_create(secp256k1_context_sign, (secp256k1_keypair*)m_keydata->data(), UCharCast(key.data())))
+    {
+        m_keydata = nullptr;
+    }
+}
+
+bool KeyPair::ApplyTapTweak(const uint256& merkle_root)
+{
+    if (!m_keydata) return false;
+    secp256k1_xonly_pubkey pubkey;
+    if (!secp256k1_keypair_xonly_pub(secp256k1_context_sign, &pubkey, nullptr, (secp256k1_keypair*)m_keydata->data())) {
+        return false;
+    }
+    unsigned char pubkey_bytes[32];
+    if (!secp256k1_xonly_pubkey_serialize(secp256k1_context_sign, pubkey_bytes, &pubkey)) {
+        return false;
+    }
+    uint256 tweak = XOnlyPubKey(pubkey_bytes).ComputeTapTweakHash(merkle_root.IsNull() ? nullptr : &merkle_root);
+    if (!secp256k1_keypair_xonly_tweak_add(secp256k1_context_static, (secp256k1_keypair*)m_keydata->data(), tweak.data())) {
+        return false;
+    }
+    return true;
+}
+
+bool KeyPair::GetKey(CKey& key) const
+{
+    if (!m_keydata) return false;
+    unsigned char tweaked_secret_bytes[32];
+    if (!secp256k1_keypair_sec(secp256k1_context_sign, tweaked_secret_bytes, (secp256k1_keypair*)m_keydata->data())) {
+        return false;
+    }
+    key.Set(std::begin(tweaked_secret_bytes), std::end(tweaked_secret_bytes), true);
+    memory_cleanse(tweaked_secret_bytes, sizeof(tweaked_secret_bytes));
+    return true;
+}
+
+bool KeyPair::SignSchnorr(const uint256& hash, Span<unsigned char> sig, const uint256& aux) const
+{
+    assert(sig.size() == 64);
+    bool ret;
+    ret = secp256k1_schnorrsig_sign32(secp256k1_context_sign, sig.data(), hash.data(), (secp256k1_keypair*)m_keydata->data(), aux.data());
+    if (ret) {
+        // Additional verification step to prevent using a potentially corrupted signature
+        secp256k1_xonly_pubkey pubkey_verify;
+        ret = secp256k1_keypair_xonly_pub(secp256k1_context_static, &pubkey_verify, nullptr, (secp256k1_keypair*)m_keydata->data());
+        ret &= secp256k1_schnorrsig_verify(secp256k1_context_static, sig.data(), hash.begin(), 32, &pubkey_verify);
+    }
+    if (!ret) memory_cleanse(sig.data(), sig.size());
+    return ret;
+}
+
 bool ECC_InitSanityCheck() {
     CKey key = GenerateRandomKey();
     CPubKey pubkey = key.GetPubKey();
diff --git a/src/key.h b/src/key.h
index 53acd179ba8f71..77c562d80e832f 100644
--- a/src/key.h
+++ b/src/key.h
@@ -236,6 +236,18 @@ struct CExtKey {
     void SetSeed(Span<const std::byte> seed);
 };
 
+class KeyPair
+{
+public:
+    KeyPair(const CKey& key);
+    [[nodiscard]] bool ApplyTapTweak(const uint256& merkle_root);
+    [[nodiscard]] bool GetKey(CKey& key) const;
+    [[nodiscard]] bool SignSchnorr(const uint256& hash, Span<unsigned char> sig, const uint256& aux) const;
+private:
+    using KeyType = std::array<unsigned char, 96>;
+    secure_unique_ptr<KeyType> m_keydata;
+};
+
 /** Initialize the elliptic curve support. May not be called twice without calling ECC_Stop first. */
 void ECC_Start();