From 39934ea3aee3b8acb37bc56864b63a80c6068614 Mon Sep 17 00:00:00 2001
From: Duncan Ogilvie <mr.exodia.tpodt@gmail.com>
Date: Fri, 9 Sep 2022 22:56:39 +0200
Subject: [PATCH] Implement safe stepping with DBG_REPLY_LATER

---
 GleeBug/Debugger.Loop.cpp         | 84 +++++++++++++++++++++++++++++++
 GleeBug/Debugger.h                |  1 +
 TitanEngineEmulator/Emulator.h    |  3 ++
 TitanEngineEmulator/TitanEngine.h |  3 ++
 4 files changed, 91 insertions(+)

diff --git a/GleeBug/Debugger.Loop.cpp b/GleeBug/Debugger.Loop.cpp
index 41bf3ca..12010df 100644
--- a/GleeBug/Debugger.Loop.cpp
+++ b/GleeBug/Debugger.Loop.cpp
@@ -1,6 +1,10 @@
 #include "Debugger.h"
 #include "Debugger.Thread.Registers.h"
 
+#ifndef DBG_REPLY_LATER
+#define DBG_REPLY_LATER ((NTSTATUS)0x40010001L)
+#endif // DBG_REPLY_LATER
+
 namespace GleeBug
 {
     void Debugger::Start()
@@ -24,6 +28,18 @@ namespace GleeBug
             return;
         }
 
+        DWORD ThreadBeingProcessed = 0;
+        std::unordered_map<DWORD, HANDLE> SuspendedThreads;
+        bool IsDbgReplyLaterSupported = false;
+
+        // Check if DBG_REPLY_LATER is supported based on Windows version (Windows 10, version 1507 or above)
+        // https://www.gaijin.at/en/infos/windows-version-numbers
+        const uint32_t NtBuildNumber = *(uint32_t*)(0x7FFE0000 + 0x260);
+        if (NtBuildNumber != 0 && NtBuildNumber >= 10240)
+        {
+            IsDbgReplyLaterSupported = mSafeStep;
+        }
+
         while(!mBreakDebugger)
         {
             //wait for a debug event
@@ -53,6 +69,38 @@ namespace GleeBug
                     continue;
                 }
             }
+
+            // Handle safe stepping
+            if (IsDbgReplyLaterSupported)
+            {
+                if (mDebugEvent.dwDebugEventCode == EXCEPTION_DEBUG_EVENT)
+                {
+                    // Check if there is a thread processing a single step
+                    if (ThreadBeingProcessed != 0 && mDebugEvent.dwThreadId != ThreadBeingProcessed)
+                    {
+                        // Reply to the event later
+                        if (!ContinueDebugEvent(mDebugEvent.dwProcessId, mDebugEvent.dwThreadId, DBG_REPLY_LATER))
+                            break;
+
+                        // Wait for the next event
+                        continue;
+                    }
+                }
+                else if (mDebugEvent.dwDebugEventCode == EXIT_THREAD_DEBUG_EVENT)
+                {
+                    if (ThreadBeingProcessed != 0 && mDebugEvent.dwThreadId == ThreadBeingProcessed)
+                    {
+                        // Resume the other threads since the thread being processed is exiting
+                        for (auto& itr : SuspendedThreads)
+                            ResumeThread(itr.second);
+
+                        SuspendedThreads.clear();
+                        ThreadBeingProcessed = 0;
+                    }
+                }
+            }
+
+            // Signal we are currently paused
             mIsRunning = false;
 
             //set default continue status
@@ -108,6 +156,15 @@ namespace GleeBug
                 unloadDllEvent(mDebugEvent.u.UnloadDll);
                 break;
             case EXCEPTION_DEBUG_EVENT:
+                if (IsDbgReplyLaterSupported && mDebugEvent.u.Exception.ExceptionRecord.ExceptionCode == STATUS_SINGLE_STEP)
+                {
+                    // Resume the other threads since we are done processing the single step
+                    for (auto& itr : SuspendedThreads)
+                        ResumeThread(itr.second);
+
+                    SuspendedThreads.clear();
+                    ThreadBeingProcessed = 0;
+                }
                 exceptionEvent(mDebugEvent.u.Exception);
                 break;
             case OUTPUT_DEBUG_STRING_EVENT:
@@ -139,6 +196,33 @@ namespace GleeBug
                     Registers(mThread->hThread, CONTEXT_CONTROL).TrapFlag = false;
             }
 
+            // Handle safe stepping
+            if (IsDbgReplyLaterSupported && mDebugEvent.dwDebugEventCode != EXIT_THREAD_DEBUG_EVENT)
+            {
+                // If TF is set (single step), then suspend all the other threads
+                if (mThread && mThread->isInternalStepping)
+                {
+                    ThreadBeingProcessed = mDebugEvent.dwThreadId;
+
+                    for (auto& Thread : mProcess->threads)
+                    {
+                        auto dwThreadId = Thread.first;
+                        auto hThread = Thread.second->hThread;
+
+                        // Do not suspend the current thread
+                        if (ThreadBeingProcessed == dwThreadId)
+                            continue;
+
+                        // Check if the thread is already suspended
+                        if (SuspendedThreads.count(dwThreadId) != 0)
+                            continue;
+
+                        if (SuspendThread(hThread) != -1)
+                            SuspendedThreads.emplace(dwThreadId, hThread);
+                    }
+                }
+            }
+
             //continue the debug event
             if(!ContinueDebugEvent(mDebugEvent.dwProcessId, mDebugEvent.dwThreadId, mContinueStatus))
                 break;
diff --git a/GleeBug/Debugger.h b/GleeBug/Debugger.h
index 0aea4a0..f537c56 100644
--- a/GleeBug/Debugger.h
+++ b/GleeBug/Debugger.h
@@ -292,6 +292,7 @@ namespace GleeBug
         bool mDetach = false;
         bool mDetachAndBreak = false;
         bool mAttachedToProcess = false;
+        bool mSafeStep = true;
 
         /**
         \brief The current process (can be null in some cases).
diff --git a/TitanEngineEmulator/Emulator.h b/TitanEngineEmulator/Emulator.h
index 3df1729..d4ce969 100644
--- a/TitanEngineEmulator/Emulator.h
+++ b/TitanEngineEmulator/Emulator.h
@@ -191,6 +191,9 @@ public:
         case UE_ENGINE_SAFE_ATTACH:
             mSafeAttach = VariableSet;
             break;
+        case UE_ENGINE_SAFE_STEP:
+            mSafeStep = VariableSet;
+            break;
         }
     }
 
diff --git a/TitanEngineEmulator/TitanEngine.h b/TitanEngineEmulator/TitanEngine.h
index 9c0a2b4..c64fb1b 100644
--- a/TitanEngineEmulator/TitanEngine.h
+++ b/TitanEngineEmulator/TitanEngine.h
@@ -57,6 +57,9 @@
 #define UE_ENGINE_CALL_PLUGIN_DEBUG_CALLBACK 8
 #define UE_ENGINE_SET_DEBUG_PRIVILEGE 9
 #define UE_ENGINE_SAFE_ATTACH 10
+#define UE_ENGINE_MEMBP_ALT 11
+#define UE_ENGINE_DISABLE_ASLR 12
+#define UE_ENGINE_SAFE_STEP 13
 
 #define UE_OPTION_REMOVEALL 1
 #define UE_OPTION_DISABLEALL 2