Replaced PsSetLoadImageNotifyRoutine with a PatchGuard safe version.

5 files changed, 77 insertions, 34 deletions

diff --git a/IntegrationTest/IntegrationTest.vcxproj b/IntegrationTest/IntegrationTest.vcxproj
index 228c0cd..b107b09 100644
--- a/IntegrationTest/IntegrationTest.vcxproj
+++ b/IntegrationTest/IntegrationTest.vcxproj
@@ -74,6 +74,7 @@
   <PropertyGroup Label="UserMacros" />
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">
     <LinkIncremental>true</LinkIncremental>
+    <TargetName>$(ProjectName)-kmem</TargetName>
   </PropertyGroup>
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">
     <LinkIncremental>true</LinkIncremental>
diff --git a/KMemDriver/Imports.h b/KMemDriver/Imports.h
index ed70956..8a33dbb 100644
--- a/KMemDriver/Imports.h
+++ b/KMemDriver/Imports.h
@@ -138,4 +138,19 @@ NTAPI
 RtlAvlRemoveNode(
 	IN PRTL_AVL_TREE pTree,
 	IN PMMADDRESS_NODE pNode
+);
+
+__kernel_entry
+NTSTATUS
+ZwQuerySystemInformation(
+	IN int SystemInformationClass,
+	OUT PVOID SystemInformation,
+	IN ULONG SystemInformationLength,
+	OUT OPTIONAL PULONG ReturnLength
+);
+
+NTKERNELAPI
+PVOID
+PsGetProcessSectionBaseAddress(
+	IN PEPROCESS Process
 );
\ No newline at end of file
diff --git a/KMemDriver/KMemDriver.c b/KMemDriver/KMemDriver.c
index 23916ef..741f932 100644
--- a/KMemDriver/KMemDriver.c
+++ b/KMemDriver/KMemDriver.c
@@ -188,24 +188,6 @@ NTSTATUS DriverEntry(
 	return status;
 }
 
