Import public keys in order

Do public key imports in the order that they are specified in the import
or in the descriptor range.

2 files changed, 62 insertions, 29 deletions

diff --git a/src/wallet/rpcdump.cpp b/src/wallet/rpcdump.cpp
index 30555886f7..930274f8a1 100644
--- a/src/wallet/rpcdump.cpp
+++ b/src/wallet/rpcdump.cpp
@@ -967,7 +967,7 @@ static std::string RecurseImportData(const CScript& script, ImportData& import_d
     }
 }
 
-static UniValue ProcessImportLegacy(ImportData& import_data, std::map<CKeyID, CPubKey>& pubkey_map, std::map<CKeyID, CKey>& privkey_map, std::set<CScript>& script_pub_keys, bool& have_solving_data, const UniValue& data)
+static UniValue ProcessImportLegacy(ImportData& import_data, std::map<CKeyID, CPubKey>& pubkey_map, std::map<CKeyID, CKey>& privkey_map, std::set<CScript>& script_pub_keys, bool& have_solving_data, const UniValue& data, std::vector<CKeyID>& ordered_pubkeys)
 {
     UniValue warnings(UniValue::VARR);
 
@@ -1038,6 +1038,7 @@ static UniValue ProcessImportLegacy(ImportData& import_data, std::map<CKeyID, CP
             throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Pubkey \"" + str + "\" is not a valid public key");
         }
         pubkey_map.emplace(pubkey.GetID(), pubkey);
