From e6805a2be3853344c02c1421f314110de2eeb75d Mon Sep 17 00:00:00 2001 From: Akash Mozumdar Date: Mon, 4 Feb 2019 15:18:47 -0500 Subject: [PATCH] refactor, perf improvement --- GUI/CMakeLists.txt | 2 +- GUI/host/host.cpp | 60 ++++++++++++++++++++++++----------------- GUI/host/host.h | 7 ++--- GUI/host/textthread.cpp | 20 ++++++++------ GUI/host/textthread.h | 9 +++---- GUI/mainwindow.cpp | 52 +++++++++++++++++------------------ GUI/mainwindow.h | 10 +++---- 7 files changed, 88 insertions(+), 72 deletions(-) diff --git a/GUI/CMakeLists.txt b/GUI/CMakeLists.txt index 5969a92..266c054 100644 --- a/GUI/CMakeLists.txt +++ b/GUI/CMakeLists.txt @@ -17,4 +17,4 @@ set(gui_src add_executable(${PROJECT_NAME} WIN32 ${gui_src}) target_link_libraries(${PROJECT_NAME} Qt5::Widgets winhttp) -install_qt5_libs(${PROJECT_NAME}) # can be commented out for consecutive builds +#install_qt5_libs(${PROJECT_NAME}) # can be commented out for consecutive builds diff --git a/GUI/host/host.cpp b/GUI/host/host.cpp index 0cbe5ab..7e52d7e 100644 --- a/GUI/host/host.cpp +++ b/GUI/host/host.cpp @@ -10,21 +10,16 @@ namespace class ProcessRecord { public: - inline static Host::ProcessEventCallback OnConnect, OnDisconnect; - ProcessRecord(DWORD processId, HANDLE pipe) : processId(processId), pipe(pipe), mappedFile(OpenFileMappingW(FILE_MAP_READ, FALSE, (ITH_SECTION_ + std::to_wstring(processId)).c_str())), view(*(const TextHook(*)[MAX_HOOK])MapViewOfFile(mappedFile, FILE_MAP_READ, 0, 0, HOOK_SECTION_SIZE / 2)), // jichi 1/16/2015: Changed to half to hook section size viewMutex(ITH_HOOKMAN_MUTEX_ + std::to_wstring(processId)) - { - OnConnect(processId); - } + {} ~ProcessRecord() { - OnDisconnect(processId); UnmapViewOfFile(view); } @@ -59,15 +54,21 @@ namespace { return std::hash()(tp.processId + tp.addr) + std::hash()(tp.ctx + tp.ctx2); } - ThreadSafe, Functor>, std::recursive_mutex> textThreadsByParams; + ThreadSafe>, std::recursive_mutex> textThreadsByParams; ThreadSafe, std::recursive_mutex> processRecordsByIds; + Host::ProcessEventHandler OnConnect, OnDisconnect; + Host::ThreadEventHandler OnCreate, OnDestroy; + void RemoveThreads(std::function removeIf) { - std::vector> removedThreads; // delay destruction until after lock is released - auto[lock, textThreadsByParams] = ::textThreadsByParams.operator->(); - for (auto it = textThreadsByParams->begin(); it != textThreadsByParams->end(); removeIf(it->first) ? it = textThreadsByParams->erase(it) : ++it) - if (removeIf(it->first)) removedThreads.emplace_back(std::move(it->second)); + std::vector threadsToRemove; + std::for_each(textThreadsByParams->begin(), textThreadsByParams->end(), [&](auto& it) { if (removeIf(it.first)) threadsToRemove.push_back(&it.second); }); + for (auto thread : threadsToRemove) + { + OnDestroy(*thread); + textThreadsByParams->erase(thread->tp); + } } void CreatePipe() @@ -89,6 +90,7 @@ namespace DWORD bytesRead, processId; ReadFile(hookPipe, &processId, sizeof(processId), &bytesRead, nullptr); processRecordsByIds->try_emplace(processId, processId, hostPipe); + OnConnect(processId); CreatePipe(); @@ -110,13 +112,18 @@ namespace default: { auto tp = *(ThreadParam*)buffer; - if (textThreadsByParams->count(tp) == 0) textThreadsByParams->insert({ tp, std::make_unique(tp, Host::GetHookParam(tp)) }); - textThreadsByParams->at(tp)->Push(buffer + sizeof(tp), bytesRead - sizeof(tp)); + if (textThreadsByParams->count(tp) == 0) + { + TextThread& created = textThreadsByParams->try_emplace(tp, tp, Host::GetHookParam(tp)).first->second; + OnCreate(created); + } + textThreadsByParams->find(tp)->second.Push(buffer + sizeof(tp), bytesRead - sizeof(tp)); } break; } RemoveThreads([&](ThreadParam tp) { return tp.processId == processId; }); + OnDisconnect(processId); processRecordsByIds->erase(processId); }).detach(); } @@ -124,22 +131,27 @@ namespace namespace Host { - void Start(ProcessEventCallback OnConnect, ProcessEventCallback OnDisconnect, TextThread::EventCallback OnCreate, TextThread::EventCallback OnDestroy, TextThread::OutputCallback Output) + void Start(ProcessEventHandler Connect, ProcessEventHandler Disconnect, ThreadEventHandler Create, ThreadEventHandler Destroy, TextThread::OutputCallback Output) { - ProcessRecord::OnConnect = OnConnect; - ProcessRecord::OnDisconnect = OnDisconnect; - TextThread::OnCreate = OnCreate; - TextThread::OnDestroy = OnDestroy; + OnConnect = Connect; + OnDisconnect = Disconnect; + OnCreate = [Create](TextThread& thread) { Create(thread); thread.Start(); }; + OnDestroy = [Destroy](TextThread& thread) { thread.Stop(); Destroy(thread); }; TextThread::Output = Output; + processRecordsByIds->try_emplace(console.processId, console.processId, INVALID_HANDLE_VALUE); - textThreadsByParams->insert({ console, std::make_unique(console, HookParam{}, CONSOLE) }); - textThreadsByParams->insert({ Host::clipboard, std::make_unique(Host::clipboard, HookParam{}, CLIPBOARD) }); + OnConnect(console.processId); + textThreadsByParams->try_emplace(console, console, HookParam{}, CONSOLE); + OnCreate(GetThread(console)); + textThreadsByParams->try_emplace(clipboard, clipboard, HookParam{}, CLIPBOARD); + OnCreate(GetThread(clipboard)); + CreatePipe(); SetWindowsHookExW(WH_GETMESSAGE, [](int statusCode, WPARAM wParam, LPARAM lParam) { if (statusCode == HC_ACTION && wParam == PM_REMOVE && ((MSG*)lParam)->message == WM_CLIPBOARDUPDATE) - if (auto text = Util::GetClipboardText()) Host::GetThread(Host::clipboard)->AddSentence(text.value()); + if (auto text = Util::GetClipboardText()) GetThread(clipboard).AddSentence(text.value()); return CallNextHookEx(NULL, statusCode, wParam, lParam); }, NULL, GetCurrentThreadId()); } @@ -200,13 +212,13 @@ namespace Host return processRecordsByIds->at(tp.processId).GetHook(tp.addr).hp; } - TextThread* GetThread(ThreadParam tp) + TextThread& GetThread(ThreadParam tp) { - return textThreadsByParams->at(tp).get(); + return textThreadsByParams->at(tp); } void AddConsoleOutput(std::wstring text) { - GetThread(console)->AddSentence(text); + GetThread(console).AddSentence(text); } } diff --git a/GUI/host/host.h b/GUI/host/host.h index c65f2ef..e9b58c2 100644 --- a/GUI/host/host.h +++ b/GUI/host/host.h @@ -5,8 +5,9 @@ namespace Host { - using ProcessEventCallback = std::function; - void Start(ProcessEventCallback OnConnect, ProcessEventCallback OnDisconnect, TextThread::EventCallback OnCreate, TextThread::EventCallback OnDestroy, TextThread::OutputCallback Output); + using ProcessEventHandler = std::function; + using ThreadEventHandler = std::function; + void Start(ProcessEventHandler Connect, ProcessEventHandler Disconnect, ThreadEventHandler Create, ThreadEventHandler Destroy, TextThread::OutputCallback Output); bool InjectProcess(DWORD processId, DWORD timeout = 5000); void DetachProcess(DWORD processId); @@ -14,7 +15,7 @@ namespace Host HookParam GetHookParam(ThreadParam tp); - TextThread* GetThread(ThreadParam tp); + TextThread& GetThread(ThreadParam tp); void AddConsoleOutput(std::wstring text); inline int defaultCodepage = SHIFT_JIS; diff --git a/GUI/host/textthread.cpp b/GUI/host/textthread.cpp index 0eed0ba..365d23d 100644 --- a/GUI/host/textthread.cpp +++ b/GUI/host/textthread.cpp @@ -9,14 +9,16 @@ TextThread::TextThread(ThreadParam tp, HookParam hp, std::optional name(name.value_or(Util::StringToWideString(hp.name).value())), tp(tp), hp(hp) +{} + +void TextThread::Start() { CreateTimerQueueTimer(&timer, NULL, [](void* This, BOOLEAN) { ((TextThread*)This)->Flush(); }, this, 10, 10, WT_EXECUTELONGFUNCTION); - OnCreate(this); } -TextThread::~TextThread() +void TextThread::Stop() { - OnDestroy(this); + timer = NULL; } void TextThread::AddSentence(const std::wstring& sentence) @@ -40,7 +42,7 @@ void TextThread::Push(const BYTE* data, int len) lastPushTime = GetTickCount(); if (std::all_of(buffer.begin(), buffer.end(), [&](wchar_t c) { return repeatingChars.count(c) > 0; })) buffer.clear(); - if (Util::RemoveRepetition(buffer)) // repetition detected, which means the entire sentence has already been received + if (Util::RemoveRepetition(buffer)) // sentence repetition detected, which means the entire sentence has already been received { repeatingChars = std::unordered_set(buffer.begin(), buffer.end()); AddSentence(buffer); @@ -53,11 +55,13 @@ void TextThread::Flush() std::vector sentences; queuedSentences->swap(sentences); for (auto& sentence : sentences) - if (Output(this, sentence)) storage->append(sentence); + if (Output(*this, sentence)) storage->append(sentence); std::scoped_lock lock(bufferMutex); if (buffer.empty()) return; - if (buffer.size() < maxBufferSize && GetTickCount() - lastPushTime < flushDelay) return; - AddSentence(buffer); - buffer.clear(); + if (buffer.size() > maxBufferSize || GetTickCount() - lastPushTime > flushDelay) + { + AddSentence(buffer); + buffer.clear(); + } } diff --git a/GUI/host/textthread.h b/GUI/host/textthread.h index c5216fb..d9f1367 100644 --- a/GUI/host/textthread.h +++ b/GUI/host/textthread.h @@ -6,17 +6,16 @@ class TextThread { public: - using EventCallback = std::function; - using OutputCallback = std::function; - inline static EventCallback OnCreate, OnDestroy; + using OutputCallback = std::function; inline static OutputCallback Output; inline static int flushDelay = 400; // flush every 400ms by default inline static int maxBufferSize = 1000; TextThread(ThreadParam tp, HookParam hp, std::optional name = {}); - ~TextThread(); + void Start(); + void Stop(); void AddSentence(const std::wstring& sentence); void Push(const BYTE* data, int len); @@ -38,5 +37,5 @@ private: DWORD lastPushTime = 0; ThreadSafe> queuedSentences; struct TimerDeleter { void operator()(HANDLE h) { DeleteTimerQueueTimer(NULL, h, INVALID_HANDLE_VALUE); } }; - AutoHandle timer = NULL; // this needs to be last so it's destructed first + AutoHandle timer = NULL; }; diff --git a/GUI/mainwindow.cpp b/GUI/mainwindow.cpp index 333810b..112688e 100644 --- a/GUI/mainwindow.cpp +++ b/GUI/mainwindow.cpp @@ -49,9 +49,9 @@ MainWindow::MainWindow(QWidget *parent) : Host::Start( [this](DWORD processId) { ProcessConnected(processId); }, [this](DWORD processId) { ProcessDisconnected(processId); }, - [this](TextThread* thread) { ThreadAdded(thread); }, - [this](TextThread* thread) { ThreadRemoved(thread); }, - [this](TextThread* thread, std::wstring& output) { return SentenceReceived(thread, output); } + [this](TextThread& thread) { ThreadAdded(thread); }, + [this](TextThread& thread) { ThreadRemoved(thread); }, + [this](TextThread& thread, std::wstring& output) { return SentenceReceived(thread, output); } ); Host::AddConsoleOutput(ABOUT); @@ -104,7 +104,7 @@ void MainWindow::ProcessConnected(DWORD processId) auto hookList = std::find_if(allProcesses.rbegin(), allProcesses.rend(), [&](QString hookList) { return hookList.contains(process); }); if (hookList != allProcesses.rend()) for (auto hookInfo : hookList->split(" , ")) - if (auto hp = Util::ParseCode(S(hookInfo))) QMetaObject::invokeMethod(this, [processId, hp] { Host::InsertHook(processId, hp.value()); }); + if (auto hp = Util::ParseCode(S(hookInfo))) Host::InsertHook(processId, hp.value()); else swscanf_s(S(hookInfo).c_str(), L"|%I64d:%I64d:%[^\n]", &savedThreadCtx.first, &savedThreadCtx.second, savedThreadCode, ARRAYSIZE(savedThreadCode)); } @@ -116,11 +116,11 @@ void MainWindow::ProcessDisconnected(DWORD processId) }, Qt::BlockingQueuedConnection); } -void MainWindow::ThreadAdded(TextThread* thread) +void MainWindow::ThreadAdded(TextThread& thread) { - std::wstring threadCode = Util::GenerateCode(thread->hp, thread->tp.processId); - QString ttString = TextThreadString(thread) + S(thread->name) + " (" + S(threadCode) + ")"; - bool savedMatch = savedThreadCtx.first == thread->tp.ctx && savedThreadCtx.second == thread->tp.ctx2 && savedThreadCode == threadCode; + std::wstring threadCode = Util::GenerateCode(thread.hp, thread.tp.processId); + QString ttString = TextThreadString(thread) + S(thread.name) + " (" + S(threadCode) + ")"; + bool savedMatch = savedThreadCtx.first == thread.tp.ctx && savedThreadCtx.second == thread.tp.ctx2 && savedThreadCode == threadCode; if (savedMatch) savedThreadCtx.first = savedThreadCtx.second = savedThreadCode[0] = 0; QMetaObject::invokeMethod(this, [this, ttString, savedMatch] { @@ -129,7 +129,7 @@ void MainWindow::ThreadAdded(TextThread* thread) }); } -void MainWindow::ThreadRemoved(TextThread* thread) +void MainWindow::ThreadRemoved(TextThread& thread) { QString ttString = TextThreadString(thread); QMetaObject::invokeMethod(this, [this, ttString] @@ -140,7 +140,7 @@ void MainWindow::ThreadRemoved(TextThread* thread) }, Qt::BlockingQueuedConnection); } -bool MainWindow::SentenceReceived(TextThread* thread, std::wstring& sentence) +bool MainWindow::SentenceReceived(TextThread& thread, std::wstring& sentence) { if (DispatchSentenceToExtensions(sentence, GetMiscInfo(thread))) { @@ -160,14 +160,14 @@ bool MainWindow::SentenceReceived(TextThread* thread, std::wstring& sentence) return false; } -QString MainWindow::TextThreadString(TextThread* thread) +QString MainWindow::TextThreadString(TextThread& thread) { return QString("%1:%2:%3:%4:%5: ").arg( - QString::number(thread->handle, 16), - QString::number(thread->tp.processId, 16), - QString::number(thread->tp.addr, 16), - QString::number(thread->tp.ctx, 16), - QString::number(thread->tp.ctx2, 16) + QString::number(thread.handle, 16), + QString::number(thread.tp.processId, 16), + QString::number(thread.tp.addr, 16), + QString::number(thread.tp.ctx, 16), + QString::number(thread.tp.ctx2, 16) ).toUpper(); } @@ -182,16 +182,16 @@ DWORD MainWindow::GetSelectedProcessId() return ui->processCombo->currentText().split(":")[0].toULong(nullptr, 16); } -std::unordered_map MainWindow::GetMiscInfo(TextThread* thread) +std::unordered_map MainWindow::GetMiscInfo(TextThread& thread) { return { { "current select", ui->ttCombo->currentText().startsWith(TextThreadString(thread)) }, - { "text number", thread->handle }, - { "process id", thread->tp.processId }, - { "hook address", thread->tp.addr }, - { "text handle", thread->handle }, - { "text name", (int64_t)thread->name.c_str() } + { "text number", thread.handle }, + { "process id", thread.tp.processId }, + { "hook address", thread.tp.addr }, + { "text handle", thread.handle }, + { "text name", (int64_t)thread.name.c_str() } }; } @@ -280,9 +280,9 @@ void MainWindow::SaveHooks() } } auto hookInfo = QStringList() << S(processName.value()) << hookCodes.values(); - TextThread* current = Host::GetThread(ParseTextThreadString(ui->ttCombo->currentText())); - if (current->tp.processId == GetSelectedProcessId()) - hookInfo << QString("|%1:%2:%3").arg(current->tp.ctx).arg(current->tp.ctx2).arg(S(Util::GenerateCode(Host::GetHookParam(current->tp), current->tp.processId))); + TextThread& current = Host::GetThread(ParseTextThreadString(ui->ttCombo->currentText())); + if (current.tp.processId == GetSelectedProcessId()) + hookInfo << QString("|%1:%2:%3").arg(current.tp.ctx).arg(current.tp.ctx2).arg(S(Util::GenerateCode(Host::GetHookParam(current.tp), current.tp.processId))); QTextFile(HOOK_SAVE_FILE, QIODevice::WriteOnly | QIODevice::Append).write((hookInfo.join(" , ") + "\n").toUtf8()); } } @@ -328,6 +328,6 @@ void MainWindow::Extensions() void MainWindow::ViewThread(int index) { ui->ttCombo->setCurrentIndex(index); - ui->textOutput->setPlainText(S(Host::GetThread(ParseTextThreadString(ui->ttCombo->itemText(index)))->storage->c_str())); + ui->textOutput->setPlainText(S(Host::GetThread(ParseTextThreadString(ui->ttCombo->itemText(index))).storage->c_str())); ui->textOutput->moveCursor(QTextCursor::End); } diff --git a/GUI/mainwindow.h b/GUI/mainwindow.h index f919a07..c7b82ed 100644 --- a/GUI/mainwindow.h +++ b/GUI/mainwindow.h @@ -20,13 +20,13 @@ private: void closeEvent(QCloseEvent*) override; void ProcessConnected(DWORD processId); void ProcessDisconnected(DWORD processId); - void ThreadAdded(TextThread* thread); - void ThreadRemoved(TextThread* thread); - bool SentenceReceived(TextThread* thread, std::wstring& sentence); - QString TextThreadString(TextThread* thread); + void ThreadAdded(TextThread& thread); + void ThreadRemoved(TextThread& thread); + bool SentenceReceived(TextThread& thread, std::wstring& sentence); + QString TextThreadString(TextThread& thread); ThreadParam ParseTextThreadString(QString ttString); DWORD GetSelectedProcessId(); - std::unordered_map GetMiscInfo(TextThread* thread); + std::unordered_map GetMiscInfo(TextThread& thread); void AttachProcess(); void LaunchProcess(); void DetachProcess();