wallet: Add scan_utxo option to getbalances RPC

Add the option to getbalances by scanning the utxo set
BrandonOdiwuor · Nov 28, 2023 · b210d3b · b210d3b
1 parent d752349
commit b210d3b
Show file tree

Hide file tree

Showing 6 changed files with 69 additions and 4 deletions.
diff --git a/src/interfaces/chain.h b/src/interfaces/chain.h
@@ -188,6 +188,9 @@ class Chain
     //! populates the values.
     virtual void findCoins(std::map<COutPoint, Coin>& coins) = 0;
 
+    //! Scan UTXO set from coins belonging to the output_scripts
+    virtual void getCoinsByScript(std::set<CScript>& output_scripts, std::map<COutPoint, Coin>& coins) = 0;
+
     //! Estimate fraction of total transactions verified if blocks up to
     //! the specified block hash are verified.
     virtual double guessVerificationProgress(const uint256& block_hash) = 0;

diff --git a/src/node/coin.cpp b/src/node/coin.cpp
@@ -23,4 +23,21 @@ void FindCoins(const NodeContext& node, std::map<COutPoint, Coin>& coins)
         }
     }
 }
+
+void GetCoins(const NodeContext& node, std::set<CScript>& output_scripts, std::map<COutPoint, Coin>& coins)
+{
+    assert(node.chainman);
+    LOCK(cs_main);
+    std::unique_ptr<CCoinsViewCursor> cursor = node.chainman->ActiveChainstate().CoinsDB().Cursor();
+    while (cursor->Valid()) {
+        COutPoint key;
+        Coin coin;
+        if (cursor->GetKey(key) && cursor->GetValue(coin)) {
+            if (output_scripts.count(coin.out.scriptPubKey)) {
+                coins.emplace(key, coin);
+            }
+        }
+        cursor->Next();
+    }
+}
 } // namespace node
diff --git a/src/node/coin.h b/src/node/coin.h
@@ -6,9 +6,11 @@
 #define BITCOIN_NODE_COIN_H
 
 #include <map>
+#include <set>
 
 class COutPoint;
 class Coin;
+class CScript;
 
 namespace node {
 struct NodeContext;
@@ -22,6 +24,15 @@ struct NodeContext;
  * @param[in,out] coins map to fill
  */
 void FindCoins(const node::NodeContext& node, std::map<COutPoint, Coin>& coins);
+
+/**
+ * Scans the UTXO set for coins belonging to output_scripts.
+ *
+ * @param[in] node The node context to use for lookup
+ * @param[in] output_scripts The scripts to scan for coins
+ * @param[in,out] coins map to fill
+ */
+void GetCoins(const NodeContext& node, std::set<CScript>& output_scripts, std::map<COutPoint, Coin>& coins);
 } // namespace node
 
 #endif // BITCOIN_NODE_COIN_H
diff --git a/src/node/interfaces.cpp b/src/node/interfaces.cpp
@@ -610,6 +610,10 @@ class ChainImpl : public Chain
                int{FillBlock(block2, block2_out, lock, active, chainman().m_blockman)};
     }
     void findCoins(std::map<COutPoint, Coin>& coins) override { return FindCoins(m_node, coins); }
+    void getCoinsByScript(std::set<CScript>& output_scripts, std::map<COutPoint, Coin>& coins) override
+    {
+        return GetCoins(m_node, output_scripts, coins);
+    }
     double guessVerificationProgress(const uint256& block_hash) override
     {
         LOCK(::cs_main);

diff --git a/src/rpc/client.cpp b/src/rpc/client.cpp
@@ -185,6 +185,7 @@ static const CRPCConvertParam vRPCConvertParams[] =
     { "lockunspent", 0, "unlock" },
     { "lockunspent", 1, "transactions" },
     { "lockunspent", 2, "persistent" },
+    { "getbalances", 0, "scan_utxoset" },
     { "send", 0, "outputs" },
     { "send", 1, "conf_target" },
     { "send", 3, "fee_rate"},

diff --git a/src/wallet/rpc/coins.cpp b/src/wallet/rpc/coins.cpp
@@ -432,7 +432,9 @@ RPCHelpMan getbalances()
     return RPCHelpMan{
         "getbalances",
         "Returns an object with all balances in " + CURRENCY_UNIT + ".\n",
-        {},
+        {
+            {"scan_utxoset", RPCArg::Type::BOOL, RPCArg::Default{false}, "Whether to scan (true) or not to scan (false) the unspent transactions output set"},
+        },
         RPCResult{
             RPCResult::Type::OBJ, "", "",
             {
@@ -453,8 +455,10 @@ RPCHelpMan getbalances()
             }
             },
         RPCExamples{
-            HelpExampleCli("getbalances", "") +
-            HelpExampleRpc("getbalances", "")},
+            HelpExampleCli("getbalances", "")
+            + HelpExampleRpc("getbalances", "")
+            + HelpExampleCli("getbalances", "true")
+        },
         [&](const RPCHelpMan& self, const JSONRPCRequest& request) -> UniValue
 {
     const std::shared_ptr<const CWallet> rpc_wallet = GetWalletForJSONRPCRequest(request);
@@ -468,10 +472,35 @@ RPCHelpMan getbalances()
     LOCK(wallet.cs_wallet);
 
     const auto bal = GetBalance(wallet);
+
+    CAmount utxo_balance = 0;
+
+    bool scan_utxoset = false;
+    if (!request.params[0].isNull()) {
+        scan_utxoset = request.params[0].get_bool();
+    }
+    if (scan_utxoset) {
+        std::set<CScript> output_scripts;
+        for (const auto spkm : wallet.GetAllScriptPubKeyMans()) {
+            for (const auto& script : spkm->GetScriptPubKeys()) {
+                output_scripts.emplace(script);
+            }
+        }
+
+        // Get coins belonging to wallet scripts from utxo set
+        std::map<COutPoint, Coin> coins;
+        wallet.chain().getCoinsByScript(output_scripts, coins);
+
+        for (const auto& it : coins) {
+            const Coin& coin = it.second;
+            utxo_balance += coin.out.nValue;
+        }
+    }
+
     UniValue balances{UniValue::VOBJ};
     {
         UniValue balances_mine{UniValue::VOBJ};
-        balances_mine.pushKV("trusted", ValueFromAmount(bal.m_mine_trusted));
+        balances_mine.pushKV("trusted", ValueFromAmount(scan_utxoset ? utxo_balance : bal.m_mine_trusted));
         balances_mine.pushKV("untrusted_pending", ValueFromAmount(bal.m_mine_untrusted_pending));
         balances_mine.pushKV("immature", ValueFromAmount(bal.m_mine_immature));
         if (wallet.IsWalletFlagSet(WALLET_FLAG_AVOID_REUSE)) {