From 38464e6450efe0aac4292f9ae2a8204e3854b663 Mon Sep 17 00:00:00 2001
From: Changhua <lichanghua0821@gmail.com>
Date: Wed, 3 Jul 2024 18:36:52 +0800
Subject: [PATCH] Fix Hook/Unhook issue

---
 WeChatFerry/spy/receive_msg.cpp | 97 +++++++++++++++++----------------
 1 file changed, 50 insertions(+), 47 deletions(-)

diff --git a/WeChatFerry/spy/receive_msg.cpp b/WeChatFerry/spy/receive_msg.cpp
index a1834fb..8d385db 100644
--- a/WeChatFerry/spy/receive_msg.cpp
+++ b/WeChatFerry/spy/receive_msg.cpp
@@ -195,6 +195,33 @@ static void DispatchPyq(QWORD arg1, QWORD arg2, QWORD arg3)
     }
 }
 
+static MH_STATUS InitializeHook()
+{
+    if (isMH_Initialized) {
+        return MH_OK;
+    }
+    MH_STATUS status = MH_Initialize();
+    if (status == MH_OK) {
+        isMH_Initialized = true;
+    }
+    return status;
+}
+
+static MH_STATUS UninitializeHook()
+{
+    if (!isMH_Initialized) {
+        return MH_OK;
+    }
+    if (gIsLogging || gIsListening || gIsListeningPyq) {
+        return MH_OK;
+    }
+    MH_STATUS status = MH_Uninitialize();
+    if (status == MH_OK) {
+        isMH_Initialized = false;
+    }
+    return status;
+}
+
 void EnableLog()
 {
     MH_STATUS status = MH_UNKNOWN;
@@ -204,13 +231,10 @@ void EnableLog()
     }
     WxLog_t funcWxLog = (WxLog_t)(g_WeChatWinDllAddr + OS_WXLOG);
 
-    if (!isMH_Initialized) {
-        status = MH_Initialize();
-        if (status != MH_OK) {
-            LOG_ERROR("MH_Initialize failed: {}", to_string(status));
-            return;
-        }
-        isMH_Initialized = true;
+    status = InitializeHook();
+    if (status != MH_OK) {
+        LOG_ERROR("MH_Initialize failed: {}", to_string(status));
+        return;
     }
 
     status = MH_CreateHook(funcWxLog, &PrintWxLog, reinterpret_cast<LPVOID *>(&realWxLog));
@@ -241,12 +265,11 @@ void DisableLog()
     }
 
     gIsLogging = false;
-    if (isMH_Initialized && !gIsLogging && !gIsListening && !gIsListeningPyq) {
-        status = MH_Uninitialize();
-        if (status != MH_OK) {
-            LOG_ERROR("MH_Uninitialize failed: {}", to_string(status));
-            return;
-        }
+
+    status = UninitializeHook();
+    if (status != MH_OK) {
+        LOG_ERROR("MH_Uninitialize failed: {}", to_string(status));
+        return;
     }
 }
 
@@ -259,13 +282,10 @@ void ListenMessage()
     }
     funcRecvMsg = (RecvMsg_t)(g_WeChatWinDllAddr + OS_RECV_MSG_CALL);
 
-    if (!isMH_Initialized) {
-        status = MH_Initialize();
-        if (status != MH_OK) {
-            LOG_ERROR("MH_Initialize failed: {}", to_string(status));
-            return;
-        }
-        isMH_Initialized = true;
+    status = InitializeHook();
+    if (status != MH_OK) {
+        LOG_ERROR("MH_Initialize failed: {}", to_string(status));
+        return;
     }
 
     status = MH_CreateHook(funcRecvMsg, &DispatchMsg, reinterpret_cast<LPVOID *>(&realRecvMsg));
@@ -296,20 +316,13 @@ void UnListenMessage()
         return;
     }
 
-    status = MH_Uninitialize();
+    gIsListening = false;
+
+    status = UninitializeHook();
     if (status != MH_OK) {
         LOG_ERROR("MH_Uninitialize failed: {}", to_string(status));
         return;
     }
-
-    gIsListening = false;
-    if (isMH_Initialized && !gIsLogging && !gIsListening && !gIsListeningPyq) {
-        status = MH_Uninitialize();
-        if (status != MH_OK) {
-            LOG_ERROR("MH_Uninitialize failed: {}", to_string(status));
-            return;
-        }
-    }
 }
 
 void ListenPyq()
@@ -321,13 +334,10 @@ void ListenPyq()
     }
     funcRecvPyq = (RecvPyq_t)(g_WeChatWinDllAddr + OS_PYQ_MSG_CALL);
 
-    if (!isMH_Initialized) {
-        status = MH_Initialize();
-        if (status != MH_OK) {
-            LOG_ERROR("MH_Initialize failed: {}", to_string(status));
-            return;
-        }
-        isMH_Initialized = true;
+    status = InitializeHook();
+    if (status != MH_OK) {
+        LOG_ERROR("MH_Initialize failed: {}", to_string(status));
+        return;
     }
 
     status = MH_CreateHook(funcRecvPyq, &DispatchPyq, reinterpret_cast<LPVOID *>(&realRecvPyq));
@@ -358,18 +368,11 @@ void UnListenPyq()
         return;
     }
 
-    status = MH_Uninitialize();
+    gIsListeningPyq = false;
+
+    status = UninitializeHook();
     if (status != MH_OK) {
         LOG_ERROR("MH_Uninitialize failed: {}", to_string(status));
         return;
     }
-
-    gIsListeningPyq = false;
-    if (isMH_Initialized && !gIsLogging && !gIsListening && !gIsListeningPyq) {
-        status = MH_Uninitialize();
-        if (status != MH_OK) {
-            LOG_ERROR("MH_Uninitialize failed: {}", to_string(status));
-            return;
-        }
-    }
 }