From c877d9cd319f7058f8c1074fe2937200d661899f Mon Sep 17 00:00:00 2001 From: Akash Mozumdar Date: Wed, 31 Oct 2018 12:04:32 -0400 Subject: [PATCH] use raii for process records and mutexes. remove a lot of manual resource management --- GUI/extensions.cpp | 10 ++-- GUI/host/host.cc | 108 ++++++++++++++++++------------------- GUI/host/host.h | 12 ++--- GUI/mainwindow.cpp | 3 +- include/const.h | 2 +- include/types.h | 16 ++++-- vnrhook/hijack/texthook.cc | 43 ++++----------- vnrhook/hijack/texthook.h | 14 ++--- vnrhook/main.cc | 5 +- 9 files changed, 98 insertions(+), 115 deletions(-) diff --git a/GUI/extensions.cpp b/GUI/extensions.cpp index 467046e..082f3d0 100644 --- a/GUI/extensions.cpp +++ b/GUI/extensions.cpp @@ -1,4 +1,5 @@ #include "extensions.h" +#include "types.h" static std::optional LoadExtension(QString extenName) { @@ -13,13 +14,13 @@ static std::optional LoadExtension(QString extenName) void Extension::Load(QString extenName) { - std::unique_lock extenLock(extenMutex); + LOCK(extenMutex); if (auto extension = LoadExtension(extenName)) extensions.push_back(extension.value()); } void Extension::SendToBack(QString extenName) { - std::unique_lock extenLock(extenMutex); + LOCK(extenMutex); Extension* extenIter = std::find_if(extensions.begin(), extensions.end(), [&](Extension extension) { return extension.name == extenName; }); Extension extension = *extenIter; extensions.erase(extenIter); @@ -28,13 +29,14 @@ void Extension::SendToBack(QString extenName) void Extension::Unload(QString extenName) { - std::unique_lock extenLock(extenMutex); + LOCK(extenMutex); extensions.erase(std::find_if(extensions.begin(), extensions.end(), [&](Extension extension) { return extension.name == extenName; })); FreeLibrary(GetModuleHandleW(extenName.toStdWString().c_str())); } QVector Extension::GetNames() { + std::shared_lock sharedLock(extenMutex); QVector ret; for (auto extension : extensions) ret.push_back(extension.name); return ret; @@ -50,7 +52,7 @@ bool Extension::DispatchSentence(std::wstring& sentence, std::unordered_mapnext = new InfoForExtension{ i.first.c_str(), i.second, nullptr }; - std::shared_lock extenLock(extenMutex); + std::shared_lock sharedLock(extenMutex); for (auto extension : extensions) { wchar_t* nextBuffer = extension.callback(sentenceBuffer, &miscInfoLinkedList); diff --git a/GUI/host/host.cc b/GUI/host/host.cc index 3071b4e..749d27f 100644 --- a/GUI/host/host.cc +++ b/GUI/host/host.cc @@ -9,20 +9,45 @@ namespace { - struct ProcessRecord + class ProcessRecord { - HANDLE processHandle; - HANDLE sectionMutex; + public: + ProcessRecord(DWORD processId, HANDLE hostPipe) : + hostPipe(hostPipe), + section(OpenFileMappingW(FILE_MAP_READ, FALSE, (ITH_SECTION_ + std::to_wstring(processId)).c_str())), + sectionMap(MapViewOfFile(section, FILE_MAP_READ, 0, 0, HOOK_SECTION_SIZE / 2)), // jichi 1/16/2015: Changed to half to hook section size + sectionMutex(ITH_HOOKMAN_MUTEX_ + std::to_wstring(processId)) + {} + + ~ProcessRecord() + { + UnmapViewOfFile(sectionMap); + CloseHandle(section); + } + + TextHook GetHook(uint64_t addr) + { + if (sectionMap == nullptr) return {}; + LOCK(sectionMutex); + auto hooks = (const TextHook*)sectionMap; + for (int i = 0; i < MAX_HOOK; ++i) + if (hooks[i].hp.insertion_address == addr) return hooks[i]; + return {}; + } + + HANDLE hostPipe; + + private: HANDLE section; LPVOID sectionMap; - HANDLE hostPipe; + WinMutex sectionMutex; }; ThreadEventCallback OnCreate, OnRemove; ProcessEventCallback OnAttach, OnDetach; std::unordered_map> textThreadsByParams; - std::unordered_map processRecordsByIds; + std::unordered_map> processRecordsByIds; std::recursive_mutex hostMutex; @@ -51,31 +76,19 @@ namespace } } - void RegisterProcess(DWORD pid, HANDLE hostPipe) + void RegisterProcess(DWORD processId, HANDLE hostPipe) { LOCK(hostMutex); - ProcessRecord record; - record.hostPipe = hostPipe; - record.section = OpenFileMappingW(FILE_MAP_READ, FALSE, (ITH_SECTION_ + std::to_wstring(pid)).c_str()); - record.sectionMap = MapViewOfFile(record.section, FILE_MAP_READ, 0, 0, HOOK_SECTION_SIZE / 2); // jichi 1/16/2015: Changed to half to hook section size - record.processHandle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid); - record.sectionMutex = OpenMutexW(MUTEX_ALL_ACCESS, FALSE, (ITH_HOOKMAN_MUTEX_ + std::to_wstring(pid)).c_str()); - processRecordsByIds[pid] = record; - OnAttach(pid); + processRecordsByIds.insert({ processId, std::make_unique(processId, hostPipe) }); + OnAttach(processId); } - void UnregisterProcess(DWORD pid) + void UnregisterProcess(DWORD processId) { + OnDetach(processId); LOCK(hostMutex); - ProcessRecord pr = processRecordsByIds[pid]; - if (!pr.hostPipe) return; - CloseHandle(pr.sectionMutex); - UnmapViewOfFile(pr.sectionMap); - CloseHandle(pr.processHandle); - CloseHandle(pr.section); - processRecordsByIds[pid] = {}; - RemoveThreads([&](ThreadParam tp) { return tp.pid == pid; }); - OnDetach(pid); + processRecordsByIds.erase(processId); + RemoveThreads([&](ThreadParam tp) { return tp.pid == processId; }); } void StartPipe() @@ -149,7 +162,7 @@ namespace Host // Artikash 7/25/2018: This is only called when Textractor is closed, at which point Windows should free everything itself...right? #ifdef _DEBUG // Check memory leaks LOCK(hostMutex); - for (auto[pid, pr] : processRecordsByIds) UnregisterProcess(pid); + processRecordsByIds.clear(); textThreadsByParams.clear(); #endif } @@ -204,54 +217,37 @@ namespace Host void DetachProcess(DWORD processId) { + LOCK(hostMutex); auto command = HOST_COMMAND_DETACH; - WriteFile(processRecordsByIds[processId].hostPipe, &command, sizeof(command), DUMMY, nullptr); + WriteFile(processRecordsByIds.at(processId)->hostPipe, &command, sizeof(command), DUMMY, nullptr); } - void InsertHook(DWORD pid, HookParam hp, std::string name) + void InsertHook(DWORD processId, HookParam hp, std::string name) { + LOCK(hostMutex); auto command = InsertHookCmd(hp, name); - WriteFile(processRecordsByIds[pid].hostPipe, &command, sizeof(command), DUMMY, nullptr); + WriteFile(processRecordsByIds.at(processId)->hostPipe, &command, sizeof(command), DUMMY, nullptr); } - void RemoveHook(DWORD pid, uint64_t addr) + void RemoveHook(DWORD processId, uint64_t addr) { + LOCK(hostMutex); auto command = RemoveHookCmd(addr); - WriteFile(processRecordsByIds[pid].hostPipe, &command, sizeof(command), DUMMY, nullptr); + WriteFile(processRecordsByIds.at(processId)->hostPipe, &command, sizeof(command), DUMMY, nullptr); } - HookParam GetHookParam(DWORD pid, uint64_t addr) + HookParam GetHookParam(DWORD processId, uint64_t addr) { + if (processId == 0) return {}; LOCK(hostMutex); - HookParam ret = {}; - ProcessRecord pr = processRecordsByIds[pid]; - if (pr.sectionMap == nullptr) return ret; - WaitForSingleObject(pr.sectionMutex, 0); - const TextHook* hooks = (const TextHook*)pr.sectionMap; - for (int i = 0; i < MAX_HOOK; ++i) - if (hooks[i].hp.insertion_address == addr) - ret = hooks[i].hp; - ReleaseMutex(pr.sectionMutex); - return ret; + return processRecordsByIds.at(processId)->GetHook(addr).hp; } - std::wstring GetHookName(DWORD pid, uint64_t addr) + std::wstring GetHookName(DWORD processId, uint64_t addr) { - if (pid == 0) return L"Console"; + if (processId == 0) return L"Console"; LOCK(hostMutex); - std::string buffer = ""; - ProcessRecord pr = processRecordsByIds[pid]; - if (pr.sectionMap == nullptr) return L""; - WaitForSingleObject(pr.sectionMutex, 0); - const TextHook* hooks = (const TextHook*)pr.sectionMap; - for (int i = 0; i < MAX_HOOK; ++i) - if (hooks[i].hp.insertion_address == addr) - { - buffer.resize(hooks[i].name_length); - ReadProcessMemory(pr.processHandle, hooks[i].hook_name, buffer.data(), hooks[i].name_length, nullptr); - } - ReleaseMutex(pr.sectionMutex); - return StringToWideString(buffer, CP_UTF8); + return StringToWideString(processRecordsByIds.at(processId)->GetHook(addr).hookName, CP_UTF8); } std::shared_ptr GetThread(ThreadParam tp) diff --git a/GUI/host/host.h b/GUI/host/host.h index f933b7a..00233f7 100644 --- a/GUI/host/host.h +++ b/GUI/host/host.h @@ -15,15 +15,15 @@ namespace Host void Start(ProcessEventCallback onAttach, ProcessEventCallback onDetach, ThreadEventCallback onCreate, ThreadEventCallback onRemove, TextThread::OutputCallback output); void Close(); - bool InjectProcess(DWORD pid, DWORD timeout = 5000); - void DetachProcess(DWORD pid); + bool InjectProcess(DWORD processId, DWORD timeout = 5000); + void DetachProcess(DWORD processId); - void InsertHook(DWORD pid, HookParam hp, std::string name = ""); - void RemoveHook(DWORD pid, uint64_t addr); + void InsertHook(DWORD processId, HookParam hp, std::string name = ""); + void RemoveHook(DWORD processId, uint64_t addr); - HookParam GetHookParam(DWORD pid, uint64_t addr); + HookParam GetHookParam(DWORD processId, uint64_t addr); inline HookParam GetHookParam(ThreadParam tp) { return GetHookParam(tp.pid, tp.hook); } - std::wstring GetHookName(DWORD pid, uint64_t addr); + std::wstring GetHookName(DWORD processId, uint64_t addr); inline std::wstring GetHookName(ThreadParam tp) { return GetHookName(tp.pid, tp.hook); } std::shared_ptr GetThread(ThreadParam tp); diff --git a/GUI/mainwindow.cpp b/GUI/mainwindow.cpp index 24ed687..b7f6a71 100644 --- a/GUI/mainwindow.cpp +++ b/GUI/mainwindow.cpp @@ -59,8 +59,9 @@ MainWindow::~MainWindow() settings.setValue("Flush_Delay", TextThread::flushDelay); settings.setValue("Max_Buffer_Size", TextThread::maxBufferSize); settings.sync(); - Host::Close(); delete ui; + + Host::Close(); } void MainWindow::AddProcess(unsigned processId) diff --git a/include/const.h b/include/const.h index 8010f30..99c0547 100644 --- a/include/const.h +++ b/include/const.h @@ -4,7 +4,7 @@ // 8/23/2013 jichi // Branch: ITH/common.h, rev 128 -enum { MESSAGE_SIZE = 500, PIPE_BUFFER_SIZE = 0x1000, SHIFT_JIS = 932, MAX_THREAD_COUNT = 250, MAX_MODULE_SIZE = 120 }; +enum { MESSAGE_SIZE = 500, PIPE_BUFFER_SIZE = 2000, SHIFT_JIS = 932, MAX_THREAD_COUNT = 250, MAX_MODULE_SIZE = 120, HOOK_NAME_SIZE = 30 }; // jichi 375/2014: Add offset of pusha/pushad // http://faydoc.tripod.com/cpu/pushad.htm diff --git a/include/types.h b/include/types.h index 7dd1958..2f36375 100644 --- a/include/types.h +++ b/include/types.h @@ -40,12 +40,22 @@ struct ThreadParam // From hook, used internally by host as well template <> struct std::hash { size_t operator()(const ThreadParam& tp) const { return std::hash()((tp.pid + tp.hook) ^ (tp.retn + tp.spl)); } }; static bool operator==(const ThreadParam& one, const ThreadParam& two) { return one.pid == two.pid && one.hook == two.hook && one.retn == two.retn && one.spl == two.spl; } +class WinMutex +{ + HANDLE mutex; +public: + WinMutex(std::wstring name) : mutex(CreateMutexW(nullptr, false, name.c_str())) {} + ~WinMutex() { ReleaseMutex(mutex); CloseHandle(mutex); } + void lock() { WaitForSingleObject(mutex, 0); } + void unlock() { ReleaseMutex(mutex); } +}; + struct InsertHookCmd // From host { - InsertHookCmd(HookParam hp, std::string name = "") : hp(hp) { strcpy_s(this->name, name.c_str()); }; + InsertHookCmd(HookParam hp, std::string name = "") : hp(hp) { strcpy_s(this->name, name.c_str()); }; int command = HOST_COMMAND_NEW_HOOK; HookParam hp; - char name[MESSAGE_SIZE] = {}; + char name[HOOK_NAME_SIZE] = {}; }; struct RemoveHookCmd // From host @@ -69,4 +79,4 @@ struct HookRemovedNotif // From hook uint64_t address; }; -#define LOCK(mutex) std::lock_guard lock(mutex) +#define LOCK(mutex) std::lock_guard lock(mutex) diff --git a/vnrhook/hijack/texthook.cc b/vnrhook/hijack/texthook.cc index 7000332..5b8acbc 100644 --- a/vnrhook/hijack/texthook.cc +++ b/vnrhook/hijack/texthook.cc @@ -102,16 +102,14 @@ namespace { // unnamed bool TextHook::InsertHook() { - bool ret = false; //ConsoleOutput("vnrcli:InsertHook: enter"); - WaitForSingleObject(hmMutex, 0); - if (hp.type & DIRECT_READ) ret = InsertReadCode(); + LOCK(*sectionMutex); + if (hp.type & DIRECT_READ) return InsertReadCode(); #ifndef _WIN64 - else ret = InsertHookCode(); + else return InsertHookCode(); #endif - ReleaseMutex(hmMutex); //ConsoleOutput("vnrcli:InsertHook: leave"); - return ret; + return false; } #ifndef _WIN64 @@ -222,16 +220,8 @@ DWORD TextHook::UnsafeSend(DWORD dwDataBase, DWORD dwRetn) bool TextHook::InsertHookCode() { - bool ret = false; // jichi 9/17/2013: might raise 0xC0000005 AccessViolationException on win7 - __try { ret = UnsafeInsertHookCode(); } - __except (1) {}; - return ret; -} - - -bool TextHook::UnsafeInsertHookCode() -{ + // Artikash 10/30/2018: No, I think that's impossible now that I moved to minhook if (hp.type & MODULE_OFFSET) // Map hook offset to real address if (hp.type & FUNCTION_OFFSET) if (FARPROC function = GetProcAddress(GetModuleHandleW(hp.module), hp.function)) hp.insertion_address += (uint64_t)function; @@ -243,7 +233,7 @@ bool TextHook::UnsafeInsertHookCode() if (hp.codepage == 0) hp.codepage = SHIFT_JIS; // Use Shift-JIS unless custom encoding was specified BYTE* original; - insert: +insert: if (MH_STATUS err = MH_CreateHook((void*)hp.insertion_address, (void*)trampoline, (void**)&original)) if (err == MH_ERROR_ALREADY_CREATED) { @@ -334,14 +324,13 @@ bool TextHook::InsertReadCode() return true; } -void TextHook::InitHook(const HookParam &h, LPCSTR name, WORD set_flag) +void TextHook::InitHook(const HookParam &h, LPCSTR name, DWORD set_flag) { - WaitForSingleObject(hmMutex, 0); + LOCK(*sectionMutex); hp = h; hp.insertion_address = hp.address; hp.type |= set_flag; - if (name && name != hook_name) SetHookName(name); - ReleaseMutex(hmMutex); + strcpy_s(hookName, name); } void TextHook::RemoveHookCode() @@ -358,22 +347,12 @@ void TextHook::RemoveReadCode() void TextHook::ClearHook() { - WaitForSingleObject(hmMutex, 0); - if (hook_name) ConsoleOutput(("Textractor: removing hook: " + std::string(hook_name)).c_str()); + LOCK(*sectionMutex); + ConsoleOutput(("Textractor: removing hook: " + std::string(hookName)).c_str()); if (hp.type & DIRECT_READ) RemoveReadCode(); else RemoveHookCode(); NotifyHookRemove(hp.insertion_address); - if (hook_name) delete[] hook_name; memset(this, 0, sizeof(TextHook)); // jichi 11/30/2013: This is the original code of ITH - ReleaseMutex(hmMutex); -} - -void TextHook::SetHookName(LPCSTR name) -{ - name_length = strlen(name) + 1; - if (hook_name) delete[] hook_name; - hook_name = new char[name_length]; - strcpy(hook_name, name); } int TextHook::GetLength(DWORD base, DWORD in) diff --git a/vnrhook/hijack/texthook.h b/vnrhook/hijack/texthook.h index 6f3d666..15ebc5f 100644 --- a/vnrhook/hijack/texthook.h +++ b/vnrhook/hijack/texthook.h @@ -20,32 +20,28 @@ class TextHook { bool InsertHookCode(); bool InsertReadCode(); - bool UnsafeInsertHookCode(); DWORD UnsafeSend(DWORD dwDataBase, DWORD dwRetn); + int GetLength(DWORD base, DWORD in); // jichi 12/25/2013: Return 0 if failed void RemoveHookCode(); void RemoveReadCode(); - void SetHookName(LPCSTR name); public: HookParam hp; - LPSTR hook_name; - int name_length; + char hookName[HOOK_NAME_SIZE]; BYTE trampoline[120]; HANDLE readerHandle; bool InsertHook(); - void InitHook(const HookParam &hp, LPCSTR name = 0, WORD set_flag = 0); + void InitHook(const HookParam &hp, LPCSTR name, DWORD set_flag); DWORD Send(DWORD dwDataBase, DWORD dwRetn); void ClearHook(); - int GetLength(DWORD base, DWORD in); // jichi 12/25/2013: Return 0 if failed }; enum { MAX_HOOK = 300 }; enum { HOOK_SECTION_SIZE = MAX_HOOK * sizeof(TextHook) * 2, HOOK_BUFFER_SIZE = MAX_HOOK * sizeof(TextHook) }; extern TextHook *hookman; - extern bool running; - -extern HANDLE hookPipe, hmMutex; +extern HANDLE hookPipe; +extern std::unique_ptr sectionMutex; // EOF diff --git a/vnrhook/main.cc b/vnrhook/main.cc index 2f6a4a4..972333f 100644 --- a/vnrhook/main.cc +++ b/vnrhook/main.cc @@ -21,7 +21,7 @@ HANDLE hSection; bool running; int currentHook = 0, userhookCount = 0; DWORD trigger = 0; -HANDLE hmMutex; +std::unique_ptr sectionMutex; BOOL WINAPI DllMain(HINSTANCE hModule, DWORD fdwReason, LPVOID) { @@ -29,7 +29,7 @@ BOOL WINAPI DllMain(HINSTANCE hModule, DWORD fdwReason, LPVOID) { case DLL_PROCESS_ATTACH: { - ::hmMutex = CreateMutexW(nullptr, FALSE, (ITH_HOOKMAN_MUTEX_ + std::to_wstring(GetCurrentProcessId())).c_str()); + sectionMutex = std::make_unique(ITH_HOOKMAN_MUTEX_ + std::to_wstring(GetCurrentProcessId())); if (GetLastError() == ERROR_ALREADY_EXISTS) return FALSE; DisableThreadLibraryCalls(hModule); @@ -55,7 +55,6 @@ BOOL WINAPI DllMain(HINSTANCE hModule, DWORD fdwReason, LPVOID) CloseHandle(::hookPipe); CloseHandle(hSection); - CloseHandle(hmMutex); //} ITH_EXCEPT {} } break;