alloc non cached memory (dynamic mem) and find VAD ptr works

2 files changed, 65 insertions, 1 deletions

diff --git a/KMemDriver/Driver.c b/KMemDriver/Driver.c
index 60c4c42..f2b0d03 100644
--- a/KMemDriver/Driver.c
+++ b/KMemDriver/Driver.c
@@ -105,6 +105,11 @@ NTSTATUS VADFind(
 	IN ULONG_PTR address,
 	OUT PMMVAD_SHORT* pResult
 );
+NTSTATUS VADProtect(
+	IN PEPROCESS pProcess,
+	IN ULONG_PTR address,
+	IN ULONG prot
+);
 
 #pragma alloc_text(PAGE, WaitForControlProcess)
 #pragma alloc_text(PAGE, VerifyControlProcess)
@@ -122,6 +127,7 @@ NTSTATUS VADFind(
 #pragma alloc_text(PAGE, KRThread)
 #pragma alloc_text(PAGE, VADFindNodeOrParent)
 #pragma alloc_text(PAGE, VADFind)
+#pragma alloc_text(PAGE, VADProtect)
 
 static void fn_zero_text(PVOID fn_start);
 static HANDLE ctrlPID;
@@ -337,6 +343,12 @@ NTSTATUS KRThread(IN PVOID pArg)
 
 		KeLowerIrql(0);
 		KDBG("Init ..\n");
+		{
+			ULONG_PTR low, high;
+			IoGetStackLimits(&low, &high);
+			KDBG("Stack limits (high/low/total/remaining): 0x%p/0x%p/0x%X/0x%X\n",
+				low, high, high - low, IoGetRemainingStackSize());
+		}
 
 		if (mmapedBase && !hijackedDriver &&
 			NT_SUCCESS(GetDriverObject(&hijackedDriver, L"\\Driver\\ahcache")))
@@ -629,8 +641,29 @@ NTSTATUS UpdatePPEPIfRequired(
 			}
 			else {
 				PEPROCESS pep = *lastPEP;
+				PVOID base = NULL;
+				SIZE_T size = ADDRESS_AND_SIZE_TO_SPAN_PAGES(base, 4096);
+				PKAPC_STATE apc = MmAllocateNonCachedMemory(sizeof(*apc));
+				KeStackAttachProcess((PRKPROCESS)pep, apc);
+				status = ZwAllocateVirtualMemory(ZwCurrentProcess(), &base, 0, &size, MEM_COMMIT, PAGE_READWRITE);
+				if (!NT_SUCCESS(status)) {
+					KDBG("ZwAllocateVirtualMemory failed with 0x%X\n", status);
+				}
+				else {
+					*(UINT64 *)base = 0x4141414142424242;
+				}
+				KeUnstackDetachProcess(apc);
+				KDBG("VAD Test Alloc.: 0x%p (status: 0x%X)\n", base, status);
 				PMMVAD_SHORT mmvad;
-				KDBG("VAD Test: 0x%p\n", VADFind(pep, 0x5086800, &mmvad));
+				status = VADFind(pep, (ULONG_PTR)base, &mmvad);
+				KDBG("VAD Test.......: 0x%p (status: 0x%X)\n", mmvad->StartingVpn, status);
+				KeStackAttachProcess((PRKPROCESS)pep, apc);
+				if (*(UINT64 *)base != 0x4141414142424242) {
+					KDBG("VAD Test failed: 0x%p != 0x%p\n", 0x4141414142424242, base);
+				}
+				ZwFreeVirtualMemory(ZwCurrentProcess(), &base, &size, MEM_RELEASE);
+				KeUnstackDetachProcess(apc);
+				MmFreeNonCachedMemory(apc, sizeof(*apc));
 #if 0
 				PMM_AVL_TABLE avltable = (PMM_AVL_TABLE)((ULONG_PTR *)pep + VAD_TREE_1803);
 				KDBG("VAD-ROOT.....: 0x%p\n", GET_VAD_ROOT(avltable));
@@ -979,4 +1012,19 @@ NTSTATUS VADFind(
 	}
 
 	return status;
+}
+
+NTSTATUS VADProtect(
+	IN PEPROCESS pProcess,
+	IN ULONG_PTR address, IN ULONG prot
+)
+{
+	NTSTATUS status = STATUS_SUCCESS;
+	PMMVAD_SHORT pVadShort = NULL;
+
+	status = VADFind(pProcess, address, &pVadShort);
+	if (NT_SUCCESS(status))
+		pVadShort->u.VadFlags.Protection = prot;
+
+	return status;
 }
\ No newline at end of file
diff --git a/KMemDriver/Imports.h b/KMemDriver/Imports.h
index 44f53cc..e0eaa4d 100644
--- a/KMemDriver/Imports.h
+++ b/KMemDriver/Imports.h
@@ -109,4 +109,20 @@ ObReferenceObjectByName(
 	KPROCESSOR_MODE Access,
 	PVOID ParseContext,
 	PVOID* ObjectPtr
+);
+
+NTSTATUS ZwAllocateVirtualMemory(
+	_In_    HANDLE    ProcessHandle,
+	_Inout_ PVOID     *BaseAddress,
+	_In_    ULONG_PTR ZeroBits,
+	_Inout_ PSIZE_T   RegionSize,
+	_In_    ULONG     AllocationType,
+	_In_    ULONG     Protect
+);
+
+NTSTATUS ZwFreeVirtualMemory(
+	_In_    HANDLE  ProcessHandle,
+	_Inout_ PVOID   *BaseAddress,
+	_Inout_ PSIZE_T RegionSize,
+	_In_    ULONG   FreeType
 );
\ No newline at end of file