From 65f9ab76500be9e98366324c8639bfb369e7d0c5 Mon Sep 17 00:00:00 2001
From: Changhua <lichanghua0821@gmail.com>
Date: Sat, 13 Aug 2022 21:55:08 +0800
Subject: [PATCH] Impl Hook/UnHook

---
 Spy/rpc_server.cpp | 37 +++++++++++++++++++++++--------------
 Spy/spy.cpp        | 23 ++---------------------
 2 files changed, 25 insertions(+), 35 deletions(-)

diff --git a/Spy/rpc_server.cpp b/Spy/rpc_server.cpp
index ad6e268..f92f7c0 100644
--- a/Spy/rpc_server.cpp
+++ b/Spy/rpc_server.cpp
@@ -3,31 +3,35 @@
 
 #include "exec_sql.h"
 #include "get_contacts.h"
+#include "receive_msg.h"
 #include "rpc_h.h"
 #include "rpc_server.h"
 #include "sdk.h"
 #include "send_msg.h"
-#include "spy.h"
+#include "spy.h"
 #include "spy_types.h"
 #include "util.h"
 
 using namespace std;
 
-extern HANDLE g_hEvent;
-extern MsgQueue_t g_MsgQueue;
-extern const MsgTypesMap_t g_WxMsgTypes;
-extern int IsLogin(void);
+extern int IsLogin(void);                // Defined in spy.cpp
+extern HANDLE g_hEvent;                  // New message signal
+extern BOOL g_rpcKeepAlive;              // Keep RPC server thread running
+extern MsgQueue_t g_MsgQueue;            // Queue for message
+extern const MsgTypesMap_t g_WxMsgTypes; // Map of WeChat Message types
+
+static BOOL listenMsgFlag = false;
 
 int server_IsLogin() { return IsLogin(); }
 
 void server_EnableReceiveMsg()
 {
     unsigned long ulCode = 0;
-
+    ListenMessage();
     RpcTryExcept
     {
         // 调用客户端的回调函数
-        while (true) {
+        while (listenMsgFlag) {
             // 中断式，兼顾及时性和CPU使用率
             WaitForSingleObject(g_hEvent, INFINITE); // 等待消息
             while (!g_MsgQueue.empty()) {
@@ -153,9 +157,9 @@ int server_GetDbTables(const wchar_t *db, int *pNum, PPRpcTables *tbls)
     int index = 0;
     for (auto it = tables.begin(); it != tables.end(); it++) {
         PRpcTables p = (PRpcTables)midl_user_allocate(sizeof(RpcTables_t));
-        p->table     = it->table;
-        p->sql       = it->sql;
-        pp[index++]  = p;
+        p->table    = it->table;
+        p->sql      = it->sql;
+        pp[index++] = p;
     }
 
     *tbls = pp;
@@ -192,8 +196,9 @@ int RpcStartServer(HMODULE hModule)
                                   (unsigned)-1,                   // Infinite max size of incoming data blocks.
                                   SecurityCallback);              // Naive security callback.
 
-    while (1) {
-        Sleep(10000); // 休眠，释放CPU
+    listenMsgFlag = true;
+    while (g_rpcKeepAlive) {
+        Sleep(1000); // 休眠，释放CPU
     }
 
     return 0;
@@ -201,8 +206,12 @@ int RpcStartServer(HMODULE hModule)
 
 int RpcStopServer(void)
 {
-    RPC_STATUS status;
-    status = RpcMgmtStopServerListening(NULL);
+    RPC_STATUS status;
+
+    UnListenMessage();
+
+    listenMsgFlag = false;
+    status        = RpcMgmtStopServerListening(NULL);
     if (status)
         return status;
 
diff --git a/Spy/spy.cpp b/Spy/spy.cpp
index e01bead..21a73a9 100644
--- a/Spy/spy.cpp
+++ b/Spy/spy.cpp
@@ -1,20 +1,13 @@
 #include "load_calls.h"
-#include "receive_msg.h"
 #include "rpc_server.h"
 #include "spy.h"
 #include "util.h"
 
 HANDLE g_hEvent          = NULL;
+BOOL g_rpcKeepAlive      = false;
 WxCalls_t g_WxCalls      = { 0 };
 DWORD g_WeChatWinDllAddr = 0;
 
-DWORD WINAPI Monitor(HMODULE hModule)
-{
-    ListenMessage();
-
-    return TRUE;
-}
-
 void InitSpy(HMODULE hModule)
 {
     wchar_t version[16] = { 0 };
@@ -40,20 +33,8 @@ void InitSpy(HMODULE hModule)
     if (rpcThread != 0) {
         CloseHandle(rpcThread);
     }
-
-    HANDLE mThread = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)Monitor, hModule, NULL, 0);
-    if (mThread != 0) {
-        CloseHandle(mThread);
-    }
 }
 
 void DestroySpy() { RpcStopServer(); }
 
-int IsLogin(void)
-{
-    if (g_WeChatWinDllAddr == 0) {
-        return 0;
-    }
-
-    return (int)GET_DWORD(g_WeChatWinDllAddr + g_WxCalls.login);
-}
+int IsLogin(void) { return (int)GET_DWORD(g_WeChatWinDllAddr + g_WxCalls.login); }