+        ordered_pubkeys.push_back(pubkey.GetID());
     }
     for (size_t i = 0; i < keys.size(); ++i) {
         const auto& str = keys[i].get_str();
@@ -1110,7 +1111,7 @@ static UniValue ProcessImportLegacy(ImportData& import_data, std::map<CKeyID, CP
     return warnings;
 }
 
-static UniValue ProcessImportDescriptor(ImportData& import_data, std::map<CKeyID, CPubKey>& pubkey_map, std::map<CKeyID, CKey>& privkey_map, std::set<CScript>& script_pub_keys, bool& have_solving_data, const UniValue& data)
+static UniValue ProcessImportDescriptor(ImportData& import_data, std::map<CKeyID, CPubKey>& pubkey_map, std::map<CKeyID, CKey>& privkey_map, std::set<CScript>& script_pub_keys, bool& have_solving_data, const UniValue& data, std::vector<CKeyID>& ordered_pubkeys)
 {
     UniValue warnings(UniValue::VARR);
 
@@ -1144,22 +1145,25 @@ static UniValue ProcessImportDescriptor(ImportData& import_data, std::map<CKeyID
 
     const UniValue& priv_keys = data.exists("keys") ? data["keys"].get_array() : UniValue();
 
-    FlatSigningProvider out_keys;
-
     // Expand all descriptors to get public keys and scripts.
     // TODO: get private keys from descriptors too
     for (int i = range_start; i <= range_end; ++i) {
+        FlatSigningProvider out_keys;
         std::vector<CScript> scripts_temp;
         parsed_desc->Expand(i, keys, scripts_temp, out_keys);
         std::copy(scripts_temp.begin(), scripts_temp.end(), std::inserter(script_pub_keys, script_pub_keys.end()));
-    }
+        for (const auto& key_pair : out_keys.pubkeys) {
+            ordered_pubkeys.push_back(key_pair.first);
+        }
 
-    for (const auto& x : out_keys.scripts) {
-        import_data.import_scripts.emplace(x.second);
+        for (const auto& x : out_keys.scripts) {
+            import_data.import_scripts.emplace(x.second);
+        }
+
+        std::copy(out_keys.pubkeys.begin(), out_keys.pubkeys.end(), std::inserter(pubkey_map, pubkey_map.end()));
+        import_data.key_origins.insert(out_keys.origins.begin(), out_keys.origins.end());
     }
 
-    std::copy(out_keys.pubkeys.begin(), out_keys.pubkeys.end(), std::inserter(pubkey_map, pubkey_map.end()));
-    import_data.key_origins.insert(out_keys.origins.begin(), out_keys.origins.end());
     for (size_t i = 0; i < priv_keys.size(); ++i) {
         const auto& str = priv_keys[i].get_str();
         CKey key = DecodeSecret(str);
@@ -1218,14 +1222,15 @@ static UniValue ProcessImport(CWallet * const pwallet, const UniValue& data, con
         std::map<CKeyID, CPubKey> pubkey_map;
         std::map<CKeyID, CKey> privkey_map;
         std::set<CScript> script_pub_keys;
+        std::vector<CKeyID> ordered_pubkeys;
         bool have_solving_data;
 
         if (data.exists("scriptPubKey") && data.exists("desc")) {
             throw JSONRPCError(RPC_INVALID_PARAMETER, "Both a descriptor and a scriptPubKey should not be provided.");
         } else if (data.exists("scriptPubKey")) {
-            warnings = ProcessImportLegacy(import_data, pubkey_map, privkey_map, script_pub_keys, have_solving_data, data);
+            warnings = ProcessImportLegacy(import_data, pubkey_map, privkey_map, script_pub_keys, have_solving_data, data, ordered_pubkeys);
         } else if (data.exists("desc")) {
-            warnings = ProcessImportDescriptor(import_data, pubkey_map, privkey_map, script_pub_keys, have_solving_data, data);
+            warnings = ProcessImportDescriptor(import_data, pubkey_map, privkey_map, script_pub_keys, have_solving_data, data, ordered_pubkeys);
         } else {
             throw JSONRPCError(RPC_INVALID_PARAMETER, "Either a descriptor or scriptPubKey must be provided.");
         }
@@ -1247,25 +1252,28 @@ static UniValue ProcessImport(CWallet * const pwallet, const UniValue& data, con
         for (const auto& entry : import_data.import_scripts) {
             if (!pwallet->HaveCScript(CScriptID(entry)) && !pwallet->AddCScript(entry)) {
                 throw JSONRPCError(RPC_WALLET_ERROR, "Error adding script to wallet");
+             }
+         }
+         for (const auto& entry : privkey_map) {
+             const CKey& key = entry.second;
+             CPubKey pubkey = key.GetPubKey();
+             const CKeyID& id = entry.first;
+             assert(key.VerifyPubKey(pubkey));
+             pwallet->mapKeyMetadata[id].nCreateTime = timestamp;
+             // If the private key is not present in the wallet, insert it.
+             if (!pwallet->HaveKey(id) && !pwallet->AddKeyPubKey(key, pubkey)) {
+                 throw JSONRPCError(RPC_WALLET_ERROR, "Error adding key to wallet");
+             }
+             pwallet->UpdateTimeFirstKey(timestamp);
+         }
+        for (const CKeyID& id : ordered_pubkeys) {
+            auto entry = pubkey_map.find(id);
+            if (entry == pubkey_map.end()) {
+                continue;
             }
-        }
-        for (const auto& entry : privkey_map) {
-            const CKey& key = entry.second;
-            CPubKey pubkey = key.GetPubKey();
-            const CKeyID& id = entry.first;
-            assert(key.VerifyPubKey(pubkey));
-            pwallet->mapKeyMetadata[id].nCreateTime = timestamp;
-            // If the private key is not present in the wallet, insert it.
-            if (!pwallet->HaveKey(id) && !pwallet->AddKeyPubKey(key, pubkey)) {
-                throw JSONRPCError(RPC_WALLET_ERROR, "Error adding key to wallet");
-            }
-            pwallet->UpdateTimeFirstKey(timestamp);
-        }
-        for (const auto& entry : pubkey_map) {
-            const CPubKey& pubkey = entry.second;
-            const CKeyID& id = entry.first;
-            CPubKey temp;
-            if (!pwallet->GetPubKey(id, temp) && !pwallet->AddWatchOnly(GetScriptForRawPubKey(pubkey), timestamp)) {
+             const CPubKey& pubkey = entry->second;
+             CPubKey temp;
+             if (!pwallet->GetPubKey(id, temp) && !pwallet->AddWatchOnly(GetScriptForRawPubKey(pubkey), timestamp)) {
                 throw JSONRPCError(RPC_WALLET_ERROR, "Error adding address to wallet");
             }
             const auto& key_orig_it = import_data.key_origins.find(id);
diff --git a/test/functional/wallet_importmulti.py b/test/functional/wallet_importmulti.py
index 3f77c5e5fd..46e3ab77c8 100755
--- a/test/functional/wallet_importmulti.py
+++ b/test/functional/wallet_importmulti.py
@@ -777,5 +777,30 @@ class ImportMultiTest(BitcoinTestFramework):
         assert_equal(result[0]['error']['code'], -8)
         assert_equal(result[0]['error']['message'], "Keys can only be imported to the keypool when private keys are disabled")
 
+        # Make sure ranged imports import keys in order
+        self.log.info('Key ranges should be imported in order')
+        wrpc = self.nodes[1].get_wallet_rpc("noprivkeys")
+        assert_equal(wrpc.getwalletinfo()["keypoolsize"], 0)
+        assert_equal(wrpc.getwalletinfo()["private_keys_enabled"], False)
+        xpub = "tpubDAXcJ7s7ZwicqjprRaEWdPoHKrCS215qxGYxpusRLLmJuT69ZSicuGdSfyvyKpvUNYBW1s2U3NSrT6vrCYB9e6nZUEvrqnwXPF8ArTCRXMY"
+        addresses = [
+            'bcrt1qtmp74ayg7p24uslctssvjm06q5phz4yrxucgnv', # m/0'/0'/0
+            'bcrt1q8vprchan07gzagd5e6v9wd7azyucksq2xc76k8', # m/0'/0'/1
+            'bcrt1qtuqdtha7zmqgcrr26n2rqxztv5y8rafjp9lulu', # m/0'/0'/2
+            'bcrt1qau64272ymawq26t90md6an0ps99qkrse58m640', # m/0'/0'/3
+            'bcrt1qsg97266hrh6cpmutqen8s4s962aryy77jp0fg0', # m/0'/0'/4
+        ]
+        result = wrpc.importmulti(
+            [{
+                'desc': 'wpkh([80002067/0h/0h]' + xpub + '/*)',
+                'keypool': True,
+                'timestamp': 'now',
+                'range' : {'start': 0, 'end': 4}
+            }]
+        )
+        for i in range(0, 5):
+            addr = wrpc.getnewaddress('', 'bech32')
+            assert_equal(addr, addresses[i])
+
 if __name__ == '__main__':
     ImportMultiTest().main()