From 566c0beb0a555d7c9620cbd882f9d32bfe5020ab Mon Sep 17 00:00:00 2001 From: Akash Mozumdar Date: Sun, 16 Jun 2019 22:57:41 -0400 Subject: [PATCH] make thread linker asynchronous. also refactored, but in a way that makes me very concerned about the stability of host --- GUI/host/host.cpp | 20 ++++++++++---------- GUI/host/host.h | 3 ++- GUI/mainwindow.cpp | 36 ++++++++++++++++++++++-------------- extensions/threadlinker.cpp | 12 +++++------- 4 files changed, 39 insertions(+), 32 deletions(-) diff --git a/GUI/host/host.cpp b/GUI/host/host.cpp index eb9f1a0..eff9c52 100644 --- a/GUI/host/host.cpp +++ b/GUI/host/host.cpp @@ -60,7 +60,7 @@ namespace size_t HashThreadParam(ThreadParam tp) { return std::hash()(tp.processId + tp.addr) + std::hash()(tp.ctx + tp.ctx2); } Synchronized>, std::recursive_mutex> textThreadsByParams; - Synchronized, std::recursive_mutex> processRecordsByIds; + Synchronized> processRecordsByIds; Host::ProcessEventHandler OnConnect, OnDisconnect; Host::ThreadEventHandler OnCreate, OnDestroy; @@ -102,7 +102,7 @@ namespace case HOST_NOTIFICATION_FOUND_HOOK: { auto info = *(HookFoundNotif*)buffer; - auto& OnHookFound = processRecordsByIds->at(processId).OnHookFound; + auto OnHookFound = processRecordsByIds->at(processId).OnHookFound; std::wstring wide = info.text; if (wide.size() > STRING) OnHookFound(info.hp, info.text); info.hp.type &= ~USING_UNICODE; @@ -132,7 +132,7 @@ namespace auto textThread = textThreadsByParams->find(tp); if (textThread == textThreadsByParams->end()) { - try { textThread = textThreadsByParams->try_emplace(tp, tp, Host::GetHookParam(tp)).first; } + try { textThread = textThreadsByParams->try_emplace(tp, tp, processRecordsByIds->at(tp.processId).GetHook(tp.addr).hp).first; } catch (std::out_of_range) { continue; } // probably garbage data in pipe, try again OnCreate(textThread->second); } @@ -158,8 +158,6 @@ namespace Host OnDestroy = [Destroy](TextThread& thread) { thread.Stop(); Destroy(thread); }; TextThread::Output = Output; - processRecordsByIds->try_emplace(console.processId, console.processId, INVALID_HANDLE_VALUE); - OnConnect(console.processId); textThreadsByParams->try_emplace(console, console, HookParam{}, CONSOLE); OnCreate(GetThread(console)); textThreadsByParams->try_emplace(clipboard, clipboard, HookParam{}, CLIPBOARD); @@ -233,16 +231,18 @@ namespace Host processRecordsByIds->at(processId).Send(FindHookCmd(sp)); } - HookParam GetHookParam(ThreadParam tp) - { - return processRecordsByIds->at(tp.processId).GetHook(tp.addr).hp; - } - TextThread& GetThread(ThreadParam tp) { return textThreadsByParams->at(tp); } + TextThread* GetThread(int64_t handle) + { + auto textThreadsByParams = ::textThreadsByParams.Acquire(); + auto thread = std::find_if(textThreadsByParams->begin(), textThreadsByParams->end(), [&](const auto& thread) { return thread.second.handle == handle; }); + return thread != textThreadsByParams->end() ? &thread->second : nullptr; + } + void AddConsoleOutput(std::wstring text) { GetThread(console).AddSentence(std::move(text)); diff --git a/GUI/host/host.h b/GUI/host/host.h index 9a203f5..ea7caa2 100644 --- a/GUI/host/host.h +++ b/GUI/host/host.h @@ -12,11 +12,12 @@ namespace Host void InjectProcess(DWORD processId); void DetachProcess(DWORD processId); + void InsertHook(DWORD processId, HookParam hp); void RemoveHook(DWORD processId, uint64_t address); void FindHooks(DWORD processId, SearchParam sp, HookEventHandler HookFound = {}); - HookParam GetHookParam(ThreadParam tp); + TextThread* GetThread(int64_t handle); TextThread& GetThread(ThreadParam tp); void AddConsoleOutput(std::wstring text); diff --git a/GUI/mainwindow.cpp b/GUI/mainwindow.cpp index bc90ca6..e5ae208 100644 --- a/GUI/mainwindow.cpp +++ b/GUI/mainwindow.cpp @@ -234,6 +234,13 @@ DWORD MainWindow::GetSelectedProcessId() std::array MainWindow::GetMiscInfo(TextThread& thread) { + void(*AddSentence)(MainWindow*, int64_t, const wchar_t*) = [](MainWindow* This, int64_t number, const wchar_t* sentence) + { + std::wstring sentenceStr = sentence; + // pointer from Host::GetThread may not stay valid unless on main thread + QMetaObject::invokeMethod(This, [=]() mutable { if (TextThread* thread = Host::GetThread(number)) thread->AddSentence(std::move(sentenceStr)); }); + }; + return { { { "current select", &thread == current }, @@ -242,6 +249,8 @@ std::array MainWindow::GetMiscInfo(TextThread& thread) { "hook address", (int64_t)thread.tp.addr }, { "text handle", thread.handle }, { "text name", (int64_t)thread.name.c_str() }, + { "this", (int64_t)this }, + { "void (*AddSentence)(void* this, int64_t number, const wchar_t* sentence)", (int64_t)AddSentence }, { nullptr, 0 } // nullptr marks end of info array } }; } @@ -323,7 +332,7 @@ void MainWindow::RemoveHooks() for (int i = 0; i < ui->ttCombo->count(); ++i) { ThreadParam tp = ParseTextThreadString(ui->ttCombo->itemText(i)); - if (tp.processId == GetSelectedProcessId()) hooks[tp.addr] = Host::GetHookParam(tp); + if (tp.processId == GetSelectedProcessId()) hooks[tp.addr] = Host::GetThread(tp).hp; } auto hookList = new QListWidget(this); hookList->setWindowFlags(Qt::Window | Qt::WindowCloseButtonHint); @@ -346,23 +355,22 @@ void MainWindow::RemoveHooks() void MainWindow::SaveHooks() { - if (auto processName = Util::GetModuleFilename(GetSelectedProcessId())) + auto processName = Util::GetModuleFilename(GetSelectedProcessId()); + if (!processName) return; + QHash hookCodes; + for (int i = 0; i < ui->ttCombo->count(); ++i) { - QHash hookCodes; - for (int i = 0; i < ui->ttCombo->count(); ++i) + ThreadParam tp = ParseTextThreadString(ui->ttCombo->itemText(i)); + if (tp.processId == GetSelectedProcessId()) { - ThreadParam tp = ParseTextThreadString(ui->ttCombo->itemText(i)); - if (tp.processId == GetSelectedProcessId()) - { - HookParam hp = Host::GetHookParam(tp); - if (!(hp.type & HOOK_ENGINE)) hookCodes[tp.addr] = S(Util::GenerateCode(hp, tp.processId)); - } + HookParam hp = Host::GetThread(tp).hp; + if (!(hp.type & HOOK_ENGINE)) hookCodes[tp.addr] = S(Util::GenerateCode(hp, tp.processId)); } - auto hookInfo = QStringList() << S(processName.value()) << hookCodes.values(); - ThreadParam tp = current.load()->tp; - if (tp.processId == GetSelectedProcessId()) hookInfo << QString("|%1:%2:%3").arg(tp.ctx).arg(tp.ctx2).arg(S(Util::GenerateCode(Host::GetHookParam(tp), tp.processId))); - QTextFile(HOOK_SAVE_FILE, QIODevice::WriteOnly | QIODevice::Append).write((hookInfo.join(" , ") + "\n").toUtf8()); } + auto hookInfo = QStringList() << S(processName.value()) << hookCodes.values(); + ThreadParam tp = current.load()->tp; + if (tp.processId == GetSelectedProcessId()) hookInfo << QString("|%1:%2:%3").arg(tp.ctx).arg(tp.ctx2).arg(S(Util::GenerateCode(Host::GetThread(tp).hp, tp.processId))); + QTextFile(HOOK_SAVE_FILE, QIODevice::WriteOnly | QIODevice::Append).write((hookInfo.join(" , ") + "\n").toUtf8()); } void MainWindow::FindHooks() diff --git a/extensions/threadlinker.cpp b/extensions/threadlinker.cpp index a78f70f..68e0014 100644 --- a/extensions/threadlinker.cpp +++ b/extensions/threadlinker.cpp @@ -93,13 +93,11 @@ BOOL WINAPI DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved bool ProcessSentence(std::wstring& sentence, SentenceInfo sentenceInfo) { std::lock_guard l(m); - static std::unordered_map queuedWritesByHandle; - int64_t textHandle = sentenceInfo["text handle"]; + int64_t textHandle = sentenceInfo["text number"]; - for (auto linkedHandle : linkedTextHandles[textHandle]) queuedWritesByHandle[linkedHandle] += L"\n" + sentence; + for (auto linkedHandle : linkedTextHandles[textHandle]) + ((void(*)(void*, int64_t, const wchar_t*))sentenceInfo["void (*AddSentence)(void* this, int64_t number, const wchar_t* sentence)"]) + ((void*)sentenceInfo["this"], linkedHandle, sentence.c_str()); - if (queuedWritesByHandle[textHandle].empty()) return false; - sentence += queuedWritesByHandle[textHandle]; - queuedWritesByHandle[textHandle].clear(); - return true; + return false; }