Put a critical section around StepInto to make it thread-safe

x64dbg · Dec 11, 2021 · fb1babc · fb1babc
1 parent 1a76d61
commit fb1babc
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 79 deletions.
diff --git a/TitanEngine/Global.Debugger.cpp b/TitanEngine/Global.Debugger.cpp
@@ -40,6 +40,7 @@ LPVOID engineAttachedProcessDebugInfo = NULL;
 wchar_t szDebuggerName[512];
 bool DebugStepFinal = false;
 LPVOID StepOutCallBack = NULL;
+CRITICAL_SECTION engineStepActiveCr;
 
 // Global.Debugger.functions:
 long DebugLoopInSecondThread(LPVOID InputParameter)

diff --git a/TitanEngine/Global.Debugger.h b/TitanEngine/Global.Debugger.h
@@ -2,6 +2,7 @@
 #define _GLOBAL_DEBUGGER_H
 
 #include <vector>
+#include <Windows.h>
 
 extern HARDWARE_DATA DebugRegister[4];
 extern PROCESS_INFORMATION dbgProcessInformation;
@@ -39,6 +40,7 @@ extern LPVOID engineAttachedProcessDebugInfo;
 extern wchar_t szDebuggerName[512];
 extern bool DebugStepFinal;
 extern LPVOID StepOutCallBack;
+extern CRITICAL_SECTION engineStepActiveCr;
 
 long DebugLoopInSecondThread(LPVOID InputParameter);
 void DebuggerReset();

diff --git a/TitanEngine/TitanEngine.Debugger.Control.cpp b/TitanEngine/TitanEngine.Debugger.Control.cpp
@@ -37,30 +37,35 @@ __declspec(dllexport) void TITCALL ForceClose()
 
 __declspec(dllexport) void TITCALL StepInto(LPVOID StepCallBack)
 {
-    ULONG_PTR ueCurrentPosition = GetContextData(UE_CIP);
-    unsigned char instr[16];
-    MemoryReadSafe(dbgProcessInformation.hProcess, (void*)ueCurrentPosition, instr, sizeof(instr), 0);
-    char* DisassembledString = (char*)StaticDisassembleEx(ueCurrentPosition, (LPVOID)instr);
-    if(strstr(DisassembledString, "PUSHF"))
-        StepOver(StepCallBack);
-    else if(strstr(DisassembledString, "POP SS") || strstr(DisassembledString, "MOV SS")) //prevent the 'PUSH SS', 'POP SS' step trick
+    EnterCriticalSection(&engineStepActiveCr);
+    if (!engineStepActive)
     {
-        ueCurrentPosition += StaticLengthDisassemble((void*)instr);
-        SetBPX(ueCurrentPosition, UE_BREAKPOINT_TYPE_INT3 + UE_SINGLESHOOT, StepCallBack);
-    }
-    else
-    {
-        CONTEXT myDBGContext;
-        HANDLE hActiveThread = EngineOpenThread(THREAD_GETSETSUSPEND, false, DBGEvent.dwThreadId);
-        myDBGContext.ContextFlags = CONTEXT_CONTROL;
-        GetThreadContext(hActiveThread, &myDBGContext);
-        myDBGContext.EFlags |= UE_TRAP_FLAG;
-        SetThreadContext(hActiveThread, &myDBGContext);
-        EngineCloseHandle(hActiveThread);
-        engineStepActive = true;
-        engineStepCallBack = StepCallBack;
-        engineStepCount = 0;
+        ULONG_PTR ueCurrentPosition = GetContextData(UE_CIP);
+        unsigned char instr[16];
+        MemoryReadSafe(dbgProcessInformation.hProcess, (void*)ueCurrentPosition, instr, sizeof(instr), 0);
+        char* DisassembledString = (char*)StaticDisassembleEx(ueCurrentPosition, (LPVOID)instr);
+        if (strstr(DisassembledString, "PUSHF"))
+            StepOver(StepCallBack);
+        else if (strstr(DisassembledString, "POP SS") || strstr(DisassembledString, "MOV SS")) //prevent the 'PUSH SS', 'POP SS' step trick
+        {
+            ueCurrentPosition += StaticLengthDisassemble((void*)instr);
+            SetBPX(ueCurrentPosition, UE_BREAKPOINT_TYPE_INT3 + UE_SINGLESHOOT, StepCallBack);
+        }
+        else
+        {
+            CONTEXT myDBGContext;
+            HANDLE hActiveThread = EngineOpenThread(THREAD_GETSETSUSPEND, false, DBGEvent.dwThreadId);
+            myDBGContext.ContextFlags = CONTEXT_CONTROL;
+            GetThreadContext(hActiveThread, &myDBGContext);
+            myDBGContext.EFlags |= UE_TRAP_FLAG;
+            SetThreadContext(hActiveThread, &myDBGContext);
+            EngineCloseHandle(hActiveThread);
+            engineStepActive = true;
+            engineStepCallBack = StepCallBack;
+            engineStepCount = 0;
+        }
     }
+    LeaveCriticalSection(&engineStepActiveCr);
 }
 
 __declspec(dllexport) void TITCALL StepOver(LPVOID StepCallBack)

