PingThread calls function on timeout.

3 files changed, 33 insertions, 14 deletions

diff --git a/CheatEngineServer/CheatEngineServer.cpp b/CheatEngineServer/CheatEngineServer.cpp
index b06b637..9f389d4 100644
--- a/CheatEngineServer/CheatEngineServer.cpp
+++ b/CheatEngineServer/CheatEngineServer.cpp
@@ -8,6 +8,9 @@
 #include "CheatEngine.h"
 #include "CommandDispatcher.h"
 
+static SOCKET sock;
+static BOOL run_main_loop = TRUE;
+
 static SOCKET make_accept_sock(const char* servspec) {
 	const int one = 1;
 	struct addrinfo hints = {};
@@ -63,19 +66,31 @@ static void new_connection(SOCKET sock) {
 	}
 }
 
-static void accept_loop(const char* servspec) {
-	SOCKET sock = make_accept_sock(servspec);
+static int accept_loop(const char* servspec) {
+	sock = make_accept_sock(servspec);
 
 	if (sock == NULL)
 	{
-		return;
+		return 1;
 	}
 
-	for (;;) {
+	while (run_main_loop == TRUE) {
 		SOCKET new_sock = accept(sock, 0, 0);
-		std::thread t(new_connection, new_sock);
-		t.detach();
+		if (new_sock != NULL) {
+			std::thread t(new_connection, new_sock);
+			t.detach();
+		}
+		else {
+			return 1;
+		}
 	}
+	return 0;
+}
+
+static void onPingThreadTimeout(void) {
+	std::cout << "PingThread timeout, abort .." << std::endl;
+	run_main_loop = FALSE;
+	closesocket(sock);
 }
 
 int main()
@@ -92,7 +107,7 @@ int main()
 	}
 	std::cout << " Ok." << std::endl;
 
-	ki.StartPingThread();
+	ki.StartPingThread(onPingThreadTimeout);
 
 	iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
 	if (iResult != 0) {
@@ -100,5 +115,5 @@ int main()
 		return 1;
 	}
 
-	accept_loop("0.0.0.0");
+	return accept_loop("0.0.0.0");
 }
\ No newline at end of file
diff --git a/MemDriverLib/MemDriverLib.cpp b/MemDriverLib/MemDriverLib.cpp
index 8123d34..1a19552 100644
--- a/MemDriverLib/MemDriverLib.cpp
+++ b/MemDriverLib/MemDriverLib.cpp
@@ -357,18 +357,22 @@ SendRecvReturn KInterface::RecvWait(DWORD timeout)
 	return SRR_ERR_UEVENT;
 }
 
-void KInterface::PingThread(void)
+void KInterface::PingThread(void(__cdecl* onTimeout)(void))
 {
 	while (m_pingThreadStarted == true) {
 		std::this_thread::sleep_for(std::chrono::milliseconds(DEFAULT_TIMEOUT_MS));
-		MtPing();
+		if (MtPing() != true) {
+			m_pingThreadStarted = false;
+			break;
+		}
 	}
+	onTimeout();
 }
 
-void KInterface::StartPingThread(void)
+void KInterface::StartPingThread(void(__cdecl* onTimeout)(void))
 {
 	m_pingThreadStarted = true;
-	m_pingThread = std::move(std::thread(&KInterface::PingThread, this));
+	m_pingThread = std::move(std::thread(&KInterface::PingThread, this, onTimeout));
 }
 
 #pragma warning(pop)
\ No newline at end of file
diff --git a/include/KInterface.h b/include/KInterface.h
index 981ae35..19c9268 100644
--- a/include/KInterface.h
+++ b/include/KInterface.h
@@ -101,11 +101,11 @@ public:
 	UINT32 getLastNtStatus();
 
 	SendRecvReturn RecvWait(DWORD timeout = DEFAULT_TIMEOUT_MS);
-	void StartPingThread(void);
+	void StartPingThread(void (__cdecl *onTimeout)(void));
 
 private:
 	SendRecvReturn SendRecvWait(UINT32 type, DWORD timeout = DEFAULT_TIMEOUT_MS);
-	void PingThread(void);
+	void PingThread(void (__cdecl *onTimeout)(void));
 
 	PVOID m_shmem = NULL;
 	HANDLE m_kevent = NULL, m_uevent = NULL;