-void OnImageLoad(
-	PUNICODE_STRING FullImageName,
-	HANDLE ProcessId,
-	PIMAGE_INFO ImageInfo
-)
-{
-	UNREFERENCED_PARAMETER(ImageInfo);
-#if 0
-	KDBG("ProcessID: 0x%X\n", ProcessId);
-	KDBG("FullImage: %wZ\n", FullImageName);
-#endif
-	if (wcsstr(FullImageName->Buffer, CHEAT_EXE)) {
-		ctrlPID = ProcessId;
-		imageBase = ImageInfo->ImageBase;
-		KDBG("Found Target !!!\n");
-	}
-}
-
 NTSTATUS WaitForControlProcess(OUT PEPROCESS *ppEProcess)
 {
 	NTSTATUS status;
@@ -216,22 +198,34 @@ NTSTATUS WaitForControlProcess(OUT PEPROCESS *ppEProcess)
 	imageBase = NULL;
 	ctrlPID = NULL;
 
-	status = PsSetLoadImageNotifyRoutine(OnImageLoad);
-	if (!NT_SUCCESS(status)) {
-		KDBG("PsSetLoadImageNotifyRoutine failed with 0x%X\n", status);
-		return status;
-	}
+	SYSTEM_PROCESS_INFORMATION * procs = MmAllocateNonCachedMemory(1024 * sizeof(*procs));
+	ULONG mem_needed = 0;
 
-	while (!ctrlPID) {
-		LARGE_INTEGER wait = { .QuadPart = -1000000 };
-		KeDelayExecutionThread(KernelMode, TRUE, &wait);
+	if (procs == NULL) {
+		return STATUS_MEMORY_NOT_ALLOCATED;
 	}
+	while (ctrlPID == NULL) {
+		status = ZwQuerySystemInformation(0x05, (PVOID)&procs[0], 1024 * sizeof(*procs), &mem_needed);
+		if (!NT_SUCCESS(status)) {
+			KDBG("NtQuerySystemInformation failed with 0x%X\n", status);
+			return status;
+		}
 
-	status = PsRemoveLoadImageNotifyRoutine(OnImageLoad);
-	if (!NT_SUCCESS(status)) {
-		KDBG("PsRemoveLoadImageNotifyRoutine failed with 0x%X\n", status);
-		return status;
+		SYSTEM_PROCESS_INFORMATION * cur_proc = procs;
+		while (cur_proc->NextEntryOffset > 0) {
+			cur_proc = (SYSTEM_PROCESS_INFORMATION *)((PUCHAR)cur_proc + cur_proc->NextEntryOffset);
+
+			if (wcsstr(cur_proc->ImageName.Buffer, CHEAT_EXE)) {
+				KDBG("FOUND %wZ with PID 0x%X\n", cur_proc->ImageName, cur_proc->UniqueProcessId);
+				ctrlPID = cur_proc->UniqueProcessId;
+				break;
+			}
+
+			LARGE_INTEGER wait = { .QuadPart = -100000 };
+			KeDelayExecutionThread(KernelMode, FALSE, &wait);
+		}
 	}
+	MmFreeNonCachedMemory(procs, 1024 * sizeof(*procs));
 
 	status = PsLookupProcessByProcessId(ctrlPID, ppEProcess);
 	if (!NT_SUCCESS(status)) {
@@ -239,8 +233,14 @@ NTSTATUS WaitForControlProcess(OUT PEPROCESS *ppEProcess)
 		return status;
 	}
 
-	KDBG("Got Ctrl Process PID: 0x%X (%d)\n",
-		ctrlPID, ctrlPID);
+	imageBase = PsGetProcessSectionBaseAddress(*ppEProcess);
+	if (imageBase == NULL)
+	{
+		KDBG("ImageBase is NULL\n");
+	}
+
+	KDBG("Got Ctrl Process PID/ImageBase: 0x%X (%d) / %p\n",
+		ctrlPID, ctrlPID, imageBase);
 
 	return STATUS_SUCCESS;
 }
@@ -762,7 +762,7 @@ NTSTATUS GetDriverObject(
 	}
 
 	return status;
-		}
+}
 
 PHANDLE_TABLE_ENTRY ExpLookupHandleTableEntry(PVOID pHandleTable, HANDLE handle)
 {
diff --git a/KMemDriver/Native.h b/KMemDriver/Native.h
index 6d7d684..6344eed 100644
--- a/KMemDriver/Native.h
+++ b/KMemDriver/Native.h
@@ -376,5 +376,32 @@ typedef struct _HANDLE_TABLE
 	PVOID DebugInfo;
 } PHANDLE_TABLE;
 
+typedef struct _SYSTEM_PROCESS_INFORMATION {
+	ULONG NextEntryOffset;
+	ULONG NumberOfThreads;
+	UINT8 Reserved1[48];
+	UNICODE_STRING ImageName;
+	KPRIORITY BasePriority;
+	ULONG Reserved2;
+	HANDLE UniqueProcessId;
+	PVOID Reserved3;
+	ULONG HandleCount;
+	ULONG SessionId;
+	PVOID Reserved4;
+	SIZE_T PeakVirtualSize;
+	SIZE_T VirtualSize;
+	ULONG Reserved5;
+	SIZE_T PeakWorkingSetSize;
+	SIZE_T WorkingSetSize;
+	PVOID Reserved6;
+	SIZE_T QuotaPagedPoolUsage;
+	PVOID Reserved7;
+	SIZE_T QuotaNonPagedPoolUsage;
+	SIZE_T PagefileUsage;
+	SIZE_T PeakPagefileUsage;
+	SIZE_T PrivatePageCount;
+	LARGE_INTEGER Reserved8[6];
+} SYSTEM_PROCESS_INFORMATION;
+
 #pragma pack(pop)
 #pragma warning(default : 4214 4201)
\ No newline at end of file
diff --git a/KTest-Only.bat b/KTest-Only.bat
index 9dadcbe..e3c1959 100644
--- a/KTest-Only.bat
+++ b/KTest-Only.bat
@@ -1,5 +1,5 @@
 @echo off
 
 REM fsutil usn deleteJournal /D C:
-%~dp0\x64\Release\KTest.exe
+%~dp0\x64\Release\IntegrationTest-kmem.exe
 timeout /t 3
\ No newline at end of file