chenx221/Textractor_test
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

refactor

Browse Source
This commit is contained in:
Akash Mozumdar 2018-07-19 01:43:31 -04:00
parent 05a14d201e
commit e4f5f3cfd5
2 changed files with 118 additions and 147 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

197
vnr/texthook/hookman.cc
View File

@ -15,7 +15,7 @@
#include "winmutex/winmutex.h" #include "winmutex/winmutex.h"
#include <atlbase.h> #include <atlbase.h>
#define HM_LOCK CriticalSectionLocker locker(hmcs) // Synchronized scope for accessing private data #define HM_LOCK CriticalSectionLocker locker(hmCs) // Synchronized scope for accessing private data
HookManager::HookManager() : HookManager::HookManager() :
current(nullptr), current(nullptr),
@ -28,68 +28,59 @@ HookManager::HookManager() :
textThreadsByParams(), textThreadsByParams(),
processRecordsByIds() processRecordsByIds()
{ {
TextThread* consoleTextThread = textThreadsByParams[{ 0, -1UL, -1UL, -1UL }] = new TextThread({ 0, -1UL, -1UL, -1UL }, nextThreadNumber++, splitDelay); InitializeCriticalSection(&hmCs);
consoleTextThread->Status() |= USING_UNICODE;
SetCurrent(consoleTextThread);
InitializeCriticalSection(&hmcs); TextThread* consoleTextThread = textThreadsByParams[{ 0, -1UL, -1UL, -1UL }] = new TextThread({ 0, -1UL, -1UL, -1UL }, nextThreadNumber++, splitDelay);
consoleTextThread->Status() |= USING_UNICODE;
SetCurrent(consoleTextThread);
} }
HookManager::~HookManager() HookManager::~HookManager()
{ {
DeleteCriticalSection(&hmcs); HM_LOCK;
DeleteCriticalSection(&hmCs);
} }
TextThread *HookManager::FindSingle(DWORD number) TextThread *HookManager::FindSingle(DWORD number)
{ {
HM_LOCK;
for (auto i : textThreadsByParams) for (auto i : textThreadsByParams)
{ if (i.second->Number() == number)
if (i.second->Number() == number)
{
return i.second; return i.second;
}
}
return nullptr; return nullptr;
} }
void HookManager::SetCurrent(TextThread *it) void HookManager::SetCurrent(TextThread *it)
{ {
if (current) HM_LOCK;
current->Status() &= ~CURRENT_SELECT; current->Status() &= ~CURRENT_SELECT;
current = it; current = it;
if (it) it->Status() |= CURRENT_SELECT;
it->Status() |= CURRENT_SELECT;
} }
void HookManager::SelectCurrent(DWORD num) void HookManager::SelectCurrent(DWORD num)
{ {
if (TextThread *st = FindSingle(num)) { HM_LOCK;
SetCurrent(st); if (TextThread *st = FindSingle(num))
if (reset) {
reset(st); SetCurrent(st);
//st->ResetEditText(); if (reset) reset(st);
} }
} }
void HookManager::RemoveSingleHook(DWORD pid, DWORD addr) void HookManager::RemoveSingleHook(DWORD pid, DWORD addr)
{ {
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 (i.first.pid == pid && i.first.hook == addr) {
{ if (remove) remove(i.second);
if (remove) delete i.second;
{ removedThreads.push_back(i.first);
remove(i.second); }
} for (auto i : removedThreads) textThreadsByParams.erase(i);
delete i.second; SelectCurrent(0);
removedThreads.push_back(i.first);
}
}
for (auto i : removedThreads)
{
textThreadsByParams.erase(i);
}
SelectCurrent(0);
} }
void HookManager::RemoveProcessContext(DWORD pid) void HookManager::RemoveProcessContext(DWORD pid)
@ -97,127 +88,108 @@ void HookManager::RemoveProcessContext(DWORD pid)
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)
if (i.first.hook == pid)
{ {
if (remove) if (remove) remove(i.second);
{
remove(i.second);
}
delete i.second; delete i.second;
removedThreads.push_back(i.first); removedThreads.push_back(i.first);
} }
} for (auto i : removedThreads) textThreadsByParams.erase(i);
for (auto i : removedThreads)
{
textThreadsByParams.erase(i);
}
SelectCurrent(0); SelectCurrent(0);
} }
void HookManager::RegisterProcess(DWORD pid, HANDLE hostPipe) void HookManager::RegisterProcess(DWORD pid, HANDLE hostPipe)
{ {
HM_LOCK; HM_LOCK;
ProcessRecord* record = processRecordsByIds[pid] = new ProcessRecord;
ProcessRecord* record = processRecordsByIds[pid] = new ProcessRecord; record->hostPipe = hostPipe;
record->hostPipe = hostPipe; record->hookman_section = OpenFileMappingW(FILE_MAP_READ, FALSE, (ITH_SECTION_ + std::to_wstring(pid)).c_str());
record->hookman_section = OpenFileMappingW(FILE_MAP_READ, FALSE, (ITH_SECTION_ + std::to_wstring(pid)).c_str()); record->hookman_map = MapViewOfFile(record->hookman_section, FILE_MAP_READ, 0, 0, HOOK_SECTION_SIZE / 2); // jichi 1/16/2015: Changed to half to hook section size
record->hookman_map = MapViewOfFile(record->hookman_section, FILE_MAP_READ, 0, 0, HOOK_SECTION_SIZE / 2); // jichi 1/16/2015: Changed to half to hook section size record->process_handle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
record->process_handle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid); record->hookman_mutex = OpenMutexW(MUTEX_ALL_ACCESS, FALSE, (ITH_HOOKMAN_MUTEX_ + std::to_wstring(pid)).c_str());
record->hookman_mutex = OpenMutexW(MUTEX_ALL_ACCESS, FALSE, (ITH_HOOKMAN_MUTEX_ + std::to_wstring(pid)).c_str()); if (attach) attach(pid);
if (attach)
attach(pid);
} }
void HookManager::UnRegisterProcess(DWORD pid) void HookManager::UnRegisterProcess(DWORD pid)
{ {
HM_LOCK; HM_LOCK;
ProcessRecord pr = *processRecordsByIds[pid];
ProcessRecord pr = *processRecordsByIds[pid]; CloseHandle(pr.hookman_mutex);
CloseHandle(pr.hookman_mutex); UnmapViewOfFile(pr.hookman_map);
UnmapViewOfFile(pr.hookman_map); CloseHandle(pr.process_handle);
CloseHandle(pr.process_handle); CloseHandle(pr.hookman_section);
CloseHandle(pr.hookman_section); processRecordsByIds.erase(pid);
processRecordsByIds.erase(pid); RemoveProcessContext(pid);
RemoveProcessContext(pid); if (detach) detach(pid);
if (detach)
detach(pid);
} }
void HookManager::DispatchText(DWORD pid, DWORD hook, DWORD retn, DWORD spl, const BYTE *text, int len) void HookManager::DispatchText(DWORD pid, DWORD hook, DWORD retn, DWORD spl, const BYTE *text, int len)
{ {
// jichi 20/27/2013: When PID is zero, the text comes from console, which I don't need // jichi 20/27/2013: When PID is zero, the text comes from console, which I don't need
if (!text || !pid || len <= 0) if (!text || !pid || len <= 0)
return; return;
HM_LOCK; HM_LOCK;
ThreadParameter tp = {pid, hook, retn, spl}; ThreadParameter tp = { pid, hook, retn, spl };
TextThread *it; TextThread *it;
if (!(it = textThreadsByParams[tp])) if ((it = textThreadsByParams[tp]) == nullptr)
{ {
it = textThreadsByParams[tp] = new TextThread(tp, nextThreadNumber++, splitDelay); it = textThreadsByParams[tp] = new TextThread(tp, nextThreadNumber++, splitDelay);
if (create) if (create) create(it);
{ }
create(it); it->AddText(text, len);
}
}
it->AddText(text, len);
} }
void HookManager::AddConsoleOutput(LPCWSTR text) void HookManager::AddConsoleOutput(LPCWSTR text)
{ {
if (text) HM_LOCK;
{ if (text)
int len = wcslen(text) * 2; {
TextThread *console = textThreadsByParams[{ 0, -1UL, -1UL, -1UL }]; int len = wcslen(text) * 2;
console->AddSentence(std::wstring(text)); TextThread *console = textThreadsByParams[{ 0, -1UL, -1UL, -1UL }];
} console->AddSentence(std::wstring(text));
}
} }
void HookManager::ClearCurrent() void HookManager::ClearCurrent()
{ {
HM_LOCK; HM_LOCK;
if (current) { current->Reset();
current->Reset(); if (reset) reset(current);
if (reset)
reset(current);
}
} }
ProcessRecord *HookManager::GetProcessRecord(DWORD pid) ProcessRecord *HookManager::GetProcessRecord(DWORD pid)
{ {
HM_LOCK; HM_LOCK;
return processRecordsByIds[pid]; return processRecordsByIds[pid];
} }
HANDLE HookManager::GetHostPipe(DWORD pid) HANDLE HookManager::GetHostPipe(DWORD pid)
{ {
HM_LOCK; HM_LOCK;
return processRecordsByIds[pid] ? processRecordsByIds[pid]->hostPipe : nullptr; return processRecordsByIds[pid] ? processRecordsByIds[pid]->hostPipe : nullptr;
} }
HookParam HookManager::GetHookParam(DWORD pid, DWORD addr) HookParam HookManager::GetHookParam(DWORD pid, DWORD addr)
{ {
HM_LOCK;
HookParam ret = {}; HookParam ret = {};
ProcessRecord* pr = GetProcessRecord(pid); ProcessRecord* pr = GetProcessRecord(pid);
if (pr == nullptr) return ret; if (pr == nullptr) return ret;
WaitForSingleObject(pr->hookman_mutex, 0); MutexLocker locker(pr->hookman_mutex);
const Hook* hooks = (const Hook*)pr->hookman_map; const Hook* hooks = (const Hook*)pr->hookman_map;
for (int i = 0; i < MAX_HOOK; ++i) for (int i = 0; i < MAX_HOOK; ++i)
if (hooks[i].Address() == addr) if (hooks[i].Address() == addr)
ret = hooks[i].hp; ret = hooks[i].hp;
ReleaseMutex(pr->hookman_mutex);
return ret; return ret;
} }
std::wstring HookManager::GetHookName(DWORD pid, DWORD addr) std::wstring HookManager::GetHookName(DWORD pid, DWORD addr)
{ {
HM_LOCK;
std::string buffer; std::string buffer;
ProcessRecord* pr = GetProcessRecord(pid); ProcessRecord* pr = GetProcessRecord(pid);
if (pr == nullptr) return L""; if (pr == nullptr) return L"";
WaitForSingleObject(pr->hookman_mutex, 0); MutexLocker locker(pr->hookman_mutex);
USES_CONVERSION; USES_CONVERSION;
const Hook* hooks = (const Hook*)pr->hookman_map; const Hook* hooks = (const Hook*)pr->hookman_map;
for (int i = 0; i < MAX_HOOK; ++i) for (int i = 0; i < MAX_HOOK; ++i)
@ -227,8 +199,7 @@ std::wstring HookManager::GetHookName(DWORD pid, DWORD addr)
buffer.resize(hooks[i].NameLength()); buffer.resize(hooks[i].NameLength());
ReadProcessMemory(pr->process_handle, hooks[i].Name(), &buffer[0], hooks[i].NameLength(), nullptr); ReadProcessMemory(pr->process_handle, hooks[i].Name(), &buffer[0], hooks[i].NameLength(), nullptr);
} }
} }
ReleaseMutex(pr->hookman_mutex);
return std::wstring(A2W(buffer.c_str())); return std::wstring(A2W(buffer.c_str()));
} }

