Add wallet keypool self-test (#91)

Co-authored-by: Gilcimar Duarte <ti@gpm.infra>

6 files changed, 309 insertions, 6 deletions

diff --git a/src/Makefile b/src/Makefile
index 4e104c0..248e62c 100644
--- a/src/Makefile
+++ b/src/Makefile
@@ -41,14 +41,16 @@ LIBS = -L$(SECP_LIB) -lsecp256k1 \
 
 HEADERS = headers_core.h headers.h util.h main.h serialize.h uint256.h \
           key.h bignum.h script.h db.h base58.h net.h nostr.h \
-          randomx_pow.h btfaddr.h btfchan.h btfrv.h btftunnel.h market.h compat.h
+          randomx_pow.h btfaddr.h btfchan.h btfrv.h btftunnel.h market.h \
+          selftest.h compat.h
 
 IMGUI_OBJS = obj/imgui.o obj/imgui_draw.o obj/imgui_tables.o \
              obj/imgui_widgets.o obj/imgui_impl_glfw.o obj/imgui_impl_opengl3.o
 
 NODE_OBJS = obj/util.o obj/script.o obj/db.o obj/net.o obj/main.o \
             obj/market.o obj/sha.o obj/nostr.o obj/randomx_pow.o \
-            obj/btfaddr.o obj/btfchan.o obj/btfrv.o obj/btftunnel.o obj/rpc.o
+            obj/btfaddr.o obj/btfchan.o obj/btfrv.o obj/btftunnel.o obj/rpc.o \
+            obj/selftest.o
 
 GUI_OBJS = obj/gui.o obj/main_gui.o
 
@@ -126,5 +128,6 @@ clean:
 test_bip32: test_bip32.cpp bip32.cpp bip32.h bip39_english.h
         $(CXX) $(CXXFLAGS) -o $@ test_bip32.cpp bip32.cpp -lssl -lcrypto
 
-tests: test_bip32
+tests: test_bip32 bitflash-node
         ./test_bip32
+        ./bitflash-node /selftest=wallet-keypool
diff --git a/src/main.cpp b/src/main.cpp
index 68ae44c..82a98d2 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -217,6 +217,11 @@ vector<unsigned char> GetKeyFromPool()
             vector<unsigned char> vchPubKey = mi->second;
             mapKeyPool.erase(mi);
             CWalletDB().ErasePool(nIndex);
+            // Leave the wallet with a full backup window after each draw too,
+            // not only immediately before a draw. If topping up fails, the key
+            // already chosen is still safe to use -- its private half was
+            // written before it entered the pool.
+            TopUpKeyPool();
             return vchPubKey;
         }
     }
diff --git a/src/main_gui.cpp b/src/main_gui.cpp
index c9e27a9..2bcdb83 100644
--- a/src/main_gui.cpp
+++ b/src/main_gui.cpp
@@ -1,6 +1,7 @@
 // Bitflash entry point -- starts node threads then runs GUI (or headless).
 
 #include "headers_core.h"
+#include "selftest.h"
 #include <thread>          // hardware_concurrency, to sanity-check /genproclimit
 #ifndef _WIN32
 #include <csignal>
@@ -48,6 +49,7 @@ static void PrintUsage()
     printf("  /debug\n");
     printf("  /gen\n");
     printf("  /nogui or /daemon\n");
+    printf("  /selftest=wallet-keypool\n");
     printf("\n");
     printf("Mining mode:\n");
     printf("  /operator\n");
@@ -219,6 +221,10 @@ int main(int argc, char* argv[])
 
     ParseStartupArguments(argc, argv);
 
+    string strSelfTest = argval2(argc, argv, "/selftest", "-selftest");
+    if (!strSelfTest.empty())
+        return RunSelfTest(strSelfTest);
+
     // Berkeley DB reports failure by throwing, and CDB's constructor lets it
     // through. Nothing on this path caught anything, so an unreadable
     // blkindex.dat or wallet.dat -- one written by another platform's Berkeley
diff --git a/src/makefile.mingw b/src/makefile.mingw
index e4dee59..c263f4e 100644
--- a/src/makefile.mingw
+++ b/src/makefile.mingw
@@ -41,13 +41,15 @@ IMGUI_OBJS = obj/imgui.o obj/imgui_draw.o obj/imgui_tables.o \
 
 NODE_OBJS = obj/util.o obj/script.o obj/db.o obj/net.o obj/main.o \
             obj/market.o obj/sha.o obj/nostr.o obj/randomx_pow.o \
-            obj/btfaddr.o obj/btfchan.o obj/btfrv.o obj/btftunnel.o obj/rpc.o
+            obj/btfaddr.o obj/btfchan.o obj/btfrv.o obj/btftunnel.o obj/rpc.o \
+            obj/selftest.o
 
 GUI_OBJS = obj/gui.o obj/main_gui.o obj/bitflash_res.o
 
 HEADERS = headers_core.h headers.h util.h main.h serialize.h uint256.h \
           key.h bignum.h script.h db.h base58.h net.h nostr.h \
