Implement SetupGeneration for DescriptorScriptPubKeyMan

achow101 · achow101 · commit e014886a3425 · 2020-04-23T13:59:48.000-04:00
diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
@@ -1577,6 +1577,73 @@ bool DescriptorScriptPubKeyMan::AddDescriptorKeyWithDB(WalletBatch& batch, const
     }
 }
 
+bool DescriptorScriptPubKeyMan::SetupDescriptorGeneration(const CExtKey& master_key)
+{
+    LOCK(cs_desc_man);
+    assert(m_storage.IsWalletFlagSet(WALLET_FLAG_DESCRIPTORS));
+
+    // Ignore when there is already a descriptor
+    if (m_wallet_descriptor.descriptor) {
+        return false;
+    }
+
+    int64_t creation_time = GetTime();
+
+    std::string xpub = EncodeExtPubKey(master_key.Neuter());
+
+    // Build descriptor string
+    std::string desc_prefix;
+    std::string desc_suffix = "/*)";
+    switch (m_address_type) {
+    case OutputType::LEGACY: {
+        desc_prefix = "pkh(" + xpub + "/44'";
+        break;
+    }
+    case OutputType::P2SH_SEGWIT: {
+        desc_prefix = "sh(wpkh(" + xpub + "/49'";
+        desc_suffix += ")";
+        break;
+    }
+    case OutputType::BECH32: {
+        desc_prefix = "wpkh(" + xpub + "/84'";
+        break;
+    }
+    default: assert(false);
+    }
+
+    // Mainnet derives at 0', testnet and regtest derive at 1'
+    if (Params().IsTestChain()) {
+        desc_prefix += "/1'";
+    } else {
+        desc_prefix += "/0'";
+    }
+
+    std::string internal_path = m_internal ? "/1" : "/0";
+    std::string desc_str = desc_prefix + "/0'" + internal_path + desc_suffix;
+
+    // Make the descriptor
+    FlatSigningProvider keys;
+    std::string error;
+    std::unique_ptr<Descriptor> desc = Parse(desc_str, keys, error, false);
+    WalletDescriptor w_desc(std::move(desc), creation_time, 0, 0, 0);
+    m_wallet_descriptor = w_desc;
+
+    // Store the master private key, and descriptor
+    WalletBatch batch(m_storage.GetDatabase());
+    if (!AddDescriptorKeyWithDB(batch, master_key.key, master_key.key.GetPubKey())) {
+        throw std::runtime_error(std::string(__func__) + ": writing descriptor master private key failed");
+    }
+    if (!batch.WriteDescriptor(GetID(), m_wallet_descriptor)) {
+        throw std::runtime_error(std::string(__func__) + ": writing descriptor failed");
+    }
+
+    // TopUp
+    TopUp();
+
+    m_storage.UnsetBlankWalletFlag(batch);
+    return true;
+}
+
 bool DescriptorScriptPubKeyMan::IsHDEnabled() const
 {
     LOCK(cs_desc_man);
diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h
@@ -532,6 +532,9 @@ class DescriptorScriptPubKeyMan : public ScriptPubKeyMan
 
     bool IsHDEnabled() const override;
 
+    //! Setup descriptors based on the given CExtkey
+    bool SetupDescriptorGeneration(const CExtKey& master_key);
+
     bool HavePrivateKeys() const override;
 
     int64_t GetOldestKeyPoolTime() const override;
