Safer memory access

Author: Dutchman101

Client/multiplayer_sa/StdInc.cpp

@@ -1,5 +1,6 @@
 // StdInc.h
 #include "StdInc.h"
 #define ALLOC_STATS_MODULE_NAME "multiplayer_sa"
+#define MTASA_EXPORT_SHARED_UTIL
 #include "SharedUtil.hpp"
 #include "SharedUtil.MemAccess.hpp"

Shared/sdk/SharedUtil.MemAccess.h

@@ -11,6 +11,16 @@
 
 #include <SharedUtil.IntTypes.h>
 
+#ifdef _WIN32
+    #ifdef MTASA_EXPORT_SHARED_UTIL
+        #define SHARED_UTIL_API __declspec(dllexport)
+    #else
+        #define SHARED_UTIL_API __declspec(dllimport)
+    #endif
+#else
+    #define SHARED_UTIL_API
+#endif
+
 namespace SharedUtil
 {
     struct SMemWrite
@@ -20,12 +30,20 @@ namespace SharedUtil
         DWORD oldProt;
     };
 
-    void      SetInitialVirtualProtect();
-    bool      IsSlowMem(const void* pAddr, uint uiAmount);
-    SMemWrite OpenMemWrite(const void* pAddr, uint uiAmount);
-    void      CloseMemWrite(SMemWrite& hMem);
-    bool      ismemset(const void* pAddr, int cValue, uint uiAmount);
+    SHARED_UTIL_API void      SetInitialVirtualProtect();
+    SHARED_UTIL_API bool      IsSlowMem(const void* pAddr, uint uiAmount);
+    SHARED_UTIL_API SMemWrite OpenMemWrite(const void* pAddr, uint uiAmount);
+    SHARED_UTIL_API void      CloseMemWrite(SMemWrite& hMem);
+    SHARED_UTIL_API bool      ismemset(const void* pAddr, int cValue, uint uiAmount);
+
+    bool IsProtectedSlowMem(const void* pAddr);
 
-    #define DEBUG_CHECK_IS_FAST_MEM(addr,size) { dassert( !IsSlowMem( (const void*)(addr), size ) ); }
-    #define DEBUG_CHECK_IS_SLOW_MEM(addr,size) { dassert( IsSlowMem( (const void*)(addr), size ) ); }
+#define DEBUG_CHECK_IS_FAST_MEM(addr, size) \
+    { \
+        dassert(!IsSlowMem((const void*)(addr), size)); \
+    }
+#define DEBUG_CHECK_IS_SLOW_MEM(addr, size) \
+    { \
+        dassert(IsSlowMem((const void*)(addr), size)); \
+    }
 }            // namespace SharedUtil

Shared/sdk/SharedUtil.MemAccess.hpp

@@ -8,6 +8,13 @@
  *
  *****************************************************************************/
 
+#pragma once		
+#include <windows.h>
+#include <cassert>
+#include "SharedUtil.MemAccess.h"
+#include "SharedUtil.Logging.h"
+#include "SString.h"
+
 namespace SharedUtil
 {
     // Returns true if matching memset would have no affect
@@ -74,13 +81,12 @@ namespace SharedUtil
     // Temporarily unprotect slow mem area
     SMemWrite OpenMemWrite(const void* pAddr, uint uiAmount)
     {
-        SMemWrite hMem;
+        SMemWrite hMem{};
 
         // Check for incorrect use of function
         if (!IsSlowMem(pAddr, uiAmount))
         {
             dassert(0 && "Should use Mem*Fast function");
-            hMem.dwFirstPage = 0;
             return hMem;
         }
 
@@ -91,10 +97,18 @@ namespace SharedUtil
         hMem.dwFirstPage = ((DWORD)pAddr) & ~0xFFF;
         DWORD dwLastPage = (((DWORD)pAddr) + uiAmount - 1) & ~0xFFF;
         hMem.dwSize = dwLastPage - hMem.dwFirstPage + 0x1000;
-        VirtualProtect((LPVOID)hMem.dwFirstPage, 0x1000, PAGE_EXECUTE_READWRITE, &hMem.oldProt);
+
+        if (!VirtualProtect((LPVOID)hMem.dwFirstPage, 0x1000, PAGE_EXECUTE_READWRITE, &hMem.oldProt))
+        {
+            DWORD error = GetLastError();
+            OutputDebugLine(SString("MemAccess::OpenMemWrite: VirtualProtect failed at %08x, error: %d", hMem.dwFirstPage, error));
+            hMem = {};
+            assert(!"Failed to unprotect memory");
+            return hMem;
+        }
 
         // Make sure not using this slow function too much
-        OutputDebugLine(SString("[Mem] OpenMemWrite at %08x for %d bytes  (oldProt:%04x)", pAddr, uiAmount, hMem.oldProt));
+        OutputDebugLine(SString("[MemAccess] OpenMemWrite at %08x for %d bytes  (oldProt:%04x)", pAddr, uiAmount, hMem.oldProt));
 
 #ifdef MTA_DEBUG
 #if 0   // Annoying
@@ -103,14 +117,25 @@ namespace SharedUtil
             assert( hMem.oldProt == PAGE_EXECUTE_READ || hMem.oldProt == PAGE_READONLY );
         else
             assert( hMem.oldProt == PAGE_EXECUTE_READWRITE || hMem.oldProt == PAGE_EXECUTE_WRITECOPY );
-#endif
+    #endif
 #endif
 
         // Extra if more than one page
         for (uint i = 0x1000; i < hMem.dwSize; i += 0x1000)
         {
             DWORD oldProtCheck;
-            VirtualProtect((LPVOID)(hMem.dwFirstPage + i), 0x1000, PAGE_EXECUTE_READWRITE, &oldProtCheck);
+            if (!VirtualProtect((LPVOID)(hMem.dwFirstPage + i), 0x1000, PAGE_EXECUTE_READWRITE, &oldProtCheck))
+            {
+                // Try to rollback
+                DWORD temp;
+                VirtualProtect((LPVOID)hMem.dwFirstPage, i, hMem.oldProt, &temp);
+
+                DWORD error = GetLastError();
+                OutputDebugLine(SString("[MemAccess] OpenMemWrite VirtualProtect failed at %08x, error: %d", hMem.dwFirstPage + i, error));
+                hMem = {};
+                assert(!"Failed to unprotect multi-page memory region");
+                return hMem;
+            }
             dassert(hMem.oldProt == oldProtCheck);
         }
 
@@ -122,9 +147,22 @@ namespace SharedUtil
     {
         if (hMem.dwFirstPage == 0)
             return;
+
         DWORD oldProt;
-        VirtualProtect((LPVOID)hMem.dwFirstPage, hMem.dwSize, hMem.oldProt, &oldProt);
-        dassert(oldProt == PAGE_EXECUTE_READWRITE);
+        BOOL  result = VirtualProtect((LPVOID)hMem.dwFirstPage, hMem.dwSize, hMem.oldProt, &oldProt);
+
+        if (!result)
+        {
+            DWORD error = GetLastError();
+            OutputDebugLine(SString("MemAccess::CloseMemWrite: VirtualProtect failed at %08x, size %08x, error: %d", hMem.dwFirstPage, hMem.dwSize, error));
+            assert(!"Failed to restore memory protection - critical");
+        }
+        else
+        {
+            dassert(oldProt == PAGE_EXECUTE_READWRITE);
+        }
+
+        hMem.dwFirstPage = 0;            // Invalidate handle
     }
 
 }            // namespace SharedUtil