NTSTATUS LeGlobalData::InitRegistryRedirection(PREGISTRY_REDIRECTION_ENTRY64 Entry64, ULONG_PTR Count, PVOID BaseAddress)
{
NTSTATUS Status;
PLEPEB LePeb;
PREGISTRY_REDIRECTION_ENTRY Entry;
if (Count == 0)
return STATUS_NO_MORE_ENTRIES;
LePeb = this->GetLePeb();
#pragma push_macro("USTR64ToUSTR")
#undef USTR64ToUSTR
#define USTR64ToUSTR(ustr64) UNICODE_STRING({ ustr64.Length, ustr64.MaximumLength, PtrAdd(ustr64.Buffer, BaseAddress) });
REGISTRY_REDIRECTION_ENTRY LocalEntry;
FOR_EACH(Entry64, Entry64, Count)
{
ULONG_PTR LastIndex;
HANDLE OriginalKey, RedirectedKey;
UNICODE_STRING KeyFullPath;
OriginalKey = nullptr;
RedirectedKey = nullptr;
Status = Reg::OpenKey(&OriginalKey, (HANDLE)Entry64->Original.Root, KEY_QUERY_VALUE, PtrAdd(Entry64->Original.SubKey.Buffer, BaseAddress));
FAIL_CONTINUE(Status);
if (Entry64->Redirected.Root != NULL)
{
Status = Reg::OpenKey(&RedirectedKey, (HANDLE)Entry64->Redirected.Root, KEY_QUERY_VALUE, PtrAdd(Entry64->Redirected.SubKey.Buffer, BaseAddress));
if (NT_FAILED(Status))
{
Reg::CloseKeyHandle(OriginalKey);
continue;
}
}
this->RegistryRedirectionEntry.Add(LocalEntry);
LastIndex = this->RegistryRedirectionEntry.GetSize() - 1;
Entry = &this->RegistryRedirectionEntry[LastIndex];
Status = QueryRegKeyFullPath(OriginalKey, &KeyFullPath);
Reg::CloseKeyHandle(OriginalKey);
if (NT_FAILED(Status))
{
Reg::CloseKeyHandle(RedirectedKey);
this->RegistryRedirectionEntry.Remove(LastIndex);
continue;
}
Entry->Original.FullPath = KeyFullPath;
RtlFreeUnicodeString(&KeyFullPath);
if (RedirectedKey != nullptr)
{
Status = QueryRegKeyFullPath(RedirectedKey, &KeyFullPath);
Reg::CloseKeyHandle(RedirectedKey);
if (NT_FAILED(Status))
{
this->RegistryRedirectionEntry.Remove(LastIndex);
continue;
}
Entry->Redirected.FullPath = KeyFullPath;
RtlFreeUnicodeString(&KeyFullPath);
}
Entry->Original.Root = (HKEY)Entry64->Original.Root;
Entry->Original.SubKey = USTR64ToUSTR(Entry64->Original.SubKey);
Entry->Original.ValueName = USTR64ToUSTR(Entry64->Original.ValueName);
Entry->Original.DataType = Entry64->Original.DataType;
Entry->Original.Data = nullptr;
Entry->Original.DataSize = 0;
Entry->Redirected.Root = (HKEY)Entry64->Redirected.Root;
Entry->Redirected.SubKey = USTR64ToUSTR(Entry64->Redirected.SubKey);
Entry->Redirected.ValueName = USTR64ToUSTR(Entry64->Redirected.ValueName);
Entry->Redirected.DataType = Entry64->Redirected.DataType;
Entry->Redirected.Data = nullptr;
Entry->Redirected.DataSize = 0;
if (Entry64->Redirected.Data != nullptr && Entry64->Redirected.DataSize != 0)
{
Entry->Redirected.DataSize = (ULONG_PTR)Entry64->Redirected.DataSize;
Entry->Redirected.Data = AllocateMemoryP(Entry->Redirected.DataSize);
if (Entry->Redirected.Data == nullptr)
{
this->RegistryRedirectionEntry.Remove(LastIndex);
continue;
}
CopyMemory(Entry->Redirected.Data, PtrAdd(Entry64->Redirected.Data, BaseAddress), Entry->Redirected.DataSize);
}
}
bool_t ExtWriteFile(handle_t h_file,
const_void_ptr buf,
ulen_t buf_len,
ulen_t *ret_len)
{
if( buf_len<=SplitLen )
{
ushortlen_t ret_short_len;
bool_t ret=WriteFile(h_file,buf,(ushortlen_t)buf_len,&ret_short_len,0);
if( ret_len ) *ret_len=ret_short_len;
return ret;
}
else
{
ulen_t sum_len=0;
while( buf_len>2*SplitLen )
{
ushortlen_t ret_short_len;
if( WriteFile(h_file,buf,SplitLen,&ret_short_len,0) )
{
if( ret_short_len==SplitLen )
{
sum_len+=SplitLen;
buf=PtrAdd(buf,SplitLen);
buf_len-=SplitLen;
}
else
{
if( ret_len ) *ret_len=sum_len+ret_short_len;
return true;
}
}
else
{
if( ret_len ) *ret_len=0;
return false;
}
}
ushortlen_t ret_short_len;
if( WriteFile(h_file,buf,(ushortlen_t)buf_len,&ret_short_len,0) )
{
if( ret_len ) *ret_len=sum_len+ret_short_len;
return true;
}
else
{
if( ret_len ) *ret_len=0;
return false;
}
}
}
void CalculateProcessorTopology(CPUNumaNodes& out_nodes, uint32_t& out_numThreadsPerProcGroup)
{
out_nodes.clear();
out_numThreadsPerProcGroup = 0;
#if defined(_WIN32)
std::vector<KAFFINITY> threadMaskPerProcGroup;
static std::mutex m;
std::lock_guard<std::mutex> l(m);
DWORD bufSize = 0;
BOOL ret = GetLogicalProcessorInformationEx(RelationProcessorCore, nullptr, &bufSize);
SWR_ASSERT(ret == FALSE && GetLastError() == ERROR_INSUFFICIENT_BUFFER);
PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX pBufferMem = (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)malloc(bufSize);
SWR_ASSERT(pBufferMem);
ret = GetLogicalProcessorInformationEx(RelationProcessorCore, pBufferMem, &bufSize);
SWR_ASSERT(ret != FALSE, "Failed to get Processor Topology Information");
uint32_t count = bufSize / pBufferMem->Size;
PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX pBuffer = pBufferMem;
for (uint32_t i = 0; i < count; ++i)
{
SWR_ASSERT(pBuffer->Relationship == RelationProcessorCore);
for (uint32_t g = 0; g < pBuffer->Processor.GroupCount; ++g)
{
auto& gmask = pBuffer->Processor.GroupMask[g];
uint32_t threadId = 0;
uint32_t procGroup = gmask.Group;
Core* pCore = nullptr;
uint32_t numThreads = (uint32_t)_mm_popcount_sizeT(gmask.Mask);
while (BitScanForwardSizeT((unsigned long*)&threadId, gmask.Mask))
{
// clear mask
KAFFINITY threadMask = KAFFINITY(1) << threadId;
gmask.Mask &= ~threadMask;
if (procGroup >= threadMaskPerProcGroup.size())
{
threadMaskPerProcGroup.resize(procGroup + 1);
}
if (threadMaskPerProcGroup[procGroup] & threadMask)
{
// Already seen this mask. This means that we are in 32-bit mode and
// have seen more than 32 HW threads for this procGroup
// Don't use it
#if defined(_WIN64)
SWR_ASSERT(false, "Shouldn't get here in 64-bit mode");
#endif
continue;
}
threadMaskPerProcGroup[procGroup] |= (KAFFINITY(1) << threadId);
// Find Numa Node
uint32_t numaId = 0;
PROCESSOR_NUMBER procNum = {};
procNum.Group = WORD(procGroup);
procNum.Number = UCHAR(threadId);
ret = GetNumaProcessorNodeEx(&procNum, (PUSHORT)&numaId);
SWR_ASSERT(ret);
// Store data
if (out_nodes.size() <= numaId) out_nodes.resize(numaId + 1);
auto& numaNode = out_nodes[numaId];
uint32_t coreId = 0;
if (nullptr == pCore)
{
numaNode.cores.push_back(Core());
pCore = &numaNode.cores.back();
pCore->procGroup = procGroup;
}
pCore->threadIds.push_back(threadId);
if (procGroup == 0)
{
out_numThreadsPerProcGroup++;
}
}
}
pBuffer = PtrAdd(pBuffer, pBuffer->Size);
}
free(pBufferMem);
#elif defined(__linux__) || defined (__gnu_linux__)
// Parse /proc/cpuinfo to get full topology
std::ifstream input("/proc/cpuinfo");
std::string line;
char* c;
uint32_t threadId = uint32_t(-1);
uint32_t coreId = uint32_t(-1);
uint32_t numaId = uint32_t(-1);
while (std::getline(input, line))
{
if (line.find("processor") != std::string::npos)
{
if (threadId != uint32_t(-1))
{
// Save information.
if (out_nodes.size() <= numaId) out_nodes.resize(numaId + 1);
auto& numaNode = out_nodes[numaId];
if (numaNode.cores.size() <= coreId) numaNode.cores.resize(coreId + 1);
auto& core = numaNode.cores[coreId];
core.procGroup = coreId;
core.threadIds.push_back(threadId);
out_numThreadsPerProcGroup++;
}
auto data_start = line.find(": ") + 2;
threadId = std::strtoul(&line.c_str()[data_start], &c, 10);
continue;
}
if (line.find("core id") != std::string::npos)
{
auto data_start = line.find(": ") + 2;
coreId = std::strtoul(&line.c_str()[data_start], &c, 10);
continue;
}
if (line.find("physical id") != std::string::npos)
{
auto data_start = line.find(": ") + 2;
numaId = std::strtoul(&line.c_str()[data_start], &c, 10);
continue;
}
}
if (threadId != uint32_t(-1))
{
// Save information.
if (out_nodes.size() <= numaId) out_nodes.resize(numaId + 1);
auto& numaNode = out_nodes[numaId];
if (numaNode.cores.size() <= coreId) numaNode.cores.resize(coreId + 1);
auto& core = numaNode.cores[coreId];
core.procGroup = coreId;
core.threadIds.push_back(threadId);
out_numThreadsPerProcGroup++;
}
for (uint32_t node = 0; node < out_nodes.size(); node++) {
auto& numaNode = out_nodes[node];
auto it = numaNode.cores.begin();
for ( ; it != numaNode.cores.end(); ) {
if (it->threadIds.size() == 0)
numaNode.cores.erase(it);
else
++it;
}
}
#else
#error Unsupported platform
#endif
}
NTSTATUS LeGlobalData::Initialize()
{
BOOL IsLoader;
PLEPEB LePeb;
PLDR_MODULE Ntdll;
PPEB_BASE Peb;
NTSTATUS Status;
NLSTABLEINFO NlsTableInfo;
UNICODE_STRING SystemDirectory, NlsFileName, OemNlsFileName, LangFileName;
PKEY_VALUE_PARTIAL_INFORMATION IndexValue;
IsLoader = IsLeLoader();
Wow64 = Ps::IsWow64Process();
Ntdll = GetNtdllLdrModule();
LOOP_ONCE
{
LePeb = OpenOrCreateLePeb();
if (LePeb == nullptr)
{
ULONG_PTR DefaultACPLength, DefaultLCIDLength, DefaultOEMCPLength;
WCHAR DefaultACP[0x20], DefaultOEMCP[0x20], DefaultLCID[0x20];
PVOID ReloadedNtdll;
PUNICODE_STRING FullDllName;
LePeb = GetLePeb();
InitDefaultLeb(&LePeb->Leb);
FullDllName = &FindLdrModuleByHandle(&__ImageBase)->FullDllName;
CopyMemory(LePeb->LeDllFullPath, FullDllName->Buffer, FullDllName->Length + sizeof(WCHAR));
Status = LoadPeImage(Ntdll->FullDllName.Buffer, &ReloadedNtdll, nullptr, LOAD_PE_IGNORE_RELOC);
if (NT_SUCCESS(Status))
{
PVOID LdrLoadDllAddress;
LdrLoadDllAddress = LookupExportTable(ReloadedNtdll, NTDLL_LdrLoadDll);
LePeb->LdrLoadDllAddress = PtrAdd(LdrLoadDllAddress, PtrOffset(Ntdll->DllBase, ReloadedNtdll));
CopyMemory(LePeb->LdrLoadDllBackup, LdrLoadDllAddress, LDR_LOAD_DLL_BACKUP_SIZE);
LePeb->LdrLoadDllBackupSize = LDR_LOAD_DLL_BACKUP_SIZE;
UnloadPeImage(ReloadedNtdll);
}
DefaultACPLength = (swprintf(DefaultACP, L"%d", LePeb->Leb.AnsiCodePage) + 1) * sizeof(WCHAR);
DefaultOEMCPLength = (swprintf(DefaultOEMCP, L"%d", LePeb->Leb.OemCodePage) + 1) * sizeof(WCHAR);
DefaultLCIDLength = (swprintf(DefaultLCID, L"%d", LePeb->Leb.LocaleID) + 1) * sizeof(WCHAR);
REGISTRY_REDIRECTION_ENTRY64 *Entry, Entries[] =
{
{
{ (ULONG64)HKEY_LOCAL_MACHINE, USTR64(REGPATH_CODEPAGE), USTR64(REGKEY_ACP), REG_SZ, },
{ (ULONG64)HKEY_LOCAL_MACHINE, USTR64(REGPATH_CODEPAGE), USTR64(REGKEY_ACP), REG_SZ, DefaultACP, DefaultACPLength },
},
{
{ (ULONG64)HKEY_LOCAL_MACHINE, USTR64(REGPATH_CODEPAGE), USTR64(REGKEY_OEMCP), REG_SZ, },
{ (ULONG64)HKEY_LOCAL_MACHINE, USTR64(REGPATH_CODEPAGE), USTR64(REGKEY_OEMCP), REG_SZ, DefaultOEMCP, DefaultOEMCPLength },
},
{
{ (ULONG64)HKEY_LOCAL_MACHINE, USTR64(REGPATH_LANGUAGE), USTR64(REGKEY_DEFAULT_LANGUAGE), REG_SZ, },
{ (ULONG64)HKEY_LOCAL_MACHINE, USTR64(REGPATH_LANGUAGE), USTR64(REGKEY_DEFAULT_LANGUAGE), REG_SZ, DefaultLCID, DefaultLCIDLength },
},
};
Status = this->InitRegistryRedirection(Entries, countof(Entries), nullptr);
}
else
{
*GetLePeb() = *LePeb;
Status = this->InitRegistryRedirection(LePeb->Leb.RegistryReplacement, LePeb->Leb.NumberOfRegistryRedirectionEntries, &LePeb->Leb);
NtClose(LePeb->Section);
CloseLePeb(LePeb);
}
if (IsLoader)
break;
Status = this->TextMetricCache.Initialize();
FAIL_RETURN(Status);
PVOID NlsBaseAddress;
LCID DefaultLocaleID;
LARGE_INTEGER DefaultCasingTableSize;
Status = NtInitializeNlsFiles(&NlsBaseAddress, &DefaultLocaleID, &DefaultCasingTableSize);
FAIL_RETURN(Status);
this->GetLePeb()->OriginalLocaleID = DefaultLocaleID;
NtUnmapViewOfSection(CurrentProcess, NlsBaseAddress);
WriteLog(L"init leb %s", GetLePeb()->LeDllFullPath);
SystemDirectory = Ntdll->FullDllName;
SystemDirectory.Length -= Ntdll->BaseDllName.Length;
Status = RtlDuplicateUnicodeString(RTL_DUPSTR_ADD_NULL, &SystemDirectory, &this->SystemDirectory);
FAIL_RETURN(Status);
RtlInitEmptyString(&NlsFileName, nullptr, 0);
RtlInitEmptyString(&OemNlsFileName, nullptr, 0);
RtlInitEmptyString(&LangFileName, nullptr, 0);
SCOPE_EXIT
{
RtlFreeUnicodeString(&NlsFileName);
RtlFreeUnicodeString(&OemNlsFileName);
RtlFreeUnicodeString(&LangFileName);
}
SCOPE_EXIT_END;
Status = GetNlsFile(&NlsFileName, GetLeb()->AnsiCodePage, REGPATH_CODEPAGE);
FAIL_RETURN(Status);
Status = GetNlsFile(&OemNlsFileName, GetLeb()->OemCodePage, REGPATH_CODEPAGE);
FAIL_RETURN(Status);
Status = GetLangFile(&LangFileName, GetLeb()->LocaleID, REGPATH_LANGUAGE);
FAIL_RETURN(Status);
NtFileMemory AnsiFile, OemFile, LangFile;
Status = ReadFileInSystemDirectory(AnsiFile, &NlsFileName);
FAIL_RETURN(Status);
Status = ReadFileInSystemDirectory(OemFile, &OemNlsFileName);
FAIL_RETURN(Status);
Status = ReadFileInSystemDirectory(LangFile, &LangFileName);
FAIL_RETURN(Status);
AnsiCodePageOffset = 0;
OemCodePageOffset = ROUND_UP(AnsiFile.GetSize32(), 16);
UnicodeCaseTableOffset = OemCodePageOffset + ROUND_UP(OemFile.GetSize32(), 16);
Status = AllocVirtualMemory(&CodePageMapView, UnicodeCaseTableOffset + LangFile.GetSize32(), PAGE_READWRITE, MEM_COMMIT | MEM_TOP_DOWN);
FAIL_RETURN(Status);
CopyMemory(PtrAdd(CodePageMapView, AnsiCodePageOffset), AnsiFile.GetBuffer(), AnsiFile.GetSize32());
CopyMemory(PtrAdd(CodePageMapView, OemCodePageOffset), OemFile.GetBuffer(), OemFile.GetSize32());
CopyMemory(PtrAdd(CodePageMapView, UnicodeCaseTableOffset), LangFile.GetBuffer(), LangFile.GetSize32());
ProtectVirtualMemory(CodePageMapView, UnicodeCaseTableOffset + LangFile.GetSize32(), PAGE_READONLY);
RtlInitNlsTables(
(PUSHORT)PtrAdd(CodePageMapView, AnsiCodePageOffset),
(PUSHORT)PtrAdd(CodePageMapView, OemCodePageOffset),
(PUSHORT)PtrAdd(CodePageMapView, UnicodeCaseTableOffset),
&NlsTableInfo
);
RtlResetRtlTranslations(&NlsTableInfo);
WriteLog(L"reset nls");
Peb = CurrentPeb();
Peb->AnsiCodePageData = (PUSHORT)PtrAdd(CodePageMapView, AnsiCodePageOffset);
Peb->OemCodePageData = (PUSHORT)PtrAdd(CodePageMapView, OemCodePageOffset);
Peb->UnicodeCaseTableData = (PUSHORT)PtrAdd(CodePageMapView, UnicodeCaseTableOffset);
// LdrInitShimEngineDynamic(&__ImageBase);
LdrRegisterDllNotification(0,
[] (ULONG NotificationReason, PCLDR_DLL_NOTIFICATION_DATA NotificationData, PVOID Context)
{
return ((PLeGlobalData)Context)->DllNotification(NotificationReason, NotificationData);
},
this,
&DllNotificationCookie
);
}
Status = InstallHookPort();
WriteLog(L"inst hp: %08X", Status);
FAIL_RETURN(Status);
HookNtdllRoutines(Ntdll->DllBase);
WriteLog(L"hook ntdll");
if (IsLoader)
return Status;
PLDR_MODULE Kernel32Ldr;
Kernel32Ldr = GetKernel32Ldr();
if (Kernel32Ldr != nullptr)
{
Kernel32Ldr->EntryPoint = DelayInitDllEntry;
// HookKernel32Routines(Kernel32Ldr->DllBase);
}
WriteLog(L"init %p", Status);
return Status;
}
EXPORT_C const TAny* TBtreeKey::Key(const TAny* anEntry) const
{
return PtrAdd(anEntry,iKeyOffset);
}
static Identifier DeserializeEntity(
InputStreamFormatter<utf8>& formatter,
IEntityInterface& interf,
DocumentId docId)
{
using Blob = InputStreamFormatter<utf8>::Blob;
using Section = InputStreamFormatter<utf8>::InteriorSection;
utf8 tempBuffer[256];
auto beginLoc = formatter.GetLocation();
Section objType = { nullptr, nullptr };
if (!formatter.TryBeginElement(objType))
Throw(FormatException("Error in begin element in entity file", formatter.GetLocation()));
XlCopyNString(tempBuffer, objType._start, objType._end - objType._start);
auto typeId = interf.GetTypeId((const char*)tempBuffer);
std::vector<PropertyInitializer> inits;
std::vector<char> initsBuffer;
initsBuffer.reserve(256);
std::vector<Identifier> children;
for (;;) {
switch (formatter.PeekNext()) {
case Blob::BeginElement:
{
auto child = DeserializeEntity(formatter, interf, docId);
if (child.Object())
children.push_back(child);
}
break;
case Blob::AttributeName:
{
Section name, value;
if (!formatter.TryAttribute(name, value))
Throw(FormatException("Error in begin element in entity file", formatter.GetLocation()));
// parse the value and add it as a property initializer
char intermediateBuffer[64];
auto type = ImpliedTyping::Parse(
(const char*)value._start, (const char*)value._end,
intermediateBuffer, dimof(intermediateBuffer));
size_t bufferOffset = initsBuffer.size();
if (type._type == ImpliedTyping::TypeCat::Void) {
type._type = ImpliedTyping::TypeCat::UInt8;
type._arrayCount = uint16(value._end - value._start);
type._typeHint = ImpliedTyping::TypeHint::String;
initsBuffer.insert(initsBuffer.end(), value._start, value._end);
} else {
auto size = std::min(type.GetSize(), (unsigned)sizeof(intermediateBuffer));
initsBuffer.insert(initsBuffer.end(), intermediateBuffer, PtrAdd(intermediateBuffer, size));
}
XlCopyNString(tempBuffer, name._start, name._end - name._start);
auto id = interf.GetPropertyId(typeId, (const char*)tempBuffer);
PropertyInitializer i;
i._prop = id;
i._elementType = unsigned(type._type);
i._arrayCount = type._arrayCount;
i._src = (const void*)bufferOffset;
inits.push_back(i);
}
break;
case Blob::EndElement:
default:
if (!formatter.TryEndElement())
Throw(FormatException("Expecting end element in entity deserialisation", formatter.GetLocation()));
if (typeId != ~ObjectTypeId(0x0)) {
for (auto&i:inits) i._src = PtrAdd(AsPointer(initsBuffer.cbegin()), size_t(i._src));
auto id = interf.AssignObjectId(docId, typeId);
Identifier identifier(docId, id, typeId);
if (!interf.CreateObject(identifier, AsPointer(inits.cbegin()), inits.size()))
Throw(FormatException("Error while creating object in entity deserialisation", beginLoc));
for (const auto&c:children)
interf.SetParent(c, identifier, -1);
typeId = ~ObjectTypeId(0x0);
initsBuffer.clear();
return identifier;
}
return Identifier();
}
}
}
NTSTATUS LeGlobalData::Initialize()
{
PLEPEB LePeb;
PLDR_MODULE Ntdll;
PPEB_BASE Peb;
NTSTATUS Status;
NLSTABLEINFO NlsTableInfo;
UNICODE_STRING SystemDirectory, NlsFileName, OemNlsFileName, LangFileName;
PKEY_VALUE_PARTIAL_INFORMATION IndexValue;
Wow64 = Ps::IsWow64Process();
Ntdll = GetNtdllLdrModule();
LePeb = OpenOrCreateLePeb();
if (LePeb == nullptr)
{
PVOID ReloadedNtdll;
PUNICODE_STRING FullDllName;
LePeb = GetLePeb();
InitDefaultLeb(&LePeb->Leb);
FullDllName = &FindLdrModuleByHandle(&__ImageBase)->FullDllName;
CopyMemory(LePeb->LeDllFullPath, FullDllName->Buffer, FullDllName->Length + sizeof(WCHAR));
Status = LoadPeImage(Ntdll->FullDllName.Buffer, &ReloadedNtdll, nullptr, LOAD_PE_IGNORE_RELOC);
if (NT_SUCCESS(Status))
{
PVOID LdrLoadDllAddress;
LdrLoadDllAddress = EATLookupRoutineByHashPNoFix(ReloadedNtdll, NTDLL_LdrLoadDll);
LePeb->LdrLoadDllAddress = PtrAdd(LdrLoadDllAddress, PtrOffset(Ntdll->DllBase, ReloadedNtdll));
CopyMemory(LePeb->LdrLoadDllBackup, LdrLoadDllAddress, LDR_LOAD_DLL_BACKUP_SIZE);
LePeb->LdrLoadDllBackupSize = LDR_LOAD_DLL_BACKUP_SIZE;
UnloadPeImage(ReloadedNtdll);
}
}
else
{
*GetLePeb() = *LePeb;
ZwClose(LePeb->Section);
CloseLePeb(LePeb);
}
WriteLog(L"init leb %s", GetLePeb()->LeDllFullPath);
SystemDirectory = Ntdll->FullDllName;
SystemDirectory.Length -= Ntdll->BaseDllName.Length;
Status = RtlDuplicateUnicodeString(RTL_DUPSTR_ADD_NULL, &SystemDirectory, &this->SystemDirectory);
FAIL_RETURN(Status);
RtlInitEmptyUnicodeString(&NlsFileName, nullptr, 0);
RtlInitEmptyUnicodeString(&OemNlsFileName, nullptr, 0);
RtlInitEmptyUnicodeString(&LangFileName, nullptr, 0);
SCOPE_EXIT
{
RtlFreeUnicodeString(&NlsFileName);
RtlFreeUnicodeString(&OemNlsFileName);
RtlFreeUnicodeString(&LangFileName);
}
SCOPE_EXIT_END;
Status = GetNlsFile(&NlsFileName, GetLeb()->AnsiCodePage, REGPATH_CODEPAGE);
FAIL_RETURN(Status);
Status = GetNlsFile(&OemNlsFileName, GetLeb()->OemCodePage, REGPATH_CODEPAGE);
FAIL_RETURN(Status);
Status = GetLangFile(&LangFileName, GetLeb()->LocaleID, REGPATH_LANGUAGE);
FAIL_RETURN(Status);
NtFileMemory AnsiFile, OemFile, LangFile;
Status = ReadFileInSystemDirectory(AnsiFile, &NlsFileName);
FAIL_RETURN(Status);
Status = ReadFileInSystemDirectory(OemFile, &OemNlsFileName);
FAIL_RETURN(Status);
Status = ReadFileInSystemDirectory(LangFile, &LangFileName);
FAIL_RETURN(Status);
AnsiCodePageOffset = 0;
OemCodePageOffset = ROUND_UP(AnsiFile.GetSize32(), 16);
UnicodeCaseTableOffset = OemCodePageOffset + ROUND_UP(OemFile.GetSize32(), 16);
Status = AllocVirtualMemory(&CodePageMapView, UnicodeCaseTableOffset + LangFile.GetSize32(), PAGE_READWRITE, MEM_COMMIT | MEM_TOP_DOWN);
FAIL_RETURN(Status);
CopyMemory(PtrAdd(CodePageMapView, AnsiCodePageOffset), AnsiFile.GetBuffer(), AnsiFile.GetSize32());
CopyMemory(PtrAdd(CodePageMapView, OemCodePageOffset), OemFile.GetBuffer(), OemFile.GetSize32());
CopyMemory(PtrAdd(CodePageMapView, UnicodeCaseTableOffset), LangFile.GetBuffer(), LangFile.GetSize32());
ProtectVirtualMemory(CodePageMapView, UnicodeCaseTableOffset + LangFile.GetSize32(), PAGE_READONLY);
RtlInitNlsTables(
(PUSHORT)PtrAdd(CodePageMapView, AnsiCodePageOffset),
(PUSHORT)PtrAdd(CodePageMapView, OemCodePageOffset),
(PUSHORT)PtrAdd(CodePageMapView, UnicodeCaseTableOffset),
&NlsTableInfo
);
RtlResetRtlTranslations(&NlsTableInfo);
WriteLog(L"reset nls");
Peb = CurrentPeb();
Peb->AnsiCodePageData = (PUSHORT)PtrAdd(CodePageMapView, AnsiCodePageOffset);
Peb->OemCodePageData = (PUSHORT)PtrAdd(CodePageMapView, OemCodePageOffset);
Peb->UnicodeCaseTableData = (PUSHORT)PtrAdd(CodePageMapView, UnicodeCaseTableOffset);
// LdrInitShimEngineDynamic(&__ImageBase);
LdrRegisterDllNotification(0,
[] (ULONG NotificationReason, PCLDR_DLL_NOTIFICATION_DATA NotificationData, PVOID Context)
{
return ((PLeGlobalData)Context)->DllNotification(NotificationReason, NotificationData);
},
this,
&DllNotificationCookie
);
Status = InstallHookPort();
FAIL_RETURN(Status);
WriteLog(L"inst hp");
HookNtdllRoutines(Ntdll->DllBase);
WriteLog(L"hook ntdll");
PLDR_MODULE Kernel32Ldr;
Kernel32Ldr = GetKernel32Ldr();
if (Kernel32Ldr != nullptr)
{
Kernel32Ldr->EntryPoint = DelayInitDllEntry;
// HookKernel32Routines(Kernel32Ldr->DllBase);
}
WriteLog(L"init %p", Status);
return Status;
}
NTSTATUS GetRedirectFile(PUNICODE_STRING Redirected, PUNICODE_STRING Original)
{
ULONG_PTR Length;
PWSTR Buffer;
UNICODE_STRING FileName;
typedef struct
{
UNICODE_STRING SubPath;
ULONG_PTR SuffixLength;
PCWSTR NewSubPath;
} REDIRECT_ENTRY, *PDB_REDIRECT;
PDB_REDIRECT Entry;
static REDIRECT_ENTRY RedirectEntries[] =
{
{ RTL_CONSTANT_STRING(L"\\All Users\\QQ\\History.db"), CONST_STRLEN(L"History.db") * sizeof(WCHAR), GlobalHistoryDb },
{ RTL_CONSTANT_STRING(L"\\All Users\\QQ\\Registry.db"), CONST_STRLEN(L"Registry.db") * sizeof(WCHAR), GlobalRegistryDb },
//{ RTL_CONSTANT_STRING(L"QQProtect.exe"), CONST_STRLEN(L"QQProtect.exe") * sizeof(WCHAR), NULL },
};
RtlInitEmptyString(Redirected);
LOOP_ONCE
{
if (Original == nullptr)
continue;
FOR_EACH_ARRAY(Entry, RedirectEntries)
{
if (Original->Length <= Entry->SubPath.Length)
continue;
FileName = *Original;
FileName.Buffer = PtrSub(PtrAdd(FileName.Buffer, FileName.Length), Entry->SubPath.Length);
FileName.Length = Entry->SubPath.Length;
if (!RtlEqualUnicodeString(&FileName, &Entry->SubPath, TRUE))
continue;
/*
if (Entry->NewSubPath == NULL)
{
ExceptionBox(L"qqprotect");
++Entry;
}
*/
break;
}
if (Entry == &RedirectEntries[countof(RedirectEntries)])
break;
Length = Original->Length + Entry->SubPath.Length + Entry->SuffixLength + sizeof(Entry->NewSubPath);
Buffer = (PWSTR)AllocStack(Length);
FileName.MaximumLength = Length;
FileName.Buffer = Buffer;
Length = StrLengthW(Entry->NewSubPath) * sizeof(WCHAR);
RtlCopyUnicodeString(&FileName, Original);
CopyMemory(PtrSub(PtrAdd(FileName.Buffer, FileName.Length), Entry->SuffixLength), Entry->NewSubPath, Length + sizeof(WCHAR));
FileName.Length = FileName.Length - Entry->SuffixLength + Length;
RtlDuplicateUnicodeString(RTL_DUPSTR_ADD_NULL, &FileName, Redirected);
return STATUS_SUCCESS;
}
return STATUS_NOT_FOUND;
}
void CalculateProcessorTopology(CPUNumaNodes& out_nodes)
{
out_nodes.clear();
#if defined(_WIN32)
SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX buffer[KNOB_MAX_NUM_THREADS];
DWORD bufSize = sizeof(buffer);
BOOL ret = GetLogicalProcessorInformationEx(RelationProcessorCore, buffer, &bufSize);
SWR_ASSERT(ret != FALSE, "Failed to get Processor Topology Information");
uint32_t count = bufSize / buffer->Size;
PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX pBuffer = buffer;
for (uint32_t i = 0; i < count; ++i)
{
SWR_ASSERT(pBuffer->Relationship == RelationProcessorCore);
for (uint32_t g = 0; g < pBuffer->Processor.GroupCount; ++g)
{
auto& gmask = pBuffer->Processor.GroupMask[g];
uint32_t threadId = 0;
uint32_t procGroup = gmask.Group;
Core* pCore = nullptr;
uint32_t numThreads = (uint32_t)_mm_popcount_sizeT(gmask.Mask);
while (BitScanForwardSizeT((unsigned long*)&threadId, gmask.Mask))
{
// clear mask
gmask.Mask &= ~(KAFFINITY(1) << threadId);
// Find Numa Node
PROCESSOR_NUMBER procNum = {};
procNum.Group = WORD(procGroup);
procNum.Number = UCHAR(threadId);
uint32_t numaId = 0;
ret = GetNumaProcessorNodeEx(&procNum, (PUSHORT)&numaId);
SWR_ASSERT(ret);
// Store data
if (out_nodes.size() <= numaId) out_nodes.resize(numaId + 1);
auto& numaNode = out_nodes[numaId];
uint32_t coreId = 0;
if (nullptr == pCore)
{
numaNode.cores.push_back(Core());
pCore = &numaNode.cores.back();
pCore->procGroup = procGroup;
#if !defined(_WIN64)
coreId = (uint32_t)numaNode.cores.size();
if ((coreId * numThreads) >= 32)
{
// Windows doesn't return threadIds >= 32 for a processor group correctly
// when running a 32-bit application.
// Just save -1 as the threadId
threadId = uint32_t(-1);
}
#endif
}
pCore->threadIds.push_back(threadId);
}
}
pBuffer = PtrAdd(pBuffer, pBuffer->Size);
}
#elif defined(__linux__) || defined (__gnu_linux__)
// Parse /proc/cpuinfo to get full topology
std::ifstream input("/proc/cpuinfo");
std::string line;
char* c;
uint32_t threadId = uint32_t(-1);
uint32_t coreId = uint32_t(-1);
uint32_t numaId = uint32_t(-1);
while (std::getline(input, line))
{
if (line.find("processor") != std::string::npos)
{
if (threadId != uint32_t(-1))
{
// Save information.
if (out_nodes.size() <= numaId) out_nodes.resize(numaId + 1);
auto& numaNode = out_nodes[numaId];
if (numaNode.cores.size() <= coreId) numaNode.cores.resize(coreId + 1);
auto& core = numaNode.cores[coreId];
core.procGroup = coreId;
core.threadIds.push_back(threadId);
}
auto data_start = line.find(": ") + 2;
threadId = std::strtoul(&line.c_str()[data_start], &c, 10);
continue;
}
if (line.find("core id") != std::string::npos)
{
auto data_start = line.find(": ") + 2;
coreId = std::strtoul(&line.c_str()[data_start], &c, 10);
continue;
}
if (line.find("physical id") != std::string::npos)
{
auto data_start = line.find(": ") + 2;
numaId = std::strtoul(&line.c_str()[data_start], &c, 10);
continue;
}
}
if (threadId != uint32_t(-1))
{
// Save information.
if (out_nodes.size() <= numaId) out_nodes.resize(numaId + 1);
auto& numaNode = out_nodes[numaId];
if (numaNode.cores.size() <= coreId) numaNode.cores.resize(coreId + 1);
auto& core = numaNode.cores[coreId];
core.procGroup = coreId;
core.threadIds.push_back(threadId);
}
for (uint32_t node = 0; node < out_nodes.size(); node++) {
auto& numaNode = out_nodes[node];
auto it = numaNode.cores.begin();
for ( ; it != numaNode.cores.end(); ) {
if (it->threadIds.size() == 0)
numaNode.cores.erase(it);
else
++it;
}
}
#else
#error Unsupported platform
#endif
}
