From cde22aace74eb58f10f8fac98906cba948f8b952 Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Tue, 25 Oct 2016 12:03:28 +0200
Subject: [PATCH] Hook WSAStartup/WSACleanup to workaround app bugs. Refs #362,
 #403

* If the application incorrectly calls WSACleanup() before WSAStartup()
  then instead of just an error for invalid use, it will instead kill
  all of the sockets we created from our early WSAStartup().
* To fix this, we hook those functions and track the process wide ref-
  count, to prevent it from every going to 0 until we shut down our own.
* This was found on GTA 5 in particular but possibly some other
  application will do this too.
---
 renderdoc/os/win32/sys_win32_hooks.cpp | 41 ++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/renderdoc/os/win32/sys_win32_hooks.cpp b/renderdoc/os/win32/sys_win32_hooks.cpp
index f7d56ef80..94e97a4c7 100644
--- a/renderdoc/os/win32/sys_win32_hooks.cpp
+++ b/renderdoc/os/win32/sys_win32_hooks.cpp
@@ -23,6 +23,7 @@
  * THE SOFTWARE.
  ******************************************************************************/
 
+#include <winsock2.h>
 #include "api/replay/renderdoc_replay.h"
 #include "core/core.h"
 #include "hooks/hooks.h"
@@ -30,6 +31,9 @@
 
 #define DLL_NAME "kernel32.dll"
 
+typedef int(WSAAPI *PFN_WSASTARTUP)(__in WORD wVersionRequested, __out LPWSADATA lpWSAData);
+typedef int(WSAAPI *PFN_WSACLEANUP)();
+
 typedef BOOL(WINAPI *PFN_CREATE_PROCESS_A)(
     __in_opt LPCSTR lpApplicationName, __inout_opt LPSTR lpCommandLine,
     __in_opt LPSECURITY_ATTRIBUTES lpProcessAttributes,
@@ -63,6 +67,12 @@ public:
     success &= CreateProcessA.Initialize("CreateProcessA", DLL_NAME, CreateProcessA_hook);
     success &= CreateProcessW.Initialize("CreateProcessW", DLL_NAME, CreateProcessW_hook);
 
+    success &= WSAStartup.Initialize("WSAStartup", "ws2_32.dll", WSAStartup_hook);
+    success &= WSACleanup.Initialize("WSACleanup", "ws2_32.dll", WSACleanup_hook);
+
+    // we start with a refcount of 1 because we initialise WSA ourselves for our own sockets.
+    m_WSARefCount = 1;
+
     if(!success)
       return false;
 
@@ -80,10 +90,41 @@ private:
   bool m_HasHooks;
   bool m_EnabledHooks;
 
+  int m_WSARefCount;
+
   // D3DPERF api
   Hook<PFN_CREATE_PROCESS_A> CreateProcessA;
   Hook<PFN_CREATE_PROCESS_W> CreateProcessW;
 
+  Hook<PFN_WSASTARTUP> WSAStartup;
+  Hook<PFN_WSACLEANUP> WSACleanup;
+
+  static int WSAAPI WSAStartup_hook(WORD wVersionRequested, LPWSADATA lpWSAData)
+  {
+    int ret = syshooks.WSAStartup()(wVersionRequested, lpWSAData);
+
+    // only increment the refcount if the function succeeded
+    if(ret == 0)
+      syshooks.m_WSARefCount++;
+
+    return ret;
+  }
+
+  static int WSAAPI WSACleanup_hook()
+  {
+    // don't let the application murder our sockets with a mismatched WSACleanup() call
+    if(syshooks.m_WSARefCount == 1)
+    {
+      RDCLOG("WSACleanup called with (to the application) no WSAStartup! Ignoring.");
+      SetLastError(WSANOTINITIALISED);
+      return SOCKET_ERROR;
+    }
+
+    // decrement refcount and call the real thing
+    syshooks.m_WSARefCount--;
+    return syshooks.WSACleanup()();
+  }
+
   static BOOL WINAPI CreateProcessA_hook(
       __in_opt LPCSTR lpApplicationName, __inout_opt LPSTR lpCommandLine,
       __in_opt LPSECURITY_ATTRIBUTES lpProcessAttributes,