minor improvments

4 files changed, 95 insertions, 65 deletions

diff --git a/CheatEngineServer/CommandDispatcher.cpp b/CheatEngineServer/CommandDispatcher.cpp
index aac833a..096e90a 100644
--- a/CheatEngineServer/CommandDispatcher.cpp
+++ b/CheatEngineServer/CommandDispatcher.cpp
@@ -243,15 +243,23 @@ CommandReturn DispatchCommand(CEConnection& con, char command)
 				cret = CommandReturn::CR_FAIL_ALLOC;
 				break;
 			}
-			if (KInterface::getInstance().MtRPM((HANDLE)((ULONG_PTR)params.handle), (PVOID)params.address, (BYTE*)out + sizeof(*out), params.size, &krr) != true) {
-				free(out);
-				cret = CommandReturn::CR_FAIL_KMEM;
-				break;
+
+			if (params.address == NULL) {
+				std::wcout << "Got a RPM to NULL, ignore." << std::endl;
+				out->read = 0;
 			}
-			if (params.size != krr.SizeReq || params.size != krr.SizeRes || krr.StatusRes != 0) {
-				free(out);
-				cret = CommandReturn::CR_FAIL_OTHER;
-				break;
+			else {
+				if (KInterface::getInstance().MtRPM((HANDLE)((ULONG_PTR)params.handle), (PVOID)params.address, (BYTE*)&out[1], params.size, &krr) != true) {
+					free(out);
+					cret = CommandReturn::CR_FAIL_KMEM;
+					break;
+				}
+				if (params.size != krr.SizeReq || params.size != krr.SizeRes || krr.StatusRes != 0) {
+					free(out);
+					cret = CommandReturn::CR_FAIL_OTHER;
+					break;
+				}
+				out->read = (int)krr.SizeRes;
 			}
 			if (sendall(con.getSocket(), out, sizeof(*out) + params.size, 0) > 0)
 			{
@@ -334,6 +342,7 @@ CommandReturn DispatchCommand(CEConnection& con, char command)
 			else {
 				//std::wcout << "Modules NEXT for PID 0x" << std::hex << toolhelpsnapshot << std::endl;
 			}
+
 			if (con.m_cachedModules.size() > 0) {
 				MODULE_DATA md = con.m_cachedModules[0];
 				int imageNameLen = (int)strnlen(md.BaseDllName, sizeof(md.BaseDllName));
@@ -348,6 +357,7 @@ CommandReturn DispatchCommand(CEConnection& con, char command)
 				pcme->modulesize = md.SizeOfImage;
 				pcme->modulenamesize = imageNameLen;
 				pcme->result = 1;
+
 				memcpy(((BYTE*)pcme) + sizeof(*pcme), md.BaseDllName, imageNameLen);
 				if (sendall(con.getSocket(), pcme, sizeof(*pcme) + imageNameLen, 0) > 0)
 				{
@@ -430,6 +440,10 @@ CommandReturn DispatchCommand(CEConnection& con, char command)
 			}
 			cret = CommandReturn::CR_OK;
 			for (auto& page : con.m_cachedPages) {
+				if (KInterface::PageIsFreed(page) == true || KInterface::PageIsPrivateReserved(page) == true)
+				{
+					continue;
+				}
 				RegionInfo out;
 				out.baseaddress = (UINT64)page.BaseAddress;
 				out.protection = page.Protect;
@@ -447,8 +461,10 @@ CommandReturn DispatchCommand(CEConnection& con, char command)
 		break;
 	}
 
-	case CMD_VIRTUALQUERYEX:
-	case CMD_GETREGIONINFO: {
+	case CMD_GETREGIONINFO:
+		break;
+
+	case CMD_VIRTUALQUERYEX: {
 		CeVirtualQueryExInput params;
 		if (recvall(con.getSocket(), &params, sizeof(params), MSG_WAITALL) > 0) {
 			con.m_cachedPages.clear();
@@ -458,22 +474,27 @@ CommandReturn DispatchCommand(CEConnection& con, char command)
 				cret = CommandReturn::CR_FAIL_KMEM;
 				break;
 			}
+			SIZE_T i = 0;
+			for (auto& page : con.m_cachedPages) {
+				if (KInterface::PageIsFreed(page) == false && KInterface::PageIsPrivateReserved(page) == false)
+				{
+					break;
+				}
+				i++;
+			}
+			if (i == con.m_cachedPages.size()) {
+				cret = CommandReturn::CR_FAIL_KMEM;
+				break;
+			}
+			std::wcout << "---" << con.m_cachedPages[i].BaseAddress << std::endl;
 			CeVirtualQueryExOutput out;
-			out.baseaddress = (UINT64)con.m_cachedPages[0].BaseAddress;
-			out.protection = con.m_cachedPages[0].Protect;
-			out.size = con.m_cachedPages[0].RegionSize;
-			out.type = con.m_cachedPages[0].Type;
-			out.result = 1;
+			out.baseaddress = (UINT64)con.m_cachedPages[i].BaseAddress;
+			out.protection = con.m_cachedPages[i].Protect;
+			out.size = con.m_cachedPages[i].RegionSize;
+			out.type = con.m_cachedPages[i].Type;
+			out.result = sizeof(MEMORY_BASIC_INFORMATION);
 			if (sendall(con.getSocket(), &out, sizeof(out), 0) > 0) {
-				if (cmd == CMD_GETREGIONINFO) {
-					uint8_t size = 0;
-					if (sendall(con.getSocket(), &size, sizeof(size), 0) > 0) {
-						cret = CommandReturn::CR_OK;
-					}
-					else {
-						cret = CommandReturn::CR_FAIL_NETWORK;
-					}
-				}
+				cret = CommandReturn::CR_OK;
 			}
 			else {
 				cret = CommandReturn::CR_FAIL_NETWORK;
diff --git a/KMemDriver/Imports.h b/KMemDriver/Imports.h
index 8a33dbb..48bc882 100644
--- a/KMemDriver/Imports.h
+++ b/KMemDriver/Imports.h
@@ -25,7 +25,15 @@ NTSTATUS
 NTAPI
 PsLookupProcessByProcessId(
 	_In_ HANDLE ProcessId,
-	_Outptr_ PEPROCESS *Process
+	_Outptr_ PEPROCESS* Process
+);
+
+NTKERNELAPI
+NTSTATUS
+NTAPI
+PsLookupThreadByThreadId(
+	HANDLE   ThreadId,
+	PETHREAD* Thread
 );
 
 typedef struct _KAPC_STATE
@@ -35,7 +43,7 @@ typedef struct _KAPC_STATE
 	UCHAR KernelApcInProgress;
 	UCHAR KernelApcPending;
 	UCHAR UserApcPending;
-} KAPC_STATE, *PKAPC_STATE, *PRKAPC_STATE;
+} KAPC_STATE, * PKAPC_STATE, * PRKAPC_STATE;
 
 NTKERNELAPI
 VOID
@@ -113,7 +121,7 @@ ObReferenceObjectByName(
 
 NTSTATUS ZwAllocateVirtualMemory(
 	_In_    HANDLE    ProcessHandle,
-	_Inout_ PVOID     *BaseAddress,
+	_Inout_ PVOID* BaseAddress,
 	_In_    ULONG_PTR ZeroBits,
 	_Inout_ PSIZE_T   RegionSize,
 	_In_    ULONG     AllocationType,
@@ -122,7 +130,7 @@ NTSTATUS ZwAllocateVirtualMemory(
 
 NTSTATUS ZwFreeVirtualMemory(
 	_In_    HANDLE  ProcessHandle,
-	_Inout_ PVOID   *BaseAddress,
+	_Inout_ PVOID* BaseAddress,
 	_Inout_ PSIZE_T RegionSize,
 	_In_    ULONG   FreeType
 );
diff --git a/KMemDriver/KMemDriver.c b/KMemDriver/KMemDriver.c
index edcf834..7639d8f 100644
--- a/KMemDriver/KMemDriver.c
+++ b/KMemDriver/KMemDriver.c
@@ -233,16 +233,17 @@ NTSTATUS WaitForControlProcess(OUT PEPROCESS* ppEProcess)
 	imageBase = NULL;
 	ctrlPID = NULL;
 
-	SYSTEM_PROCESS_INFORMATION* procs = MmAllocateNonCachedMemory((1024 + 128) * sizeof(*procs));
+	ULONG const max_procs = 1024 + 256;
+	SYSTEM_PROCESS_INFORMATION* procs = MmAllocateNonCachedMemory(max_procs * sizeof(*procs));
 	ULONG mem_needed = 0;
 
 	if (procs == NULL) {
 		return STATUS_MEMORY_NOT_ALLOCATED;
 	}
 	while (ctrlPID == NULL) {
-		status = ZwQuerySystemInformation(SystemProcessInformation, (PVOID)&procs[0], (1024 + 128) * sizeof(*procs), &mem_needed);
+		status = ZwQuerySystemInformation(SystemProcessInformation, (PVOID)&procs[0], max_procs * sizeof(*procs), &mem_needed);
 		if (!NT_SUCCESS(status)) {
-			KDBG("ZwQuerySystemInformation(%zu,%lu) failed with 0x%X\n", 1024 * sizeof(*procs), mem_needed, status);
+			KDBG("ZwQuerySystemInformation(%zu,%lu) failed with 0x%X\n", max_procs * sizeof(*procs), mem_needed, status);
 			return status;
 		}
 
diff --git a/KMemDriver/Native.h b/KMemDriver/Native.h
index 6344eed..ed2fe1e 100644
--- a/KMemDriver/Native.h
+++ b/KMemDriver/Native.h
@@ -16,7 +16,7 @@ typedef struct _PEB_LDR_DATA
 	LIST_ENTRY InLoadOrderModuleList;
 	LIST_ENTRY InMemoryOrderModuleList;
 	LIST_ENTRY InInitializationOrderModuleList;
-} PEB_LDR_DATA, *PPEB_LDR_DATA;
+} PEB_LDR_DATA, * PPEB_LDR_DATA;
 
 typedef struct _LDR_DATA_TABLE_ENTRY
 {
@@ -33,7 +33,7 @@ typedef struct _LDR_DATA_TABLE_ENTRY
 	USHORT TlsIndex;
 	LIST_ENTRY HashLinks;
 	ULONG TimeDateStamp;
-} LDR_DATA_TABLE_ENTRY, *PLDR_DATA_TABLE_ENTRY;
+} LDR_DATA_TABLE_ENTRY, * PLDR_DATA_TABLE_ENTRY;
 
 typedef struct _PEB
 {
@@ -55,7 +55,7 @@ typedef struct _PEB
 	ULONG SystemReserved;
 	ULONG AtlThunkSListPtr32;
 	PVOID ApiSetMap;
-} PEB, *PPEB;
+} PEB, * PPEB;
 
 typedef struct _PEB_LDR_DATA32
 {
@@ -65,7 +65,7 @@ typedef struct _PEB_LDR_DATA32
 	LIST_ENTRY32 InLoadOrderModuleList;
 	LIST_ENTRY32 InMemoryOrderModuleList;
 	LIST_ENTRY32 InInitializationOrderModuleList;
-} PEB_LDR_DATA32, *PPEB_LDR_DATA32;
+} PEB_LDR_DATA32, * PPEB_LDR_DATA32;
 
 typedef struct _LDR_DATA_TABLE_ENTRY32
 {
@@ -82,7 +82,7 @@ typedef struct _LDR_DATA_TABLE_ENTRY32
 	USHORT TlsIndex;
 	LIST_ENTRY32 HashLinks;
 	ULONG TimeDateStamp;
-} LDR_DATA_TABLE_ENTRY32, *PLDR_DATA_TABLE_ENTRY32;
+} LDR_DATA_TABLE_ENTRY32, * PLDR_DATA_TABLE_ENTRY32;
 
 typedef struct _PEB32
 {
@@ -104,7 +104,7 @@ typedef struct _PEB32
 	ULONG SystemReserved;
 	ULONG AtlThunkSListPtr32;
 	ULONG ApiSetMap;
-} PEB32, *PPEB32;
+} PEB32, * PPEB32;
 
 typedef struct _MEMORY_BASIC_INFORMATION {
 	PVOID  BaseAddress;
@@ -114,7 +114,7 @@ typedef struct _MEMORY_BASIC_INFORMATION {
 	ULONG  State;
 	ULONG  Protect;
 	ULONG  Type;
-} MEMORY_BASIC_INFORMATION, *PMEMORY_BASIC_INFORMATION;
+} MEMORY_BASIC_INFORMATION, * PMEMORY_BASIC_INFORMATION;
 
 typedef struct _KLDR_DATA_TABLE_ENTRY {
 	LIST_ENTRY InLoadOrderLinks;
@@ -134,15 +134,15 @@ typedef struct _KLDR_DATA_TABLE_ENTRY {
 	ULONG CheckSum;
 	PVOID LoadedImports;
 	PVOID PatchInformation;
-} KLDR_DATA_TABLE_ENTRY, *PKLDR_DATA_TABLE_ENTRY;
+} KLDR_DATA_TABLE_ENTRY, * PKLDR_DATA_TABLE_ENTRY;
 
 #pragma warning(disable : 4214 4201)
 #pragma pack(push, 1)
 
 typedef struct _MM_AVL_NODE // Size=24
 {
-	struct _MM_AVL_NODE * LeftChild; // Size=8 Offset=0
-	struct _MM_AVL_NODE * RightChild; // Size=8 Offset=8
+	struct _MM_AVL_NODE* LeftChild; // Size=8 Offset=0
+	struct _MM_AVL_NODE* RightChild; // Size=8 Offset=8
 
 	union // Size=8
 	{
@@ -154,9 +154,9 @@ typedef struct _MM_AVL_NODE // Size=24
 		{
 			__int64  Balance : 2; // Size=8 Offset=0 BitOffset=0 BitCount=2
 		};
-		struct _MM_AVL_NODE * Parent; // Size=8 Offset=0
+		struct _MM_AVL_NODE* Parent; // Size=8 Offset=0
 	};
-} MM_AVL_NODE, *PMM_AVL_NODE, *PMMADDRESS_NODE;
+} MM_AVL_NODE, * PMM_AVL_NODE, * PMMADDRESS_NODE;
 
 union _EX_PUSH_LOCK // Size=8
 {
@@ -169,7 +169,7 @@ union _EX_PUSH_LOCK // Size=8
 		unsigned __int64 Shared : 60; // Size=8 Offset=0 BitOffset=4 BitCount=60
 	};
 	unsigned __int64 Value; // Size=8 Offset=0
-	void * Ptr; // Size=8 Offset=0
+	void* Ptr; // Size=8 Offset=0
 };
 
 struct _MMVAD_FLAGS // Size=4