68
vnr/texthook/hookman.h
View File

@ -10,13 +10,13 @@
#include <string> #include <string>
#include "vnrhook/include/types.h" #include "vnrhook/include/types.h"
struct ProcessRecord struct ProcessRecord
{ {
HANDLE process_handle; HANDLE process_handle;
HANDLE hookman_mutex; HANDLE hookman_mutex;
HANDLE hookman_section; HANDLE hookman_section;
LPVOID hookman_map; LPVOID hookman_map;
HANDLE hostPipe; HANDLE hostPipe;
}; };
typedef void(*ProcessEventCallback)(DWORD pid); typedef void(*ProcessEventCallback)(DWORD pid);
@ -33,45 +33,45 @@ struct ThreadParameterHasher
class __declspec(dllexport) HookManager class __declspec(dllexport) HookManager
{ {
public: public:
HookManager(); HookManager();
~HookManager(); ~HookManager();
TextThread *FindSingle(DWORD number); TextThread *FindSingle(DWORD number);
ProcessRecord *GetProcessRecord(DWORD pid); ProcessRecord *GetProcessRecord(DWORD pid);
HANDLE GetHostPipe(DWORD pid); HANDLE GetHostPipe(DWORD pid);
void ClearCurrent(); void ClearCurrent();
void SelectCurrent(DWORD num); void SelectCurrent(DWORD num);
void SetCurrent(TextThread *it); void SetCurrent(TextThread *it);
void AddConsoleOutput(LPCWSTR text); void AddConsoleOutput(LPCWSTR 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 RemoveProcessContext(DWORD pid); // private
void RemoveSingleHook(DWORD pid, DWORD addr); 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);
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);
void RegisterThreadCreateCallback(ThreadEventCallback cf) { create = cf; } void RegisterThreadCreateCallback(ThreadEventCallback cf) { create = cf; }
void RegisterThreadRemoveCallback(ThreadEventCallback cf) { remove = cf; } void RegisterThreadRemoveCallback(ThreadEventCallback cf) { remove = cf; }
void RegisterThreadResetCallback(ThreadEventCallback cf) { reset = cf; } void RegisterThreadResetCallback(ThreadEventCallback cf) { reset = cf; }
void RegisterProcessAttachCallback(ProcessEventCallback cf) { attach = cf; } void RegisterProcessAttachCallback(ProcessEventCallback cf) { attach = cf; }
void RegisterProcessDetachCallback(ProcessEventCallback cf) { detach = cf; } void RegisterProcessDetachCallback(ProcessEventCallback cf) { detach = cf; }
void SetSplitInterval(unsigned int splitDelay) { this->splitDelay = splitDelay; } void SetSplitInterval(unsigned int splitDelay) { this->splitDelay = splitDelay; }
private: private:
std::unordered_map<ThreadParameter, TextThread*, ThreadParameterHasher> textThreadsByParams; std::unordered_map<ThreadParameter, TextThread*, ThreadParameterHasher> textThreadsByParams;
std::unordered_map<DWORD, ProcessRecord*> processRecordsByIds; std::unordered_map<DWORD, ProcessRecord*> processRecordsByIds;
CRITICAL_SECTION hmcs; CRITICAL_SECTION hmCs;
TextThread *current; TextThread *current;
ThreadEventCallback create, remove, reset; ThreadEventCallback create, remove, reset;
ProcessEventCallback attach, detach; ProcessEventCallback attach, detach;
WORD nextThreadNumber; WORD nextThreadNumber;
unsigned int splitDelay; unsigned int splitDelay;
}; };
// EOF // EOF