diff --git a/TitanEngine/TitanEngine.Debugger.DebugLoop.cpp b/TitanEngine/TitanEngine.Debugger.DebugLoop.cpp
@@ -12,6 +12,33 @@
 #define UE_MODULEx86 0x2000;
 #define UE_MODULEx64 0x2000;
 
+static void engineStep()
+{
+    EnterCriticalSection(&engineStepActiveCr);
+    if (engineStepActive)
+    {
+        DBGCode = DBG_CONTINUE;
+        if (engineStepCount == 0)
+        {
+            typedef void(TITCALL* fCustomBreakPoint)(void);
+            auto cbStep = fCustomBreakPoint(engineStepCallBack);
+            engineStepActive = false;
+            engineStepCallBack = NULL;
+            LeaveCriticalSection(&engineStepActiveCr);
+            cbStep();
+        }
+        else
+        {
+            SingleStep(engineStepCount, engineStepCallBack);
+            LeaveCriticalSection(&engineStepActiveCr);
+        }
+    }
+    else
+    {
+        LeaveCriticalSection(&engineStepActiveCr);
+    }
+}
+
 __declspec(dllexport) void TITCALL DebugLoop()
 {
     bool FirstBPX = true;
@@ -642,20 +669,7 @@ __declspec(dllexport) void TITCALL DebugLoop()
                             EnableBPX(ResetBPXAddressTo);
                             ResetBPXAddressTo = NULL;
                             ResetBPX = false;
-                            if(engineStepActive)
-                            {
-                                if(engineStepCount == 0)
-                                {
-                                    myCustomBreakPoint = (fCustomBreakPoint)(engineStepCallBack);
-                                    engineStepActive = false;
-                                    engineStepCallBack = NULL;
-                                    myCustomBreakPoint();
-                                }
-                                else
-                                {
-                                    SingleStep(engineStepCount, engineStepCallBack);
-                                }
-                            }
+                            engineStep();
                         }
                         else
                         {
@@ -671,20 +685,7 @@ __declspec(dllexport) void TITCALL DebugLoop()
                     {
                         ResetHwBPX = false;
                         SetHardwareBreakPoint(DebugRegisterX.DrxBreakAddress, DebugRegisterXId, DebugRegisterX.DrxBreakPointType, DebugRegisterX.DrxBreakPointSize, (LPVOID)DebugRegisterX.DrxCallBack);
-                        if(engineStepActive)
-                        {
-                            if(engineStepCount == 0)
-                            {
-                                myCustomBreakPoint = (fCustomBreakPoint)(engineStepCallBack);
-                                engineStepActive = false;
-                                engineStepCallBack = NULL;
-                                myCustomBreakPoint();
-                            }
-                            else
-                            {
-                                SingleStep(engineStepCount, engineStepCallBack);
-                            }
-                        }
+                        engineStep();
                     }
                     if(ResetMemBPX) //restore memory breakpoint
                     {
@@ -719,20 +720,7 @@ __declspec(dllexport) void TITCALL DebugLoop()
                             VirtualProtectEx(dbgProcessInformation.hProcess, (LPVOID)ResetMemBPXAddress, ResetMemBPXSize, NewProtect, &OldProtect);
                         }
 
-                        if(engineStepActive)
-                        {
-                            if(engineStepCount == 0)
-                            {
-                                myCustomBreakPoint = (fCustomBreakPoint)(engineStepCallBack);
-                                engineStepActive = false;
-                                engineStepCallBack = NULL;
-                                myCustomBreakPoint();
-                            }
-                            else
-                            {
-                                SingleStep(engineStepCount, engineStepCallBack);
-                            }
-                        }
+                        engineStep();
                     }
                 }
                 else //no resetting needed (debugger reached hardware breakpoint or the user stepped)
@@ -867,21 +855,7 @@ __declspec(dllexport) void TITCALL DebugLoop()
                         if(strstr(DisassembledString, "PUSHF"))
                             PushfBPX = true;
                     }
-                    if(engineStepActive)
-                    {
-                        DBGCode = DBG_CONTINUE;
-                        if(engineStepCount == 0)
-                        {
-                            myCustomBreakPoint = (fCustomBreakPoint)(engineStepCallBack);
-                            engineStepActive = false;
-                            engineStepCallBack = NULL;
-                            myCustomBreakPoint();
-                        }
-                        else
-                        {
-                            SingleStep(engineStepCount, engineStepCallBack);
-                        }
-                    }
+                    engineStep();
                 }
                 if(DBGCode == DBG_EXCEPTION_NOT_HANDLED) //NOTE: only call the chSingleStep callback when the debuggee generated the exception
                 {

diff --git a/TitanEngine/TitanEngine.cpp b/TitanEngine/TitanEngine.cpp
@@ -5,6 +5,7 @@
 #include "Global.Injector.h"
 #include "Global.Engine.Extension.h"
 #include "Global.Engine.Threading.h"
+#include "Global.Debugger.h"
 
 // Global.Engine.Entry:
 BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved)
@@ -13,6 +14,7 @@ BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved)
     {
     case DLL_PROCESS_ATTACH:
         engineHandle = hinstDLL;
+        InitializeCriticalSection(&engineStepActiveCr);
         EngineInit();
         EmptyGarbage();
         for(int i = 0; i < UE_MAX_RESERVED_MEMORY_LEFT; i++)