AuroraRuntime/Source/Process/AuProcessMap.NT.cpp
Jamie Reece Wilson ca2c0462ab [*] Continue to refactor allocations of synchronization primitive away
[*] NT: Fix missing CoTaskMemFree leak on startup
[*] Fix ConsoleStd restart bug
2023-06-28 10:33:12 +01:00

319 lines
8.4 KiB
C++

/***
Copyright (C) 2022 J Reece Wilson (a/k/a "Reece"). All rights reserved.
File: UtilProcessMap.NT.cpp
Date: 2022-1-24
Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "AuProcessMap.NT.hpp"
#include "AuProcessMap.hpp"
namespace Aurora::Process
{
static AuThreadPrimitives::Mutex gMutex;
static AuBST<AuUInt, AuString> gPathCache;
static AuBST<AuUInt, AuString> gModNameCache;
AuString ModuleToSomething(HMODULE handle, bool path)
{
try
{
std::wstring file;
if (!handle)
{
return {};
}
if (gMutex)
{
AU_LOCK_GUARD(gMutex);
if (path)
{
auto itr = gPathCache.find(reinterpret_cast<AuUInt>(handle));
if (itr != gPathCache.end()) return itr->second;
}
else
{
auto itr = gModNameCache.find(reinterpret_cast<AuUInt>(handle));
if (itr != gModNameCache.end()) return itr->second;
}
}
if (!AuTryResize(file, 16 * 1024))
{
return {};
}
auto length = GetModuleFileNameW(handle, &file[0], DWORD(file.size()));
if (!length)
{
return {};
}
file.resize(length);
if (!path)
{
auto idx = file.find_last_of('\\');
if (idx != std::wstring::npos)
{
file = file.substr(idx + 1);
}
}
auto ret = Locale::ConvertFromWChar(file.c_str(), file.length());
if (gMutex)
{
AU_LOCK_GUARD(gMutex);
if (path)
{
AuTryInsert(gPathCache, AuMakePair(reinterpret_cast<AuUInt>(handle), ret));
}
else
{
AuTryInsert(gModNameCache, AuMakePair(reinterpret_cast<AuUInt>(handle), ret));
}
}
return ret;
}
catch (...)
{
return "";
}
}
AuString ModuleToName(HMODULE handle)
{
return ModuleToSomething(handle, false);
}
AuString ModuleToPath(HMODULE handle)
{
return ModuleToSomething(handle, true);
}
static AuUInt GetModuleBaseAddressFromPathGrossYesThisIsCached(HMODULE mod, const AuString &path)
{
if (path.empty())
{
return 0;
}
auto dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(mod);
auto nt = &(reinterpret_cast<PIMAGE_NT_HEADERS>(reinterpret_cast<AuUInt>(mod) + dosHeader->e_lfanew)->OptionalHeader.ImageBase);
auto offset = static_cast<AuUInt32>(reinterpret_cast<const AuUInt8 *>(nt) - reinterpret_cast<const AuUInt8 *>(mod));
using Value_t = AuRemovePointer_t<decltype(nt)>;
Value_t value {};
auto handle = CreateFileW(Locale::ConvertFromUTF8(path).c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, NULL, NULL);
if (handle == INVALID_HANDLE_VALUE)
{
return 0;
}
DWORD bytesRead {};
if (!SetFilePointer(handle, offset, NULL, 0))
{
return {};
}
if (!ReadFile(handle, &value, sizeof(Value_t), &bytesRead, NULL))
{
return {};
}
if (bytesRead != sizeof(Value_t))
{
return {};
}
CloseHandle(handle);
return value;
}
static void FetchModuleSections(Sections &sections, AuUInt offset, HMODULE mod)
{
auto dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(mod);
auto nt = reinterpret_cast<PIMAGE_NT_HEADERS>(reinterpret_cast<AuUInt>(mod) + dosHeader->e_lfanew);
auto pSection = IMAGE_FIRST_SECTION(nt);
auto pSectionCur = pSection;
auto imageBase = nt->OptionalHeader.ImageBase;
for (auto i = 0; i < nt->FileHeader.NumberOfSections; i++)
{
auto cur = pSectionCur++;
Section seg;
seg.origVa = cur->VirtualAddress + offset;
seg.baseVa = cur->VirtualAddress + imageBase;
seg.size = cur->SizeOfRawData;
seg.fsOff = cur->PointerToRawData;
seg.pt.readable = cur->Characteristics & (IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_EXECUTE);
seg.pt.NX = !(cur->Characteristics & IMAGE_SCN_MEM_EXECUTE);
seg.pt.writable = cur->Characteristics & IMAGE_SCN_MEM_WRITE;
seg.name = AuString(cur->Name, cur->Name + strnlen(reinterpret_cast<const char *>(cur->Name), sizeof(cur->Name)));
AuTryInsert(sections, seg);
}
}
static AuSPtr<PublicModule> HandleToPublicModule(HMODULE h)
{
auto pub = AuMakeShared<PublicModule>();
if (!pub) return {};
pub->moduleMeta = AuMakeShared<ModuleMeta>();
if (!pub->moduleMeta) return {};
pub->moduleMeta->moduleBase = AuUInt(h);
pub->moduleMeta->moduleName = ModuleToName(h);
pub->moduleMeta->modulePath = ModuleToPath(h);
pub->moduleMeta->origVa = GetModuleBaseAddressFromPathGrossYesThisIsCached(h, pub->moduleMeta->modulePath);
FetchModuleSections(pub->sections, pub->moduleMeta->origVa, h);
return pub;
}
void InvaildateModule(HMODULE hmod)
{
if (!gMutex)
{
return;
}
try
{
// acquire gmutex
{
AU_LOCK_GUARD(gMutex);
auto itr1 = gPathCache.find(reinterpret_cast<AuUInt>(hmod));
if (itr1 != gPathCache.end()) gPathCache.erase(itr1);
auto itr2 = gModNameCache.find(reinterpret_cast<AuUInt>(hmod));
if (itr2 != gModNameCache.end()) gModNameCache.erase(itr2);
}
// release lock
{
RemoveModuleCache({"", reinterpret_cast<AuUInt>(hmod)});
}
}
catch (...)
{
}
}
bool MakeAwarePtr(AuUInt pointer)
{
HMODULE handle;
if (!GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, reinterpret_cast<LPCWSTR>(pointer), &handle))
{
return false;
}
MakeAware(handle);
return true;
}
AuOptional<Section> LookupArbitrarySection(AuUInt address)
{
MEMORY_BASIC_INFORMATION info;
Section section;
if (!VirtualQuery((LPCVOID)address, &info, sizeof(info)))
{
return {};
}
section.baseVa = reinterpret_cast<AuUInt>(info.BaseAddress);
section.size = info.RegionSize;
section.pt.NX = (info.Protect & (PAGE_EXECUTE | PAGE_EXECUTE_READ)) == 0;
section.pt.writable = (info.Protect & (PAGE_READWRITE | PAGE_WRITECOPY | PAGE_EXECUTE_WRITECOPY)) != 0;
section.pt.readable = (!section.pt.NX) && ((info.Protect & (PAGE_NOACCESS)) != 0);
section.pt.acSanity = info.Protect == PAGE_EXECUTE_WRITECOPY;
if (info.Type & (MEM_MAPPED | MEM_IMAGE))
{
section.name = kSectionNameFile;
}
else if (info.Type & (MEM_PRIVATE))
{
section.name = kSectionNameHeap;
}
return section;
}
void MakeAware(HMODULE hmod)
{
ModuleBasePair handle {"", reinterpret_cast<AuUInt>(hmod)};
if (!gMutex)
{
return;
}
if (hmod == INVALID_HANDLE_VALUE)
{
SysPushErrorArg();
return;
}
if (!hmod)
{
SysPushErrorArg();
return;
}
try
{
if (IsInModuleCache(handle))
{
return;
}
auto ptr = HandleToPublicModule(hmod);
if (!ptr)
{
return;
}
InsertModuleCache(handle, ptr);
}
catch (...)
{
}
}
PublicModule GetExecutableRoot()
{
try
{
auto handle = GetModuleHandleW(NULL);
MakeAware(handle);
return GetFromModuleCache(reinterpret_cast<AuUInt>(handle));
}
catch (...)
{
return {};
}
}
void InitProcessMapNt()
{
}
void DeinitProcessMapNt()
{
}
}