From 2da21065a6037a6fde702af74c1c5d7e29c0cd7e Mon Sep 17 00:00:00 2001
From: mkrsym1 <mkrsym1@gmail.com>
Date: Sun, 7 Jan 2024 18:34:03 +0200
Subject: [PATCH] Refactored protected memory writes

---
 game_payload/include/utils.h |  2 ++
 game_payload/src/ace.c       | 10 ++--------
 game_payload/src/hsr/hsr.c   | 22 ++++++++--------------
 game_payload/src/utils.c     |  9 +++++++++
 4 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/game_payload/include/utils.h b/game_payload/include/utils.h
index edd812e..18c8536 100644
--- a/game_payload/include/utils.h
+++ b/game_payload/include/utils.h
@@ -12,3 +12,5 @@ void utils_create_parent_dirs(const wchar_t *path);
 void utils_save_to_file(const wchar_t *filePath, const void *buf, size_t length);
 
 char utils_env_enabled(const char *env);
+
+void utils_write_protected_memory(void *addr, void *buf, size_t size);
diff --git a/game_payload/src/ace.c b/game_payload/src/ace.c
index 8675f2e..ae87a5f 100644
--- a/game_payload/src/ace.c
+++ b/game_payload/src/ace.c
@@ -1,6 +1,7 @@
 #include <ntdll.h>
 #include <pe.h>
 #include <msg.h>
+#include <utils.h>
 
 #include <ace.h>
 
@@ -23,14 +24,7 @@ static void _dll_notification(ULONG reason, const PLDR_DLL_NOTIFICATION_DATA dat
         0xB8, 0x01, 0x00, 0x00, 0x00,   // mov eax, 1
         0xC3                            // ret
     };
-
-    DWORD oldProtect;
-    VirtualProtect(entryPoint, sizeof(ENTRY_POINT_STUB), PAGE_EXECUTE_READWRITE, &oldProtect);
-
-    memcpy(entryPoint, ENTRY_POINT_STUB, sizeof(ENTRY_POINT_STUB));
-
-    VirtualProtect(entryPoint, sizeof(ENTRY_POINT_STUB), oldProtect, &oldProtect);
-    
+    utils_write_protected_memory(entryPoint, ENTRY_POINT_STUB, sizeof(ENTRY_POINT_STUB));
 }
 
 static void _create_driver_file(const char *path) {
diff --git a/game_payload/src/hsr/hsr.c b/game_payload/src/hsr/hsr.c
index a669656..2273d1b 100644
--- a/game_payload/src/hsr/hsr.c
+++ b/game_payload/src/hsr/hsr.c
@@ -38,13 +38,8 @@ char *wtsud_patch_addr;
 
 static void _wtsud_stub() {
     // Recover original bytes
-    DWORD oldProtect;
-    VirtualProtect(wtsud_patch_addr, JUMP_SIZE, PAGE_EXECUTE_READWRITE, &oldProtect);
+    utils_write_protected_memory(wtsud_patch_addr, wtsud_original_bytes, JUMP_SIZE);
 
-    memcpy(wtsud_patch_addr, wtsud_original_bytes, JUMP_SIZE);
-
-    VirtualProtect(wtsud_patch_addr, JUMP_SIZE, oldProtect, &oldProtect);
-    
     unload_ctr_dec();
 }
 
@@ -59,21 +54,20 @@ static void _unityplayer_callback(HMODULE unityModule) {
 
     wtsud_patch_addr = ((char*)unityModule) + WTSUD_PATCH_OFFSET;
 
-    DWORD oldProtect;
-    VirtualProtect(wtsud_patch_addr, JUMP_SIZE, PAGE_EXECUTE_READWRITE, &oldProtect);
-
     // Save original bytes
     memcpy(wtsud_original_bytes, wtsud_patch_addr, JUMP_SIZE);
 
-    // Write jump
+    // Prepare payload
+    char payload[JUMP_SIZE];
+
     const char JUMP_INST[] = { 0xFF, 0x25, 0x00, 0x00, 0x00, 0x00 }; // jmp [$ + 6]
-    memcpy(wtsud_patch_addr, JUMP_INST, sizeof(JUMP_INST));
+    memcpy(payload, JUMP_INST, sizeof(JUMP_INST));
 
-    // Write destination address
     void *destAddr = &_wtsud_stub;
-    memcpy(wtsud_patch_addr + sizeof(JUMP_INST), &destAddr, sizeof(destAddr));
+    memcpy(payload + sizeof(JUMP_INST), &destAddr, sizeof(destAddr));
 
-    VirtualProtect(wtsud_patch_addr, JUMP_SIZE, oldProtect, &oldProtect);
+    // Write payload
+    utils_write_protected_memory(wtsud_patch_addr, payload, sizeof(payload));
 }
 
 void hsr_fill_data(struct game_data *buf) {
diff --git a/game_payload/src/utils.c b/game_payload/src/utils.c
index 9ae826e..6b41e3d 100644
--- a/game_payload/src/utils.c
+++ b/game_payload/src/utils.c
@@ -64,3 +64,12 @@ char utils_env_enabled(const char *env) {
     char *envText = getenv(env);
     return envText && *envText;
 }
+
+void utils_write_protected_memory(void *addr, void *buf, size_t size) {
+    DWORD oldProtect;
+    VirtualProtect(addr, size, PAGE_READWRITE, &oldProtect);
+
+    memcpy(addr, buf, size);
+    
+    VirtualProtect(addr, size, oldProtect, &oldProtect);
+}