@@ -209,7 +209,7 @@ typedef struct _MMVAD_SHORT // Size=64
 	union
 	{
 		struct _RTL_BALANCED_NODE VadNode; // Size=24 Offset=0
-		struct _MMVAD_SHORT * NextVad; // Size=8 Offset=0
+		struct _MMVAD_SHORT* NextVad; // Size=8 Offset=0
 	};
 	unsigned long StartingVpn; // Size=4 Offset=24
 	unsigned long EndingVpn; // Size=4 Offset=28
@@ -221,8 +221,8 @@ typedef struct _MMVAD_SHORT // Size=64
 	union _EX_PUSH_LOCK PushLock; // Size=8 Offset=40
 	union MMVAD_SHORT_u1 u; // Size=4 Offset=48
 	union MMVAD_SHORT_u2 u1; // Size=4 Offset=52
-	struct _MI_VAD_EVENT_BLOCK * EventList; // Size=8 Offset=56
-} MMVAD_SHORT, *PMMVAD_SHORT;
+	struct _MI_VAD_EVENT_BLOCK* EventList; // Size=8 Offset=56
+} MMVAD_SHORT, * PMMVAD_SHORT;
 
 struct _MMVAD_FLAGS2 // Size=4
 {
@@ -250,22 +250,22 @@ union ___unnamed2047 // Size=4
 union ___unnamed2048 // Size=8
 {
 	struct _MI_VAD_SEQUENTIAL_INFO SequentialVa; // Size=8 Offset=0
-	struct _MMEXTEND_INFO * ExtendedInfo; // Size=8 Offset=0
+	struct _MMEXTEND_INFO* ExtendedInfo; // Size=8 Offset=0
 };
 
 typedef union _EX_FAST_REF // Size=8
 {
-	void * Object;
+	void* Object;
 	struct
 	{
 		unsigned __int64 RefCnt : 4;
 	};
 	unsigned __int64 Value;
-} EX_FAST_REF, *PEX_FAST_REF;
+} EX_FAST_REF, * PEX_FAST_REF;
 
 typedef struct _CONTROL_AREA // Size=120
 {
-	struct _SEGMENT * Segment;
+	struct _SEGMENT* Segment;
 	struct _LIST_ENTRY ListHead;
 	unsigned __int64 NumberOfSectionReferences;
 	unsigned __int64 NumberOfPfnReferences;
@@ -275,27 +275,27 @@ typedef struct _CONTROL_AREA // Size=120
 	unsigned long f2;
 	EX_FAST_REF FilePointer;
 	// Other fields
-} CONTROL_AREA, *PCONTROL_AREA;
+} CONTROL_AREA, * PCONTROL_AREA;
 
 typedef struct _SUBSECTION // Size=56
 {
 	PCONTROL_AREA ControlArea;
 	// Other fields
-} SUBSECTION, *PSUBSECTION;
+} SUBSECTION, * PSUBSECTION;
 
 typedef struct _MMVAD // Size=128
 {
 	struct _MMVAD_SHORT Core; // Size=64 Offset=0
 	union ___unnamed2047 u2; // Size=4 Offset=64
 	unsigned long pad0;  // Size=4 Offset=68
-	struct _SUBSECTION * Subsection; // Size=8 Offset=72
-	struct _MMPTE * FirstPrototypePte; // Size=8 Offset=80
-	struct _MMPTE * LastContiguousPte; // Size=8 Offset=88
+	struct _SUBSECTION* Subsection; // Size=8 Offset=72
+	struct _MMPTE* FirstPrototypePte; // Size=8 Offset=80
+	struct _MMPTE* LastContiguousPte; // Size=8 Offset=88
 	struct _LIST_ENTRY ViewLinks; // Size=16 Offset=96
-	struct _EPROCESS * VadsProcess; // Size=8 Offset=112
+	struct _EPROCESS* VadsProcess; // Size=8 Offset=112
 	union ___unnamed2048 u4; // Size=8 Offset=120
-	struct _FILE_OBJECT * FileObject; // Size=8 Offset=128
-} MMVAD, *PMMVAD;
+	struct _FILE_OBJECT* FileObject; // Size=8 Offset=128
+} MMVAD, * PMMVAD;
 
 typedef enum _MI_VAD_TYPE
 {
@@ -307,19 +307,19 @@ typedef enum _MI_VAD_TYPE
 	VadLargePages,
 	VadRotatePhysical,
 	VadLargePageSection
-} MI_VAD_TYPE, *PMI_VAD_TYPE;
+} MI_VAD_TYPE, * PMI_VAD_TYPE;
 
 typedef struct _RTL_AVL_TREE // Size=8
 {
 	PMM_AVL_NODE BalancedRoot;
-	void * NodeHint;
+	void* NodeHint;
 	UINT64 NumberGenericTableElements;
-} RTL_AVL_TREE, *PRTL_AVL_TREE, MM_AVL_TABLE, *PMM_AVL_TABLE;
+} RTL_AVL_TREE, * PRTL_AVL_TREE, MM_AVL_TABLE, * PMM_AVL_TABLE;
 
 typedef struct _HANDLE_TABLE_ENTRY_INFO {
 	UINT32 AuditMask;
 	UINT32 MaxRelativeAccessMask;
-} HANDLE_TABLE_ENTRY_INFO, *PHANDLE_TABLE_ENTRY_INFO;
+} HANDLE_TABLE_ENTRY_INFO, * PHANDLE_TABLE_ENTRY_INFO;
 
 typedef struct _HANDLE_TABLE_ENTRY
 {
@@ -340,7 +340,7 @@ typedef struct _HANDLE_TABLE_ENTRY
 		};
 		LONG NextFreeTableEntry;
 	};
-} HANDLE_TABLE_ENTRY, *PHANDLE_TABLE_ENTRY;
+} HANDLE_TABLE_ENTRY, * PHANDLE_TABLE_ENTRY;
 
 typedef struct _HANDLE_TABLE_FREE_LIST
 {
@@ -349,7 +349,7 @@ typedef struct _HANDLE_TABLE_FREE_LIST
 	PHANDLE_TABLE_ENTRY LastFreeHandleEntry;
 	UINT32 HandleCount;
 	UINT32 HighWaterMark;
-} HANDLE_TABLE_FREE_LIST, *PHANDLE_TABLE_FREE_LIST;
+} HANDLE_TABLE_FREE_LIST, * PHANDLE_TABLE_FREE_LIST;
 
 typedef struct _HANDLE_TABLE
 {