Add private key derivation functions to descriptors

2 files changed, 48 insertions, 8 deletions

diff --git a/src/script/descriptor.cpp b/src/script/descriptor.cpp
index 9be87fabb0..50119ba184 100644
--- a/src/script/descriptor.cpp
+++ b/src/script/descriptor.cpp
@@ -164,6 +164,9 @@ struct PubkeyProvider
 
     /** Get the descriptor string form including private data (if available in arg). */
     virtual bool ToPrivateString(const SigningProvider& arg, std::string& out) const = 0;
+
+    /** Derive a private key, if private data is available in arg. */
+    virtual bool GetPrivKey(int pos, const SigningProvider& arg, CKey& key) const = 0;
 };
 
 class OriginPubkeyProvider final : public PubkeyProvider
@@ -195,6 +198,10 @@ public:
         ret = "[" + OriginString() + "]" + std::move(sub);
         return true;
     }
+    bool GetPrivKey(int pos, const SigningProvider& arg, CKey& key) const override
+    {
+        return m_provider->GetPrivKey(pos, arg, key);
+    }
 };
 
 /** An object representing a parsed constant public key in a descriptor. */
@@ -222,6 +229,10 @@ public:
         ret = EncodeSecret(key);
         return true;
     }
+    bool GetPrivKey(int pos, const SigningProvider& arg, CKey& key) const override
+    {
+        return arg.GetKey(m_pubkey.GetID(), key);
+    }
 };
 
 enum class DeriveType {
@@ -266,14 +277,9 @@ public:
     {
         if (key) {
             if (IsHardened()) {
-                CExtKey extkey;
-                if (!GetExtKey(arg, extkey)) return false;
-                for (auto entry : m_path) {
-                    extkey.Derive(extkey, entry);
-                }
-                if (m_derive == DeriveType::UNHARDENED) extkey.Derive(extkey, pos);
-                if (m_derive == DeriveType::HARDENED) extkey.Derive(extkey, pos | 0x80000000UL);
-                *key = extkey.Neuter().pubkey;
+                CKey priv_key;
+                if (!GetPrivKey(pos, arg, priv_key)) return false;
+                *key = priv_key.GetPubKey();
             } else {
                 // TODO: optimize by caching
                 CExtPubKey extkey = m_extkey;
@@ -312,6 +318,18 @@ public:
         }
         return true;
     }
+    bool GetPrivKey(int pos, const SigningProvider& arg, CKey& key) const override
+    {
+        CExtKey extkey;
+        if (!GetExtKey(arg, extkey)) return false;
+        for (auto entry : m_path) {
+            extkey.Derive(extkey, entry);
+        }
+        if (m_derive == DeriveType::UNHARDENED) extkey.Derive(extkey, pos);
+        if (m_derive == DeriveType::HARDENED) extkey.Derive(extkey, pos | 0x80000000UL);
+        key = extkey.key;
+        return true;
+    }
 };
 
 /** Base class for all Descriptor implementations. */
@@ -462,6 +480,20 @@ public:
         Span<const unsigned char> span = MakeSpan(cache);
         return ExpandHelper(pos, DUMMY_SIGNING_PROVIDER, &span, output_scripts, out, nullptr) && span.size() == 0;
     }
+
+    void ExpandPrivate(int pos, const SigningProvider& provider, FlatSigningProvider& out) const final
+    {
+        for (const auto& p : m_pubkey_args) {
+            CKey key;
+            if (!p->GetPrivKey(pos, provider, key)) continue;
+            out.keys.emplace(key.GetPubKey().GetID(), key);
+        }
+        if (m_script_arg) {
+            FlatSigningProvider subprovider;
+            m_script_arg->ExpandPrivate(pos, provider, subprovider);
+            out = Merge(out, subprovider);
+        }
+    }
 };
 
 /** Construct a vector with one element, which is moved into it. */
diff --git a/src/script/descriptor.h b/src/script/descriptor.h
index 907a102284..af7ae229ca 100644
--- a/src/script/descriptor.h
+++ b/src/script/descriptor.h
@@ -60,6 +60,14 @@ struct Descriptor {
      * out: scripts and public keys necessary for solving the expanded scriptPubKeys will be put here (may be equal to provider).
      */
     virtual bool ExpandFromCache(int pos, const std::vector<unsigned char>& cache, std::vector<CScript>& output_scripts, FlatSigningProvider& out) const = 0;
+
+    /** Expand the private key for a descriptor at a specified position, if possible.
+     *
+     * pos: the position at which to expand the descriptor. If IsRange() is false, this is ignored.
+     * provider: the provider to query for the private keys.
+     * out: any private keys available for the specified pos will be placed here.
+     */
+    virtual void ExpandPrivate(int pos, const SigningProvider& provider, FlatSigningProvider& out) const = 0;
 };
 
 /** Parse a descriptor string. Included private keys are put in out.