This commit is contained in:
Akash Mozumdar 2018-08-28 17:21:20 -04:00
parent d1007097ad
commit 189691a562
3 changed files with 29 additions and 28 deletions

View File

@ -33,7 +33,7 @@ namespace
void DispatchText(ThreadParam tp, const BYTE* text, int len) void DispatchText(ThreadParam tp, const BYTE* text, int len)
{ {
if (!text || len <= 0) return; if (!text || len <= 0) return;
LOCK hostLock(hostMutex); LOCK(hostMutex);
TextThread *it; TextThread *it;
if ((it = textThreadsByParams[tp]) == nullptr) if ((it = textThreadsByParams[tp]) == nullptr)
OnCreate(it = textThreadsByParams[tp] = new TextThread(tp, Host::GetHookParam(tp).type)); OnCreate(it = textThreadsByParams[tp] = new TextThread(tp, Host::GetHookParam(tp).type));
@ -42,7 +42,7 @@ namespace
void RemoveThreads(std::function<bool(ThreadParam)> removeIf) void RemoveThreads(std::function<bool(ThreadParam)> removeIf)
{ {
LOCK hostLock(hostMutex); LOCK(hostMutex);
std::vector<ThreadParam> removedThreads; std::vector<ThreadParam> removedThreads;
for (auto i : textThreadsByParams) for (auto i : textThreadsByParams)
if (removeIf(i.first)) if (removeIf(i.first))
@ -56,7 +56,7 @@ namespace
void RegisterProcess(DWORD pid, HANDLE hostPipe) void RegisterProcess(DWORD pid, HANDLE hostPipe)
{ {
LOCK hostLock(hostMutex); LOCK(hostMutex);
ProcessRecord record; ProcessRecord record;
record.hostPipe = hostPipe; record.hostPipe = hostPipe;
record.section = OpenFileMappingW(FILE_MAP_READ, FALSE, (ITH_SECTION_ + std::to_wstring(pid)).c_str()); record.section = OpenFileMappingW(FILE_MAP_READ, FALSE, (ITH_SECTION_ + std::to_wstring(pid)).c_str());
@ -69,7 +69,7 @@ namespace
void UnregisterProcess(DWORD pid) void UnregisterProcess(DWORD pid)
{ {
LOCK hostLock(hostMutex); LOCK(hostMutex);
ProcessRecord pr = processRecordsByIds[pid]; ProcessRecord pr = processRecordsByIds[pid];
if (!pr.hostPipe) return; if (!pr.hostPipe) return;
CloseHandle(pr.sectionMutex); CloseHandle(pr.sectionMutex);
@ -148,7 +148,7 @@ namespace Host
{ {
// Artikash 7/25/2018: This is only called when NextHooker is closed, at which point Windows should free everything itself...right? // Artikash 7/25/2018: This is only called when NextHooker is closed, at which point Windows should free everything itself...right?
#ifdef _DEBUG // Check memory leaks #ifdef _DEBUG // Check memory leaks
LOCK hostLock(hostMutex); LOCK(hostMutex);
OnRemove = [](TextThread* textThread) { delete textThread; }; OnRemove = [](TextThread* textThread) { delete textThread; };
for (auto i : processRecordsByIds) UnregisterProcess(i.first); for (auto i : processRecordsByIds) UnregisterProcess(i.first);
delete textThreadsByParams[CONSOLE]; delete textThreadsByParams[CONSOLE];
@ -223,7 +223,7 @@ namespace Host
HookParam GetHookParam(DWORD pid, unsigned __int64 addr) HookParam GetHookParam(DWORD pid, unsigned __int64 addr)
{ {
LOCK hostLock(hostMutex); LOCK(hostMutex);
HookParam ret = {}; HookParam ret = {};
ProcessRecord pr = processRecordsByIds[pid]; ProcessRecord pr = processRecordsByIds[pid];
if (pr.sectionMap == nullptr) return ret; if (pr.sectionMap == nullptr) return ret;
@ -241,7 +241,7 @@ namespace Host
std::wstring GetHookName(DWORD pid, unsigned __int64 addr) std::wstring GetHookName(DWORD pid, unsigned __int64 addr)
{ {
if (pid == 0) return L"Console"; if (pid == 0) return L"Console";
LOCK hostLock(hostMutex); LOCK(hostMutex);
std::string buffer = ""; std::string buffer = "";
ProcessRecord pr = processRecordsByIds[pid]; ProcessRecord pr = processRecordsByIds[pid];
if (pr.sectionMap == nullptr) return L""; if (pr.sectionMap == nullptr) return L"";
@ -260,14 +260,13 @@ namespace Host
TextThread* GetThread(ThreadParam tp) TextThread* GetThread(ThreadParam tp)
{ {
LOCK hostLock(hostMutex); LOCK(hostMutex);
return textThreadsByParams[tp]; return textThreadsByParams[tp];
} }
void AddConsoleOutput(std::wstring text) void AddConsoleOutput(std::wstring text)
{ {
LOCK hostLock(hostMutex); GetThread(CONSOLE)->AddSentence(std::wstring(text));
textThreadsByParams[CONSOLE]->AddSentence(std::wstring(text));
} }
} }

View File

@ -16,42 +16,44 @@ TextThread::~TextThread()
std::wstring TextThread::GetStore() std::wstring TextThread::GetStore()
{ {
LOCK ttLock(ttMutex); LOCK(ttMutex);
return storage; return storage;
} }
void TextThread::Flush() void TextThread::Flush()
{ {
LOCK ttLock(ttMutex);
if (buffer.size() < 400 && (timestamp - GetTickCount() < 250 || buffer.size() == 0)) return; // TODO: let user change delay before sentence is flushed
std::wstring sentence; std::wstring sentence;
if (status & USING_UNICODE)
{ {
sentence = std::wstring((wchar_t*)buffer.data(), buffer.size() / 2); LOCK(ttMutex);
if (buffer.size() < 400 && (timestamp - GetTickCount() < 250 || buffer.size() == 0)) return; // TODO: let user change delay before sentence is flushed
std::wstring sentence;
if (status & USING_UNICODE)
{
sentence = std::wstring((wchar_t*)buffer.data(), buffer.size() / 2);
}
else
{
wchar_t* converted = new wchar_t[buffer.size()];
sentence = std::wstring(converted, MultiByteToWideChar(status & USING_UTF8 ? CP_UTF8 : 932, 0, buffer.data(), buffer.size(), converted, buffer.size()));
delete[] converted;
}
memset(buffer.data(), 0, buffer.size());
buffer.clear();
} }
else
{
wchar_t* converted = new wchar_t[buffer.size()];
sentence = std::wstring(converted, MultiByteToWideChar(status & USING_UTF8 ? CP_UTF8 : 932, 0, buffer.data(), buffer.size(), converted, buffer.size()));
delete[] converted;
}
ttMutex.unlock();
AddSentence(sentence); AddSentence(sentence);
ttMutex.lock();
memset(buffer.data(), 0, buffer.size());
buffer.clear();
} }
void TextThread::AddSentence(std::wstring sentence) void TextThread::AddSentence(std::wstring sentence)
{ {
// Dispatch to extensions occurs here. Don't hold mutex! Extensions might take a while!
if (Output) sentence = Output(this, sentence); if (Output) sentence = Output(this, sentence);
LOCK ttLock(ttMutex); LOCK(ttMutex);
storage.append(sentence); storage.append(sentence);
} }
void TextThread::AddText(const BYTE *con, int len) void TextThread::AddText(const BYTE *con, int len)
{ {
LOCK ttLock(ttMutex); LOCK(ttMutex);
// Artikash 8/27/2018: add repetition filter // Artikash 8/27/2018: add repetition filter
if (len > 6 && buffer.data() && (strstr(buffer.data(), (const char*)con) || wcsstr((const wchar_t*)buffer.data(), (const wchar_t*)con))) return; if (len > 6 && buffer.data() && (strstr(buffer.data(), (const char*)con) || wcsstr((const wchar_t*)buffer.data(), (const wchar_t*)con))) return;
buffer.insert(buffer.end(), con, con + len); buffer.insert(buffer.end(), con, con + len);

View File

@ -68,4 +68,4 @@ struct HookRemovedNotif // From hook
unsigned __int64 address; unsigned __int64 address;
}; };
typedef std::lock_guard<std::recursive_mutex> LOCK; #define LOCK(mutex) std::lock_guard<std::recursive_mutex> lock(mutex)