Implement GetSolvingProvider for DescriptorScriptPubKeyMan

Internally, a GetSigningProvider function is introduced which allows for
some private keys to be optionally included. This can be called with a
script as the argument (i.e. a scriptPubKey from our wallet when we are
signing) or with a pubkey. In order to know what index to expand the
private keys for that pubkey, we need to also cache all of the pubkeys
involved when we expand the descriptor. So SetCache and TopUp are
updated to do this too.

2 files changed, 75 insertions, 1 deletions

diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
index cf2882ed61..67063ad23c 100644
--- a/src/wallet/scriptpubkeyman.cpp
+++ b/src/wallet/scriptpubkeyman.cpp
@@ -1680,6 +1680,15 @@ bool DescriptorScriptPubKeyMan::TopUp(unsigned int size)
         for (const CScript& script : scripts_temp) {
             m_map_script_pub_keys[script] = i;
         }
+        for (const auto& pk_pair : out_keys.pubkeys) {
+            const CPubKey& pubkey = pk_pair.second;
+            if (m_map_pubkeys.count(pubkey) != 0) {
+                // We don't need to give an error here.
+                // It doesn't matter which of many valid indexes the pubkey has, we just need an index where we can derive it and it's private key
+                continue;
+            }
+            m_map_pubkeys[pubkey] = i;
+        }
         // Write the cache
         for (const auto& parent_xpub_pair : temp_cache.GetCachedParentExtPubKeys()) {
             CExtPubKey xpub;
@@ -1876,9 +1885,55 @@ int64_t DescriptorScriptPubKeyMan::GetTimeFirstKey() const
     return m_wallet_descriptor.creation_time;
 }
 
+std::unique_ptr<FlatSigningProvider> DescriptorScriptPubKeyMan::GetSigningProvider(const CScript& script, bool include_private) const
+{
+    LOCK(cs_desc_man);
+
+    // Find the index of the script
+    auto it = m_map_script_pub_keys.find(script);
+    if (it == m_map_script_pub_keys.end()) {
+        return nullptr;
+    }
+    int32_t index = it->second;
+
+    return GetSigningProvider(index, include_private);
+}
+
+std::unique_ptr<FlatSigningProvider> DescriptorScriptPubKeyMan::GetSigningProvider(const CPubKey& pubkey) const
+{
+    LOCK(cs_desc_man);
+
+    // Find index of the pubkey
+    auto it = m_map_pubkeys.find(pubkey);
+    if (it == m_map_pubkeys.end()) {
+        return nullptr;
+    }
+    int32_t index = it->second;
+
+    // Always try to get the signing provider with private keys. This function should only be called during signing anyways
+    return GetSigningProvider(index, true);
+}
+
+std::unique_ptr<FlatSigningProvider> DescriptorScriptPubKeyMan::GetSigningProvider(int32_t index, bool include_private) const
+{
+    AssertLockHeld(cs_desc_man);
+    // Get the scripts, keys, and key origins for this script
+    std::unique_ptr<FlatSigningProvider> out_keys = MakeUnique<FlatSigningProvider>();
+    std::vector<CScript> scripts_temp;
+    if (!m_wallet_descriptor.descriptor->ExpandFromCache(index, m_wallet_descriptor.cache, scripts_temp, *out_keys)) return nullptr;
+
+    if (HavePrivateKeys() && include_private) {
+        FlatSigningProvider master_provider;
+        master_provider.keys = GetKeys();
+        m_wallet_descriptor.descriptor->ExpandPrivate(index, master_provider, *out_keys);
+    }
+
+    return out_keys;
+}
+
 std::unique_ptr<SigningProvider> DescriptorScriptPubKeyMan::GetSolvingProvider(const CScript& script) const
 {
-    return nullptr;
+    return GetSigningProvider(script, false);
 }
 
 bool DescriptorScriptPubKeyMan::CanProvide(const CScript& script, SignatureData& sigdata)
@@ -1938,6 +1993,15 @@ void DescriptorScriptPubKeyMan::SetCache(const DescriptorCache& cache)
             }
             m_map_script_pub_keys[script] = i;
         }
+        for (const auto& pk_pair : out_keys.pubkeys) {
+            const CPubKey& pubkey = pk_pair.second;
+            if (m_map_pubkeys.count(pubkey) != 0) {
+                // We don't need to give an error here.
+                // It doesn't matter which of many valid indexes the pubkey has, we just need an index where we can derive it and it's private key
+                continue;
+            }
+            m_map_pubkeys[pubkey] = i;
+        }
         m_max_cached_index++;
     }
 }
diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h
index 7522dce9c5..0932957419 100644
--- a/src/wallet/scriptpubkeyman.h
+++ b/src/wallet/scriptpubkeyman.h
@@ -488,10 +488,12 @@ private:
     WalletDescriptor m_wallet_descriptor GUARDED_BY(cs_desc_man);
 
     using ScriptPubKeyMap = std::map<CScript, int32_t>; // Map of scripts to descriptor range index
+    using PubKeyMap = std::map<CPubKey, int32_t>; // Map of pubkeys involved in scripts to descriptor range index
     using CryptedKeyMap = std::map<CKeyID, std::pair<CPubKey, std::vector<unsigned char>>>;
     using KeyMap = std::map<CKeyID, CKey>;
 
     ScriptPubKeyMap m_map_script_pub_keys GUARDED_BY(cs_desc_man);
+    PubKeyMap m_map_pubkeys GUARDED_BY(cs_desc_man);
     int32_t m_max_cached_index = -1;
 
     OutputType m_address_type;
@@ -508,6 +510,14 @@ private:
     bool AddDescriptorKeyWithDB(WalletBatch& batch, const CKey& key, const CPubKey &pubkey);
 
     KeyMap GetKeys() const EXCLUSIVE_LOCKS_REQUIRED(cs_desc_man);
+
+    // Fetch the SigningProvider for the given script and optionally include private keys
+    std::unique_ptr<FlatSigningProvider> GetSigningProvider(const CScript& script, bool include_private = false) const;
+    // Fetch the SigningProvider for the given pubkey and always include private keys. This should only be called by signing code.
+    std::unique_ptr<FlatSigningProvider> GetSigningProvider(const CPubKey& pubkey) const;
+    // Fetch the SigningProvider for a given index and optionally include private keys. Called by the above functions.
+    std::unique_ptr<FlatSigningProvider> GetSigningProvider(int32_t index, bool include_private = false) const EXCLUSIVE_LOCKS_REQUIRED(cs_desc_man);
+
 public:
     DescriptorScriptPubKeyMan(WalletStorage& storage, WalletDescriptor& descriptor)
         :   ScriptPubKeyMan(storage),