Stop using pre_split_keypool for MWEB keys in upgraded wallets, and support recovering coins sent to stealth addresses generated from pre_split_keypool

(cherry picked from commit 4b45fdf7f3cb8e113c1c04970e7e33751b7d7473)

src/libmw/include/mw/models/crypto/PublicKey.h

@@ -3,6 +3,8 @@
 #include <mw/common/Traits.h>
 #include <mw/models/crypto/BigInteger.h>
 #include <mw/models/crypto/SecretKey.h>
+#include <pubkey.h>
+
 #include <boost/functional/hash.hpp>
 
 class PublicKey :
@@ -45,6 +47,7 @@ public:
     static PublicKey Random();
 
     const BigInt<33>& GetBigInt() const { return m_compressed; }
+    CKeyID GetID() const { return CPubKey(vec()).GetID(); }
     std::array<uint8_t, 33> array() const { return m_compressed.ToArray(); }
     const std::vector<uint8_t>& vec() const { return m_compressed.vec(); }
     const uint8_t& operator[](const size_t x) const { return m_compressed[x]; }

src/libmw/include/mw/models/wallet/Coin.h

@@ -20,6 +20,13 @@ static constexpr uint32_t CHANGE_INDEX{0};
 /// </summary>
 static constexpr uint32_t PEGIN_INDEX{1};
 
+/// <summary>
+/// Outputs sent to a stealth address whose spend key was not generated using the MWEB
+/// keychain won't have an address index. We use 0xfffffffe to represent this.
+/// In that case, we must lookup the secret key in the wallet DB, rather than the MWEB keychain.
+/// </summary>
+static constexpr uint32_t CUSTOM_KEY{std::numeric_limits<uint32_t>::max() - 1};
+
 /// <summary>
 /// Outputs sent to others will be marked with an address_index of 0xffffffff.
 /// </summary>
@@ -35,7 +42,8 @@ struct Coin : public Traits::ISerializable {
     uint32_t address_index{UNKNOWN_INDEX};
 
     // The private key needed in order to spend the coin.
-    // May be empty for watch-only wallets.
+    // Will be empty for watch-only wallets.
+    // May be empty for locked wallets. Upon unlock, spend_key will get populated.
     boost::optional<SecretKey> spend_key;
 
     // The blinding factor of the coin's output.

src/libmw/include/mw/wallet/Keychain.h

@@ -7,7 +7,7 @@
 #include <memory>
 
 // Forward Declarations
-class ScriptPubKeyMan;
+class LegacyScriptPubKeyMan;
 
 MW_NAMESPACE
 
@@ -16,7 +16,7 @@ class Keychain
 public:
     using Ptr = std::shared_ptr<Keychain>;
 
-    Keychain(const ScriptPubKeyMan& spk_man, SecretKey scan_secret, SecretKey spend_secret)
+    Keychain(const LegacyScriptPubKeyMan& spk_man, SecretKey scan_secret, SecretKey spend_secret)
         : m_spk_man(spk_man),
         m_scanSecret(std::move(scan_secret)),
         m_spendSecret(std::move(spend_secret)) { }
@@ -28,9 +28,11 @@ public:
     // used to calculate the spend key when the wallet becomes unlocked.
     bool RewindOutput(const Output& output, mw::Coin& coin) const;
 
-    // Calculates the output spend key for the address index and shared secret.
-    // Requires that keychain be unlocked and not watch-only.
-    SecretKey CalculateOutputKey(const uint32_t index, const SecretKey& shared_secret) const;
+    // Calculates the output secret key for the given coin.
+    // If the address index is known, it calculates from the keychain's master spend key.
+    // If not, it attempts to lookup the spend key in the database.
+    // Returns boost::empty when keychain is locked or watch-only.
+    boost::optional<SecretKey> CalculateOutputKey(const mw::Coin& coin) const;
 
     // Calculates the StealthAddress at the given index.
     // Requires that keychain be unlocked and not watch-only.
@@ -42,6 +44,8 @@ public:
     const SecretKey& GetScanSecret() const noexcept { return m_scanSecret; }
     const SecretKey& GetSpendSecret() const noexcept { return m_spendSecret; }
 
+    bool HasSpendSecret() const noexcept { return !m_spendSecret.IsNull(); }
+
 	// Clears the spend secret from memory, effectively making this a watch-only keychain.
     void Lock() { m_spendSecret = SecretKey::Null(); }
 	
@@ -49,7 +53,7 @@ public:
     void Unlock(const SecretKey& spend_secret) { m_spendSecret = spend_secret; }
     
 private:
-    const ScriptPubKeyMan& m_spk_man;
+    const LegacyScriptPubKeyMan& m_spk_man;
     SecretKey m_scanSecret;
     SecretKey m_spendSecret;
 };

