fix hook for windows

idealvin · Oct 21, 2023 · d9cece6 · d9cece6
1 parent 80a27d1
commit d9cece6
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 25 deletions.
diff --git a/src/co/hook.h b/src/co/hook.h
@@ -21,13 +21,6 @@ void hook_sleep(bool x);
 #define __sys_api(x)        _sys_##x
 #define _CO_DEC_SYS_API(x)  extern x##_fp_t __sys_api(x)
 
-struct HookInitializer {
-    HookInitializer();
-    ~HookInitializer();
-};
-
-static HookInitializer g_hook_initializer;
-
 #ifdef _WIN32
 #include <WinSock2.h>
 #include <ws2tcpip.h> // for inet_ntop...
@@ -123,6 +116,13 @@ _CO_DEC_SYS_API(GetQueuedCompletionStatusEx);
 
 #define _CO_DEF_SYS_API(x)  x##_fp_t __sys_api(x) = 0
 
+struct HookInitializer {
+    HookInitializer();
+    ~HookInitializer();
+};
+
+static HookInitializer g_hook_initializer;
+
 namespace co {
 
 // deduce type of the second parameter of ioctl

diff --git a/src/co/hook_win.cc b/src/co/hook_win.cc
@@ -3,6 +3,7 @@
 
 #ifdef _CO_DISABLE_HOOK
 namespace co {
+void init_hook() {}
 void hook_sleep(bool) {}
 } // co
 
@@ -1327,7 +1328,7 @@ WSARecvMsg_fp_t get_WSARecvMsg_fp() {
     CHECK_EQ(r, 0) << "get WSARecvMsg failed: " << co::strerror();
     CHECK(fp != NULL) << "pointer to WSARecvMsg is NULL..";
 
-    __sys_api(closesocket)(fd);
+    ::closesocket(fd);
     return fp;
 }
 
@@ -1351,7 +1352,7 @@ WSASendMsg_fp_t get_WSASendMsg_fp() {
     CHECK_EQ(r, 0) << "get WSASendMsg failed: " << co::strerror();
     CHECK(fp != NULL) << "pointer to WSASendMsg is NULL..";
 
-    __sys_api(closesocket)(fd);
+    ::closesocket(fd);
     return fp;
 }
 
@@ -1372,7 +1373,12 @@ inline void detour_detach(PVOID* ppbReal, PVOID pbMine, PCHAR psz) {
 #define attach_hook(x)  detour_attach(&(PVOID&)__sys_api(x), (PVOID)hook_##x, #x)
 #define detach_hook(x)  detour_detach(&(PVOID&)__sys_api(x), (PVOID)hook_##x, #x)
 
+// it will be called at initialization of coroutine schedulers
 void init_hook() {
+    g_hook = co::_make_static<co::Hook>();
+    __sys_api(WSARecvMsg) = get_WSARecvMsg_fp();
+    __sys_api(WSASendMsg) = get_WSASendMsg_fp();
+
     DetourTransactionBegin();
     DetourUpdateThread(GetCurrentThread());
     attach_hook(Sleep);
@@ -1452,19 +1458,6 @@ void hook_sleep(bool x) {
 
 } // co
 
-static int g_nifty_counter;
-HookInitializer::HookInitializer() {
-    if (g_nifty_counter++ == 0) {
-        __sys_api(WSARecvMsg) = get_WSARecvMsg_fp();
-        __sys_api(WSASendMsg) = get_WSASendMsg_fp();
-        g_hook = co::_make_static<co::Hook>();
-        co::init_hook();
-    }
-}
-
-HookInitializer::~HookInitializer() {}
-
-
 #undef attach_hook
 #undef detach_hook
 #undef do_hard_hook

diff --git a/src/co/sock_win.cc b/src/co/sock_win.cc
@@ -371,6 +371,8 @@ bool _can_skip_iocp_on_success() {
     return true;
 }
 
+// call init_sock() first, then init_hook()
+extern void init_hook();
 void init_sock() {
     WSADATA x;
     WSAStartup(MAKEWORD(2, 2), &x);
@@ -411,12 +413,12 @@ void init_sock() {
 
     ::closesocket(fd);
     can_skip_iocp_on_success = _can_skip_iocp_on_success();
-}
 
-void cleanup_sock() {
-    WSACleanup();
+    init_hook();
 }
 
+void cleanup_sock() { WSACleanup(); }
+
 } // co
 
 #if defined(_MSC_VER) && defined(BUILDING_CO_SHARED)