AuroraRuntime/Source/Process/ProcessMap.cpp
2022-03-23 16:14:53 +00:00

238 lines
5.8 KiB
C++

/***
Copyright (C) 2022 J Reece Wilson (a/k/a "Reece"). All rights reserved.
File: ProcessMap.cpp
Date: 2022-1-23
Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "ProcessMap.hpp"
#if defined(AURORA_PLATFORM_WIN32)
#include "ProcessMap.Win32.hpp"
#endif
#if defined(AURORA_IS_MODERNNT_DERIVED)
#include "ProcessMap.NT.hpp"
#endif
namespace Aurora::Process
{
struct ModuleBasePairHash
{
AuUInt operator()(const ModuleBasePair &in) const
{
return in.modBase;
}
};
struct ModuleBasePairEq
{
constexpr bool operator()(const ModuleBasePair &lhs, const AuString &rhs) const
{
return lhs.module == rhs;
}
constexpr bool operator()(const ModuleBasePair &lhs, const AuUInt &rhs) const
{
return lhs.modBase == rhs;
}
constexpr bool operator()(const ModuleBasePair &lhs, const ModuleBasePair &rhs) const
{
return rhs.modBase ? lhs.modBase == rhs.modBase : (rhs.module.size() ? rhs.module == rhs.module : false);
}
};
struct ModuleLookup : Section
{
ModuleLookup()
{}
ModuleLookup(const Section &s) : Section(s)
{}
};
static AuBST<AuUInt, ModuleLookup> gModulePtrMap;
static const auto kMinPageAlignment = 4096;
static const auto kPageBufferPad = 20;
static AuThreadPrimitives::MutexUnique_t gMutexUnique;
static AuHashMap<ModuleBasePair, AuSPtr<PublicModule>, ModuleBasePairHash, ModuleBasePairEq> gModuleMap;
static AuUInt ToLowestPageAlignment(AuUInt in)
{
return in & ~(kMinPageAlignment - 1);
}
static AuSPtr<PublicModule> GetModuleFromSectionCache(AuUInt pointer)
{
AU_LOCK_GUARD(gMutexUnique);
auto itr = gModulePtrMap.find(pointer);
if (itr == gModulePtrMap.end()) return {};
return itr->second.moduleMeta.lock();
}
static AuOptional<Section> GetSectionCache(AuUInt pointer)
{
AU_LOCK_GUARD(gMutexUnique);
auto itr = gModulePtrMap.find(pointer);
if (itr == gModulePtrMap.end()) return {};
return itr->second;
}
bool IsInModuleCache(const ModuleBasePair &pair)
{
AU_LOCK_GUARD(gMutexUnique);
return gModuleMap.find(pair) != gModuleMap.end();
}
void InsertModuleCache(const ModuleBasePair &pair, const AuSPtr<PublicModule> &mod)
{
for (auto &section : mod->sections)
{
AU_LOCK_GUARD(gMutexUnique);
section.moduleMeta = mod;
for (AuUInt i = section.baseVa; i < section.baseVa + section.size; i += (kMinPageAlignment * kPageBufferPad))
{
ModuleLookup a(section);
gModulePtrMap[i] = a;
}
}
{
AU_LOCK_GUARD(gMutexUnique);
gModuleMap[pair] = mod;
}
}
void RemoveModuleCache(const ModuleBasePair &eitherOr)
{
AU_LOCK_GUARD(gMutexUnique);
auto itr = gModuleMap.find(eitherOr);
if (itr == gModuleMap.end()) return;
auto mod = itr->second;
for (const auto &section : mod->sections)
{
for (AuUInt i = section.baseVa; i < section.baseVa + section.size; i += (kMinPageAlignment * kPageBufferPad))
{
auto itr = gModulePtrMap.find(i);
if (itr != gModulePtrMap.end())
{
gModulePtrMap.erase(itr);
}
}
}
gModuleMap.erase(itr);
}
PublicModule GetFromModuleCache(AuUInt handle)
{
AU_LOCK_GUARD(gMutexUnique);
auto itr = gModuleMap.find({"", handle});
if (itr == gModuleMap.end()) return {};
return *itr->second;
}
static AuOptional<Section> FindInCache(AuUInt pointer)
{
auto curPtr = ToLowestPageAlignment(pointer);
auto temp = GetSectionCache(curPtr);
if (temp.has_value()) return temp;
for (int i = 0; i < kPageBufferPad + 1; i++)
{
curPtr -= kMinPageAlignment;
temp = GetSectionCache(curPtr); // TODO: i dont want to start from the top of the tree, thats stupid
if (temp.has_value()) return temp;
}
return AuOptional<Section>{};
}
void InitProcessMap()
{
gMutexUnique = AuThreadPrimitives::MutexUnique();
#if defined(AURORA_IS_MODERNNT_DERIVED)
InitProcessMapNt();
#endif
TryRescanSlow();
}
AUKN_SYM void TryRescanSlow()
{
#if defined(AURORA_PLATFORM_WIN32)
MakeToolHelp32Snapshot();
#endif
}
void DeinitProcessMap()
{
gMutexUnique.reset();
DeinitProcessMapNt();
}
AUKN_SYM AuOptional<Section> GetSection(AuUInt pointer)
{
try
{
auto ceg = FindInCache(pointer);
if (ceg.has_value())
{
return ceg;
}
#if defined(AURORA_IS_MODERNNT_DERIVED)
if (MakeAwarePtr(pointer))
{
return FindInCache(pointer);
}
else
{
return LookupArbitrarySection(pointer);
}
#endif
return {};
}
catch (...)
{
return {};
}
}
AUKN_SYM PublicModule DumpExecutableRoot()
{
try
{
return GetExecutableRoot();
}
catch (...)
{
return {};
}
}
AUKN_SYM Sections DumpExecutableAll()
{
try
{
Sections ret;
for (const auto &[meta, ptr] : gModuleMap)
{
ret.insert(ret.end(), ptr->sections.begin(), ptr->sections.end());
}
return ret;
}
catch (...)
{
return {};
}
}
}