diff --git a/CMakeLists.txt b/CMakeLists.txt index 49bfad6..863826a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -22,6 +22,7 @@ set (PROGRAM_SOURCE src/checksum.cpp src/WIF.cpp src/key_search.cpp + src/key_search_helpers.cpp src/main.cpp ) @@ -38,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_mt.cpp) endif (Threads_FOUND) endif (USE_THREADS) diff --git a/src/key_search.cpp b/src/key_search.cpp index 9286ddf..4201588 100644 --- a/src/key_search.cpp +++ b/src/key_search.cpp @@ -1,46 +1,21 @@ -#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); -} - -void key_search(strlist_t word_list, size_t n) { +void key_search(const strlist_t& word_list, size_t n) { size_t count = 0; struct ec_keypair pair; while (count < n) { - std::string pubstr; + std::string word; ec_generate_key(&pair); - pubstr = wif_pub_encode(pair.pub); - strtolower(pubstr); - - for(auto const& word: word_list) { - if (pubstr.find(word) != std::string::npos) { - key_result(word, &pair); - count++; - } + if (key_contains_word(&pair, word_list, word)) { + key_search_result(word, &pair); + count++; } } } diff --git a/src/key_search.h b/src/key_search.h index 742195c..9676053 100644 --- a/src/key_search.h +++ b/src/key_search.h @@ -25,7 +25,12 @@ #define KEY_SEARCH_H #include "string.h" +#include "ec.h" -void key_search(strlist_t word_list, size_t n); +void key_search(const strlist_t& word_list, size_t n); + +#ifdef HAVE_THREADS +void key_search_mt(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_helpers.cpp b/src/key_search_helpers.cpp new file mode 100644 index 0000000..82d625e --- /dev/null +++ b/src/key_search_helpers.cpp @@ -0,0 +1,25 @@ + +#include +#include "WIF.h" +#include "key_search_helpers.h" + +void key_search_result(const std::string& word, const struct ec_keypair* pair) { + + std::cout << "----" << std::endl; + std::cout << "Found: " << word << std::endl; + wif_print_key(pair); +} + +bool key_contains_word(const struct ec_keypair* key, const strlist_t& word_list, std::string& word) { + + std::string pubstr = wif_pub_encode(key->pub); + strtolower(pubstr); + + for(auto const& w: word_list) { + if (pubstr.find(w) != std::string::npos) { + word = w; + return true; + } + } + return false; +} diff --git a/src/key_search_helpers.h b/src/key_search_helpers.h new file mode 100644 index 0000000..66d4291 --- /dev/null +++ b/src/key_search_helpers.h @@ -0,0 +1,36 @@ +/** + * MIT License + * + * Copyright (c) 2019-2020 EOS Sw/eden + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ +#ifndef KEY_SEARCH_HELPSER_H +#define KEY_SEARCH_HELPERS_H + +#include "string.h" +#include "ec.h" + +void key_search_result(const std::string& word, const struct ec_keypair* pair); + +// Check if any word in appears in 's public key. +// returns true if a word was found (stored in ), false otherwise. +bool key_contains_word(const struct ec_keypair* key, const strlist_t& word_list, std::string& word); + +#endif /* KEY_SEARCH_HELPERS_H */ diff --git a/src/key_search_mt.cpp b/src/key_search_mt.cpp new file mode 100644 index 0000000..bfb2584 --- /dev/null +++ b/src/key_search_mt.cpp @@ -0,0 +1,75 @@ + +#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 _mt_search(const strlist_t& word_list) { + + struct ec_keypair pair; + + while (g_count < g_max) { + std::string word; + + ec_generate_key(&pair); + if (key_contains_word(&pair, word_list, word)) { + + // 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_mt(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(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(_mt_search, word_list); + } + + // Use main thread for 1 search + _mt_search(word_list); + + // Wait for all threads to compelete. + for(int i = 0; i < t.size(); i++) { + t[i].join(); + } +} diff --git a/src/main.cpp b/src/main.cpp index 84c78b7..dec65e9 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -21,6 +21,9 @@ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE * SOFTWARE. */ +#ifdef HAVE_THREADS +#include +#endif /* HAVE_THREADS */ #include #include #include @@ -31,58 +34,27 @@ #include "key_search.h" #ifdef HAVE_THREADS -#include -#include - -// Minium number of threads. -#define MIN_THREADS 2 - -// Number of threads to use. -static int n_threads = std::thread::hardware_concurrency(); - -#define search_func thread_search -static void thread_search(const strlist_t& words, int n) { - - std::vector t; - int d, m; - - // We can use all threads - if (n >= n_threads) { - - // create n_threads - 1 as we use main process also. - t.resize(n_threads - 1); - // divide the number of results for all threads. - d = n / n_threads; - // Also calculate the reminder (will be assigned to the main thread) - m = n % n_threads; +#define n_thread_decl int n_threads = std::thread::hardware_concurrency() +#define n_thread_argv \ + if (argc > 2) { \ + n_threads = atoi(argv[2]); \ + if (n_threads < 2) { \ + n_threads = 2; \ + } \ } - // not enough results to use all threads. - else { - t.resize(n); - d = 1; - m = 0; - } - - // Launch threads. - for(int i = 0; i < t.size(); i++) { - t[i] = std::thread(key_search, words, d); - } - - // Use main thread for 1 search - key_search(words, d + m); - - // Wait for all threads to compelete. - for(int i = 0; i < t.size(); i++) { - t[i].join(); - } -} +#define n_thread_outp << ", Using: " << n_threads << " threads" +#define call_search key_search_mt(words, n, n_threads) #else -#define search_func key_search -#endif +#define n_thread_decl +#define n_thread_argv +#define n_thread_outp +#define call_search key_search(words, n) +#endif /* HAVE_THREADS */ void cmd_search(int argc, char **argv) { int n = 100; + n_thread_decl; std::string search(argv[0]); strlist_t words = strsplitwords(strtolower(search)); @@ -93,24 +65,14 @@ void cmd_search(int argc, char **argv) { } } -#ifdef HAVE_THREADS - if (argc > 2) { - n_threads = atoi(argv[2]); - // Make sure we never go under min threads. - if (n_threads < MIN_THREADS) { - n_threads = MIN_THREADS; - } - } -# endif /* HAVE_THREADS */ + n_thread_argv; std::cout << "Searching for " << n - << " keys containing: " << search -#ifdef HAVE_THREADS - << ", Using: " << n_threads << " threads" -#endif /* HAVE_THREADS */ - << std::endl; + << " keys containing: " << search + n_thread_outp + << std::endl; - search_func(words, n); + call_search; } void usage(const char *name) {