optimize performance

This commit is contained in:
Akash Mozumdar 2019-10-03 16:00:19 -04:00
parent 822c2ae233
commit 415ceb4694
3 changed files with 20 additions and 18 deletions

View File

@ -22,18 +22,16 @@ namespace
{ {
AutoHandle<> hookPipe = INVALID_HANDLE_VALUE, mappedFile = INVALID_HANDLE_VALUE; AutoHandle<> hookPipe = INVALID_HANDLE_VALUE, mappedFile = INVALID_HANDLE_VALUE;
TextHook (*hooks)[MAX_HOOK]; TextHook (*hooks)[MAX_HOOK];
bool running;
int currentHook = 0; int currentHook = 0;
} }
DWORD WINAPI Pipe(LPVOID) DWORD WINAPI Pipe(LPVOID)
{ {
while (running) for (bool running = true; running; hookPipe = INVALID_HANDLE_VALUE)
{ {
DWORD count = 0; DWORD count = 0;
BYTE buffer[PIPE_BUFFER_SIZE] = {}; BYTE buffer[PIPE_BUFFER_SIZE] = {};
AutoHandle<> hostPipe = INVALID_HANDLE_VALUE; AutoHandle<> hostPipe = INVALID_HANDLE_VALUE;
hookPipe = INVALID_HANDLE_VALUE;
while (!hostPipe || !hookPipe) while (!hostPipe || !hookPipe)
{ {
@ -82,10 +80,10 @@ DWORD WINAPI Pipe(LPVOID)
break; break;
} }
} }
hookPipe = INVALID_HANDLE_VALUE;
for (auto& hook : *hooks) if (hook.address) hook.Clear(); MH_Uninitialize();
for (auto& hook : *hooks) hook.Clear();
FreeLibraryAndExitThread(GetModuleHandleW(ITH_DLL), 0); FreeLibraryAndExitThread(GetModuleHandleW(ITH_DLL), 0);
return 0;
} }
void TextOutput(ThreadParam tp, BYTE* text, int len) void TextOutput(ThreadParam tp, BYTE* text, int len)
@ -136,16 +134,14 @@ BOOL WINAPI DllMain(HINSTANCE hModule, DWORD fdwReason, LPVOID)
memset(hooks, 0, HOOK_BUFFER_SIZE); memset(hooks, 0, HOOK_BUFFER_SIZE);
MH_Initialize(); MH_Initialize();
running = true;
CloseHandle(CreateThread(nullptr, 0, Pipe, nullptr, 0, nullptr)); // Using std::thread here = deadlock CloseHandle(CreateThread(nullptr, 0, Pipe, nullptr, 0, nullptr)); // Using std::thread here = deadlock
} }
break; break;
case DLL_PROCESS_DETACH: case DLL_PROCESS_DETACH:
{ {
running = false;
UnmapViewOfFile(hooks);
MH_Uninitialize(); MH_Uninitialize();
UnmapViewOfFile(hooks);
} }
break; break;
} }

View File