-          randomx_pow.h btfaddr.h btfchan.h btfrv.h btftunnel.h market.h
+          randomx_pow.h btfaddr.h btfchan.h btfrv.h btftunnel.h market.h \
+          selftest.h
 
 all: bitflash.exe
 
@@ -98,5 +100,6 @@ test_bip32.exe: test_bip32.cpp bip32.cpp bip32.h bip39_english.h
         $(CXX) $(CXXFLAGS) -o $@ test_bip32.cpp bip32.cpp -lssl -lcrypto -lcrypt32
         @echo "Built: test_bip32.exe"
 
-tests: test_bip32.exe
+tests: test_bip32.exe bitflash.exe
         ./test_bip32.exe
+        ./bitflash.exe /selftest=wallet-keypool
diff --git a/src/selftest.cpp b/src/selftest.cpp
new file mode 100644
index 0000000..d649b23
--- /dev/null
+++ b/src/selftest.cpp
@@ -0,0 +1,275 @@
+// Copyright (c) 2026 Bitflash developers
+// Distributed under the MIT/X11 software license.
+//
+// Integration self-tests that need the real wallet/database code linked into
+// the node. These are intentionally run behind an explicit command-line flag
+// and against a temporary data directory.
+
+#include "headers_core.h"
+#include "selftest.h"
+
+#include <mutex>
+#include <stdexcept>
+#include <thread>
+
+#ifdef _WIN32
+#include <direct.h>
+#else
+#include <dirent.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+#endif
+
+static std::string HexStrLocal(const std::vector<unsigned char>& v)
+{
+    static const char* hex = "0123456789abcdef";
+    std::string out;
+    out.reserve(v.size() * 2);
+    for (unsigned char c : v)
+    {
+        out.push_back(hex[c >> 4]);
+        out.push_back(hex[c & 0x0f]);
+    }
+    return out;
+}
+
+static bool MakeTempDir(std::string& out)
+{
+#ifdef _WIN32
+    char base[MAX_PATH + 1] = {0};
+    if (!GetTempPathA(MAX_PATH, base))
+        return false;
+    for (int i = 0; i < 100; i++)
+    {
+        std::string path = strprintf("%sbitflash-selftest-%lu-%lld-%d",
+            base, (unsigned long)GetCurrentProcessId(), (long long)GetTime(), i);
+        if (CreateDirectoryA(path.c_str(), NULL))
+        {
+            out = path;
+            return true;
+        }
+    }
+    return false;
+#else
+    char tmpl[256];
+    snprintf(tmpl, sizeof(tmpl), "/tmp/bitflash-selftest-%ld-XXXXXX", (long)getpid());
+    char* p = mkdtemp(tmpl);
+    if (!p)
+        return false;
+    out = p;
+    return true;
+#endif
+}
+
+static bool GetCurrentDir(std::string& out)
+{
+#ifdef _WIN32
+    char buf[MAX_PATH + 1] = {0};
+    DWORD n = GetCurrentDirectoryA(MAX_PATH, buf);
+    if (n == 0 || n > MAX_PATH)
+        return false;
+    out = buf;
+    return true;
+#else
+    char buf[4096];
+    if (!getcwd(buf, sizeof(buf)))
+        return false;
+    out = buf;
+    return true;
+#endif
+}
+
+static bool SetCurrentDir(const std::string& path)
+{
+#ifdef _WIN32
+    return SetCurrentDirectoryA(path.c_str()) != 0;
+#else
+    return chdir(path.c_str()) == 0;
+#endif
+}
+
+static void RemoveTree(const std::string& path)
+{
+#ifdef _WIN32
+    DWORD attr = GetFileAttributesA(path.c_str());
+    if (attr == INVALID_FILE_ATTRIBUTES)
+        return;
+    if (!(attr & FILE_ATTRIBUTE_DIRECTORY))
+    {
+        DeleteFileA(path.c_str());
+        return;
+    }
+
+    WIN32_FIND_DATAA findData;
+    std::string pattern = path + "\\*";
+    HANDLE hFind = FindFirstFileA(pattern.c_str(), &findData);
+    if (hFind != INVALID_HANDLE_VALUE)
+    {
+        do
+        {
+            std::string name = findData.cFileName;
+            if (name == "." || name == "..")
+                continue;
+            RemoveTree(path + "\\" + name);
+        }
+        while (FindNextFileA(hFind, &findData));
+        FindClose(hFind);
+    }
+    RemoveDirectoryA(path.c_str());
+#else
+    struct stat st;
+    if (lstat(path.c_str(), &st) != 0)
+        return;
+    if (!S_ISDIR(st.st_mode))
+    {
+        unlink(path.c_str());
+        return;
+    }
+
+    DIR* dir = opendir(path.c_str());
+    if (dir)
+    {
+        struct dirent* ent;
+        while ((ent = readdir(dir)) != NULL)
+        {
+            std::string name = ent->d_name;
+            if (name == "." || name == "..")
+                continue;
+            RemoveTree(path + "/" + name);
+        }
+        closedir(dir);
+    }
+    rmdir(path.c_str());
+#endif
+}
+
+static bool WalletHasPrivateKey(const std::vector<unsigned char>& vchPubKey)
+{
+    CPrivKey priv;
+    return CWalletDB("r").ReadKey(vchPubKey, priv) && !priv.empty();
+}
+
+static bool Check(bool condition, const char* message)
+{
+    if (condition)
+    {
+        printf("  ok   %s\n", message);
+        return true;
+    }
+    printf("  FAIL %s\n", message);
+    return false;
+}
+
+static int RunWalletKeyPoolSelfTest()
+{
+    printf("wallet-keypool self-test\n");
+
+    std::string tmp;
+    std::string cwd;
+    if (!MakeTempDir(tmp))
+    {
+        printf("  FAIL could not create a temporary data directory\n");
+        return 1;
+    }
+    if (!GetCurrentDir(cwd) || !SetCurrentDir(tmp))
+    {
+        printf("  FAIL could not move into the temporary data directory\n");
+        RemoveTree(tmp);
+        return 1;
+    }
+
+    int nFail = 0;
+    strSetDataDir = tmp;
+    printf("  temp datadir: %s\n", tmp.c_str());
+
+    try
+    {
+        if (!LoadWallet())
+            throw std::runtime_error("LoadWallet failed");
+
+        TopUpKeyPool();
+        int nInitialPoolSize = 0;
+        CRITICAL_BLOCK(cs_keyPool)
+            nInitialPoolSize = (int)mapKeyPool.size();
+        nFail += Check(nInitialPoolSize == KEYPOOL_SIZE, "TopUpKeyPool fills the pool") ? 0 : 1;
+
+        bool fPooledKeysStored = true;
+        CRITICAL_BLOCK(cs_keyPool)
+        {
+            for (map<int64, vector<unsigned char> >::const_iterator mi = mapKeyPool.begin();
+                 mi != mapKeyPool.end(); ++mi)
+                if (!WalletHasPrivateKey(mi->second))
+                    fPooledKeysStored = false;
+        }
+        nFail += Check(fPooledKeysStored, "every pooled public key has a stored private key") ? 0 : 1;
+
+        std::vector<std::vector<unsigned char> > drawn;
+        std::mutex drawnMutex;
+        const int nThreads = 8;
+        const int nDrawsPerThread = 32;
+        std::vector<std::thread> threads;
+
+        for (int t = 0; t < nThreads; t++)
+        {
+            threads.emplace_back([&drawn, &drawnMutex, nDrawsPerThread]() {
+                for (int i = 0; i < nDrawsPerThread; i++)
+                {
+                    std::vector<unsigned char> key = GetKeyFromPool();
+                    std::lock_guard<std::mutex> lock(drawnMutex);
+                    drawn.push_back(key);
+                }
+            });
+        }
+        for (std::thread& t : threads)
+            t.join();
+
+        nFail += Check((int)drawn.size() == nThreads * nDrawsPerThread,
+                       "all concurrent draws returned a key") ? 0 : 1;
+
+        std::set<std::string> seen;
+        bool fUnique = true;
+        bool fAllStored = true;
+        for (const std::vector<unsigned char>& key : drawn)
+        {
+            if (!seen.insert(HexStrLocal(key)).second)
+                fUnique = false;
+            if (!WalletHasPrivateKey(key))
+                fAllStored = false;
+        }
+
+        nFail += Check(fUnique, "concurrent draws never return the same key") ? 0 : 1;
+        nFail += Check(fAllStored, "each drawn key was already stored in wallet.dat") ? 0 : 1;
+        int nPoolSize = 0;
+        CRITICAL_BLOCK(cs_keyPool)
+            nPoolSize = (int)mapKeyPool.size();
+        nFail += Check(nPoolSize == KEYPOOL_SIZE, "drawing leaves the key pool topped up") ? 0 : 1;
+    }
+    catch (const std::exception& e)
+    {
+        printf("  FAIL exception: %s\n", e.what());
+        nFail++;
+    }
+    catch (...)
+    {
+        printf("  FAIL unknown exception\n");
+        nFail++;
+    }
+
+    DBFlush(true);
+    SetCurrentDir(cwd);
+    RemoveTree(tmp);
+    printf("%s (%d failure%s)\n", nFail == 0 ? "ALL TESTS PASSED" : "TESTS FAILED",
+           nFail, nFail == 1 ? "" : "s");
+    return nFail == 0 ? 0 : 1;
+}
+
+int RunSelfTest(const std::string& name)
+{
+    if (name == "wallet-keypool")
+        return RunWalletKeyPoolSelfTest();
+
+    printf("Unknown self-test '%s'\n", name.c_str());
+    printf("Known self-tests: wallet-keypool\n");
+    return 1;
+}
diff --git a/src/selftest.h b/src/selftest.h
new file mode 100644
index 0000000..d3699c6
--- /dev/null
+++ b/src/selftest.h
@@ -0,0 +1,11 @@
+// Copyright (c) 2026 Bitflash developers
+// Distributed under the MIT/X11 software license.
+
+#ifndef BITFLASH_SELFTEST_H
+#define BITFLASH_SELFTEST_H
+
+#include <string>
+
+int RunSelfTest(const std::string& name);
+
+#endif