src/libmw/src/wallet/Keychain.cpp

@@ -3,6 +3,7 @@
 #include <mw/crypto/SecretKeys.h>
 #include <mw/models/tx/OutputMask.h>
 #include <wallet/scriptpubkeyman.h>
+#include <key_io.h>
 
 MW_NAMESPACE
 
@@ -25,12 +26,10 @@ bool Keychain::RewindOutput(const Output& output, mw::Coin& coin) const
     // Check if B_i belongs to wallet
     StealthAddress address(B_i.Mul(m_scanSecret), B_i);
     auto pMetadata = m_spk_man.GetMetadata(address);
-    if (!pMetadata || !pMetadata->mweb_index) {
+    if (!pMetadata) {
         return false;
     }
 
-    uint32_t index = *pMetadata->mweb_index;
-
     // Calc blinding factor and unmask nonce and amount.
     OutputMask mask = OutputMask::FromShared(t);
     uint64_t value = mask.MaskValue(output.GetMaskedValue());
@@ -51,28 +50,53 @@ bool Keychain::RewindOutput(const Output& output, mw::Coin& coin) const
         return false;
     }
 
-    // Spend secret will be null for locked or watch-only wallets.
-    if (!GetSpendSecret().IsNull()) {
-        coin.spend_key = boost::make_optional(CalculateOutputKey(index, t));
-    }
-
-    coin.address_index = index;
+    // v0.21.2 incorrectly generated MWEB keys from the pre-split keypool for upgraded wallets.
+    // These keys will not have an mweb_index, so we set the address_index as CUSTOM_KEY.
+    coin.address_index = pMetadata->mweb_index.get_value_or(CUSTOM_KEY);
     coin.blind = boost::make_optional(mask.GetRawBlind());
     coin.amount = value;
     coin.output_id = output.GetOutputID();
     coin.address = address;
     coin.shared_secret = boost::make_optional(std::move(t));
+    coin.spend_key = CalculateOutputKey(coin);
 
     return true;
 }
 
-SecretKey Keychain::CalculateOutputKey(const uint32_t index, const SecretKey& shared_secret) const
+boost::optional<SecretKey> Keychain::CalculateOutputKey(const mw::Coin& coin) const
 {
-    assert(!m_spendSecret.IsNull());
+    // If we already calculated the spend key, there's no need to calculate it again.
+    if (coin.HasSpendKey()) {
+        return coin.spend_key;
+    }
+
+    // Watch-only or locked wallets will not have the spend secret.
+    if (!HasSpendSecret() || !coin.HasSharedSecret() || !coin.IsMine()) {
+        return boost::none;
+    }
+
+    auto derive_output_key = [](const SecretKey& spend_key, const SecretKey& shared_secret) -> SecretKey {
+        return SecretKeys::From(spend_key)
+            .Mul(Hashed(EHashTag::OUT_KEY, shared_secret))
+            .Total();
+    };
+
+    // An address_index of CUSTOM_KEY means the spend key was not generated from the MWEB keychain.
+    // We should lookup the secret key in the wallet DB, instead of calculating by index.
+    if (coin.address_index == CUSTOM_KEY) {
+        if (!coin.HasAddress()) {
+            return boost::none;
+        }
+
+        CKey key;
+        if (!m_spk_man.GetKey(coin.address->B().GetID(), key)) {
+            return boost::none;
+        }
+
+        return derive_output_key(SecretKey(key.begin()), *coin.shared_secret);
+    }
 
-    return SecretKeys::From(GetSpendKey(index))
-        .Mul(Hashed(EHashTag::OUT_KEY, shared_secret))
-        .Total();
+    return derive_output_key(GetSpendKey(coin.address_index), *coin.shared_secret);
 }
 
 StealthAddress Keychain::GetStealthAddress(const uint32_t index) const

src/mweb/mweb_models.h

@@ -36,6 +36,11 @@ struct Block {
         return IsNull() ? 0 : m_block->GetSupplyChange();
     }
 
