diff --git a/CMakeLists.txt b/CMakeLists.txt index 7bdf10b..0222841 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -39,6 +39,7 @@ if (USE_THREADS) if (Threads_FOUND) # Add preprocessor flag add_definitions( "-DHAVE_THREADS=1" ) + set (PROGRAM_SOURCE ${PROGRAM_SOURCE} src/key_search_nt.cpp) endif (Threads_FOUND) endif (USE_THREADS) diff --git a/src/key_search.cpp b/src/key_search.cpp index c18c015..8dfe10e 100644 --- a/src/key_search.cpp +++ b/src/key_search.cpp @@ -1,30 +1,10 @@ -#ifdef HAVE_THREADS -#include -#endif /* HAVE_THREADS */ #include #include -#include "ec.h" #include "WIF.h" +#include "key_search_helpers.h" #include "key_search.h" -#ifdef HAVE_THREADS -// Guards result output. -std::mutex search_mutex; -#endif /* HAVE_THREADS */ - -static void key_result(const std::string& word, const struct ec_keypair* pair) { - -#ifdef HAVE_THREADS - // Guard output with mutex, so we don't get interrupted mid write. - const std::lock_guard lock(search_mutex); -#endif /* HAVE_THREADS */ - - std::cout << "----" << std::endl; - std::cout << "Found: " << word << std::endl; - wif_print_key(pair); -} - bool key_search(struct ec_keypair* key, std::string& word, const strlist_t& word_list) { std::string pubstr; @@ -50,7 +30,7 @@ void key_search_n(const strlist_t& word_list, size_t n) { while (count < n) { std::string word; if (key_search(&pair, word, word_list)) { - key_result(word, &pair); + key_search_result(word, &pair); count++; } } diff --git a/src/key_search.h b/src/key_search.h index 3f73e16..b147b92 100644 --- a/src/key_search.h +++ b/src/key_search.h @@ -31,4 +31,8 @@ bool key_search(struct ec_keypair* out, std::string& word, const strlist_t& word void key_search_n(const strlist_t& word_list, size_t n); +#ifdef HAVE_THREADS +void key_search_nt(const strlist_t& word_list, size_t n, size_t n_threads = 0); +#endif /* HAVE_THREADS */ + #endif /* KEY_SEARCH_H */ diff --git a/src/key_search_nt.cpp b/src/key_search_nt.cpp new file mode 100644 index 0000000..73738d7 --- /dev/null +++ b/src/key_search_nt.cpp @@ -0,0 +1,74 @@ + +#include +#include +#include +#include "key_search_helpers.h" +#include "key_search.h" + +// Max keys to search for, +unsigned int g_max; + +// How many keys we have found so far. +unsigned int g_count; + +// Mutex guard for g_count. +std::mutex g_count_mtx; + +// Thread process. +static void thr_proc(const strlist_t& word_list) { + + struct ec_keypair pair; + + while (g_count < g_max) { + std::string word; + + if (key_search(&pair, word, word_list)) { + + // Guard output with mutex, so we don't get + // interrupted mid write and can write to g_count safely. + const std::lock_guard lock(g_count_mtx); + + // It is possible g_count was updated by another thread + // after we checked it in the while loop. + // So while we have the lock, we need to check it again. + if (g_count >= g_max) { + return; + } + + // Update count and print result. + g_count++; + key_search_result(word, &pair); + } + } +} + +void key_search_nt(const strlist_t& word_list, size_t n, size_t n_threads) { + + std::vector t; + + // Not enough threads passed in by caller. + if (n_threads < 2) { + // Just use linear function. + key_search_n(word_list, n); + return; + } + + t.resize(n_threads - 1); + + // Setup global variables for the threads. + g_max = n; + g_count = 0; + + // Launch them. + for(int i = 0; i < t.size(); i++) { + t[i] = std::thread(thr_proc, word_list); + } + + // Use main thread for 1 search + thr_proc(word_list); + + // Wait for all threads to compelete. + for(int i = 0; i < t.size(); i++) { + t[i].join(); + } +}