@ -105,13 +105,13 @@ void SetTrigger()
// - TextHook methods - // - TextHook methods -
bool TextHook::Insert(HookParam h, DWORD set_flag) bool TextHook::Insert(HookParam hp, DWORD set_flag)
{ {
std::scoped_lock lock(viewMutex); std::scoped_lock lock(viewMutex);
hp = h;
address = hp.address;
hp.type |= set_flag; hp.type |= set_flag;
if (hp.type & USING_UTF8) hp.codepage = CP_UTF8; if (hp.type & USING_UTF8) hp.codepage = CP_UTF8;
this->hp = hp;
address = hp.address;
if (hp.type & DIRECT_READ) return InsertReadCode(); if (hp.type & DIRECT_READ) return InsertReadCode();
else return InsertHookCode(); else return InsertHookCode();
} }
@ -120,6 +120,7 @@ bool TextHook::Insert(HookParam h, DWORD set_flag)
// - dwDataBase: the stack address // - dwDataBase: the stack address
void TextHook::Send(uintptr_t dwDataBase) void TextHook::Send(uintptr_t dwDataBase)
{ {
_InterlockedIncrement(&useCount);
__try __try
{ {
if (trigger) trigger = Engine::InsertDynamicHook(location, *(DWORD *)(dwDataBase - 0x1c), *(DWORD *)(dwDataBase - 0x18)); if (trigger) trigger = Engine::InsertDynamicHook(location, *(DWORD *)(dwDataBase - 0x1c), *(DWORD *)(dwDataBase - 0x18));
@ -134,7 +135,7 @@ void TextHook::Send(uintptr_t dwDataBase)
// jichi 10/24/2014: generic hook function // jichi 10/24/2014: generic hook function
if (hp.hook_fun && !hp.hook_fun(dwDataBase, &hp)) hp.hook_fun = nullptr; if (hp.hook_fun && !hp.hook_fun(dwDataBase, &hp)) hp.hook_fun = nullptr;
if (hp.type & HOOK_EMPTY) return; // jichi 10/24/2014: dummy hook only for dynamic hook if (hp.type & HOOK_EMPTY) goto done; // jichi 10/24/2014: dummy hook only for dynamic hook
if (hp.text_fun) { if (hp.text_fun) {
hp.text_fun(dwDataBase, &hp, 0, &dwDataIn, &dwSplit, &dwCount); hp.text_fun(dwDataBase, &hp, 0, &dwDataIn, &dwSplit, &dwCount);
@ -150,7 +151,7 @@ void TextHook::Send(uintptr_t dwDataBase)
dwCount = GetLength(dwDataBase, dwDataIn); dwCount = GetLength(dwDataBase, dwDataIn);
} }
if (dwCount == 0) return; if (dwCount == 0) goto done;
if (dwCount > TEXT_BUFFER_SIZE) dwCount = TEXT_BUFFER_SIZE; if (dwCount > TEXT_BUFFER_SIZE) dwCount = TEXT_BUFFER_SIZE;
BYTE pbData[TEXT_BUFFER_SIZE]; BYTE pbData[TEXT_BUFFER_SIZE];
if (hp.length_offset == 1) { if (hp.length_offset == 1) {
@ -161,13 +162,13 @@ void TextHook::Send(uintptr_t dwDataBase)
} }
else ::memcpy(pbData, (void*)dwDataIn, dwCount); else ::memcpy(pbData, (void*)dwDataIn, dwCount);
if (hp.filter_fun && !hp.filter_fun(pbData, &dwCount, &hp, 0) || dwCount <= 0) return; if (hp.filter_fun && !hp.filter_fun(pbData, &dwCount, &hp, 0) || dwCount <= 0) goto done;
if (hp.type & (NO_CONTEXT | FIXING_SPLIT)) dwRetn = 0; if (hp.type & (NO_CONTEXT | FIXING_SPLIT)) dwRetn = 0;
TextOutput({ GetCurrentProcessId(), address, dwRetn, dwSplit }, pbData, dwCount); TextOutput({ GetCurrentProcessId(), address, dwRetn, dwSplit }, pbData, dwCount);
#else // _WIN32 #else // _WIN32
if (hp.type & HOOK_EMPTY) return; // jichi 10/24/2014: dummy hook only for dynamic hook if (hp.type & HOOK_EMPTY) goto done; // jichi 10/24/2014: dummy hook only for dynamic hook
int count = 0; int count = 0;
ThreadParam tp = { GetCurrentProcessId(), address, *(uintptr_t*)dwDataBase, 0 }; // first value on stack (if hooked start of function, this is return address) ThreadParam tp = { GetCurrentProcessId(), address, *(uintptr_t*)dwDataBase, 0 }; // first value on stack (if hooked start of function, this is return address)
uintptr_t data = *(uintptr_t*)(dwDataBase + hp.offset); // default value uintptr_t data = *(uintptr_t*)(dwDataBase + hp.offset); // default value
@ -181,7 +182,7 @@ void TextHook::Send(uintptr_t dwDataBase)
data += hp.padding; data += hp.padding;
count = GetLength(dwDataBase, data); count = GetLength(dwDataBase, data);
if (count == 0) return; if (count == 0) goto done;
if (count > TEXT_BUFFER_SIZE) count = TEXT_BUFFER_SIZE; if (count > TEXT_BUFFER_SIZE) count = TEXT_BUFFER_SIZE;
BYTE pbData[TEXT_BUFFER_SIZE]; BYTE pbData[TEXT_BUFFER_SIZE];
if (hp.length_offset == 1) if (hp.length_offset == 1)
@ -202,10 +203,13 @@ void TextHook::Send(uintptr_t dwDataBase)
{ {
if (!err) if (!err)
{ {
ConsoleOutput(SEND_ERROR);
err = true; err = true;
ConsoleOutput(SEND_ERROR);
ConsoleOutput("Textractor: in %s", hp.name);
} }
} }
done:
_InterlockedDecrement(&useCount);
} }
bool TextHook::InsertHookCode() bool TextHook::InsertHookCode()
@ -278,6 +282,7 @@ bool TextHook::InsertReadCode()
void TextHook::RemoveHookCode() void TextHook::RemoveHookCode()
{ {
MH_DisableHook(location); MH_DisableHook(location);
while (useCount != 0);
MH_RemoveHook(location); MH_RemoveHook(location);
} }

View File

@ -38,6 +38,7 @@ private:
void RemoveHookCode(); void RemoveHookCode();
void RemoveReadCode(); void RemoveReadCode();
volatile DWORD useCount;
HANDLE readerThread, readerEvent; HANDLE readerThread, readerEvent;
bool err; bool err;
BYTE trampoline[x64 ? 140 : 40]; BYTE trampoline[x64 ? 140 : 40];