CheattEngineServer: CMD_OPENPROCESS, CMD_PROCESS32FIRST, CMD_PROCESS32NEXT, CMD_CLOSEHANDLE

 * KMemDriver integration
 * generic MT-Support achieved by synchronized wrapper functions
 * ability to spawn a Ping-Only Thread

9 files changed, 206 insertions, 39 deletions

diff --git a/CheatEngineServer/CheatEngine.h b/CheatEngineServer/CheatEngine.h
index d2f4d41..ad75b59 100644
--- a/CheatEngineServer/CheatEngine.h
+++ b/CheatEngineServer/CheatEngine.h
@@ -1,7 +1,10 @@
 #pragma once
 
+#include "KInterface.h"
 #include <winsock.h>
 
+#include <vector>
+
 #define CE_PORT "52736"
 #define MSG_WAITALL 0x8
 
@@ -84,9 +87,11 @@ typedef struct {
 
 class CEConnection {
 public:
-	explicit CEConnection(SOCKET s) : sock(s) {}
-	SOCKET getSocket(void) { return sock; }
-	void closeSocket(void) { closesocket(sock); }
+	explicit CEConnection(SOCKET s) : m_sock(s) {}
+	SOCKET getSocket(void) { return m_sock; }
+	void closeSocket(void) { closesocket(m_sock); }
+
+	std::vector<PROCESS_DATA> m_cachedProcesses;
 private:
-	SOCKET sock;
+	SOCKET m_sock;
 };
\ No newline at end of file
diff --git a/CheatEngineServer/CheatEngineServer.cpp b/CheatEngineServer/CheatEngineServer.cpp
index 2a9f54d..b06b637 100644
--- a/CheatEngineServer/CheatEngineServer.cpp
+++ b/CheatEngineServer/CheatEngineServer.cpp
@@ -7,7 +7,6 @@
 
 #include "CheatEngine.h"
 #include "CommandDispatcher.h"
-#include "KInterface.h"
 
 static SOCKET make_accept_sock(const char* servspec) {
 	const int one = 1;
@@ -74,9 +73,8 @@ static void accept_loop(const char* servspec) {
 
 	for (;;) {
 		SOCKET new_sock = accept(sock, 0, 0);
-		new_connection(new_sock);
-		//std::thread t(new_connection, new_sock);
-		//t.detach();
+		std::thread t(new_connection, new_sock);
+		t.detach();
 	}
 }
 
@@ -86,8 +84,15 @@ int main()
 	DWORD iResult;
 	KInterface& ki = KInterface::getInstance();
 
-	ki.Init();
-	ki.Handshake();
+	std::cout << "KMemDriver Init/Handshake.";
+	if (ki.Init() == false || ki.Handshake() == false) {
+		std::cout << " Failed. [PRESS RETURN TO EXIT]" << std::endl;
+		getchar();
+		return 1;
+	}
+	std::cout << " Ok." << std::endl;
+
+	ki.StartPingThread();
 
 	iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
 	if (iResult != 0) {
diff --git a/CheatEngineServer/CheatEngineServer.vcxproj b/CheatEngineServer/CheatEngineServer.vcxproj
index 26dc9fd..3beba75 100644
--- a/CheatEngineServer/CheatEngineServer.vcxproj
+++ b/CheatEngineServer/CheatEngineServer.vcxproj
@@ -74,12 +74,14 @@
   <PropertyGroup Label="UserMacros" />
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'">
     <LinkIncremental>false</LinkIncremental>
+    <TargetName>$(ProjectName)-kmem</TargetName>
   </PropertyGroup>
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">
     <LinkIncremental>true</LinkIncremental>
   </PropertyGroup>
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">
     <LinkIncremental>true</LinkIncremental>
+    <TargetName>$(ProjectName)-kmem</TargetName>
   </PropertyGroup>
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">
     <LinkIncremental>false</LinkIncremental>
diff --git a/CheatEngineServer/CommandDispatcher.cpp b/CheatEngineServer/CommandDispatcher.cpp
index 386e2c0..8c04129 100644
--- a/CheatEngineServer/CommandDispatcher.cpp
+++ b/CheatEngineServer/CommandDispatcher.cpp
@@ -5,11 +5,11 @@
 #include <iostream>
 
 
-static int recvall(SOCKET s, void *buf, int size, int flags)
+static int recvall(SOCKET s, void* buf, int size, int flags)
 {
 	int totalreceived = 0;
 	int sizeleft = size;
-	char *buffer = (char*)buf;
+	char* buffer = (char*)buf;
 
 	flags = flags | MSG_WAITALL;
 	while (sizeleft > 0)
@@ -40,11 +40,11 @@ static int recvall(SOCKET s, void *buf, int size, int flags)
 	return totalreceived;
 }
 
-static int sendall(SOCKET s, void *buf, int size, int flags)
+static int sendall(SOCKET s, void* buf, int size, int flags)
 {
 	int totalsent = 0;
 	int sizeleft = size;
-	char *buffer = (char*)buf;
+	char* buffer = (char*)buf;
 
 	while (sizeleft > 0)
 	{
@@ -72,7 +72,7 @@ static int sendall(SOCKET s, void *buf, int size, int flags)
 	return totalsent;
 }
 
-int DispatchCommand(CEConnection & con, char command)
+int DispatchCommand(CEConnection& con, char command)
 {
 	enum ce_command cmd = (enum ce_command)command;
 
@@ -86,29 +86,89 @@ int DispatchCommand(CEConnection & con, char command)
 		break;
 	case CMD_TERMINATESERVER:
 		break;
-	case CMD_OPENPROCESS:
+	case CMD_OPENPROCESS: {
+		int pid = 0;
+
+		if (recvall(con.getSocket(), &pid, sizeof(pid), MSG_WAITALL) > 0)
+		{
+			if (sendall(con.getSocket(), &pid, sizeof(pid), 0) > 0) {
+				return 0;
+			}
+		}
 		break;
+	}
+
 	case CMD_CREATETOOLHELP32SNAPSHOT: {
-		HANDLE result = (HANDLE)((ULONG_PTR)0x1);
+		UINT32 result = 0x1;
 		CeCreateToolhelp32Snapshot params;
 
 		if (recvall(con.getSocket(), &params, sizeof(CeCreateToolhelp32Snapshot), MSG_WAITALL) > 0)
 		{
+#if 0
 			std::cout << "Calling CreateToolhelp32Snapshot with flags 0x" << std::hex << params.dwFlags
 				<< " for PID 0x" << std::hex << params.th32ProcessID << std::endl;
-		}
-		if (sendall(con.getSocket(), &result, sizeof(result), 0) == sizeof(result))
-		{
-			return 0;
+#endif
+			if (sendall(con.getSocket(), &result, sizeof(result), 0) > 0)
+			{
+				return 0;
+			}
 		}
 		break;
 	}
+
 	case CMD_PROCESS32FIRST:
+		con.m_cachedProcesses.clear();
+		KInterface::getInstance().MtProcesses(con.m_cachedProcesses);
+	case CMD_PROCESS32NEXT: {
+		UINT32 toolhelpsnapshot;
+
+		if (recvall(con.getSocket(), &toolhelpsnapshot, sizeof(toolhelpsnapshot), MSG_WAITALL) > 0)
+		{
+			if (con.m_cachedProcesses.size() > 0) {
+				PROCESS_DATA pd = con.m_cachedProcesses[0];
+				int imageNameLen = (int)strnlen(pd.ImageName, sizeof(pd.ImageName));
+				CeProcessEntry* pcpe = (CeProcessEntry*)malloc(sizeof(*pcpe) + imageNameLen);
+
+				con.m_cachedProcesses.erase(con.m_cachedProcesses.begin());
+				if (pcpe == NULL) {
+					return 1;
+				}
+				pcpe->pid = (int)((ULONG_PTR)pd.UniqueProcessId);
+				pcpe->processnamesize = imageNameLen;
+				memcpy(((BYTE*)pcpe) + sizeof(*pcpe), pd.ImageName, imageNameLen);
+				pcpe->result = 1;
+				if (sendall(con.getSocket(), pcpe, sizeof(*pcpe) + imageNameLen, 0) > 0)
+				{
+					free(pcpe);
+					return 0;
+				}
+				free(pcpe);
+			}
+			else {
+				CeProcessEntry cpe;
+				cpe.pid = 0;
+				cpe.processnamesize = 0;
+				cpe.result = 0;
+				if (sendall(con.getSocket(), &cpe, sizeof(cpe), 0) > 0)
+				{
+					return 0;
+				}
+			}
+		}
 		break;
-	case CMD_PROCESS32NEXT:
-		break;
-	case CMD_CLOSEHANDLE:
+	}
+
+	case CMD_CLOSEHANDLE: {
+		UINT32 handle;
+		if (recvall(con.getSocket(), &handle, sizeof(handle), MSG_WAITALL) > 0)
+		{
+			UINT32 r = 1;
+			sendall(con.getSocket(), &r, sizeof(r), 0);
+			return 0;
+		}
 		break;
+	}
+
 	case CMD_VIRTUALQUERYEX:
 		break;
 	case CMD_READPROCESSMEMORY:
@@ -135,8 +195,24 @@ int DispatchCommand(CEConnection & con, char command)
 		break;
 	case CMD_SETTHREADCONTEXT:
 		break;
-	case CMD_GETARCHITECTURE:
-		break;
+	case CMD_GETARCHITECTURE: {
+		unsigned char arch;
+#ifdef __i386__
+		arch = 0;
+#endif
+#ifdef __x86_64__
+		arch = 1;
+#endif
+#ifdef __arm__
+		arch = 2;
+#endif
+#ifdef __aarch64__
+		arch = 3;
+#endif
+		sendall(con.getSocket(), &arch, sizeof(arch), 0);
+		return 0;
+	}
+
 	case CMD_MODULE32FIRST:
 		break;
 	case CMD_MODULE32NEXT:
@@ -168,7 +244,7 @@ int DispatchCommand(CEConnection & con, char command)
 	return 1;
 }
 
-int CheckForAndDispatchCommand(CEConnection & con)
+int CheckForAndDispatchCommand(CEConnection& con)
 {
 	int r;
 	char command;
diff --git a/IntegrationTest/IntegrationTest.cpp b/IntegrationTest/IntegrationTest.cpp
index b4f852e..5503345 100644
--- a/IntegrationTest/IntegrationTest.cpp
+++ b/IntegrationTest/IntegrationTest.cpp
@@ -138,8 +138,8 @@ int main()
 	try {
 		KM_ASSERT_EQUAL(true, ki.Init(), "Kernel Interface Init");
 		KM_ASSERT_EQUAL(true, ki.Handshake(), "Kernel Interface Handshake");
-		KM_ASSERT_EQUAL(true, ki.getBuffer() != NULL, "Kernel Interface Buffer != NULL");
 		KM_ASSERT_EQUAL(true, ki.getKHandle() != ki.getUHandle() && ki.getKHandle() != NULL && ki.getUHandle() != NULL, "Kernel Interface Handles");
+		KM_ASSERT_EQUAL(true, ki.getBuffer() != NULL, "Kernel Interface Buffer != NULL");
 		KM_ASSERT_EQUAL(SRR_TIMEOUT, ki.RecvWait(), "Kernel Interface Receive Wait");
 		KM_ASSERT_EQUAL(true, ki.Ping(), "Kernel Interface PING - PONG #1");
 		KM_ASSERT_EQUAL(true, ki.Ping(), "Kernel Interface PING - PONG #2");
diff --git a/KMemDriver/KMemDriver.c b/KMemDriver/KMemDriver.c
index 4c7592f..f574e17 100644
--- a/KMemDriver/KMemDriver.c
+++ b/KMemDriver/KMemDriver.c
@@ -8,6 +8,7 @@
 
 #define CHEAT_EXE L"kmem"
 #define WAIT_OBJECT_0 ((STATUS_WAIT_0 ) + 0 )
+#define SystemProcessInformation 0x05
 
 DRIVER_INITIALIZE DriverEntry;
 #pragma alloc_text(INIT, DriverEntry)
@@ -172,7 +173,7 @@ NTSTATUS GetProcesses(OUT PROCESS_DATA* procs, IN OUT SIZE_T* psiz)
 	SIZE_T const max_siz = *psiz;
 	ULONG mem_needed = 0;
 
-	NTSTATUS status = ZwQuerySystemInformation(0x05, NULL, 0, &mem_needed);
+	NTSTATUS status = ZwQuerySystemInformation(SystemProcessInformation, NULL, 0, &mem_needed);
 	if (!NT_SUCCESS(status) && mem_needed == 0) {
 		KDBG("NtQuerySystemInformation(ReturnLength: %lu) failed with 0x%X\n", mem_needed, status);
 		return status;
@@ -188,7 +189,7 @@ NTSTATUS GetProcesses(OUT PROCESS_DATA* procs, IN OUT SIZE_T* psiz)
 		return STATUS_NO_MEMORY;
 	}
 
-	status = ZwQuerySystemInformation(0x05, sysinfo, mem_needed, NULL);
+	status = ZwQuerySystemInformation(SystemProcessInformation, sysinfo, mem_needed, NULL);
 	if (!NT_SUCCESS(status)) {
 		KDBG("NtQuerySystemInformation(SystemInformationLength: %lu) failed with 0x%X\n", mem_needed, status);
 		goto free_memory;
@@ -231,16 +232,16 @@ NTSTATUS WaitForControlProcess(OUT PEPROCESS* ppEProcess)
 	imageBase = NULL;
 	ctrlPID = NULL;
 
-	SYSTEM_PROCESS_INFORMATION* procs = MmAllocateNonCachedMemory(1024 * sizeof(*procs));
+	SYSTEM_PROCESS_INFORMATION* procs = MmAllocateNonCachedMemory((1024 + 128) * sizeof(*procs));
 	ULONG mem_needed = 0;
 
 	if (procs == NULL) {
 		return STATUS_MEMORY_NOT_ALLOCATED;
 	}
 	while (ctrlPID == NULL) {
-		status = ZwQuerySystemInformation(0x05, (PVOID)&procs[0], 1024 * sizeof(*procs), &mem_needed);
+		status = ZwQuerySystemInformation(SystemProcessInformation, (PVOID)&procs[0], (1024 + 128) * sizeof(*procs), &mem_needed);
 		if (!NT_SUCCESS(status)) {
-			KDBG("NtQuerySystemInformation failed with 0x%X\n", status);
+			KDBG("ZwQuerySystemInformation(%zu,%lu) failed with 0x%X\n", 1024 * sizeof(*procs), mem_needed, status);
 			return status;
 		}
 
@@ -448,7 +449,7 @@ NTSTATUS KRThread(IN PVOID pArg)
 		INT running = 1;
 		SIZE_T maxWaits = 20;
 		do {
-			LARGE_INTEGER wait = { .QuadPart = -10000000 };
+			LARGE_INTEGER wait = { .QuadPart = -(KRNL_WAIT_TIME_US * 10) };
 			status = KeWaitForSingleObject(pk_kevent, Executive, UserMode, FALSE, &wait);
 			if (NT_SUCCESS(status) && status == WAIT_OBJECT_0) {
 				maxWaits = 20;
@@ -464,7 +465,7 @@ NTSTATUS KRThread(IN PVOID pArg)
 						break;
 					case MEM_PING: {
 						PKERNEL_PING ping = (PKERNEL_PING)shm_buf;
-						KDBG("Got a PING with rng 0x%X, sending PONG !\n",
+						KDBG("Got a PING with rng 0x%08X, sending PONG !\n",
 							ping->rnd_user);
 						ping->rnd_kern = ping->rnd_user;
 
diff --git a/MemDriverLib/MemDriverLib.cpp b/MemDriverLib/MemDriverLib.cpp
index c79a01a..8123d34 100644
--- a/MemDriverLib/MemDriverLib.cpp
+++ b/MemDriverLib/MemDriverLib.cpp
@@ -154,6 +154,10 @@ bool KInterface::Modules(HANDLE targetPID,
 
 bool KInterface::Exit()
 {
+	if (m_pingThreadStarted == true) {
+		m_pingThreadStarted = false;
+		m_pingThread.join();
+	}
 	m_last_ntstatus = INVALID_NTSTATUS;
 	return SendRecvWait(MEM_EXIT, INFINITE) == SRR_SIGNALED;
 }
@@ -353,4 +357,18 @@ SendRecvReturn KInterface::RecvWait(DWORD timeout)
 	return SRR_ERR_UEVENT;
 }
 
+void KInterface::PingThread(void)
+{
+	while (m_pingThreadStarted == true) {
+		std::this_thread::sleep_for(std::chrono::milliseconds(DEFAULT_TIMEOUT_MS));
+		MtPing();
+	}
+}
+
+void KInterface::StartPingThread(void)
+{
+	m_pingThreadStarted = true;
+	m_pingThread = std::move(std::thread(&KInterface::PingThread, this));
+}
+
 #pragma warning(pop)
\ No newline at end of file
diff --git a/include/KInterface.h b/include/KInterface.h
index 3fee236..981ae35 100644
--- a/include/KInterface.h
+++ b/include/KInterface.h
@@ -2,11 +2,13 @@
 
 #include "KMemDriver.h"
 
+#include <Windows.h>
+
+#include <mutex>
 #include <stdexcept>
 #include <vector>
-#include <Windows.h>
 
-#define DEFAULT_TIMEOUT 2500
+#define DEFAULT_TIMEOUT_MS ((KRNL_WAIT_TIME_US / 1000) * (KRNL_MAX_WAITS - 1))
 #define INVALID_NTSTATUS (UINT32)-1
 
 typedef enum SendRecvReturn {
@@ -43,21 +45,77 @@ public:
 	bool VFree(HANDLE targetPID, PVOID address, SIZE_T size);
 	bool VUnlink(HANDLE targetPID, PVOID address);
 
+	bool MtInit() {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return Init();
+	}
+	bool MtHandshake() {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return Handshake();
+	}
+	bool MtPing() {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return Ping();
+	}
+	bool MtProcesses(std::vector<PROCESS_DATA>& dest) {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return Processes(dest);
+	}
+	bool MtPages(HANDLE targetPID, std::vector<MEMORY_BASIC_INFORMATION>& dest, PVOID start_address = NULL) {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return Pages(targetPID, dest, start_address);
+	}
+	bool MtModules(HANDLE targetPID, std::vector<MODULE_DATA>& dest) {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return Modules(targetPID, dest);
+	}
+	bool MtExit() {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return Exit();
+	}
+	bool MtRPM(HANDLE targetPID, PVOID address, BYTE* buf, SIZE_T size, PKERNEL_READ_REQUEST result) {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return RPM(targetPID, address, buf, size, result);
+	}
+	bool MtWPM(HANDLE targetPID, PVOID address, BYTE* buf, SIZE_T size, PKERNEL_WRITE_REQUEST result) {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return WPM(targetPID, address, buf, size, result);
+	}
+	bool MtVAlloc(HANDLE targetPID, PVOID* address, SIZE_T* size, ULONG protection) {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return VAlloc(targetPID, address, size, protection);
+	}
+	bool MtVFree(HANDLE targetPID, PVOID address, SIZE_T size) {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return VFree(targetPID, address, size);
+	}
+	bool MtVUnlink(HANDLE targetPID, PVOID address) {
+		std::unique_lock<std::mutex> lck(m_jobLock);
+		return VUnlink(targetPID, address);
+	}
+
 	PVOID getBuffer();
 	HANDLE getKHandle();
 	HANDLE getUHandle();
 	UINT32 getLastPingValue();
 	UINT32 getLastNtStatus();
-	SendRecvReturn RecvWait(DWORD timeout = DEFAULT_TIMEOUT);
+
+	SendRecvReturn RecvWait(DWORD timeout = DEFAULT_TIMEOUT_MS);
+	void StartPingThread(void);
 
 private:
-	SendRecvReturn SendRecvWait(UINT32 type, DWORD timeout = DEFAULT_TIMEOUT);
+	SendRecvReturn SendRecvWait(UINT32 type, DWORD timeout = DEFAULT_TIMEOUT_MS);
+	void PingThread(void);
 
 	PVOID m_shmem = NULL;
 	HANDLE m_kevent = NULL, m_uevent = NULL;
 
 	UINT32 m_last_ping_value = 0;
 	UINT32 m_last_ntstatus = INVALID_NTSTATUS;
+
+	bool m_pingThreadStarted = false;
+	std::thread m_pingThread;
+	std::mutex m_jobLock;
 };
 
 class KMemory
diff --git a/include/KMemDriver.h b/include/KMemDriver.h
index bc50d61..778006c 100644
--- a/include/KMemDriver.h
+++ b/include/KMemDriver.h
@@ -18,6 +18,8 @@ typedef _Return_type_success_(return >= 0) LONG NTSTATUS;
 #define SHMEM_ADDR				0x60000000
 #define SHMEM_SIZE				8192*8*2
 #define INVALID_REQUEST			(UINT32)-1
+#define KRNL_MAX_WAITS			20
+#define KRNL_WAIT_TIME_US		1000000
 
 #define MEM_HANDSHAKE			0x800
 #define MEM_PING				0x801