+    mw::Hash GetHash() const noexcept
+    {
+        return IsNull() ? mw::Hash{} : m_block->GetHeader()->GetHash();
+    }
+
     mw::Header::CPtr GetMWEBHeader() const noexcept
     {
         return IsNull() ? mw::Header::CPtr{nullptr} : m_block->GetHeader();

src/mweb/mweb_wallet.cpp

@@ -8,7 +8,7 @@ using namespace MWEB;
 bool Wallet::UpgradeCoins()
 {
     mw::Keychain::Ptr keychain = GetKeychain();
-    if (!keychain || keychain->GetSpendSecret().IsNull()) {
+    if (!keychain || !keychain->HasSpendSecret()) {
         return false;
     }
 
@@ -19,14 +19,17 @@ bool Wallet::UpgradeCoins()
 
         if (wtx->mweb_wtx_info && wtx->mweb_wtx_info->received_coin) {
             mw::Coin& coin = *wtx->mweb_wtx_info->received_coin;
-            if (!coin.HasSpendKey() && coin.HasSharedSecret()) {
-                coin.spend_key = keychain->CalculateOutputKey(coin.address_index, *coin.shared_secret);
+            if (!coin.HasSpendKey()) {
+                coin.spend_key = keychain->CalculateOutputKey(coin);
 
-                m_coins[coin.output_id] = coin;
+                // If spend key was populated, update the database and m_coins map.
+                if (coin.HasSpendKey()) {
+                    m_coins[coin.output_id] = coin;
 
-                WalletBatch batch(m_pWallet->GetDatabase());
-                batch.WriteMWEBCoin(coin);
-                batch.WriteTx(*wtx);
+                    WalletBatch batch(m_pWallet->GetDatabase());
+                    batch.WriteMWEBCoin(coin);
+                    batch.WriteTx(*wtx);
+                }
             }
         }
     }
@@ -57,7 +60,7 @@ bool Wallet::RewindOutput(const Output& output, mw::Coin& coin)
     if (GetCoin(output.GetOutputID(), coin) && coin.IsMine()) {
         // If the coin has the spend key, it's fully rewound.
         // If not, try rewinding further if we have the master spend key (i.e. wallet is unlocked).
-        if (coin.HasSpendKey() || !keychain || keychain->GetSpendSecret().IsNull()) {
+        if (coin.HasSpendKey() || !keychain || !keychain->HasSpendSecret()) {
             return true;
         }
     }
@@ -90,7 +93,7 @@ bool Wallet::GetStealthAddress(const mw::Coin& coin, StealthAddress& address) co
 bool Wallet::GetStealthAddress(const uint32_t index, StealthAddress& address) const
 {
     mw::Keychain::Ptr keychain = GetKeychain();
-    if (!keychain || index == mw::UNKNOWN_INDEX) {
+    if (!keychain || index == mw::UNKNOWN_INDEX || index == mw::CUSTOM_KEY) {
         return false;
     }
 

src/script/signingprovider.cpp

@@ -196,7 +196,7 @@ CKeyID GetKeyForDestination(const SigningProvider& store, const CTxDestination&
         }
     }
     if (auto stealth_address = boost::get<StealthAddress>(&dest)) {
-        return CPubKey(stealth_address->B().vec()).GetID();
+        return stealth_address->B().GetID();
     }
     return CKeyID();
 }

src/wallet/scriptpubkeyman.cpp

@@ -244,8 +244,7 @@ isminetype LegacyScriptPubKeyMan::IsMine(const DestinationAddr& script) const
             return ISMINE_NO;
         }
 
-        CPubKey pubkey(mweb_address.GetSpendPubKey().vec());
-        return HaveKey(pubkey.GetID()) ? ISMINE_SPENDABLE : ISMINE_NO;
+        return HaveKey(mweb_address.GetSpendPubKey().GetID()) ? ISMINE_SPENDABLE : ISMINE_NO;
     }
 
     switch (IsMineInner(*this, script.GetScript(), IsMineSigVersion::TOP)) {
@@ -1444,11 +1443,11 @@ bool LegacyScriptPubKeyMan::ReserveKeyFromKeyPool(int64_t& nIndex, CKeyPool& key
         bool fMWEB = (purpose == KeyPurpose::MWEB) && IsHDEnabled() && m_storage.CanSupportFeature(FEATURE_HD_SPLIT);
 
         auto fn_get_keypool = [this](const bool internal, const bool mweb) -> std::set<int64_t>& {
-            if (!set_pre_split_keypool.empty()) {
-                return set_pre_split_keypool;
-            } else if (mweb) {
+             if (mweb) {
                 return set_mweb_keypool;
-            } else if (internal) {
+             } else if (!set_pre_split_keypool.empty()) {
+                 return set_pre_split_keypool;
+             } else if (internal) {
                 return setInternalKeyPool;
             }
 
@@ -1474,8 +1473,11 @@ bool LegacyScriptPubKeyMan::ReserveKeyFromKeyPool(int64_t& nIndex, CKeyPool& key
         if (!GetPubKey(keypool.vchPubKey.GetID(), pk)) {
             throw std::runtime_error(std::string(__func__) + ": unknown key in key pool");
         }
+        if (keypool.fMWEB != fMWEB) {
+            throw std::runtime_error(std::string(__func__) + ": keypool entry misclassified");
+        }
         // If the key was pre-split keypool, we don't care about what type it is
-        if (set_pre_split_keypool.empty() && (keypool.fInternal != fReturningInternal || keypool.fMWEB != fMWEB)) {
+        if (set_pre_split_keypool.empty() && keypool.fInternal != fReturningInternal) {
             throw std::runtime_error(std::string(__func__) + ": keypool entry misclassified");
         }
         if (!keypool.vchPubKey.IsValid()) {
@@ -1548,8 +1550,7 @@ void LegacyScriptPubKeyMan::MarkReserveKeysAsUsed(int64_t keypool_id)
 std::vector<CKeyID> GetAffectedKeys(const DestinationAddr& spk, const SigningProvider& provider)
 {
     if (spk.IsMWEB()) {
-        CPubKey spend_pubkey(spk.GetMWEBAddress().GetSpendPubKey().vec());
-        return std::vector<CKeyID>{spend_pubkey.GetID()};
+        return std::vector<CKeyID>{spk.GetMWEBAddress().GetSpendPubKey().GetID()};
     }
 
     std::vector<DestinationAddr> dummy;

src/wallet/test/scriptpubkeyman_tests.cpp

@@ -66,24 +66,21 @@ BOOST_AUTO_TEST_CASE(StealthAddresses)
     StealthAddress change_address = mweb_keychain->GetStealthAddress(0);
     BOOST_CHECK(EncodeDestination(change_address) == "ltcmweb1qq20e2arnhvxw97katjkmsd35agw3capxjkrkh7dk8d30rczm8ypxuq329nwh2twmchhqn3jqh7ua4ps539f6aazh79jy76urqht4qa59ts3at6gf");
     BOOST_CHECK(keyman.IsMine(change_address) == ISMINE_SPENDABLE);
-    CPubKey change_pubkey(change_address.B().vec());
-    BOOST_CHECK(keyman.GetAllReserveKeys().find(change_pubkey.GetID()) == keyman.GetAllReserveKeys().end());
+    BOOST_CHECK(keyman.GetAllReserveKeys().find(change_address.B().GetID()) == keyman.GetAllReserveKeys().end());
     BOOST_CHECK(*keyman.GetMetadata(change_address)->mweb_index == 0);
 
     // Check "peg-in" (idx=1) address is USED
     StealthAddress pegin_address = mweb_keychain->GetStealthAddress(1);
     BOOST_CHECK(EncodeDestination(pegin_address) == "ltcmweb1qqg5hddkl4uhspjwg9tkmatxa4s6gswdaq9swl8vsg5xxznmye7phcqatzc62mzkg788tsrfcuegxe9q3agf5cplw7ztqdusqf7x3n2tl55x4gvyt");
     BOOST_CHECK(keyman.IsMine(pegin_address) == ISMINE_SPENDABLE);
-    CPubKey pegin_pubkey(pegin_address.B().vec());
-    BOOST_CHECK(keyman.GetAllReserveKeys().find(pegin_pubkey.GetID()) == keyman.GetAllReserveKeys().end());
+    BOOST_CHECK(keyman.GetAllReserveKeys().find(pegin_address.B().GetID()) == keyman.GetAllReserveKeys().end());
     BOOST_CHECK(*keyman.GetMetadata(pegin_address)->mweb_index == 1);
 
     // Check first receive (idx=2) address is UNUSED
     StealthAddress receive_address = mweb_keychain->GetStealthAddress(2);
     BOOST_CHECK(EncodeDestination(receive_address) == "ltcmweb1qq0yq03ewm830ugmkkvrvjmyyeslcpwk8ayd7k27qx63sryy6kx3ksqm3k6jd24ld3r5dp5lzx7rm7uyxfujf8sn7v4nlxeqwrcq6k6xxwqdc6tl3");
     BOOST_CHECK(keyman.IsMine(receive_address) == ISMINE_SPENDABLE);
-    CPubKey receive_pubkey(receive_address.B().vec());
-    BOOST_CHECK(keyman.GetAllReserveKeys().find(receive_pubkey.GetID()) != keyman.GetAllReserveKeys().end());
+    BOOST_CHECK(keyman.GetAllReserveKeys().find(receive_address.B().GetID()) != keyman.GetAllReserveKeys().end());
     BOOST_CHECK(*keyman.GetMetadata(receive_address)->mweb_index == 2);
 
     BOOST_CHECK(keyman.GetHDChain().nMWEBIndexCounter == 1002);