refactoring and bugfixes

This commit is contained in:
Akash Mozumdar 2018-07-22 18:10:00 -07:00
parent 89f8cb9d9b
commit 3d4ee3d698
4 changed files with 15 additions and 26 deletions

View File

@ -35,11 +35,14 @@ HookManager::HookManager() :
HookManager::~HookManager() HookManager::~HookManager()
{ {
HM_LOCK; EnterCriticalSection(&hmCs);
RemoveThreads([](auto one, auto two) { return true; }, {});
for (auto i : processRecordsByIds) UnRegisterProcess(i.first);
LeaveCriticalSection(&hmCs);
DeleteCriticalSection(&hmCs); DeleteCriticalSection(&hmCs);
} }
TextThread *HookManager::FindSingle(DWORD number) TextThread* HookManager::FindSingle(DWORD number)
{ {
HM_LOCK; HM_LOCK;
for (auto i : textThreadsByParams) for (auto i : textThreadsByParams)
@ -48,26 +51,12 @@ TextThread *HookManager::FindSingle(DWORD number)
return nullptr; return nullptr;
} }
void HookManager::RemoveSingleHook(DWORD pid, DWORD addr) void HookManager::RemoveThreads(bool(*RemoveIf)(ThreadParameter, ThreadParameter), ThreadParameter cmp)
{ {
HM_LOCK; HM_LOCK;
std::vector<ThreadParameter> removedThreads; std::vector<ThreadParameter> removedThreads;
for (auto i : textThreadsByParams) for (auto i : textThreadsByParams)
if (i.first.pid == pid && i.first.hook == addr) if (RemoveIf(i.first, cmp))
{
if (remove) remove(i.second);
delete i.second;
removedThreads.push_back(i.first);
}
for (auto i : removedThreads) textThreadsByParams.erase(i);
}
void HookManager::RemoveProcessContext(DWORD pid)
{
HM_LOCK;
std::vector<ThreadParameter> removedThreads;
for (auto i : textThreadsByParams)
if (i.first.pid == pid)
{ {
if (remove) remove(i.second); if (remove) remove(i.second);
delete i.second; delete i.second;
@ -98,7 +87,7 @@ void HookManager::UnRegisterProcess(DWORD pid)
CloseHandle(pr.process_handle); CloseHandle(pr.process_handle);
CloseHandle(pr.hookman_section); CloseHandle(pr.hookman_section);
processRecordsByIds.erase(pid); processRecordsByIds.erase(pid);
RemoveProcessContext(pid); RemoveThreads([](auto one, auto two) { return one.pid == two.pid; }, { pid, 0, 0, 0 });
if (detach) detach(pid); if (detach) detach(pid);
} }
@ -149,7 +138,7 @@ HookParam HookManager::GetHookParam(DWORD pid, DWORD addr)
std::wstring HookManager::GetHookName(DWORD pid, DWORD addr) std::wstring HookManager::GetHookName(DWORD pid, DWORD addr)
{ {
HM_LOCK; HM_LOCK;
std::string buffer; std::string buffer = "";
ProcessRecord pr = processRecordsByIds[pid]; ProcessRecord pr = processRecordsByIds[pid];
if (pr.hookman_map == nullptr) return L""; if (pr.hookman_map == nullptr) return L"";
MutexLocker locker(pr.hookman_mutex); MutexLocker locker(pr.hookman_mutex);

View File

@ -37,14 +37,13 @@ public:
HookManager(); HookManager();
~HookManager(); ~HookManager();
TextThread *FindSingle(DWORD number); TextThread* FindSingle(DWORD number);
HANDLE GetHostPipe(DWORD pid);
void AddConsoleOutput(std::wstring text); void AddConsoleOutput(std::wstring text);
void DispatchText(DWORD pid, DWORD hook, DWORD retn, DWORD split, const BYTE *text, int len); void DispatchText(DWORD pid, DWORD hook, DWORD retn, DWORD split, const BYTE *text, int len);
void RemoveProcessContext(DWORD pid); // private void RemoveThreads(bool(*RemoveIf)(ThreadParameter, ThreadParameter), ThreadParameter cmp);
void RemoveSingleHook(DWORD pid, DWORD addr);
void RegisterProcess(DWORD pid, HANDLE hostPipe); void RegisterProcess(DWORD pid, HANDLE hostPipe);
void UnRegisterProcess(DWORD pid); void UnRegisterProcess(DWORD pid);
HANDLE GetHostPipe(DWORD pid);
HookParam GetHookParam(DWORD pid, DWORD addr); HookParam GetHookParam(DWORD pid, DWORD addr);
std::wstring GetHookName(DWORD pid, DWORD addr); std::wstring GetHookName(DWORD pid, DWORD addr);

View File

@ -151,7 +151,7 @@ DLLEXPORT bool RemoveHook(DWORD pid, DWORD addr)
WaitForSingleObject(hookRemovalEvent, 1000); WaitForSingleObject(hookRemovalEvent, 1000);
CloseHandle(hookRemovalEvent); CloseHandle(hookRemovalEvent);
man->RemoveSingleHook(pid, addr); man->RemoveThreads([](auto one, auto two) { return one.pid == two.pid && one.hook == two.hook; }, { pid, addr, 0, 0 });
return true; return true;
} }

View File

@ -29,7 +29,8 @@ TextThread::TextThread(ThreadParameter tp, unsigned int threadNumber, unsigned i
TextThread::~TextThread() TextThread::~TextThread()
{ {
TT_LOCK; EnterCriticalSection(&ttCs);
LeaveCriticalSection(&ttCs);
DeleteCriticalSection(&ttCs); DeleteCriticalSection(&ttCs);
} }