diff --git a/src/key_search.cpp b/src/key_search.cpp index 4201588..408caa0 100644 --- a/src/key_search.cpp +++ b/src/key_search.cpp @@ -5,7 +5,32 @@ #include "key_search_helpers.h" #include "key_search.h" -void key_search(const strlist_t& word_list, size_t n) { +void KeySearch::addWord(const std::string& str) +{ + std::string tmp = str; + base58_strip(tmp); + strtolower(tmp); + m_words.push_back(tmp); +} + +void KeySearch::addList(const strlist_t& list) +{ + for(const std::string& item : list) { + addWord(item); + } +} + +const strlist_t& KeySearch::getList() +{ + return m_words; +} + +void KeySearch::clear() +{ + m_words.clear(); +} + +void KeySearch::_search_linear(size_t n) { size_t count = 0; struct ec_keypair pair; @@ -13,9 +38,22 @@ void key_search(const strlist_t& word_list, size_t n) { while (count < n) { std::string word; ec_generate_key(&pair); - if (key_contains_word(&pair, word_list, word)) { + if (key_contains_word(&pair, m_words, word)) { key_search_result(word, &pair); count++; } } } + +void KeySearch::find(size_t num_results) { + +#ifdef HAVE_THREADS + // Only do multithread if number of threads makes sense. + if (m_threads >= 2) { + _search_mt(num_results); + return; + } +#endif /* HAVE_THREADS */ + + _search_linear(num_results); +} diff --git a/src/key_search.h b/src/key_search.h index 9676053..3a92639 100644 --- a/src/key_search.h +++ b/src/key_search.h @@ -27,10 +27,45 @@ #include "string.h" #include "ec.h" -void key_search(const strlist_t& word_list, size_t n); +class KeySearch +{ +public : + // Add a word to search for. + void addWord(const std::string& str); + + // Add a list of words to search for. + void addList(const strlist_t& list); + + // get the list of words to search for. + const strlist_t& getList(); + + // Clears the search list. + void clear(); #ifdef HAVE_THREADS -void key_search_mt(const strlist_t& word_list, size_t n, size_t n_threads = 0); + // Set the number of threads to use while searching. + void setThreadCount(size_t num); #endif /* HAVE_THREADS */ + // Perform a search. + void find(size_t num_results); + +protected : + +#ifdef HAVE_THREADS + void _search_mt(size_t n); +#endif /* HAVE_THREADS */ + + void _search_linear(size_t n); + +protected : + // List of words to search for. + strlist_t m_words; + +#ifdef HAVE_THREADS + // Number of threads to use. + size_t m_threads; +#endif /* HAVE_THREADS */ +}; + #endif /* KEY_SEARCH_H */ diff --git a/src/key_search_mt.cpp b/src/key_search_mt.cpp index bfb2584..2c30fc0 100644 --- a/src/key_search_mt.cpp +++ b/src/key_search_mt.cpp @@ -15,7 +15,7 @@ unsigned int g_count; std::mutex g_count_mtx; // Thread process. -static void _mt_search(const strlist_t& word_list) { +static void _thr_proc(const strlist_t& word_list) { struct ec_keypair pair; @@ -43,18 +43,16 @@ static void _mt_search(const strlist_t& word_list) { } } -void key_search_mt(const strlist_t& word_list, size_t n, size_t n_threads) { +void KeySearch::setThreadCount(size_t num) +{ + m_threads = num; +} +void KeySearch::_search_mt(size_t n) +{ 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); + t.resize(m_threads - 1); // Setup global variables for the threads. g_max = n; @@ -62,11 +60,11 @@ void key_search_mt(const strlist_t& word_list, size_t n, size_t n_threads) { // Launch them. for(int i = 0; i < t.size(); i++) { - t[i] = std::thread(_mt_search, word_list); + t[i] = std::thread(_thr_proc, m_words); } // Use main thread for 1 search - _mt_search(word_list); + _thr_proc(m_words); // Wait for all threads to compelete. for(int i = 0; i < t.size(); i++) { diff --git a/src/main.cpp b/src/main.cpp index e6ea752..432549c 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -44,33 +44,29 @@ bool option_l33t = false; if (n_threads < 2) { \ n_threads = 2; \ } \ - } + } \ + ks.setThreadCount(n_threads); #define n_thread_outp << ", Using: " << n_threads << " threads" -#define call_search key_search_mt(words, n, n_threads) #else #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; + std::string input(argv[0]); + KeySearch ks; if (option_l33t) { - strlist_t tmp = strsplitwords(search); + strlist_t tmp = strsplitwords(input); for(int i = 0; i < tmp.size(); i++) { - strlist_t list = l33twords(base58_strip(tmp[i])); - words.reserve(words.size() + list.size()); - words.insert(words.end(), list.begin(), list.end()); + ks.addList(l33twords(tmp[i])); } } else { - words = strsplitwords(strtolower(search)); - base58_strip(words); + ks.addList(strsplitwords(input)); } if (argc > 1) { @@ -83,11 +79,11 @@ void cmd_search(int argc, char **argv) { n_thread_argv; std::cout << "Searching for " << n - << " keys containing: " << strjoin(words, ",") + << " keys containing: " << strjoin(ks.getList(), ",") n_thread_outp << std::endl; - call_search; + ks.find(n); } void usage(const char *name) {