From deb791c7ba057a3765d09b12bf3e55547a5298e4 Mon Sep 17 00:00:00 2001
From: Andrew Chow <achow101-github@achow101.com>
Date: Wed, 26 Feb 2020 16:52:03 -0500
Subject: [PATCH] Only cache xpubs that have a hardened last step

Also adds tests for this:
For ranged descriptors with unhardened derivation, we expect to
find parent keys in the cache but no child keys.

For descriptors containing an xpub but do not have unhardened derivation
(i.e. hardened derivation or single xpub with or without derivation),
we expect to find all of the keys in the cache, and the same
number of keys in the cache as in the SigningProvider.

For everything else (no xpub), nothing should be cached at all.
---
 src/script/descriptor.cpp     |  3 ++-
 src/test/descriptor_tests.cpp | 48 +++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/src/script/descriptor.cpp b/src/script/descriptor.cpp
index 30384fbd06..588e583339 100644
--- a/src/script/descriptor.cpp
+++ b/src/script/descriptor.cpp
@@ -345,10 +345,11 @@ public:
         key_out = final_extkey.pubkey;
 
         if (write_cache) {
-            write_cache->CacheDerivedExtPubKey(m_expr_index, pos, final_extkey);
             // Only cache parent if there is any unhardened derivation
             if (m_derive != DeriveType::HARDENED) {
                 write_cache->CacheParentExtPubKey(m_expr_index, parent_extkey);
+            } else if (final_info_out.path.size() > 0) {
+                write_cache->CacheDerivedExtPubKey(m_expr_index, pos, final_extkey);
             }
         }
 
diff --git a/src/test/descriptor_tests.cpp b/src/test/descriptor_tests.cpp
index 7f20a5c022..3154c619d2 100644
--- a/src/test/descriptor_tests.cpp
+++ b/src/test/descriptor_tests.cpp
@@ -149,6 +149,54 @@ void DoCheck(const std::string& prv, const std::string& pub, int flags, const st
             BOOST_CHECK(script_provider.scripts == script_provider_cached.scripts);
             BOOST_CHECK(script_provider.origins == script_provider_cached.origins);
 
+            // Check whether keys are in the cache
+            const auto& der_xpub_cache = desc_cache.GetCachedDerivedExtPubKeys();
+            const auto& parent_xpub_cache = desc_cache.GetCachedParentExtPubKeys();
+            if ((flags & RANGE) && !(flags & DERIVE_HARDENED)) {
+                // For ranged, unhardened derivation, None of the keys in origins should appear in the cache but the cache should have parent keys
+                // But we can derive one level from each of those parent keys and find them all
+                BOOST_CHECK(der_xpub_cache.empty());
+                BOOST_CHECK(parent_xpub_cache.size() > 0);
+                std::set<CPubKey> pubkeys;
+                for (const auto& xpub_pair : parent_xpub_cache) {
+                    const CExtPubKey& xpub = xpub_pair.second;
+                    CExtPubKey der;
+                    xpub.Derive(der, i);
+                    pubkeys.insert(der.pubkey);
+                }
+                for (const auto& origin_pair : script_provider_cached.origins) {
+                    const CPubKey& pk = origin_pair.second.first;
+                    BOOST_CHECK(pubkeys.count(pk) > 0);
+                }
+            } else if (pub1.find("xpub") != std::string::npos) {
+                // For ranged, hardened derivation, or not ranged, but has an xpub, all of the keys should appear in the cache
+                BOOST_CHECK(der_xpub_cache.size() + parent_xpub_cache.size() == script_provider_cached.origins.size());
+                // Get all of the derived pubkeys
+                std::set<CPubKey> pubkeys;
+                for (const auto& xpub_map_pair : der_xpub_cache) {
+                    for (const auto& xpub_pair : xpub_map_pair.second) {
+                        const CExtPubKey& xpub = xpub_pair.second;
+                        pubkeys.insert(xpub.pubkey);
+                    }
+                }
+                // Derive one level from all of the parents
+                for (const auto& xpub_pair : parent_xpub_cache) {
+                    const CExtPubKey& xpub = xpub_pair.second;
+                    pubkeys.insert(xpub.pubkey);
+                    CExtPubKey der;
+                    xpub.Derive(der, i);
+                    pubkeys.insert(der.pubkey);
+                }
+                for (const auto& origin_pair : script_provider_cached.origins) {
+                    const CPubKey& pk = origin_pair.second.first;
+                    BOOST_CHECK(pubkeys.count(pk) > 0);
+                }
+            } else {
+                // No xpub, nothing should be cached
+                BOOST_CHECK(der_xpub_cache.empty());
+                BOOST_CHECK(parent_xpub_cache.empty());
+            }
+
             // Make sure we can expand using cached xpubs for unhardened derivation
             if (!(flags & DERIVE_HARDENED)) {
                 // Evaluate the descriptor at i + 1